Require Import Coq.ZArith.ZArith.
Require Import Coq.Strings.String.
Require Import Coq.Lists.List.
Require Import Coq.micromega.Lia.
Require Import bedrock2.Array.
Require Import bedrock2.Scalars.
Require Import bedrock2.Syntax.
Require Import bedrock2.ProgramLogic.
Require Import bedrock2.Map.Separation.
Require Import bedrock2.Map.SeparationLogic.
Require Import bedrock2.WeakestPreconditionProperties.
Require Import coqutil.Word.Interface coqutil.Word.Properties.
Require Import coqutil.Map.Interface coqutil.Map.Properties.
Require Import coqutil.Datatypes.List.
Require Import coqutil.Datatypes.PropSet.
Require Import Crypto.AbstractInterpretation.AbstractInterpretation.
Require Import Crypto.Bedrock.Field.Common.Types.
Require Import Crypto.Bedrock.Field.Common.Tactics.
Require Import Crypto.Bedrock.Field.Common.Util.
Require Import Crypto.Bedrock.Field.Translation.Proofs.Equivalence.
Require Import Crypto.Bedrock.Field.Translation.Proofs.EquivalenceProperties.
Require Import Crypto.Bedrock.Field.Translation.Proofs.Flatten.
Require Import Crypto.Bedrock.Field.Translation.Proofs.UsedVarnames.
Require Import Crypto.Bedrock.Field.Translation.Proofs.VarnameSet.
Require Import Crypto.Bedrock.Field.Translation.Flatten.
Require Import Crypto.Bedrock.Field.Translation.LoadStoreList.
Require Import Crypto.Language.API.
Require Import Crypto.Util.ListUtil.
Require Import Crypto.Util.ZUtil.Tactics.RewriteModSmall.
Import ListNotations.

Import AbstractInterpretation.Compilers.
Import API.Compilers.
Import Types.Notations.

Section LoadStoreList.
  Context {p : parameters} {p_ok : @ok p}.

  Local Existing Instance rep.Z.
  Local Instance sem_ok : Semantics.parameters_ok semantics
    := semantics_ok.

  (* states that access sizes are within machine integer width *)
  Fixpoint base_access_sizes_good {t}
    : base_access_sizes (listZ:=rep.listZ_mem) t -> Prop :=
    match t with
    | base.type.prod a b =>
      fun s => base_access_sizes_good (fst s)
               /\ base_access_sizes_good (snd s)
    | base_listZ =>
      fun s =>
        (Z.of_nat (Memory.bytes_per
                     (width:=Semantics.width) s) * 8 <= Semantics.width)%Z
    | _ => fun _ => True
    end.
  Definition access_sizes_good {t}
    : access_sizes (listZ:=rep.listZ_mem) t -> Prop :=
    match t with
    | type.base _ => base_access_sizes_good
    | _ => fun _ => True
    end.
  Fixpoint access_sizes_good_args {t}
    : type.for_each_lhs_of_arrow
        (access_sizes (listZ:=rep.listZ_mem)) t -> Prop :=
    match t with
    | type.base _ => fun _ : unit => True
    | type.arrow _ _ =>
      fun s => access_sizes_good (fst s) /\ access_sizes_good_args (snd s)
    end.

  Lemma access_sizes_good_in_range size :
    base_access_sizes_good (t:=base_listZ) size ->
    (2 ^ (Z.of_nat (Memory.bytes_per
                      (width:=Semantics.width) size) * 8)
     <= 2 ^ Semantics.width)%Z.
  Proof.
    cbv [base_access_sizes_good]; intros.
    apply Z.pow_le_mono_r; lia.
  Qed.

  (* states that each list Z is within its access size *)
  Fixpoint within_base_access_sizes {t}
    : base.interp t ->
      base_access_sizes (listZ:=rep.listZ_mem) t -> Prop :=
      match t with
      | base.type.prod a b =>
        fun x s =>
          within_base_access_sizes (fst x) (fst s) /\
          within_base_access_sizes (snd x) (snd s)
      | base_listZ =>
        fun x s =>
          let bytes :=
              Z.of_nat (Memory.bytes_per (width:=Semantics.width) s) in
          Forall (fun z => 0 <= z < 2 ^ (bytes * 8))%Z x
      | _ => fun _ _ => True
      end.

  Fixpoint within_access_sizes_args {t}
    : type.for_each_lhs_of_arrow API.interp_type t ->
      type.for_each_lhs_of_arrow
        (access_sizes (listZ:=rep.listZ_mem)) t -> Prop :=
    match t with
    | type.base _ => fun _ _ => True
    | type.arrow (type.base s) _ =>
      fun x s =>
        within_base_access_sizes (fst x) (fst s) /\
        within_access_sizes_args (snd x) (snd s)
    | _ => fun _ _ => False
    end.

  (* separation-logic relation that says space exists in memory for lists
     (other values are ignored) *)
  Fixpoint lists_reserved {t}
    : base_listonly nat t ->
      base_rtype (listZ:=rep.listZ_mem) t ->
      base_access_sizes (listZ:=rep.listZ_mem) t ->
      Semantics.locals ->
      Semantics.mem ->
      Prop :=
    match t as t0 return
          base_listonly nat t0 -> base_rtype t0 ->
          base_access_sizes (listZ:=rep.listZ_mem) t0 -> _ with
    | base.type.prod a b =>
      fun x y s locals =>
        sep (lists_reserved (fst x) (fst y) (fst s) locals)
            (lists_reserved (snd x) (snd y) (snd s) locals)
    | base_listZ =>
      fun n loc sz locals =>
        (Lift1Prop.ex1
           (fun oldvalues : list Z =>
              sep (emp (length oldvalues = n))
                  (rep.equiv (rep:=rep.listZ_mem)
                              oldvalues loc sz locals)))
    | base_Z => fun _ _ _ _ => emp True
    |  _ => fun _ _ _ _ => emp False
    end.

  Fixpoint list_lengths_from_value {t}
    : base.interp t -> base_listonly nat t :=
    match t as t0 return base.interp t0 -> base_listonly nat t0 with
    | base.type.prod a b =>
      fun x : base.interp a * base.interp b =>
        (list_lengths_from_value (fst x),
         list_lengths_from_value (snd x))
    | base_listZ => fun x : list Z => length x
    | _ => fun _ => tt
    end.

  Fixpoint list_lengths_from_args {t}
    : type.for_each_lhs_of_arrow API.interp_type t ->
      type.for_each_lhs_of_arrow list_lengths t :=
    match t with
    | type.base b => fun _ => tt
    | type.arrow (type.base a) b =>
      fun x =>
        (list_lengths_from_value (fst x), list_lengths_from_args (snd x))
    | type.arrow a b =>
      fun x => (tt, list_lengths_from_args (snd x))
    end.

  Definition lists_reserved_with_initial_context {t}
             (lengths : base_listonly nat (type.final_codomain t))
             (argnames : type.for_each_lhs_of_arrow
                           (ltype (listZ:=rep.listZ_mem)) t)
             (retnames : base_ltype (type.final_codomain t))
             (retsizes : base_access_sizes (type.final_codomain t))
             (flat_args : list Semantics.word)
    : Semantics.mem -> Prop :=
    let outptrs :=
        flatten_listonly_base_ltype (fst (extract_listnames retnames)) in
    match map.of_list_zip
            (outptrs ++ (flatten_argnames argnames))
            flat_args with
    | Some init_locals =>
      lists_reserved
        lengths (base_rtype_of_ltype retnames) retsizes init_locals
    | None => emp True
    end.

  Lemma lists_reserved_0 {listZ:rep.rep base_listZ} start sizes locals:
    Lift1Prop.iff1
      (lists_reserved (t:=base_listZ) 0%nat start sizes locals)
      (emp (exists x, locally_equivalent_base (t:=base_Z) x start locals)).
  Proof.
    cbv [lists_reserved].
    split; intros;
      repeat match goal with
             | _ => progress (sepsimpl; subst)
             | _ => progress cbn [locally_equivalent_base
                                    equivalent_base rep.equiv rep.Z] in *
             | H : Datatypes.length _ = 0%nat |- _ =>
               apply length0_nil in H; subst
             | H : rep.equiv _ _ _ _ _ |- _ => apply equiv_nil_iff1 in H
             | |- Lift1Prop.ex1 _ _ => eexists
             | |- exists _, _ => eexists
             | |- _ /\ _ => split
             | |- rep.equiv [] _ _ _ _ => apply equiv_nil_iff1
             | _ => solve [eauto using List.length_nil]
             | _ => eexists; split; solve [eauto]
             end.
  Qed.

  Lemma lists_reserved_only_differ_undef t :
    forall lengths locs s locals locals' vset,
      map.only_differ locals vset locals' ->
      map.undef_on locals vset ->
      Lift1Prop.impl1
        (lists_reserved (t:=t) lengths locs s locals)
        (lists_reserved lengths locs s locals').
  Proof.
    induction t; cbn [lists_reserved];
      break_match; intros; try reflexivity; [ | ].
    { rewrite IHt1, IHt2; try reflexivity; eauto. }
    { repeat intro; sepsimpl; eexists; sepsimpl; eauto.
      eapply equiv_listZ_only_differ_undef_mem; eauto. }
  Qed.

  Lemma equivalent_extract_listnames {t} :
    forall (x : base.interp t)
           (names : base_ltype t)
           (locals : Semantics.locals)
           (sizes : base_access_sizes t)
           (R : Semantics.mem -> Prop),
    Lift1Prop.iff1
      (sep (equivalent_base x (base_rtype_of_ltype names) sizes locals) R)
      (sep (equivalent_listexcl
              x (map_listexcl (@base_rtype_of_ltype _ rep.listZ_mem)
                              (snd (extract_listnames names))) sizes locals)
           (sep (equivalent_listonly
                   x (map_listonly
                        (base_rtype_of_ltype)
                        (fst (extract_listnames names))) sizes locals) R)).
  Proof.
    induction t;
      cbn [fst snd equivalent_base extract_listnames
               base_rtype_of_ltype rep.rtype_of_ltype
               equivalent_listexcl map_listexcl
               equivalent_listonly map_listonly];
      break_match; intros;
        try match goal with
              |- context [emp False] =>
              split; intros; sepsimpl
            end; [ solve [cancel] | | ].
    { rewrite sep_assoc. (* needed for rewrites for some reason *)
      rewrite IHt1. rewrite IHt2.
      rewrite !sep_assoc.
      cancel.
      (* TODO : this should go by associativity/commutativity; is there a better
      way than manually selecting indices? *)
      cancel_seps_at_indices 0 0; [ reflexivity | ].
      cancel_seps_at_indices 0 1; [ reflexivity | ].
      cancel_seps_at_indices 0 0; [ reflexivity | ].
      reflexivity. }
    { cancel. reflexivity. }
  Qed.

  Lemma load_list_item_correct
        (size : access_size)
        (name : base_ltype (listZ:=rep.listZ_mem) base_listZ)
        i l :
    forall mem locals R,
      (i < length l)%nat ->
      base_access_sizes_good (t:=base_listZ) size ->
      sep (rep.equiv (t:=base_listZ)
                     l (base_rtype_of_ltype name) size locals) R mem ->
      WeakestPrecondition.dexpr
        mem locals
        (load_list_item size (expr.var name) i)
        (word.of_Z (hd 0%Z (skipn i l))).
  Proof.
    cbv [load_list_item];
      cbn [rep.equiv rep.listZ_mem rep.Z base_rtype_of_ltype
                     rep.rtype_of_ltype rep.listZ_mem ]; intros.
    cbn [WeakestPrecondition.dexpr
           WeakestPrecondition.expr WeakestPrecondition.expr_body] in *.
    sepsimpl; subst.
    match goal with
    | H : context[array] |- _ =>
      eapply Proper_sep_iff1 in H;
        [ | symmetry; apply array_index_nat_inbounds
                        with (n:=i) (default:=0%Z);
            rewrite !map_length in *; lia
          | reflexivity]
    end.
    match goal with
    | H : context[array] |- _ =>
      rewrite !word.ring_morph_mul, !word.of_Z_unsigned in H;
        rewrite <-!word.ring_morph_mul in H
    end.
    straightline.
    eapply Proper_get; [ repeat intro |  eassumption ].
    match goal with H : word.unsigned _ = word.unsigned _ |- _ =>
                    apply word.unsigned_inj in H
    end.
    subst. rewrite Nat2Z.inj_mul.
    eexists; split.
    { cbv [Memory.load].
      erewrite load_Z_of_sep; [ reflexivity | ].
      ecancel_assumption. }
    { unfold truncate_Z.
      rewrite Z.land_ones by auto with zarith.
      rewrite Z.mod_small; [ reflexivity | ].
      rewrite <-hd_skipn_nth_default.
      match goal with H : base_access_sizes_good _ |- _ =>
                      apply access_sizes_good_in_range in H end.
      apply Forall_nth_default; auto with zarith. }
  Qed.

  Lemma load_list_correct rem l :
    forall (i nextn : nat)
           (name : base_ltype (listZ:=rep.listZ_mem) base_listZ)
           (tr : Semantics.trace)
           (sizes : base_access_sizes
                      (listZ:=rep.listZ_mem) base_listZ)
           (locals : Semantics.locals)
           (mem : Semantics.mem)
           (functions : list _)
           (R : Semantics.mem -> Prop),
      (forall (n : nat) (H : nextn <= n), name <> varname_gen n) ->
      base_access_sizes_good sizes ->
      length l = i + rem ->
      sep (rep.equiv (t:=base_listZ)
                     l (base_rtype_of_ltype name) sizes locals) R mem ->
      (* load_list returns # vars used, variable names, cmd *)
      let out := load_list sizes (expr.var name) i rem nextn in
      let nvars := fst (fst out) in
      let names' : base_ltype (listZ:=rep.listZ_local) base_listZ
          := snd (fst out) in
      WeakestPrecondition.cmd
        (WeakestPrecondition.call functions)
        (snd out)
        tr mem locals
        (fun tr' mem' locals' =>
           tr = tr' /\
           mem = mem' /\
           Interface.map.only_differ
             locals (used_varnames nextn nvars) locals' /\
           (forall n,
               (nextn + nvars <= n)%nat ->
               ~ varname_set_base names' (varname_gen n)) /\
           locally_equivalent_base
             (t:=base_listZ) (skipn i l)
             (base_rtype_of_ltype names') locals').
  Proof.
    induction rem; cbn [fst snd load_list]; intros.
    { straightline.
      cbv [locally_equivalent_base].
      cbn [rep.Z rep.listZ_local fold_right map
                 equivalent_base rep.equiv
                 base_rtype_of_ltype rep.rtype_of_ltype
                 varname_set rep.varname_set].
        repeat split.
      { eapply only_differ_zero. }
      { intros. cbv [empty_set]. tauto. }
      { rewrite skipn_all by lia. econstructor. } }
    { cbn [WeakestPrecondition.cmd WeakestPrecondition.cmd_body].
      eexists; split; cbv [dlet.dlet].
      { eapply load_list_item_correct; try eassumption.
        cbn [Datatypes.length]; lia. }
      eapply Proper_cmd; [ solve [apply Proper_call] | | ].
      2:{
        eapply IHrem; eauto with lia.
        eapply Proper_sep_iff1; [ | reflexivity | eassumption ].
        eapply equiv_listZ_only_differ_mem_iff1;
          eauto using only_differ_put.
        cbn [varname_set rep.varname_set rep.listZ_mem rep.Z].
        apply disjoint_singleton_singleton; eauto using string_dec. }
      repeat intro.
      cbv beta in *; cleanup; subst.
      repeat match goal with |- _ /\ _ => split end;
        eauto using only_differ_succ.
      { cbn [varname_set_base
               rep.varname_set rep.listZ_local rep.Z fold_right] in *.
        intros.
        apply not_union_iff; split; eauto with lia.
        cbv [singleton_set].
        rewrite varname_gen_unique; lia. }
      { cbv [locally_equivalent] in *.
        cbn [map base_rtype_of_ltype
                 Compilers.base_interp
                 rep.rtype_of_ltype rep.listZ_local
                 equivalent rep.equiv ] in *.
        rewrite skipn_nth_default with (d:=0%Z) by lia.
        eapply Forall2_cons; eauto.
        eapply (equiv_Z_only_differ_iff1 (listZ:=rep.listZ_mem)); eauto.
        { eauto using only_differ_sym, only_differ_put. }
        { eapply disjoint_used_varnames_singleton.
          lia. }
        { assert (2 ^ (Z.of_nat (Memory.bytes_per (width:=Semantics.width) sizes) * 8)
                  <= 2 ^ Semantics.width)%Z
            by (cbv [base_access_sizes_good] in *;
                apply Z.pow_le_mono_r; lia).
          cbn [rep.equiv rep.listZ_mem rep.Z] in *.
          sepsimpl_hyps.
          match goal with H : word.unsigned _ = word.unsigned _ |- _ =>
                          apply word.unsigned_inj in H
          end.
          subst. eexists.
          split;[reflexivity|].
          split.
          { rewrite <-word.unsigned_of_Z_0, map_nth_default_always.
            reflexivity. }
          { eexists; split; [ | reflexivity ].
            rewrite map.get_put_same, hd_skipn_nth_default.
            rewrite <-word.unsigned_of_Z_0.
            rewrite skipn_map, hd_map.
            rewrite !word.of_Z_unsigned.
            reflexivity. } } } }
  Qed.

  Lemma load_all_lists_correct {t} :
    forall (argnames : base_ltype t)
           (argsizes : base_access_sizes t)
           (args : base.interp t)
           (functions : list _)
           (tr : Semantics.trace)
           (locals : Semantics.locals)
           (mem : Semantics.mem)
           (nextn : nat)
           (R : Semantics.mem -> Prop),
        (* lengths of any lists in the arguments *)
        let arglengths := list_lengths_from_value args in
        (* look up variables in argnames *)
        let argvalues := base_rtype_of_ltype argnames in
        (* argnames don't contain variables we could later overwrite *)
        (forall n,
            (nextn <= n)%nat ->
            ~ varname_set_base argnames (varname_gen n)) ->
        (* argument sizes are <= integer width *)
        base_access_sizes_good argsizes ->
        (* argument values are equivalent *)
        sep (equivalent_base args argvalues argsizes locals) R mem ->
        (* load_all_lists returns triple : # fresh variables used,
           new argnames with local lists, and cmd *)
        let out := load_all_lists nextn argnames arglengths argsizes in
        let nvars := fst (fst out) in
        (* extract loaded argument values *)
        let argnames' := snd (fst out) in
        let argvalues' := base_rtype_of_ltype argnames' in
        (* translated function produces equivalent results *)
        WeakestPrecondition.cmd
          (WeakestPrecondition.call functions)
          (snd out)
          tr mem locals
          (fun tr' mem' locals' =>
             tr = tr' /\
             mem = mem' /\
             Interface.map.only_differ
               locals (used_varnames nextn nvars) locals' /\
             (forall n,
                 (nextn + nvars <= n)%nat ->
                 ~ varname_set_base argnames' (varname_gen n)) /\
             locally_equivalent_base args argvalues' locals').
  Proof.
    (* TODO: lots of repeated steps in this proof; automate *)
    induction t;
      cbn [fst snd map load_all_lists varname_set
               locally_equivalent_base equivalent_base rep.equiv
               flatten_base_ltype equivalent_flat_base
               base_access_sizes_good
               base_rtype_of_ltype rep.varname_set rep.Z
          ]; break_match; intros; sepsimpl; [ | | ].
    { (* base_Z *)
      repeat straightline. cbv [emp].
      repeat match goal with
               | |- Lift1Prop.ex1 _ _ => eexists
               | |- _ /\ _ => split end;
        eauto using only_differ_zero with lia. }
    { (* product *)
      fold (@base_access_sizes p rep.listZ_mem).
      straightline.
      eapply Proper_cmd; [ solve [apply Proper_call] | | ].
      2:{ eapply IHt1; eauto.
          { intros.
            match goal with
            | H : _ |- _ =>
              setoid_rewrite not_union_iff in H;
                apply H; eauto
            end. }
          { cbn [base_rtype_of_ltype fst snd] in *.
            ecancel_assumption. } }
      repeat intro. cleanup; subst.
      eapply Proper_cmd; [ solve [apply Proper_call] | | ].
      2:{ eapply IHt2; eauto.
          { intros.
            match goal with
            | H : _ |- _ =>
              setoid_rewrite not_union_iff in H;
                apply H; lia
            end. }
          { cbn [base_rtype_of_ltype fst snd] in *.
            eapply Proper_sep_iff1; [ | reflexivity | ].
            { eapply equivalent_only_differ_iff1;
                eauto using equiv_listZ_only_differ_mem, only_differ_sym.
              eapply disjoint_used_varnames_lt.
              match goal with
              | H : _ |- _ =>
                setoid_rewrite not_union_iff in H;
                  apply H; eauto
              end. }
            ecancel_assumption. } }
      cbn [list_lengths_from_value fst snd] in *.
      repeat intro. cleanup; subst.
      repeat match goal with |- _ /\ _ => split end;
        eauto using only_differ_step.
      { intros; apply not_union_iff; split;
          match goal with H : _ |- _ => apply H; lia end. }
      { cbv [locally_equivalent locally_equivalent_args] in *.
        cbn [equivalent_args fst rtype_of_ltype].
        apply sep_empty_iff; split; eauto.
        eapply equivalent_only_differ_iff1;
          eauto using equiv_listZ_only_differ_local, only_differ_sym.
        intros.
        eapply disjoint_used_varnames_lt.
        cleanup. subst.
        match goal with H : _ |- _ => apply H; lia end. } }
  { (* base_listZ *)
    clear IHt.
     cbn [flatten_base_rtype base_rtype_of_ltype] in *.
    eapply Proper_cmd;
      [ solve [apply Proper_call]
      | | eapply load_list_correct; eauto; reflexivity ].
    repeat intro. cleanup; subst.
    repeat match goal with |- _ /\ _ => split end;
      eauto using only_differ_step. }
  Qed.

  Lemma load_arguments_correct {t} :
    forall (argnames : type.for_each_lhs_of_arrow ltype t)
           (argsizes : type.for_each_lhs_of_arrow access_sizes t)
           (args : type.for_each_lhs_of_arrow API.interp_type t)
           (functions : list _)
           (tr : Semantics.trace)
           (locals : Semantics.locals)
           (mem : Semantics.mem)
           (nextn : nat),
        (* lengths of any lists in the arguments *)
        let arglengths := list_lengths_from_args args in
        (* look up variables in argnames *)
        let argvalues :=
            type.map_for_each_lhs_of_arrow rtype_of_ltype argnames in
        (* argnames don't contain variables we could later overwrite *)
        (forall n,
            (nextn <= n)%nat ->
            ~ varname_set_args argnames (varname_gen n)) ->
        (* argument sizes are <= width *)
        access_sizes_good_args argsizes ->
        (* argument values are equivalent *)
        equivalent_args args argvalues argsizes locals mem ->
        (* load_arguments returns triple : # fresh variables used,
           new argnames with local lists, and cmd *)
        let out := load_arguments nextn argnames arglengths argsizes in
        let nvars := fst (fst out) in
        (* extract loaded argument values *)
        let argnames' := snd (fst out) in
        let argvalues' :=
            type.map_for_each_lhs_of_arrow rtype_of_ltype argnames' in
        (* translated function produces equivalent results *)
        WeakestPrecondition.cmd
          (WeakestPrecondition.call functions)
          (snd out)
          tr mem locals
          (fun tr' mem' locals' =>
             tr = tr' /\
             mem = mem' /\
             Interface.map.only_differ
               locals (used_varnames nextn nvars) locals' /\
             (forall n,
                 (nextn + nvars <= n)%nat ->
                 ~ varname_set_args argnames' (varname_gen n)) /\
             locally_equivalent_args args argvalues' locals').
  Proof.
    cbv zeta.
    induction t;
      destruct argnames;
      cbn [fst snd load_arguments access_sizes_good_args
               list_lengths_from_args varname_set_args
               type.for_each_lhs_of_arrow equivalent_args
               type.map_for_each_lhs_of_arrow flatten_args
               ]; break_match; cbn [fst snd]; intros.
    { (* base type *)
      repeat straightline.
      repeat match goal with
             | |- _ /\ _ => split end;
        eauto using only_differ_zero; reflexivity. }
    { (* arrow (base _) _ *)
      cleanup. straightline.
      eapply Proper_cmd.
      3:{
        eapply load_all_lists_correct; eauto; [ ].
        intros.
        match goal with
        | H : _ |- _ =>
          setoid_rewrite not_union_iff in H;
            apply H; eauto
        end. }
      { apply Proper_call. }
      { repeat intro; cleanup; subst.
        eapply Proper_cmd; [ solve [apply Proper_call] | | ].
        2:{
          eapply IHt2; eauto.
        { intros.
          match goal with
            | H : _ |- _ =>
          setoid_rewrite not_union_iff in H;
            apply H; lia
          end. }
        { eapply equivalent_args_only_differ;
            eauto using only_differ_sym with equiv; [ ].
          intros.
          eapply disjoint_used_varnames_lt.
          cleanup. subst.
          match goal with
            | H : _ |- _ =>
          setoid_rewrite not_union_iff in H;
            apply H; eauto
          end. } }
        { repeat intro. cleanup. subst.
          repeat match goal with |- _ /\ _ => split end;
            eauto using only_differ_step.
          { intros; apply not_union_iff; split;
              match goal with H : _ |- _ => apply H; lia end. }
          { cbv [locally_equivalent locally_equivalent_args] in *.
            cbn [equivalent_args fst rtype_of_ltype].
            split; [ | solve [eauto] ].
            exists (emp True). sepsimpl; eauto.
            eapply equivalent_only_differ;
              eauto using only_differ_sym with equiv; [ ].
            eapply disjoint_used_varnames_lt.
            intros; cleanup. subst.
            eauto with lia. } } } }
    { (* arrow (arrow _ _) _ *)
      sepsimpl. }
  Qed.

  Lemma store_list_correct :
    forall (size : access_size)
           (start : string)
           (functions : list _)
           (value_names2 value_names1 : list string)
           (rets1 rets2 rets : base.interp base_listZ)
           (i : nat)
           (tr : Semantics.trace)
           (locals : Semantics.locals)
           (mem : Semantics.mem)
           (R : Semantics.mem -> Prop),
      let retlengths :=
          list_lengths_from_value (t:=base_listZ) rets2 in
      let loc := expr.var start in
      let nbytes := Memory.bytes_per (width:=Semantics.width) size in
      let offset := expr.literal (Z.of_nat (nbytes * i)) in
      let loc' := expr.op bopname.add loc offset in
      let values1 := map expr.var value_names1 in
      let values2 := map expr.var value_names2 in
      rets = rets1 ++ rets2 ->
      length rets1 = length value_names1 ->
      length rets2 = length value_names2 ->
      length value_names1 = i ->
      (* bounds are OK *)
      Forall (fun x =>
                0 <= x < 2 ^ (Z.of_nat nbytes * 8))%Z rets2 ->
      (* yet-to-be-stored values *)
      locally_equivalent_base
        (listZ:=rep.listZ_local) rets2 values2 locals ->
      base_access_sizes_good (t:=base_listZ) size ->
      (* already-stored values *)
      sep (map:=Semantics.mem)
          (sep (map:=Semantics.mem)
               (rep.equiv (rep:=rep.listZ_mem)
                          rets1 (expr.var start) size locals)
               (lists_reserved retlengths loc' size locals)) R mem ->
      WeakestPrecondition.cmd
        (WeakestPrecondition.call functions)
        (store_list size (expr.var start) values2 i) tr mem locals
        (fun tr' mem' locals' =>
           tr = tr' /\
           locals = locals' /\
           sep (map:=Semantics.mem)
               (rep.equiv (rep:=rep.listZ_mem)
                          rets (expr.var start) size locals') R mem').
  Proof.
    cbv zeta.
    induction value_names2 as [|n0 ?]; destruct rets2 as [|r0 rets2];
      cbn [store_list Datatypes.length base_access_sizes_good
                      locally_equivalent_base equivalent_base
                      rep.equiv rep.listZ_local rep.Z rep.listZ_mem];
      intros; subst; try lia; sepsimpl.
    { repeat split; try reflexivity.
      repeat match goal with
             | _ => progress sepsimpl
             | H : _ |- _ =>
               rewrite List.firstn_all, List.skipn_all_exact in H
             | _ => rewrite app_nil_r
             | |- Lift1Prop.ex1 _ _ => eexists
             | |- _ /\ _ => split
             | _ => eassumption
             end.
      match goal with
        H : sep _ _ _ |- _ =>
        eapply Proper_sep_iff1 in H;
          [ | rewrite ?(lists_reserved_0 (listZ:=rep.listZ_local));
              reflexivity .. ]
      end.
      sepsimpl.
      ecancel_assumption. }
    { repeat match goal with
             | H : _ |- _ => rewrite word.of_Z_unsigned in H end.
      subst; eexists; split.
      { cbn [WeakestPrecondition.dexpr
               WeakestPrecondition.expr
               WeakestPrecondition.expr_body] in *.
        eapply Proper_get; [ | eassumption ].
        repeat intro; subst.
        reflexivity. }
      cbn [map] in *.
      match goal with
      | H : Forall2 _ (_ :: _) (_ :: _) |- _ => inversion H; subst; clear H
      end.
      sepsimpl.
      eexists; split.
      { cbn [WeakestPrecondition.dexpr
               WeakestPrecondition.expr
               WeakestPrecondition.expr_body] in *.
        eassumption. }
      cbn [lists_reserved
             list_lengths_from_value
             equivalent
             rep.equiv rep.Z rep.listZ_mem] in *.
      sepsimpl.
      match goal with
        H : Datatypes.length ?x = Datatypes.length (_ :: _) |- _ =>
        destruct x; cbn [Datatypes.length] in H; [ lia | ]
      end.
      cbn [map array Semantics.interp_binop] in *.

      (* we now have a Z in context that should be equivalent to the offsetted
         location; destruct WeakestPrecondition.dexpr to expose that equivalence *)
      match goal with
        H : WeakestPrecondition.dexpr _ _ (expr.op _ _ _) ?v |- _ =>
          cbn [WeakestPrecondition.dexpr
                 WeakestPrecondition.expr
                 WeakestPrecondition.expr_body
                 Semantics.interp_binop] in H;
          cbv [dlet.dlet WeakestPrecondition.literal
                         WeakestPrecondition.get] in H;
          rewrite ?word.of_Z_unsigned in H;
          cleanup; subst
      end.
      match goal with
      | H1 : (WeakestPrecondition.dexpr _ ?l (expr.var ?x) ?v1),
             H2 : (map.get ?l ?x = Some ?v2) |- _ =>
        replace v2 with v1 in *
          by (cbn [WeakestPrecondition.dexpr
                     WeakestPrecondition.expr
                     WeakestPrecondition.expr_body] in H1;
              cbv [WeakestPrecondition.get] in H1;
              cleanup; congruence)
      end.
      repeat match goal with
             | H: word.unsigned _ = word.unsigned _ |- _ =>
               apply word.unsigned_inj in H; subst
             end.
      repeat match goal with
             | H : map word.unsigned ?x = _ :: _ |- _ =>
               destruct x; cbn [map] in H;
                 [ congruence | inversion H; clear H ]
             end.
      cbn [map array] in *.
      eapply store_Z_of_sep; [ ecancel_assumption | ].
      intros. sepsimpl.
      repeat match goal with
             | H : Forall _ (_ :: _) |- _ => inversion H; subst; clear H
             end.

      eapply Proper_cmd; [ solve [apply Proper_call] | repeat intro | ].
      2:{
        let s :=
            lazymatch goal with
              H : map.get locals start = Some ?x |- _ => x end in
        let H := lazymatch goal with
                 | H : sep _ _ ?m |- context [?m] => H end in
        let r1 :=
            lazymatch type of H with
                context [array _ _ s ?xs] => xs end in
        let r2 :=
            lazymatch type of H with
                context [array _ _ (word.add _ _) ?xs] => xs end in
        let r :=
            lazymatch type of H with
              context [truncated_scalar _ _ ?x] => x end in
        apply IHvalue_names2 with
            (rets1:= r1 ++ [r])
            (rets2:=rets2) (R:=R)
            (value_names1:=value_names1 ++ [n0]);
        rewrite ?app_length; cbn [length]; eauto; try lia; [ ].
        { intros.
          (* look at goal and plug in evars very carefully *)
          sepsimpl.
          match goal with H: ?P (expr.var start) ?x |- _ =>
                          exists x end.
          sepsimpl; [ solve [eauto using Forall_snoc] .. | ].
          match goal with
            |- context [map word.unsigned _ =
                        map word.unsigned ?xs ++ [word.unsigned ?x] ] =>
            exists (xs ++ [x])
          end.
          sepsimpl; [ rewrite map_app; reflexivity
                    | solve [eauto using Forall_snoc] | ].
          eexists; sepsimpl; [ reflexivity | solve [eauto] | ].
          match goal with
            H : sep _ _ ?m |- context [?m] =>
            match type of H with
              context [array _ _ (word.add ?a ?b) ?x] =>
              exists x;
                sepsimpl; [solve [auto] .. | ];
                  exists (word.add a b)
            end
          end.
          sepsimpl; [ ].
          eexists; sepsimpl; [ reflexivity | solve [eauto] | ].
          eexists; sepsimpl; [ reflexivity | | ].
          { eexists; split; [ eassumption | ].
            cbn [WeakestPrecondition.dexpr
                   WeakestPrecondition.expr
                   WeakestPrecondition.expr_body
                   Semantics.interp_binop].
            cbv [WeakestPrecondition.literal dlet.dlet].
            match goal with
            | |- context[word.add ?x (word.of_Z ?y)] =>
              progress match x with
                       | word.of_Z _ => idtac
                       | word.add _ _ => idtac
                       | _ => rewrite <-(word.of_Z_unsigned x)
                       end
            end.
            rewrite <-!word.ring_morph_add.
            f_equal. clear; lia. }
          { match goal with
              H : context [array _ _ (word.add (word.add _ _) _) _] |- _ =>
              refine (Lift1Prop.subrelation_iff1_impl1 _ _ _ _ _ H); clear H
            end.
            rewrite map_app, array_append. cbn [map array].
            cbn [Compilers.base_interp] in *. (* simplifies implicit types *)
            match goal with
            | |- Lift1Prop.iff1 ?L ?R =>
              let la := match L with
                          context [truncated_scalar _ ?la] => la end in
              let ra := match R with
                          context [truncated_scalar _ ?ra] => ra end in
              replace ra with la
                by (apply word.unsigned_inj;
                    rewrite !word.unsigned_add, !word.unsigned_of_Z;
                    cbv [word.wrap]; Modulo.pull_Zmod; f_equal; lia)
            end.
            cancel. } } }
      cbv beta in *; cleanup; subst.
      repeat match goal with |- _ /\ _ => split end;
        eauto.
      rewrite app_cons_app_app.
      eassumption. }
  Qed.

  Lemma store_return_values_correct {t} init_locals :
    forall (retnames_local : base_ltype t)
           (retnames_mem : base_ltype t)
           (retsizes : base_access_sizes t)
           (vset : set string)
           (rets : base.interp t)
           (functions : list _)
           (tr : Semantics.trace)
           (locals : Semantics.locals)
           (mem : Semantics.mem)
           (R : Semantics.mem -> Prop),
      let out := store_return_values retnames_local retnames_mem retsizes in
      let retlengths := list_lengths_from_value rets in
      sep (lists_reserved retlengths
                          (base_rtype_of_ltype retnames_mem)
                          retsizes init_locals) R mem ->
      map.undef_on
        init_locals
        (union vset (varname_set_listexcl retnames_mem)) ->
      map.only_differ init_locals vset locals ->
      (* rets are stored in local retnames *)
      locally_equivalent_base
        rets (base_rtype_of_ltype retnames_local) locals ->
      (* retnames are disjoint *)
      disjoint (varname_set_base retnames_local)
               (varname_set_base retnames_mem) ->
      (* retnames don't contain duplicates *)
      NoDup (flatten_base_ltype retnames_mem) ->
      (* rets bounds obey access sizes *)
      within_base_access_sizes rets retsizes ->
      (* access sizes not larger than integer width *)
      base_access_sizes_good retsizes ->
      (* translated function produces equivalent results *)
      WeakestPrecondition.cmd
        (WeakestPrecondition.call functions)
        (snd out)
        tr mem locals
        (fun tr' mem' locals' =>
           tr = tr' /\
           retlengths = fst out /\
           map.only_differ
             locals (varname_set_listexcl retnames_mem) locals' /\
           sep (equivalent_base
                  rets (base_rtype_of_ltype retnames_mem) retsizes locals')
               R mem').
  Proof.
    cbv zeta.
    induction t;
      cbn [fst snd store_return_values lists_reserved
               locally_equivalent_base equivalent_base
               varname_set_base varname_set_listexcl
               base_rtype_of_ltype rep.rtype_of_ltype
               rep.equiv rep.varname_set rep.Z
               base_access_sizes_good within_base_access_sizes
               flatten_base_ltype list_lengths_from_value
          ]; break_match; intros; sepsimpl; [ | | ].
    { (* base_Z *)
      repeat straightline.
      eexists; split; [eassumption|].
      cbv [dlet.dlet]; split; [reflexivity|].
      split; [reflexivity|].
      split;
        [ solve [eauto using only_differ_sym, only_differ_put] | ].
      sepsimpl. eexists; sepsimpl; try eassumption; [ ].
      cbv [WeakestPrecondition.dexpr
             WeakestPrecondition.expr
             WeakestPrecondition.expr_body
             WeakestPrecondition.get] in *.
      cleanup; subst.
      rewrite map.get_put_same.
      eexists; split; reflexivity. }
    { (* prod *)
      repeat straightline.
      match goal with H : _ |- _ =>
                      pose proof H;
                      apply NoDup_app_iff in H; cleanup end.
      repeat match goal with
             | _ => progress cleanup
             | H : disjoint (union _ _) _ |- _ =>
               apply disjoint_union_l_iff in H
             | H : disjoint _ (union _ _) |- _ =>
               apply disjoint_union_r_iff in H
             | H : map.undef_on _ (union _ _) |- _ =>
               apply undef_on_union_iff in H
             end.
      eapply Proper_cmd;
        [ solve [apply Proper_call] | repeat intro
          | eapply IHt1; try ecancel_assumption;
            try (apply undef_on_union_iff; split);
            solve [eauto] ].
      cbv beta in *. cleanup; subst.
      eapply Proper_cmd;
        [ solve [apply Proper_call] | repeat intro | ].
      2:{ eapply IHt2 with
              (vset:=union vset (varname_set_listexcl (fst retnames_mem)));
          try ecancel_assumption; eauto.
          { apply undef_on_union_iff; split; eauto.
            apply undef_on_union_iff; split; eauto. }
          { rewrite union_comm.
            eapply only_differ_trans; eauto. }
          { cbv [locally_equivalent].
            eapply equivalent_only_differ; eauto with equiv.
            apply disjoint_sym.
            eapply subset_disjoint_r;
              eauto using varname_set_listexcl_subset. } }
      cbv beta in *. cleanup; subst.
      repeat match goal with H : list_lengths_from_value _ = _ |- _ =>
                             rewrite H end.
      repeat match goal with |- _ /\ _ => split end;
        eauto using only_differ_sym, only_differ_trans.
      use_sep_assumption. cancel.
      (* TODO: why doesn't cancel handle the below? *)
      cancel_seps_at_indices 0 1; [ reflexivity | ].
      eapply equivalent_only_differ_iff1; eauto with equiv.
      eapply disjoint_sym.
      eapply subset_disjoint_r;
        eauto using varname_set_listexcl_subset; [ ].
      rewrite !varname_set_flatten.
      eapply NoDup_disjoint; eauto using string_dec. }
    { (* base_listZ *)
      repeat straightline.
      repeat match goal with p := _ |- _ => subst p end.
      repeat match goal with
             | H : map.undef_on _ (union _ _) |- _ =>
               apply undef_on_union_iff in H
             end.
      cbn [rep.Z rep.listZ_mem rep.equiv rep.rtype_of_ltype] in *.
      sepsimpl.
      eapply Proper_cmd;
        [ solve [apply Proper_call] | repeat intro | ].
      2:{
        eapply store_list_correct with
            (rets:=rets) (rets1:=[]) (rets2:=rets)
            (value_names1:=nil); auto using List.length_nil; [ | ].
        { cbn [rep.equiv rep.rtype_of_ltype rep.listZ_local] in *.
          eapply Forall.eq_length_Forall2, Forall.Forall2_map_r_iff.
          eauto. }
        { apply sep_assoc.
          eapply Proper_sep_iff1;
            [ apply equiv_nil_iff1 | reflexivity | ].
          sepsimpl. eexists.
          apply sep_ex1_l.
          eexists; sepsimpl;
            [ | eapply expr_only_differ_undef; solve [eauto] | ];
            [ reflexivity | ].
          cbn [lists_reserved equivalent rep.equiv rep.listZ_mem rep.Z].
          sepsimpl. eexists.
          sepsimpl; eauto.
          eexists; sepsimpl; try ecancel_assumption; eauto.
          cbn [WeakestPrecondition.dexpr
                 WeakestPrecondition.expr
                 WeakestPrecondition.expr_body] in *.
          eexists; sepsimpl; try eassumption; [ ].
          eexists; sepsimpl; try eassumption; [ ].
          eapply WP_get_only_differ_undef; eauto.
          eapply Proper_get; [ repeat intro | eassumption ].
          cbn [Semantics.interp_binop].
          cbv [WeakestPrecondition.literal dlet.dlet].
          subst.
          apply word.unsigned_inj.
          rewrite word.unsigned_add, word.unsigned_of_Z.
          rewrite Nat2Z.inj_mul. change (Z.of_nat 0) with 0%Z.
          autorewrite with zsimplify_fast.
          cbv [word.wrap]. rewrite word.wrap_unsigned.
          reflexivity. } }
      cbv beta in *. cleanup; subst.
      cbn [rep.rtype_of_ltype rep.listZ_local rep.Z] in *.
      match goal with
      | H : rep.equiv ?x ?y _ _ _ |- _ =>
        assert  (length x = length y)
        by (eapply Forall.eq_length_Forall2; apply H)
      end.
      rewrite !map_length in *.
      repeat match goal with |- _ /\ _ => split end;
        eauto using only_differ_sym, only_differ_put,
        only_differ_empty. }
  Qed.
End LoadStoreList.


Inductive MS : Set :=
  | X : MS -> MS
  | Y : MS -> MS.

Type (fun p : MS => match p return nat with
                    | X x => 0
                    end).


Require Import Coqlib.
Require Import Asm.
Require Import Integers.
Require Import PeekTactics.
Require Import PeepsLib.
Require Import PregTactics.
Require Import StepIn.
Require Import AsmBits.
Require Import Values.
Require Import ValEq.
Require Import Integers.
Require Import PeepsTactics.

Definition peep_extended_shift_example :=
  Psal_ri EAX Int.zero ::
          Pshr_ri EAX Int.zero ::
          Por_rr EAX EAX :: nil.

Section EXTENDED_SHIFT.

  Variable concrete : code.
  Variable r1 r2 : ireg.
  Variable c1 c2 : int.
  (* tweak this to maybe also include range facts  *)
  Hypothesis eq : Int.add c1 c2 = Int.repr (32).
  Variable r1_r2_neq : r1 <> r2.

  Definition peep_extended_shift_defs : rewrite_defs :=
    {|
      fnd :=
        Psal_ri r1 c1 ::
                Pshr_ri r2 c2 ::
                Por_rr r1 r2 :: nil
      ; rpl :=
          Pshld_ri r1 r2 c1 :: Pnop :: Pnop :: nil
      ; lv_in := PC :: IR r1 :: IR r2 :: nil
      ; lv_out := PC :: IR r1 :: nil
      ; clobbered := IR r2 :: flags
    |}.

Lemma peep_extended_shift_selr :
  StepEquiv.step_through_equiv_live (fnd peep_extended_shift_defs) (rpl peep_extended_shift_defs) (lv_in peep_extended_shift_defs) (lv_out peep_extended_shift_defs).
Proof.
    prep_l.
    step_l.
    step_l.
    step_l.
    prep_r.
    step_r.
    step_r.
    step_r.
    finish_r.
    prep_eq.
    split.
    2: eq_mem_tac.
    intros.
    simpl_and_clear.
    repeat break_or_reg.
    - (*PC*)
      preg_simpl.
      repeat find_rewrite_goal.
      simpl.
      reflexivity.
    - preg_simpl.
      eapply val_eq_vor.
      subst_max.
      preg_simpl.
      eapply val_eq_shl.
      assumption.
      simpl. auto.
      subst_max.
      preg_simpl.
      eapply val_eq_shru.
      assumption.
      simpl.
      f_equal.
      clear -eq.
      unfold Int.iwordsize in *.
      unfold Int.zwordsize in *.
      unfold Int.wordsize in *.
      unfold Wordsize_32.wordsize in *.
      simpl.
      rewrite <- eq.
      ring.
  Qed.

  Definition peep_extended_shift_proofs : rewrite_proofs :=
    {|
      defs := peep_extended_shift_defs
      ; selr := peep_extended_shift_selr
    |}.

  Definition peep_extended_shift :
    concrete = fnd peep_extended_shift_defs ->
    StepEquiv.rewrite.
  Proof.
    intros.
    peep_tac_mk_rewrite peep_extended_shift_defs peep_extended_shift_proofs.
  Qed.

End EXTENDED_SHIFT.

Definition peep_extended_shift_rewrite (c : code) : option StepEquiv.rewrite.
  name peep_extended_shift p.
  unfold peep_extended_shift_defs in p.
  simpl in p.
  specialize (p c).
  do 3 set_code_cons c.
  set_code_nil c.
  set_instr_eq i 0%nat peep_extended_shift_example.
  set_instr_eq i0 1%nat peep_extended_shift_example.
  set_instr_eq i1 2%nat peep_extended_shift_example.
  set_ireg_eq rd1 rd.
  set_ireg_eq rd0 r1.
  set_ireg_neq rd r1.
  destruct (Int.eq (Int.add n n0) (Int.repr 32)) eqn:?.
  2: exact None.
  apply PtrEquiv.int_eq_true in Heqb.
  specialize (p _ _ _ _ Heqb n1 eq_refl). exact (Some p).
Defined.

Definition extended_shift (c : code) : list StepEquiv.rewrite :=
  collect (map peep_extended_shift_rewrite (ParamSplit.matched_pat peep_extended_shift_example c)).


From Hammer Require Import Hammer.

Require Import DecidableTypeEx.
Require Export MSetInterface.
Set Implicit Arguments.
Unset Strict Implicit.

Module WFactsOn (Import E : DecidableType)(Import M : WSetsOn E).

Notation eq_dec := E.eq_dec.
Definition eqb x y := if eq_dec x y then true else false.

Section ImplSpec.
Variable s s' : t.
Variable x y : elt.

Lemma In_1 : E.eq x y -> In x s -> In y s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.In_1".   intros E; rewrite E; auto. Qed.

Lemma mem_1 : In x s -> mem x s = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.mem_1".   intros; apply <- mem_spec; auto. Qed.
Lemma mem_2 : mem x s = true -> In x s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.mem_2".   intros; apply -> mem_spec; auto. Qed.

Lemma equal_1 : Equal s s' -> equal s s' = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.equal_1".   intros; apply <- equal_spec; auto. Qed.
Lemma equal_2 : equal s s' = true -> Equal s s'.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.equal_2".   intros; apply -> equal_spec; auto. Qed.

Lemma subset_1 : Subset s s' -> subset s s' = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.subset_1".   intros; apply <- subset_spec; auto. Qed.
Lemma subset_2 : subset s s' = true -> Subset s s'.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.subset_2".   intros; apply -> subset_spec; auto. Qed.

Lemma is_empty_1 : Empty s -> is_empty s = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.is_empty_1".   intros; apply <- is_empty_spec; auto. Qed.
Lemma is_empty_2 : is_empty s = true -> Empty s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.is_empty_2".   intros; apply -> is_empty_spec; auto. Qed.

Lemma add_1 : E.eq x y -> In y (add x s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_1".   intros; apply <- add_spec. auto with relations. Qed.
Lemma add_2 : In y s -> In y (add x s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_2".   intros; apply <- add_spec; auto. Qed.
Lemma add_3 : ~ E.eq x y -> In y (add x s) -> In y s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_3".   rewrite add_spec. intros H [H'|H']; auto. elim H; auto with relations. Qed.

Lemma remove_1 : E.eq x y -> ~ In y (remove x s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_1".   intros; rewrite remove_spec; intuition. Qed.
Lemma remove_2 : ~ E.eq x y -> In y s -> In y (remove x s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_2".   intros; apply <- remove_spec; auto with relations. Qed.
Lemma remove_3 : In y (remove x s) -> In y s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_3".   rewrite remove_spec; intuition. Qed.

Lemma singleton_1 : In y (singleton x) -> E.eq x y.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.singleton_1".   rewrite singleton_spec; auto with relations. Qed.
Lemma singleton_2 : E.eq x y -> In y (singleton x).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.singleton_2".   rewrite singleton_spec; auto with relations. Qed.

Lemma union_1 : In x (union s s') -> In x s \/ In x s'.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.union_1".   rewrite union_spec; auto. Qed.
Lemma union_2 : In x s -> In x (union s s').
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.union_2".   rewrite union_spec; auto. Qed.
Lemma union_3 : In x s' -> In x (union s s').
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.union_3".   rewrite union_spec; auto. Qed.

Lemma inter_1 : In x (inter s s') -> In x s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.inter_1".   rewrite inter_spec; intuition. Qed.
Lemma inter_2 : In x (inter s s') -> In x s'.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.inter_2".   rewrite inter_spec; intuition. Qed.
Lemma inter_3 : In x s -> In x s' -> In x (inter s s').
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.inter_3".   rewrite inter_spec; intuition. Qed.

Lemma diff_1 : In x (diff s s') -> In x s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.diff_1".   rewrite diff_spec; intuition. Qed.
Lemma diff_2 : In x (diff s s') -> ~ In x s'.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.diff_2".   rewrite diff_spec; intuition. Qed.
Lemma diff_3 : In x s -> ~ In x s' -> In x (diff s s').
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.diff_3".   rewrite diff_spec; auto. Qed.

Variable f : elt -> bool.
Notation compatb := (Proper (E.eq==>Logic.eq)) (only parsing).

Lemma filter_1 : compatb f -> In x (filter f s) -> In x s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.filter_1".   intros P; rewrite filter_spec; intuition. Qed.
Lemma filter_2 : compatb f -> In x (filter f s) -> f x = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.filter_2".   intros P; rewrite filter_spec; intuition. Qed.
Lemma filter_3 : compatb f -> In x s -> f x = true -> In x (filter f s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.filter_3".   intros P; rewrite filter_spec; intuition. Qed.

Lemma for_all_1 : compatb f ->
For_all (fun x => f x = true) s -> for_all f s = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.for_all_1".   intros; apply <- for_all_spec; auto. Qed.
Lemma for_all_2 : compatb f ->
for_all f s = true -> For_all (fun x => f x = true) s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.for_all_2".   intros; apply -> for_all_spec; auto. Qed.

Lemma exists_1 : compatb f ->
Exists (fun x => f x = true) s -> exists_ f s = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.exists_1".   intros; apply <- exists_spec; auto. Qed.

Lemma exists_2 : compatb f ->
exists_ f s = true -> Exists (fun x => f x = true) s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.exists_2".   intros; apply -> exists_spec; auto. Qed.

Lemma elements_1 : In x s -> InA E.eq x (elements s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.elements_1".   intros; apply <- elements_spec1; auto. Qed.
Lemma elements_2 : InA E.eq x (elements s) -> In x s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.elements_2".   intros; apply -> elements_spec1; auto. Qed.

End ImplSpec.

Notation empty_1 := empty_spec (only parsing).
Notation fold_1 := fold_spec (only parsing).
Notation cardinal_1 := cardinal_spec (only parsing).
Notation partition_1 := partition_spec1 (only parsing).
Notation partition_2 := partition_spec2 (only parsing).
Notation choose_1 := choose_spec1 (only parsing).
Notation choose_2 := choose_spec2 (only parsing).
Notation elements_3w := elements_spec2w (only parsing).

Hint Resolve mem_1 equal_1 subset_1 empty_1
is_empty_1 choose_1 choose_2 add_1 add_2 remove_1
remove_2 singleton_2 union_1 union_2 union_3
inter_3 diff_3 fold_1 filter_3 for_all_1 exists_1
partition_1 partition_2 elements_1 elements_3w
: set.
Hint Immediate In_1 mem_2 equal_2 subset_2 is_empty_2 add_3
remove_3 singleton_1 inter_1 inter_2 diff_1 diff_2
filter_1 filter_2 for_all_2 exists_2 elements_2
: set.

Section IffSpec.
Variable s s' s'' : t.
Variable x y z : elt.

Lemma In_eq_iff : E.eq x y -> (In x s <-> In y s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.In_eq_iff".
intros E; rewrite E; intuition.
Qed.

Lemma mem_iff : In x s <-> mem x s = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.mem_iff".   apply iff_sym, mem_spec. Qed.

Lemma not_mem_iff : ~In x s <-> mem x s = false.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.not_mem_iff".
rewrite <-mem_spec; destruct (mem x s); intuition.
Qed.

Lemma equal_iff : s[=]s' <-> equal s s' = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.equal_iff".   apply iff_sym, equal_spec. Qed.

Lemma subset_iff : s[<=]s' <-> subset s s' = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.subset_iff".   apply iff_sym, subset_spec. Qed.

Lemma empty_iff : In x empty <-> False.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.empty_iff".   intuition; apply (empty_spec H). Qed.

Lemma is_empty_iff : Empty s <-> is_empty s = true.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.is_empty_iff".   apply iff_sym, is_empty_spec. Qed.

Lemma singleton_iff : In y (singleton x) <-> E.eq x y.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.singleton_iff".   rewrite singleton_spec; intuition. Qed.

Lemma add_iff : In y (add x s) <-> E.eq x y \/ In y s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_iff".   rewrite add_spec; intuition. Qed.

Lemma add_neq_iff : ~ E.eq x y -> (In y (add x s)  <-> In y s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_neq_iff".   rewrite add_spec; intuition. elim H; auto with relations. Qed.

Lemma remove_iff : In y (remove x s) <-> In y s /\ ~E.eq x y.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_iff".   rewrite remove_spec; intuition. Qed.

Lemma remove_neq_iff : ~ E.eq x y -> (In y (remove x s) <-> In y s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_neq_iff".   rewrite remove_spec; intuition. Qed.

Variable f : elt -> bool.

Lemma for_all_iff : Proper (E.eq==>Logic.eq) f ->
(For_all (fun x => f x = true) s <-> for_all f s = true).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.for_all_iff".   intros; apply iff_sym, for_all_spec; auto. Qed.

Lemma exists_iff : Proper (E.eq==>Logic.eq) f ->
(Exists (fun x => f x = true) s <-> exists_ f s = true).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.exists_iff".   intros; apply iff_sym, exists_spec; auto. Qed.

Lemma elements_iff : In x s <-> InA E.eq x (elements s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.elements_iff".   apply iff_sym, elements_spec1. Qed.

End IffSpec.

Notation union_iff := union_spec (only parsing).
Notation inter_iff := inter_spec (only parsing).
Notation diff_iff := diff_spec (only parsing).
Notation filter_iff := filter_spec (only parsing).

Ltac set_iff :=
repeat (progress (
rewrite add_iff || rewrite remove_iff || rewrite singleton_iff
|| rewrite union_iff || rewrite inter_iff || rewrite diff_iff
|| rewrite empty_iff)).

Section BoolSpec.
Variable s s' s'' : t.
Variable x y z : elt.

Lemma mem_b : E.eq x y -> mem x s = mem y s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.mem_b".
intros.
generalize (mem_iff s x) (mem_iff s y)(In_eq_iff s H).
destruct (mem x s); destruct (mem y s); intuition.
Qed.

Lemma empty_b : mem y empty = false.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.empty_b".
generalize (empty_iff y)(mem_iff empty y).
destruct (mem y empty); intuition.
Qed.

Lemma add_b : mem y (add x s) = eqb x y || mem y s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_b".
generalize (mem_iff (add x s) y)(mem_iff s y)(add_iff s x y); unfold eqb.
destruct (eq_dec x y); destruct (mem y s); destruct (mem y (add x s)); intuition.
Qed.

Lemma add_neq_b : ~ E.eq x y -> mem y (add x s) = mem y s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_neq_b".
intros; generalize (mem_iff (add x s) y)(mem_iff s y)(add_neq_iff s H).
destruct (mem y s); destruct (mem y (add x s)); intuition.
Qed.

Lemma remove_b : mem y (remove x s) = mem y s && negb (eqb x y).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_b".
generalize (mem_iff (remove x s) y)(mem_iff s y)(remove_iff s x y); unfold eqb.
destruct (eq_dec x y); destruct (mem y s); destruct (mem y (remove x s)); simpl; intuition.
Qed.

Lemma remove_neq_b : ~ E.eq x y -> mem y (remove x s) = mem y s.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_neq_b".
intros; generalize (mem_iff (remove x s) y)(mem_iff s y)(remove_neq_iff s H).
destruct (mem y s); destruct (mem y (remove x s)); intuition.
Qed.

Lemma singleton_b : mem y (singleton x) = eqb x y.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.singleton_b".
generalize (mem_iff (singleton x) y)(singleton_iff x y); unfold eqb.
destruct (eq_dec x y); destruct (mem y (singleton x)); intuition.
Qed.

Lemma union_b : mem x (union s s') = mem x s || mem x s'.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.union_b".
generalize (mem_iff (union s s') x)(mem_iff s x)(mem_iff s' x)(union_iff s s' x).
destruct (mem x s); destruct (mem x s'); destruct (mem x (union s s')); intuition.
Qed.

Lemma inter_b : mem x (inter s s') = mem x s && mem x s'.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.inter_b".
generalize (mem_iff (inter s s') x)(mem_iff s x)(mem_iff s' x)(inter_iff s s' x).
destruct (mem x s); destruct (mem x s'); destruct (mem x (inter s s')); intuition.
Qed.

Lemma diff_b : mem x (diff s s') = mem x s && negb (mem x s').
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.diff_b".
generalize (mem_iff (diff s s') x)(mem_iff s x)(mem_iff s' x)(diff_iff s s' x).
destruct (mem x s); destruct (mem x s'); destruct (mem x (diff s s')); simpl; intuition.
Qed.

Lemma elements_b : mem x s = existsb (eqb x) (elements s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.elements_b".
generalize (mem_iff s x)(elements_iff s x)(existsb_exists (eqb x) (elements s)).
rewrite InA_alt.
destruct (mem x s); destruct (existsb (eqb x) (elements s)); auto; intros.
symmetry.
rewrite H1.
destruct H0 as (H0,_).
destruct H0 as (a,(Ha1,Ha2)); [ intuition |].
exists a; intuition.
unfold eqb; destruct (eq_dec x a); auto.
rewrite <- H.
rewrite H0.
destruct H1 as (H1,_).
destruct H1 as (a,(Ha1,Ha2)); [intuition|].
exists a; intuition.
unfold eqb in *; destruct (eq_dec x a); auto; discriminate.
Qed.

Variable f : elt->bool.

Lemma filter_b : Proper (E.eq==>Logic.eq) f -> mem x (filter f s) = mem x s && f x.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.filter_b".
intros.
generalize (mem_iff (filter f s) x)(mem_iff s x)(filter_iff s x H).
destruct (mem x s); destruct (mem x (filter f s)); destruct (f x); simpl; intuition.
Qed.

Lemma for_all_b : Proper (E.eq==>Logic.eq) f ->
for_all f s = forallb f (elements s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.for_all_b".
intros.
generalize (forallb_forall f (elements s))(for_all_iff s H)(elements_iff s).
unfold For_all.
destruct (forallb f (elements s)); destruct (for_all f s); auto; intros.
rewrite <- H1; intros.
destruct H0 as (H0,_).
rewrite (H2 x0) in H3.
rewrite (InA_alt E.eq x0 (elements s)) in H3.
destruct H3 as (a,(Ha1,Ha2)).
rewrite (H _ _ Ha1).
apply H0; auto.
symmetry.
rewrite H0; intros.
destruct H1 as (_,H1).
apply H1; auto.
rewrite H2.
rewrite InA_alt. exists x0; split; auto with relations.
Qed.

Lemma exists_b : Proper (E.eq==>Logic.eq) f ->
exists_ f s = existsb f (elements s).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.exists_b".
intros.
generalize (existsb_exists f (elements s))(exists_iff s H)(elements_iff s).
unfold Exists.
destruct (existsb f (elements s)); destruct (exists_ f s); auto; intros.
rewrite <- H1; intros.
destruct H0 as (H0,_).
destruct H0 as (a,(Ha1,Ha2)); auto.
exists a; split; auto.
rewrite H2; rewrite InA_alt; exists a; auto with relations.
symmetry.
rewrite H0.
destruct H1 as (_,H1).
destruct H1 as (a,(Ha1,Ha2)); auto.
rewrite (H2 a) in Ha1.
rewrite (InA_alt E.eq a (elements s)) in Ha1.
destruct Ha1 as (b,(Hb1,Hb2)).
exists b; auto.
rewrite <- (H _ _ Hb1); auto.
Qed.

End BoolSpec.

Instance In_m : Proper (E.eq==>Equal==>iff) In.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.In_m".
unfold Equal; intros x y H s s' H0.
rewrite (In_eq_iff s H); auto.
Qed.

Instance Empty_m : Proper (Equal==>iff) Empty.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.Empty_m".
repeat red; unfold Empty; intros s s' E.
setoid_rewrite E; auto.
Qed.

Instance is_empty_m : Proper (Equal==>Logic.eq) is_empty.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.is_empty_m".
intros s s' H.
generalize (is_empty_iff s). rewrite H at 1. rewrite is_empty_iff.
destruct (is_empty s); destruct (is_empty s'); intuition.
Qed.

Instance mem_m : Proper (E.eq==>Equal==>Logic.eq) mem.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.mem_m".
intros x x' Hx s s' Hs.
generalize (mem_iff s x). rewrite Hs, Hx at 1; rewrite mem_iff.
destruct (mem x s), (mem x' s'); intuition.
Qed.

Instance singleton_m : Proper (E.eq==>Equal) singleton.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.singleton_m".
intros x y H a. rewrite !singleton_iff, H; intuition.
Qed.

Instance add_m : Proper (E.eq==>Equal==>Equal) add.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_m".
intros x x' Hx s s' Hs a. rewrite !add_iff, Hx, Hs; intuition.
Qed.

Instance remove_m : Proper (E.eq==>Equal==>Equal) remove.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_m".
intros x x' Hx s s' Hs a. rewrite !remove_iff, Hx, Hs; intuition.
Qed.

Instance union_m : Proper (Equal==>Equal==>Equal) union.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.union_m".
intros s1 s1' Hs1 s2 s2' Hs2 a. rewrite !union_iff, Hs1, Hs2; intuition.
Qed.

Instance inter_m : Proper (Equal==>Equal==>Equal) inter.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.inter_m".
intros s1 s1' Hs1 s2 s2' Hs2 a. rewrite !inter_iff, Hs1, Hs2; intuition.
Qed.

Instance diff_m : Proper (Equal==>Equal==>Equal) diff.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.diff_m".
intros s1 s1' Hs1 s2 s2' Hs2 a. rewrite !diff_iff, Hs1, Hs2; intuition.
Qed.

Instance Subset_m : Proper (Equal==>Equal==>iff) Subset.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.Subset_m".
unfold Equal, Subset; firstorder.
Qed.

Instance subset_m : Proper (Equal==>Equal==>Logic.eq) subset.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.subset_m".
intros s1 s1' Hs1 s2 s2' Hs2.
generalize (subset_iff s1 s2). rewrite Hs1, Hs2 at 1. rewrite subset_iff.
destruct (subset s1 s2); destruct (subset s1' s2'); intuition.
Qed.

Instance equal_m : Proper (Equal==>Equal==>Logic.eq) equal.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.equal_m".
intros s1 s1' Hs1 s2 s2' Hs2.
generalize (equal_iff s1 s2). rewrite Hs1,Hs2 at 1. rewrite equal_iff.
destruct (equal s1 s2); destruct (equal s1' s2'); intuition.
Qed.

Instance SubsetSetoid : PreOrder Subset.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.SubsetSetoid".   firstorder. Qed.

Definition Subset_refl := @PreOrder_Reflexive _ _ SubsetSetoid.
Definition Subset_trans := @PreOrder_Transitive _ _ SubsetSetoid.

Instance In_s_m : Morphisms.Proper (E.eq ==> Subset ++> impl) In | 1.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.In_s_m".
simpl_relation. eauto with set.
Qed.

Instance Empty_s_m : Proper (Subset-->impl) Empty.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.Empty_s_m".   firstorder. Qed.

Instance add_s_m : Proper (E.eq==>Subset++>Subset) add.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.add_s_m".
intros x x' Hx s s' Hs a. rewrite !add_iff, Hx; intuition.
Qed.

Instance remove_s_m : Proper (E.eq==>Subset++>Subset) remove.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.remove_s_m".
intros x x' Hx s s' Hs a. rewrite !remove_iff, Hx; intuition.
Qed.

Instance union_s_m : Proper (Subset++>Subset++>Subset) union.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.union_s_m".
intros s1 s1' Hs1 s2 s2' Hs2 a. rewrite !union_iff, Hs1, Hs2; intuition.
Qed.

Instance inter_s_m : Proper (Subset++>Subset++>Subset) inter.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.inter_s_m".
intros s1 s1' Hs1 s2 s2' Hs2 a. rewrite !inter_iff, Hs1, Hs2; intuition.
Qed.

Instance diff_s_m : Proper (Subset++>Subset-->Subset) diff.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.diff_s_m".
intros s1 s1' Hs1 s2 s2' Hs2 a. rewrite !diff_iff, Hs1, Hs2; intuition.
Qed.

Generalizable Variables f.

Instance filter_equal : forall `(Proper _ (E.eq==>Logic.eq) f),
Proper (Equal==>Equal) (filter f).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.filter_equal".
intros f Hf s s' Hs a. rewrite !filter_iff, Hs by auto; intuition.
Qed.

Instance filter_subset : forall `(Proper _ (E.eq==>Logic.eq) f),
Proper (Subset==>Subset) (filter f).
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.filter_subset".
intros f Hf s s' Hs a. rewrite !filter_iff, Hs by auto; intuition.
Qed.

Lemma filter_ext : forall f f', Proper (E.eq==>Logic.eq) f -> (forall x, f x = f' x) ->
forall s s', s[=]s' -> filter f s [=] filter f' s'.
Proof. hammer_hook "MSetFacts" "MSetFacts.WFactsOn.filter_ext".
intros f f' Hf Hff' s s' Hss' x. rewrite 2 filter_iff; auto.
rewrite Hff', Hss'; intuition.
red; red; intros; rewrite <- 2 Hff'; auto.
Qed.

End WFactsOn.

Module WFacts (M:WSets) := WFactsOn M.E M.
Module Facts := WFacts.


Require Export PredMonad.Unordered.Monad.

(***
 *** Monads with State Effects
 ***)

Section MonadState.
  Context (S : Type) {R_S} `{LR_S:@LR S R_S} (M : Type -> Type) `{MonadOps M}.

  (* State effects = get and put *)
  Class MonadStateOps : Type :=
  { getM : M S
  ; putM : S -> M unit
  }.

  Class MonadState `{MonadStateOps} : Prop :=
  {
    monad_state_monad :> Monad M;
    monad_proper_get :> Proper lr_leq getM;
    monad_proper_put :> Proper lr_leq putM;
    monad_state_get :
      forall {A} `{LR A} (m : M A),
        Proper lr_leq m ->
        bindM getM (fun _ => m) ~~ m ;

    monad_state_get_put :
      forall {A} `{LR A} (f : S -> unit -> M A),
        Proper lr_leq f ->
        bindM getM (fun s => bindM (putM s) (f s)) ~~ bindM getM (fun s => f s tt) ;

    monad_state_put_get :
      forall {A} `{LR A} s (f : unit -> S -> M A),
        Proper lr_leq s -> Proper lr_leq f ->
        bindM (putM s) (fun u => bindM getM (f u))
              ~~ bindM (putM s) (fun u => f u s) ;

    monad_state_put_put :
      forall {A} `{LR A} s1 s2 (f : unit -> unit -> M A),
        Proper lr_leq s1 -> Proper lr_leq s2 -> Proper lr_leq f ->
        bindM (putM s1) (fun u => bindM (putM s2) (f u)) ~~ bindM (putM s2) (f tt)
  }.

End MonadState.

(***
 *** The State Monad Transformer
 ***)

Section StateT.

Context (S:Type) {R_S:LR_Op S} `{@LR S R_S} (M:Type -> Type) `{Monad M}.

Definition StateT (X:Type) := S -> M (S * X).

Global Instance StateT_MonadOps : MonadOps StateT :=
  {returnM :=
     fun A x => fun s => returnM (s, x);
   bindM :=
     fun A B m f =>
       fun s => do s_x <- m s; f (snd s_x) (fst s_x);
   lrM := fun {A} _ => LR_Op_fun }.

(* The Monad instance for StateT *)
Global Instance StateT_Monad : Monad (StateT).
Proof.
  constructor; intros; unfold StateT, returnM, bindM, lrM, StateT_MonadOps.
  - auto with typeclass_instances.
  - prove_lr_proper.
  - prove_lr_proper.
  - intros R1 R2 subR; apply LRFun_Proper_subrelation; apply monad_proper_lrM;
      apply LRPair_Proper_subrelation; try assumption; reflexivity.
  - prove_lr.
  - transitivity (fun s => do s_x <- m s; returnM s_x).
    + split; build_lr_fun; apply monad_proper_bind; prove_lr.
    + prove_lr; autorewrite with LR; prove_lr.
  - prove_lr.
Qed.

Global Instance StateT_MonadStateOps : MonadStateOps S StateT :=
  { getM := fun s => returnM (s, s)
  ; putM := fun s _ => returnM (s, tt)
  }.

Global Instance StateT_MonadState : MonadState S StateT.
Proof.
  constructor; intros;
    unfold StateT, returnM, bindM, lrM, getM, putM,
    StateT_MonadOps, StateT_MonadStateOps.
  - auto with typeclass_instances.
  - prove_lr_proper.
  - prove_lr_proper.
  - prove_lr.
  - prove_lr.
  - prove_lr.
  - prove_lr.
Qed.

End StateT.


Ltac exfalso := elimtype False. (* Coq 8.3 feature. *)

Lemma if_cc:
  forall A B Z W (f:Z->W) (a:{A}+{B}) (b c : Z),
    f (if a then b else c) = if a then f b else f c.
Proof.
 intros.
 destruct a; auto.
Qed.

Lemma if_irrelevant:
  forall A B Z (a:{A}+{B}) (b : Z),
    (if a then b else b) = b.
Proof.
 intros.
 destruct a; auto.
Qed.

Definition fmap A B (f:A->B) x :=
  match x with
  | None => None
  | Some x => Some (f x)
  end.

(* Not list-related. Relocate. *)
Definition uncurry A B C (f:A->B->C) (xy:A*B) :=
    let (x, y) return C := xy in f x y.

Definition idy A (x:A) := x.

(* Not list-related. Relocate. *)
Lemma inhabits_cut A B (f : A -> B) : inhabited A -> inhabited B.
 intros H.
 destruct H as [X].
 exact (inhabits (f X)).
Qed.