Require Import Coq.ZArith.ZArith Coq.micromega.Lia.
Require Import Coq.Lists.List. Import ListNotations.
Require Import Rewriter.Language.Wf. Import Language.Wf.Compilers.
Require Import Rewriter.Language.Inversion. Import Language.Inversion.Compilers.
Require Import Crypto.Language.API. Import Language.API.Compilers.
Import Compilers.API.
Require Import Crypto.CastLemmas.
Require Import Crypto.Fancy.Spec.
Require Import Crypto.Util.Option.
Require Import Crypto.Util.Prod.
Require Import Crypto.Util.Tactics.BreakMatch.
Require Import Crypto.Util.Tactics.DestructHead.
Require Import Crypto.Util.Tuple.
Require Import Crypto.Util.ZRange.
Require Import Crypto.Util.ZUtil.Zselect.
Require Import Crypto.Util.ZUtil.Tactics.LtbToLt.
Require Import Crypto.Util.ZUtil.AddModulo.
Require Import Crypto.Util.ZUtil.Definitions.
Require Import Crypto.Util.ZUtil.Div.
Require Import Crypto.Util.ZUtil.Notations.
Require Import Crypto.Util.Notations.
Local Open Scope Z_scope.

(* TODO: organize this file *)
Section of_prefancy.
  (* We would like to write [Local Notation cexpr := (@API.expr).] but
     we can't, in part because implicit status doesn't work through
     nested notations, and in part because
     COQBUG(https://github.com/coq/coq/issues/10803) *)
  (* Set Printing All. Print API.expr. *)
  Local Notation cexpr := (@expr.expr (Language.Compilers.base.type.type Compilers.base) IdentifiersBasicGENERATED.Compilers.ident).
  Local Notation tZ := (base.type.type_base base.type.Z).
  Local Notation tzrange := (base.type.type_base base.type.zrange).
  Local Notation cstZ r
    := (expr.App
          (expr.Ident ident.Z_cast)
          (expr.Ident (@ident.Literal base.type.zrange r%zrange))).
  Local Notation cstZZ r1 r2
    := (expr.App
          (expr.Ident ident.Z_cast2)
          (expr.App
             (expr.App
                (expr.Ident ident.pair)
                (expr.Ident (@ident.Literal base.type.zrange r1%zrange)))
             (expr.Ident (@ident.Literal base.type.zrange r2%zrange)))).
  Local Notation LetInAppIdentZ S D r eidc x f
    := (expr.LetIn
          (A:=type.base tZ)
          (B:=type.base D)
          (expr.App
             (s:=type.base tZ)
             (d:=type.base tZ)
             (cstZ r)
             (expr.App
                (s:=type.base S)
                (d:=type.base tZ)
                eidc
                x))
          f).
  Local Notation LetInAppIdentZZ S D r1 r2 eidc x f
    := (expr.LetIn
          (A:=type.base (base.type.prod tZ tZ))
          (B:=type.base D)
          (expr.App
             (s:=type.base (base.type.prod tZ tZ))
             (d:=type.base (base.type.prod tZ tZ))
             (cstZZ r1 r2)
             (expr.App
                (s:=type.base S)
                (d:=type.base (base.type.prod tZ tZ))
                eidc
                x))
          f).
  Local Notation pair_LiteralZ B v e :=
    (expr.App
       (s:=type.base B)
       (expr.App
          (expr.Ident (@ident.pair _ _))
          (expr.Ident (ident.Literal (t:=base.type.Z) v)))
       e).
  Local Notation pair_LiteralZZ B v1 v2 e :=
    (expr.App
       (s:=type.base B)
       (expr.App
          (expr.Ident (@ident.pair _ _))
          (expr.App
             (expr.App
                (expr.Ident (@ident.pair _ _))
                (expr.Ident (ident.Literal (t:=base.type.Z) v1)))
             (expr.Ident (ident.Literal (t:=base.type.Z) v2))))
       e).

  Context (name : Type) (name_succ : name -> name) (error : name) (consts : Z -> option name).

  Fixpoint base_var (t : base.type) : Type :=
    match t with
    | tZ => name
    | base.type.prod a b => base_var a * base_var b
    | _ => unit
    end.
  Fixpoint var (t : type.type base.type) : Type :=
    match t with
    | type.base t => base_var t
    | type.arrow s d => var s -> var d
    end.
  Fixpoint base_error {t} : base_var t
    := match t with
       | tZ => error
       | base.type.prod A B => (@base_error A, @base_error B)
       | _ => tt
       end.
  Fixpoint make_error {t} : var t
    := match t with
       | type.base _ => base_error
       | type.arrow s d => fun _ => @make_error d
       end.

  Fixpoint of_prefancy_scalar {t} (s : @cexpr var t) : var t
    := match s in expr.expr t return var t with
       | Compilers.expr.Var t v => v
       | expr.App s d f x => @of_prefancy_scalar _ f (@of_prefancy_scalar _ x)
       | expr.Ident t idc
         => match idc in ident.ident t return var t with
            | ident.Literal base.type.Z v => match consts v with
                                             | Some n => n
                                             | None => error
                                             end
            | ident.pair A B => fun a b => (a, b)%core
            | ident.fst A B => fun v => fst v
            | ident.snd A B => fun v => snd v
            | ident.Z_cast => fun r v => v
            | ident.Z_cast2 => fun r v => v
            | ident.Z_land => fun x y => x
            | _ => make_error
            end
       | expr.Abs s d f => make_error
       | expr.LetIn A B x f => make_error
       end%expr_pat%etype.

  (* Note : some argument orders are reversed for MUL128LU, MUL128UL, SELC, SELM, and SELL *)
  Definition of_prefancy_ident {s d : base.type} (idc : ident.ident (s -> d))
    : @cexpr var s -> option {i : instruction & tuple name i.(num_source_regs) } :=
    match idc in ident.ident t return match t return Type with
                                      | type.arrow (type.base s) (type.base d)
                                        => @cexpr var s
                                      | _ => unit
                                      end
                                      -> option {i : instruction & tuple name i.(num_source_regs) }
    with
    | ident.fancy_selc
      => fun args : @cexpr var (tZ * tZ * tZ) => Some (existT _ SELC (of_prefancy_scalar ((#ident.snd @ args), (#ident.snd @ (#ident.fst @ args)))))
    | ident.fancy_sell
      => fun args : @cexpr var (tZ * tZ * tZ) => Some (existT _ SELL (of_prefancy_scalar ((#ident.snd @ args), (#ident.snd @ (#ident.fst @ args)))))
    | ident.fancy_addm
      => fun args : @cexpr var (tZ * tZ * tZ) => Some (existT _ ADDM (of_prefancy_scalar args))
    | ident.fancy_mulll
      => fun args : @cexpr var (tZ * (tZ * tZ)) =>
           match args with
           | pair_LiteralZ (tZ * tZ) log2wordmax args =>
             if Z.eqb log2wordmax 256
             then Some (existT _ MUL128LL (of_prefancy_scalar args))
             else None
           | _ => None
           end
    | ident.fancy_mullh
      => fun args : @cexpr var (tZ * (tZ * tZ)) =>
           match args with
           | pair_LiteralZ (tZ * tZ) log2wordmax args =>
             if Z.eqb log2wordmax 256
             then Some (existT _ MUL128LU (of_prefancy_scalar ((#ident.snd @ args), (#ident.fst @ args))))
             else None
           | _ => None
           end
    | ident.fancy_mulhl
      => fun args : @cexpr var (tZ * (tZ * tZ)) =>
           match args with
           | pair_LiteralZ (tZ * tZ) log2wordmax args =>
             if Z.eqb log2wordmax 256
             then Some (existT _ MUL128UL (of_prefancy_scalar ((#ident.snd @ args), (#ident.fst @ args))))
             else None
           | _ => None
           end
    | ident.fancy_mulhh
      => fun args : @cexpr var (tZ * (tZ * tZ)) =>
           match args with
           | pair_LiteralZ (tZ * tZ) log2wordmax args =>
             if Z.eqb log2wordmax 256
             then Some (existT _ MUL128UU (of_prefancy_scalar args))
             else None
           | _ => None
           end
    | ident.fancy_selm
      => fun args : @cexpr var (tZ * (tZ * tZ * tZ)) =>
           match args with
           | pair_LiteralZ (tZ * tZ * tZ) log2wordmax args =>
             if Z.eqb log2wordmax 256
             then Some (existT _ SELM (of_prefancy_scalar ((#ident.snd @ args), (#ident.snd @ (#ident.fst @ args)))))
             else None
           | _ => None
           end
    | ident.fancy_add
      => fun args : @cexpr var ((tZ * tZ) * (tZ * tZ)) =>
           match args with
           | pair_LiteralZZ (tZ * tZ) log2wordmax imm args =>
             if Z.eqb log2wordmax 256
             then Some (existT _ (ADD imm) (of_prefancy_scalar args))
             else None
           | _ => None
           end
    | ident.fancy_addc
      => fun args : @cexpr var ((tZ * tZ) * (tZ * tZ * tZ)) =>
           match args with
           | pair_LiteralZZ (tZ * tZ * tZ) log2wordmax imm args =>
               if Z.eqb log2wordmax 256
               then Some (existT _ (ADDC imm) (of_prefancy_scalar ((#ident.snd @ (#ident.fst @ args)), (#ident.snd @ args))))
               else None
           | _ => None
           end
    | ident.fancy_sub
      => fun args : @cexpr var ((tZ * tZ) * (tZ * tZ)) =>
           match args with
           | pair_LiteralZZ (tZ * tZ) log2wordmax imm args =>
               if Z.eqb log2wordmax 256
               then Some (existT _ (SUB imm) (of_prefancy_scalar args))
               else None
           | _ => None
           end
    | ident.fancy_subb
      => fun args : @cexpr var ((tZ * tZ) * (tZ * tZ * tZ)) =>
           match args with
           | pair_LiteralZZ (tZ * tZ * tZ) log2wordmax imm args =>
               if Z.eqb log2wordmax 256
               then Some (existT _ (SUBC imm) (of_prefancy_scalar ((#ident.snd @ (#ident.fst @ args)), (#ident.snd @ args))))
               else None
           | _ => None
           end
    | ident.fancy_rshi
      => fun args : @cexpr var ((tZ * tZ) * (tZ * tZ)) =>
           match args with
           | pair_LiteralZZ (tZ * tZ) log2wordmax imm args =>
               if Z.eqb log2wordmax 256
               then Some (existT _ (RSHI imm) (of_prefancy_scalar args))
               else None
           | _ => None
           end
    | _ => fun _ => None
    end%option.

  Local Notation "x <- y ; f" := (match y with Some x => f | None => Ret error end).
  Definition of_prefancy_step
             (of_prefancy : forall (next_name : name) {t} (e : @cexpr var t), @expr name)
             (next_name : name) {t} (e : @cexpr var t) : @expr name
    := let default _ := (e' <- type.try_transport (@cexpr var) t tZ e;
                           Ret (of_prefancy_scalar e')) in
       match e with
       | LetInAppIdentZ s d r eidc x f
         => idc <- invert_expr.invert_Ident eidc;
              instr_args <- @of_prefancy_ident s tZ idc x;
              let i : instruction := projT1 instr_args in
              let args : tuple name i.(num_source_regs) := projT2 instr_args in
              Instr i next_name args (@of_prefancy (name_succ next_name) _ (f next_name))
       | LetInAppIdentZZ s d r1 r2 eidc x f
         => idc <- invert_expr.invert_Ident eidc;
              instr_args <- @of_prefancy_ident s (tZ * tZ) idc x;
              let i : instruction := projT1 instr_args in
              let args : tuple name i.(num_source_regs) := projT2 instr_args in
              Instr i next_name args (@of_prefancy (name_succ next_name) _ (f (next_name, next_name))) (* the second argument is for the carry, and it will not be read from directly. *)
       | _  => default tt
       end.
  Fixpoint of_prefancy (next_name : name) {t} (e : @cexpr var t) : @expr name
    := @of_prefancy_step of_prefancy next_name t e.

  Section Proofs.
    Context (name_eqb : name -> name -> bool).
    Context (name_lt : name -> name -> Prop)
            (name_lt_trans : forall n1 n2 n3,
                name_lt n1 n2 -> name_lt n2 n3 -> name_lt n1 n3)
            (name_lt_irr : forall n, ~ name_lt n n)
            (name_lt_succ : forall n, name_lt n (name_succ n))
            (name_eqb_eq : forall n1 n2, name_eqb n1 n2 = true -> n1 = n2)
            (name_eqb_neq : forall n1 n2, name_eqb n1 n2 = false -> n1 <> n2).
    Local Notation wordmax := (2^256).
    Local Notation interp := (interp name_eqb wordmax cc_spec).
    Local Notation uint256 := r[0~>wordmax-1]%zrange.
    Local Notation uint128 := r[0~>(2 ^ (Z.log2 wordmax / 2) - 1)]%zrange.
    Local Notation "'existZ' x" := (existT _ (type.base tZ) x) (at level 200).
    Local Notation "'existZZ' x" := (existT _ (type.base (tZ * tZ)%etype) x) (at level 200).
    Local Notation cinterp := API.interp.
    Definition interp_if_Z {t} (e : cexpr t) : option Z :=
      option_map (API.interp (t:=tZ))
                 (type.try_transport
                    _ _ tZ e).

    Lemma interp_if_Z_Some {t} e r :
      @interp_if_Z t e = Some r ->
      exists e',
        (type.try_transport
           _ _ tZ e) = Some e' /\
        API.interp (t:=tZ) e' = r.
    Proof using Type.
      clear. cbv [interp_if_Z option_map].
      break_match; inversion 1; intros.
      subst; eexists. tauto.
    Qed.

    Inductive valid_scalar
      : @cexpr var tZ -> Prop :=
    | valid_scalar_literal :
        forall v n,
          consts v = Some n ->
          valid_scalar (expr.Ident (@ident.Literal base.type.Z v))
    | valid_scalar_Var :
        forall v,
          valid_scalar (expr.App (cstZ uint256) (expr.Var v))
    | valid_scalar_fst :
        forall v r2,
          valid_scalar
            (expr.App (cstZ uint256)
                      (expr.App (expr.Ident (@ident.fst tZ
                                                        tZ))
                                (expr.App (cstZZ uint256 r2) (expr.Var v))))
    .
    Inductive valid_carry
      :  @cexpr var tZ -> Prop :=
    | valid_carry_0 : consts 0 <> None -> valid_carry (expr.Ident (@ident.Literal base.type.Z 0))
    | valid_carry_1 : consts 1 <> None -> valid_carry (expr.Ident (@ident.Literal base.type.Z 1))
    | valid_carry_snd :
        forall v r2,
          valid_carry
            (expr.App (cstZ r[0~>1])
                      (expr.App (expr.Ident (@ident.snd tZ
                                                        tZ))
                                (expr.App (cstZZ r2 r[0~>1]) (expr.Var v))))
    .

    Definition interp_base (ctx : name -> Z) (cctx : name -> bool) {t}
      : base_var t -> base.interp t :=
      match t as t0 return base_var t0 -> base.interp t0 with
      | tZ => fun n => ctx n
      | (tZ * tZ)%etype =>
        fun v => (ctx (fst v), Z.b2z (cctx (snd v)))
      | (_a * _b)%etype =>
        fun _ => DefaultValue.type.base.default
      | _ => fun _ : unit =>
               DefaultValue.type.base.default
      end.

    Definition new_write {d} : var d -> name :=
      match d with
      | type.base tZ => fun r => r
      | type.base (tZ * tZ)%etype => fst
      | _ => fun _ => error
      end.
    Definition new_cc_to_name (old_cc_to_name : CC.code -> name) (i : instruction)
               {d} (new_r : var d) (x : CC.code) : name :=
      if (in_dec CC.code_dec x (writes_conditions i))
      then new_write new_r
      else old_cc_to_name x.

    Inductive valid_ident
      : forall {s d},
        (CC.code -> name) -> (* last variables that wrote to each flag *)
        (var d -> CC.code -> name) -> (* new last variables that wrote to each flag *)
        ident.ident (s->d) -> @cexpr var s -> Prop :=
    | valid_fancy_add :
        forall r imm x y,
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r (ADD imm)) ident.fancy_add ((##256, ##imm), (x, y))%expr
    | valid_fancy_addc :
        forall r imm c x y,
          (of_prefancy_scalar (t:= tZ) c = r CC.C) ->
          valid_carry c ->
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r (ADDC imm)) ident.fancy_addc ((##256, ##imm), (c, x, y))%expr
    | valid_fancy_sub :
        forall r imm x y,
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r (SUB imm)) ident.fancy_sub ((##256, ##imm), (x, y))%expr
    | valid_fancy_subb :
        forall r imm c x y,
          (of_prefancy_scalar (t:= tZ) c = r CC.C) ->
          valid_carry c ->
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r (SUBC imm)) ident.fancy_subb ((##256, ##imm), (c, x, y))%expr
    | valid_fancy_mulll :
        forall r x y,
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r MUL128LL) ident.fancy_mulll ((##256), (x, y))%expr
    | valid_fancy_mullh :
        forall r x y,
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r MUL128LU) ident.fancy_mullh ((##256), (x, y))%expr
    | valid_fancy_mulhl :
        forall r x y,
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r MUL128UL) ident.fancy_mulhl ((##256), (x, y))%expr
    | valid_fancy_mulhh :
        forall r x y,
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r MUL128UU) ident.fancy_mulhh ((##256), (x, y))%expr
    | valid_fancy_rshi :
        forall r imm x y,
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r (RSHI imm)) ident.fancy_rshi ((##256, ##imm), (x, y))%expr
    | valid_fancy_selc :
        forall r c x y,
          (of_prefancy_scalar (t:= tZ) c = r CC.C) ->
          valid_carry c ->
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r SELC) ident.fancy_selc (c, x, y)%expr
    | valid_fancy_selm :
        forall r c x y,
          (of_prefancy_scalar (t:= tZ) c = r CC.M) ->
          valid_scalar c ->
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r SELM) ident.fancy_selm ((##256), (c, x, y))%expr
    | valid_fancy_sell :
        forall r c x y,
          (of_prefancy_scalar (t:= tZ) c = r CC.L) ->
          valid_scalar c ->
          valid_scalar x ->
          valid_scalar y ->
          valid_ident r (new_cc_to_name r SELL) ident.fancy_sell (c, x, y)%expr
    | valid_fancy_addm :
        forall r x y m,
          valid_scalar x ->
          valid_scalar y ->
          valid_scalar m ->
          valid_ident r (new_cc_to_name r ADDM) ident.fancy_addm (x, y, m)%expr
    .

    Inductive valid_expr
      : forall t,
        (CC.code -> name) -> (* the last variables that wrote to each flag *)
        @cexpr var t -> Prop :=
    | valid_LetInZ_loosen :
        forall s d idc r rf x f u ia,
          valid_ident r rf idc x ->
          0 < u < wordmax ->
          (forall x, valid_expr _ (rf x) (f x)) ->
          of_prefancy_ident idc x = Some ia ->
          (forall cc ctx,
              (forall n v, consts v = Some n -> ctx n = v) ->
              (forall n, ctx n mod wordmax = ctx n) ->
              let args := Tuple.map ctx (projT2 ia) in
              spec (projT1 ia) args cc mod wordmax = spec (projT1 ia) args cc mod (u+1)) ->
          valid_expr _ r (LetInAppIdentZ s d r[0~>u] (expr.Ident idc) x f)
    | valid_LetInZ :
        forall s d idc r rf x f,
          valid_ident r rf idc x ->
          (forall x, valid_expr _ (rf x) (f x)) ->
          valid_expr _ r (LetInAppIdentZ s d uint256 (expr.Ident idc) x f)
    | valid_LetInZZ :
        forall s d idc r rf x f,
          valid_ident r rf idc x ->
          (forall x : var (type.base (tZ * tZ)%etype),
              fst x = snd x ->
              valid_expr _ (rf x) (f x)) ->
          valid_expr _ r (LetInAppIdentZZ s d uint256 r[0~>1] (expr.Ident idc) x f)
    | valid_Ret :
        forall r x,
          valid_scalar x ->
          valid_expr _ r x
    .

    Lemma wordmax_nonneg : 0 <= wordmax.
    Proof using Type. cbv; congruence. Qed.

    Lemma of_prefancy_scalar_correct'
          (e1 : @cexpr var (type.base tZ))
          (e2 : cexpr (type.base tZ))
          G (ctx : name -> Z) (cctx : name -> bool) :
      valid_scalar e1 ->
      API.wf G e1 e2 ->
      (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
      (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
      (forall v1 v2, In (existZ (v1, v2)) G -> ctx v1 = v2) -> (* implied by above *)
      (forall n, ctx n mod wordmax = ctx n) ->
      (forall v1 v2, In (existZZ (v1, v2)) G -> ctx (fst v1) = fst v2) ->
      (forall v1 v2, In (existZZ (v1, v2)) G -> Z.b2z (cctx (snd v1)) = snd v2) ->
      ctx (of_prefancy_scalar e1) = cinterp e2.
    Proof using Type.
      inversion 1; inversion 1;
        cbv [interp_if_Z option_map];
        cbn [of_prefancy_scalar interp_base]; intros.
      all: repeat first [
                    progress subst
                  | exfalso; assumption
                  | progress inversion_sigma
                  | progress inversion_option
                  | progress inversion_prod
                  | progress destruct_head'_sig
                  | progress destruct_head'_and
                  | progress inversion_type
                  | progress cbn [fst snd upper lower fst snd eq_rect projT1 projT2 eq_rect type.decode f_equal f_equal2] in *
                  | progress cbn [expr.interp ident.interp type.interp base.interp base.base_interp] in *
                  | progress cbv [ident.cast2 ident.literal] in *
                  | progress destruct_head'_prod
                  | progress expr.inversion_expr
                  | progress expr.inversion_wf_one_constr
                  | progress expr.invert_subst
                  | break_innermost_match_hyps_step
                  | progress break_innermost_match
                  | match goal with H : context [_ = cinterp _] |- context [cinterp _] =>
                                    rewrite <-H by eauto; try reflexivity end
                  | solve [eauto using (f_equal2 pair), cast_oor_id, wordmax_nonneg]
                  | rewrite ident.cast_out_of_bounds_simple_0_mod
                  | rewrite Z.mod_mod by lia
                  | rewrite cast_oor_mod by (cbv; congruence)
                  | lia
                  | match goal with
                      H : context[ ?x mod _ = ?x ] |- _ => rewrite H end
                  | match goal with
                    | H : context [In _ _ -> _ = _] |- _ => erewrite H by eauto end
                  | match goal with
                    | H : forall v1 v2, In _ _ -> ?ctx v1 = v2 |- ?x = ?x mod ?m =>
                      replace m with wordmax by ring; erewrite <-(H _ x)  by eauto; solve [eauto]
                    end
                  | match goal with
                    | H : forall v1 v2, In _ _ -> ?ctx (fst v1) = fst v2,
                      H' : In (existZZ (_,(?x,?y))) _ |- ?x = ?x mod ?m =>
                    replace m with wordmax by ring;
                    specialize (H _ _ H'); cbn [fst] in H; rewrite <-H; solve [eauto] end
                  ].
    Qed.

    Lemma of_prefancy_scalar_correct
          (e1 : @cexpr var (type.base tZ))
          (e2 : cexpr (type.base tZ))
          G (ctx : name -> Z) cc :
      valid_scalar e1 ->
      API.wf G e1 e2 ->
      (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
      (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cc v1 = v2) ->
      (forall n, ctx n mod wordmax = ctx n) ->
      ctx (of_prefancy_scalar e1) = cinterp e2.
    Proof using Type.
      intros; match goal with H : context [interp_base _ _ _ = _] |- _ =>
                              pose proof (H tZ);
                                pose proof (H (tZ * tZ)%etype); cbn [interp_base] in *
              end.
      eapply of_prefancy_scalar_correct'; eauto;
        match goal with
        | H : forall _ _, In _ _ -> (_, _) = _ |- _ =>
          let v1 := fresh "v" in
          let v2 := fresh "v" in
          intros v1 v2 ?; rewrite <-(H v1 v2) by auto
        end; reflexivity.
    Qed.

    Lemma of_prefancy_ident_Some_gen {s d} idc r rf x:
      @valid_ident s d r rf idc x ->
      match s, d return ident (s -> d) -> cexpr s -> Prop with
      | type.base s, type.base d
        => fun idc x => of_prefancy_ident idc x <> None
      | _, _ => fun idc x => False
      end idc x.
    Proof using Type.
      destruct 1; intros;
        repeat first [
                 progress HProp.eliminate_hprop_eq
               | progress subst
               | progress cbn [GallinaReify.base.reify ident.ident_Literal ident.buildIdent] in *
               | progress cbn [projT1 projT2 eq_rect of_prefancy_ident option_map] in *
               | progress cbv [ident.literal] in *
               | progress inversion_sigma
               | progress Z.ltb_to_lt
               | progress break_innermost_match
               | congruence
               ].
    Qed.

    Lemma of_prefancy_ident_Some {s d} idc r rf x:
      @valid_ident (type.base s) (type.base d) r rf idc x ->
      of_prefancy_ident idc x <> None.
    Proof using Type. exact (@of_prefancy_ident_Some_gen _ _ idc r rf x). Qed.

    Ltac name_eqb_to_eq :=
      repeat match goal with
             | H : name_eqb _ _ = true |- _ => apply name_eqb_eq in H
             | H : name_eqb _ _ = false |- _ => apply name_eqb_neq in H
             end.
    Ltac inversion_of_prefancy_ident :=
        match goal with
        | H : of_prefancy_ident _ _ = None |- _ =>
          eapply of_prefancy_ident_Some in H;
          [ contradiction | eassumption]
        end.

    Local Ltac hammer :=
      repeat first [
                    progress subst
                  | progress inversion_sigma
                  | progress inversion_option
                  | progress inversion_of_prefancy_ident
                  | progress inversion_prod
                  | progress cbv [id]
                  | progress cbn [eq_rect projT1 projT2 expr.interp ident.interp interp_base Coercions.base Coercions.type_base interp interp_if_Z option_map] in *
                  | progress cbv [ident.cast2 ident.literal] in *
                  | progress ident.fancy.cbv_fancy_in_all
                  | progress cbn [invert_expr.invert_Ident] in * (* N.B. Must be above [break_innermost_match] for proofs below to work *)
                  | progress Language.Inversion.Compilers.type_beq_to_eq
                  | progress name_eqb_to_eq
                  | progress Language.Inversion.Compilers.rewrite_type_transport_correct
                  | progress HProp.eliminate_hprop_eq
                  | progress break_innermost_match_hyps
                  | progress break_innermost_match
                  | progress inversion_type
                  | progress expr.invert_subst
                  | progress Language.Inversion.Compilers.expr.inversion_expr
                  | solve [auto]
                  | contradiction
             ].
    Ltac prove_Ret :=
      repeat match goal with
             | H : valid_scalar (expr.LetIn _ _) |- _ =>
               inversion H
             | _ => progress cbn [id of_prefancy of_prefancy_step of_prefancy_scalar]
             | _ => progress hammer
             | H : valid_scalar (expr.Ident _) |- _ =>
               inversion H; clear H
             | |- _ = cinterp ?f (cinterp ?x) =>
               transitivity
                 (cinterp (f @ x)%expr);
               [ | reflexivity ];
               erewrite <-of_prefancy_scalar_correct by (try reflexivity; eassumption)
             end.

    Lemma cast_mod u v :
      0 <= u ->
      ident.cast r[0~>u] v = v mod (u + 1).
    Proof using Type. apply ident.cast_out_of_bounds_simple_0_mod. Qed.

    Lemma cc_spec_c v :
      Z.b2z (cc_spec CC.C v) = (v / wordmax) mod 2.
    Proof using Type. cbv [cc_spec]; apply Z.testbit_spec'. lia. Qed.

    Lemma cc_m_zselect x z nz :
      x mod wordmax = x ->
      (if (if cc_spec CC.M x then 1 else 0) =? 1 then nz else z) =
      Z.zselect (x >> 255) z nz.
    Proof using Type.
      intro Hx_small.
      transitivity (if (Z.b2z (cc_spec CC.M x) =? 1) then nz else z); [ reflexivity | ].
      cbv [cc_spec Z.zselect].
      rewrite Z.testbit_spec', Z.shiftr_div_pow2 by lia. rewrite <-Hx_small.
      rewrite Z.div_between_0_if by (try replace (2 * (2 ^ 255)) with wordmax by reflexivity;
                                         auto with zarith).
      break_innermost_match; Z.ltb_to_lt; try rewrite Z.mod_small in * by lia; congruence.
    Qed.

    Lemma cc_l_zselect x z nz :
      (if (if cc_spec CC.L x then 1 else 0) =? 1 then nz else z) = Z.zselect (x &' 1) z nz.
    Proof using Type.
      transitivity (if (Z.b2z (cc_spec CC.L x) =? 1) then nz else z); [ reflexivity | ].
      transitivity (Z.zselect (x &' Z.ones 1) z nz); [ | reflexivity ].
      cbv [cc_spec Z.zselect]. rewrite Z.testbit_spec', Z.land_ones by lia.
      autorewrite with zsimplify_fast. rewrite Zmod_even.
      break_innermost_match; Z.ltb_to_lt; congruence.
    Qed.

    Lemma b2z_range b : 0<= Z.b2z b < 2.
    Proof using Type. cbv [Z.b2z]. break_match; lia. Qed.

    Lemma of_prefancy_scalar_carry
          (c : @cexpr var (type.base tZ))
          (e : cexpr (type.base tZ))
          G (ctx : name -> Z) cctx :
      valid_carry c ->
      API.wf G c e ->
      (forall n0, consts 0 = Some n0 -> cctx n0 = false) ->
      (forall n1, consts 1 = Some n1 -> cctx n1 = true) ->
      (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
      Z.b2z (cctx (of_prefancy_scalar c)) = cinterp e.
    Proof using Type.
      destruct 1; intros; hammer; cbn;
      repeat match goal with
             | H : context [ _ = false] |- Z.b2z _ = 0 => rewrite H; reflexivity
             | H : context [ _ = true] |- Z.b2z _ = 1 => rewrite H; reflexivity
             | _ => progress Compilers.expr.inversion_wf_one_constr
             | _ => progress cbn [fst snd]
             | _ => progress destruct_head'_sig
             | _ => progress destruct_head'_and
             | _ => progress hammer
             | _ => progress Language.Inversion.Compilers.expr.invert_subst
             | _ => rewrite cast_mod by (cbv; congruence)
             | _ => rewrite Z.mod_mod by lia
             | _ => rewrite Z.mod_small by apply b2z_range
             | H : (forall _ _ _, In _ _ -> interp_base _ _ _ = _),
                   H' : In (existZZ (?v, _)) _ |- context [cctx (snd ?v)] =>
               specialize (H _ _ _ H'); cbn in H
             end.
    Qed.

    Ltac hammer_wf :=
      repeat first [ progress subst
                   | progress cbn [eq_rect fst snd projT1 projT2] in *
                   | progress destruct_head'_False
                   | progress expr.inversion_wf_one_constr
                   | progress expr.invert_subst
                   | progress destruct_head'_and
                   | progress destruct_head'_sig
                   | break_innermost_match_hyps_step ].

    Ltac simplify_ident :=
      repeat match goal with
             | _ => progress inversion_option
             | _ => progress hammer_wf
             | _ => progress intros
             | _ => progress cbn [fst snd of_prefancy_ident] in *
             | H : upper _ = _ |- _ => rewrite H
             | _ => rewrite cc_spec_c by auto
             | _ => rewrite cast_mod by (cbv; congruence)
             | H : ?P, H' : ?P |- _ => clear H'
             | _ => progress hammer
             end.

    (* TODO: zero flag is a little tricky, since the value
      depends both on the stored variable and the carry if there
      is one. For now, since Barrett doesn't use it, we're just
      pretending it doesn't exist. *)
    Definition cc_good cc cctx ctx r :=
        CC.cc_c cc = cctx (r CC.C) /\
        CC.cc_m cc = cc_spec CC.M (ctx (r CC.M)) /\
        CC.cc_l cc = cc_spec CC.L (ctx (r CC.L)) /\
        (forall n0 : name, consts 0 = Some n0 -> cctx n0 = false) /\
        (forall n1 : name, consts 1 = Some n1 -> cctx n1 = true).

    Lemma of_prefancy_identZ_loosen_correct_gen {s d} idc:
      forall (x : @cexpr var _) i ctx G cc cctx x2 r rf f u,
        @valid_ident s d r rf idc x ->
        match s, d return ident (s -> d) -> cexpr s -> (var d -> _ -> _) -> Prop with
        | type.base s, type.base tZ
          => fun idc x rf
             => API.wf G (#idc @ x)%expr x2 ->
                API.wf G (cstZ r[0~>u]) f ->
                0 < u < wordmax ->
                cc_good cc cctx ctx r ->
                (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
                (forall t (v1 : var (type.base t)) v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
                (forall n, ctx n mod wordmax = ctx n) ->
                of_prefancy_ident idc x = Some i ->
                (spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod wordmax = spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod (u+1)) ->
                spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod wordmax = (cinterp f (cinterp x2))
        | _, _ => fun _ _ _ => True
        end idc x rf.
    Proof using Type.
      destruct 1; try exact I; cbn [of_prefancy_ident]; cbn [GallinaReify.base.reify ident.ident_Literal ident.buildIdent] in *; hammer; intros; inversion_option; (simplify_ident; [ ]).
      all:
        rewrite cast_mod by lia;
        match goal with
                 | H : context [spec _ _ _ mod _ = _] |- ?x mod wordmax = _ mod ?m =>
                   replace (x mod wordmax) with (x mod m) by auto
        end.
      all:  cbn - [Z.shiftl wordmax]; cbv [cc_good] in *; destruct_head'_and;
          repeat match goal with
                 | H : CC.cc_c _ = _ |- _ => rewrite H
                 | H : CC.cc_m _ = _ |- _ => rewrite H
                 | H : CC.cc_l _ = _ |- _ => rewrite H
                 | H : CC.cc_z _ = _ |- _ => rewrite H
                 | H: of_prefancy_scalar _ = ?r ?c |- _ => rewrite <-H
                 | _ => progress rewrite ?cc_m_zselect, ?cc_l_zselect by auto
                 | _ => progress rewrite ?Z.add_modulo_correct, ?Z.geb_leb by auto
                 | |- context [cinterp ?x] =>
                   erewrite of_prefancy_scalar_correct with (e2:=x) by eauto
                 | |- context [cinterp ?x] =>
                   erewrite <-of_prefancy_scalar_carry with (e:=x) by eauto
                 | |- context [if _ (of_prefancy_scalar _) then _ else _ ] =>
                   cbv [Z.zselect Z.b2z];
                     break_innermost_match; Z.ltb_to_lt; try reflexivity;
                       congruence
                 end; try reflexivity.

      { (* RSHI case *)
        cbv [Z.rshi].
        rewrite Z.land_ones, Z.shiftl_mul_pow2 by (cbv; congruence).
        change (2 ^ Z.log2 wordmax) with wordmax.
        break_innermost_match; try congruence; [ ]. autorewrite with zsimplify_fast.
        repeat (f_equal; try ring). }
    Qed.
    Lemma of_prefancy_identZ_loosen_correct {s} idc:
      forall (x : @cexpr var _) i ctx G cc cctx x2 r rf f u,
        @valid_ident (type.base s) tZ r rf idc x ->
        API.wf G (#idc @ x)%expr x2 ->
        API.wf G (cstZ r[0~>u]) f ->
        0 < u < wordmax ->
        cc_good cc cctx ctx r ->
        (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
        (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
        (forall n, ctx n mod wordmax = ctx n) ->
        of_prefancy_ident idc x = Some i ->
        (spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod wordmax = spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod (u+1)) ->
        spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod wordmax = (cinterp f (cinterp x2)).
    Proof using Type. apply (@of_prefancy_identZ_loosen_correct_gen (type.base s) (type.base tZ)). Qed.
    Lemma of_prefancy_identZ_correct {s} idc:
      forall (x : @cexpr var _) i ctx G cc cctx x2 r rf f,
        @valid_ident (type.base s) tZ r rf idc x ->
        API.wf G (#idc @ x)%expr x2 ->
        API.wf G (cstZ uint256) f ->
        cc_good cc cctx ctx r ->
        (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
        (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
        (forall n, ctx n mod wordmax = ctx n) ->
        of_prefancy_ident idc x = Some i ->
        spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod wordmax = (cinterp f (cinterp x2)).
    Proof using Type.
      intros; eapply of_prefancy_identZ_loosen_correct; try eassumption; [ | ].
      { cbn; lia. } { intros; f_equal; ring. }
    Qed.
    Lemma of_prefancy_identZZ_correct' {s} idc:
      forall (x : @cexpr var _) i ctx G cc cctx x2 r rf f,
        @valid_ident (type.base s) (tZ * tZ) r rf idc x ->
        API.wf G (#idc @ x)%expr x2 ->
        API.wf G (cstZZ uint256 r[0~>1]) f ->
        cc_good cc cctx ctx r ->
        (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
        (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
        (forall n, ctx n mod wordmax = ctx n) ->
        of_prefancy_ident idc x = Some i ->
        spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod wordmax = fst (cinterp f (cinterp x2)) /\
        Z.b2z (cc_spec CC.C (spec (projT1 i) (Tuple.map ctx (projT2 i)) cc)) = snd (cinterp f (cinterp x2)).
    Proof using Type.
      inversion 1; inversion 1; cbn [of_prefancy_ident]; intros; hammer; (simplify_ident; [ ]);
        cbn - [Z.div Z.modulo]; cbv [Z.sub_with_borrow Z.add_with_carry];
          cbv [cc_good] in *; destruct_head'_and; autorewrite with zsimplify_fast.
      all: repeat match goal with
               | H : CC.cc_c _ = _ |- _ => rewrite H
               | H: of_prefancy_scalar _ = ?r ?c |- _ => rewrite <-H
               | H : API.wf _ ?x ?e |- context [cinterp ?e] =>
                 erewrite <-of_prefancy_scalar_correct with (e1:=x) (e2:=e) by eauto
               | H : API.wf _ ?x ?e2 |- context [cinterp ?e2] =>
                 erewrite <-of_prefancy_scalar_carry with (c:=x) (e:=e2) by eauto
                  end.
      all: match goal with |- context [(?x << ?n) mod ?m] =>
                           pose proof (Z.mod_pos_bound (x << n) m ltac:(lia)) end.
      all:repeat match goal with
               | |- context [if _ (of_prefancy_scalar _) then _ else _ ] =>
                 cbv [Z.zselect Z.b2z]; break_innermost_match; Z.ltb_to_lt; try congruence; [ | ]
               | _ => rewrite Z.add_opp_r
               | _ => rewrite Z.div_sub_small by auto with zarith
               | H : forall n, ?ctx n mod wordmax = ?ctx n |- context [?ctx ?m - _] => rewrite <-(H m)
               |  |- ((?x - ?y - ?c) / _) mod _ = - ((- ?c + ?x - ?y) / _) mod _ =>
                  replace (-c + x - y) with (x - (y + c)) by ring; replace (x - y - c) with (x - (y + c)) by ring
               | _ => split
               | _ => try apply (f_equal2 Z.modulo); try apply (f_equal2 Z.div); ring
               | _ => break_innermost_match; reflexivity
               end.
    Qed.
    Lemma of_prefancy_identZZ_correct {s} idc:
      forall (x : @cexpr var _) i ctx G cc cctx x2 r rf f,
        @valid_ident (type.base s) (tZ * tZ) r rf idc x ->
        API.wf G (#idc @ x)%expr x2 ->
        API.wf G (cstZZ uint256 r[0~>1]) f ->
        cc_good cc cctx ctx r ->
        (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
        (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
        (forall n, ctx n mod wordmax = ctx n) ->
        of_prefancy_ident idc x = Some i ->
        spec (projT1 i) (Tuple.map ctx (projT2 i)) cc mod wordmax = fst (cinterp f (cinterp x2)).
    Proof using Type. apply of_prefancy_identZZ_correct'. Qed.
    Lemma of_prefancy_identZZ_correct_carry {s} idc:
      forall (x : @cexpr var _) i ctx G cc cctx x2 r rf f,
        @valid_ident (type.base s) (tZ * tZ) r rf idc x ->
        API.wf G (#idc @ x)%expr x2 ->
        API.wf G (cstZZ uint256 r[0~>1]) f ->
        cc_good cc cctx ctx r ->
        (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
        (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
        (forall n, ctx n mod wordmax = ctx n) ->
        of_prefancy_ident idc x = Some i ->
        Z.b2z (cc_spec CC.C (spec (projT1 i) (Tuple.map ctx (projT2 i)) cc)) = snd (cinterp f (cinterp x2)).
    Proof using Type. apply of_prefancy_identZZ_correct'. Qed.

    Lemma identZZ_writes {s} idc r rf x:
      @valid_ident (type.base s) (tZ * tZ) r rf idc x ->
      forall i, of_prefancy_ident idc x = Some i ->
                  In CC.C (writes_conditions (projT1 i)).
    Proof using Type.
      inversion 1;
        repeat match goal with
               | _ => progress intros
               | _ => progress cbn [of_prefancy_ident writes_conditions ADD ADDC SUB SUBC In] in *
               | _ => progress hammer; Z.ltb_to_lt
               | _ => congruence
               end.
    Qed.

    (* Common side conditions for cases in of_prefancy_correct *)
    Local Ltac side_cond :=
      repeat match goal with
             | _ => progress intros
             | _ => progress cbn [In fst snd] in *
             | H : _ \/ _ |- _ => destruct H
             | [H : forall _ _, In _ ?l -> _, H' : In _ ?l |- _] =>
               let H'' := fresh in
               pose proof H'; apply H in H''; clear H
             | H : name_lt ?n ?n |- _ =>
               specialize (name_lt_irr n); contradiction
             | _ => progress hammer
             | _ => solve [eauto]
             end.

    Lemma interp_base_helper G next_name ctx cctx :
      (forall n v2, In (existZ (n, v2)) G -> name_lt n next_name) ->
      (forall n v2, In (existZZ (n, v2)) G -> name_lt (fst n) next_name) ->
      (forall n v2, In (existZZ (n, v2)) G -> fst n = snd n) ->
      (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
      (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G ->
                       t = tZ
                       \/ t = (tZ * tZ)%etype) ->
      forall t v1 v2 x xc,
        In (existT (fun t : type => (var t * type.interp base.interp t)%type) (type.base t) (v1, v2)%zrange)
           ((existZ (next_name, x)%zrange) :: G) ->
        interp_base (fun n : name => if name_eqb n next_name then x else ctx n)
                    (fun n : name => if name_eqb n next_name then xc else cctx n) v1 = v2.
    Proof using name_eqb_eq name_eqb_neq name_lt_irr.
      intros.
      repeat match goal with
             | H: In _ (_ :: _) |- _ => cbn [In] in H; destruct H; [ solve [side_cond] | ]
             | H : (forall t _ _, In _ ?G -> (t = _ \/ t = _)), H' : In _ ?G |- _ =>
               destruct (H _ _ _ H'); subst t
             | H : forall _ _ _, In _ ?G -> interp_base _ _ _ = _, H' : In _ G |- _ => specialize (H _ _ _ H')
      end; side_cond.
    Qed.

    Lemma name_eqb_refl n : name_eqb n n = true.
    Proof using name_eqb_eq name_eqb_neq. case_eq (name_eqb n n); intros; name_eqb_to_eq; auto. Qed.

    Lemma valid_ident_new_cc_to_name s d r rf idc x y n :
      @valid_ident (type.base s) (type.base d) r rf idc x ->
      of_prefancy_ident idc x = Some y ->
      rf n = new_cc_to_name r (projT1 y) n.
    Proof using Type. inversion 1; intros; hammer; simplify_ident. Qed.

    Lemma new_cc_to_name_Z_cases r i n x :
      new_cc_to_name (d:=tZ) r i n x
      = if in_dec CC.code_dec x (writes_conditions i)
        then n else r x.
    Proof using Type. reflexivity. Qed.
    Lemma new_cc_to_name_ZZ_cases r i n x :
      new_cc_to_name (d:=tZ * tZ) r i n x
      = if in_dec CC.code_dec x (writes_conditions i)
        then fst n else r x.
    Proof using Type. reflexivity. Qed.

    Lemma cc_good_helper cc cctx ctx r i x next_name :
      (forall c, name_lt (r c) next_name) ->
      (forall n v, consts v = Some n -> name_lt n next_name) ->
      cc_good cc cctx ctx r ->
      cc_good (CC.update (writes_conditions i) x cc_spec cc)
              (fun n : name =>
                 if name_eqb n next_name
                 then CC.cc_c (CC.update (writes_conditions i) x cc_spec cc)
                 else cctx n)
              (fun n : name => if name_eqb n next_name then x mod wordmax else ctx n)
              (new_cc_to_name (d:=tZ) r i next_name).
    Proof using name_eqb_eq name_eqb_neq name_lt_irr.
      cbv [cc_good]; intros; destruct_head'_and.
      rewrite !new_cc_to_name_Z_cases.
      cbv [CC.update CC.cc_c CC.cc_m CC.cc_l CC.cc_z].
      repeat match goal with
             | _ => split; intros
             | _ => progress hammer
             | H : forall c, name_lt (r c) (r ?c2) |- _ => specialize (H c2)
             | H : (forall n v, consts v = Some n -> name_lt _ _),
                   H' : consts _ = Some _ |- _ => specialize (H _ _ H')
             | H : name_lt ?n ?n |- _ => apply name_lt_irr in H; contradiction
             | _ => cbv [cc_spec]; rewrite Z.mod_pow2_bits_low by lia
             | _ => congruence
             end.
    Qed.

    Lemma of_prefancy_correct
          {t} (e1 : @cexpr var t) (e2 : @cexpr _ t) r :
      valid_expr _ r e1 ->
      forall G,
      API.wf G e1 e2 ->
      forall ctx cc cctx,
        cc_good cc cctx ctx r ->
        (forall n v, consts v = Some n -> In (existZ (n, v)) G) ->
        (forall n v2, In (existZZ (n, v2)) G -> fst n = snd n) ->
        (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G -> interp_base ctx cctx v1 = v2) ->
        (forall t v1 v2, In (existT _ (type.base t) (v1, v2)) G ->
                         t = tZ
                         \/ t = (tZ * tZ)%etype) ->
        (forall n, ctx n mod wordmax = ctx n) ->
        forall next_name result,
          (forall c : CC.code, name_lt (r c) next_name) ->
          (forall n v2, In (existZ (n, v2)) G -> name_lt n next_name) ->
          (forall n v2, In (existZZ (n, v2)) G -> name_lt (fst n) next_name) ->
          (interp_if_Z e2 = Some result) ->
          interp (@of_prefancy next_name t e1) cc ctx = result.
    Proof using name_eqb_eq name_eqb_neq name_lt_irr name_lt_succ name_lt_trans.
      induction 1; inversion 1; cbv [interp_if_Z];
        cbn [of_prefancy of_prefancy_step]; intros;
          match goal with H : context [interp_base _ _ _ = _] |- _ =>
                          pose proof (H tZ) end;
          try solve [prove_Ret]; [ | | ]; hammer;
            match goal with
            | H : context [interp (of_prefancy _ _) _ _ = _]
              |- interp _ ?cc' ?ctx' = _ =>
              match goal with
              | _ : context [LetInAppIdentZ _ _ _ _ _ _] |-  _=>
                erewrite H with
                    (G := (existZ (next_name, ctx' next_name)) :: G)
                    (e2 := _ (ctx' next_name))
                    (cctx := (fun n => if name_eqb n next_name then CC.cc_c cc' else cctx n))
              | _ : context [LetInAppIdentZZ _ _ _ _ _ _ _] |-  _=>
                erewrite H with
                    (G := (existZZ ((next_name, next_name), (ctx' next_name, Z.b2z (CC.cc_c cc')))) :: G)
                    (e2 := _ (ctx' next_name, Z.b2z (CC.cc_c cc')))
                    (cctx := (fun n => if name_eqb n next_name then CC.cc_c cc' else cctx n))
              end
            end;
            repeat match goal with
                   | _ => progress intros
                   | _ => rewrite name_eqb_refl in *
                   | _ => rewrite Z.testbit_spec' in *
                   | _ =>  erewrite valid_ident_new_cc_to_name by eassumption
                   | _ => rewrite new_cc_to_name_Z_cases
                   | _ => rewrite new_cc_to_name_ZZ_cases
                   | _ => solve [intros; eapply interp_base_helper; side_cond]
                   | _ => solve [intros; apply cc_good_helper; eauto]
                   | _ => reflexivity
                   | _ => solve [eauto using Z.mod_small, b2z_range]
                   | _ => progress autorewrite with zsimplify_fast
                   | _ => progress side_cond
                   end; [ | | ].
      { cbn - [cc_spec]; cbv [id]; cbn - [cc_spec].
        inversion wf_x; hammer.
        erewrite @of_prefancy_identZ_loosen_correct by eauto.
        reflexivity. }
      { cbn - [cc_spec]; cbv [id]; cbn - [cc_spec].
        inversion wf_x; hammer.
        erewrite of_prefancy_identZ_correct by eassumption.
        reflexivity. }
      { cbn - [cc_spec]; cbv [id]; cbn - [cc_spec].
        match goal with H : _ |- _ => pose proof H; eapply identZZ_writes in H; [ | eassumption] end.
        inversion wf_x; hammer.
        erewrite of_prefancy_identZZ_correct by eassumption.
        erewrite of_prefancy_identZZ_correct_carry by eassumption.
        rewrite <-surjective_pairing. reflexivity. }
    Qed.
  End Proofs.
End of_prefancy.

Section allocate_registers.
  Context (reg name : Type) (name_eqb : name -> name -> bool) (error : reg).
  Fixpoint allocate (e : @expr name) (reg_list : list reg) (name_to_reg : name -> reg) : @expr reg :=
    match e with
    | Ret n => Ret (name_to_reg n)
    | Instr i rd args cont =>
      match reg_list with
      | r :: reg_list' => Instr i r (Tuple.map name_to_reg args) (allocate cont reg_list' (fun n => if name_eqb n rd then r else name_to_reg n))
      | nil => Ret error
      end
    end.
End allocate_registers.

Definition test_prog : @expr positive :=
  Instr (ADD (128)) 3%positive (1, 2)%positive
        (Instr (ADDC 0) 4%positive (3,1)%positive
               (Ret 4%positive)).

Definition x1 := 2^256 - 1.
Definition x2 := 2^128 - 1.
Definition wordmax := 2^256.
Definition expected :=
  let r3' := (x1 + (x2 << 128)) in
  let r3 := r3' mod wordmax in
  let c := r3' / wordmax in
  let r4' := (r3 + x1 + c) in
  r4' mod wordmax.
Definition actual :=
  interp Pos.eqb
         (2^256) cc_spec test_prog {|CC.cc_c:=false; CC.cc_m:=false; CC.cc_l:=false; CC.cc_z:=false|}
         (fun n => if Pos.eqb n 1%positive
                   then x1
                   else if Pos.eqb n 2%positive
                        then x2
                        else 0).
Lemma test_prog_ok : expected = actual.
Proof. reflexivity. Qed.

Definition of_Expr {t} next_name (consts : Z -> option positive)
           (e : expr.Expr t)
           (x : type.for_each_lhs_of_arrow (var positive) t)
  : positive -> @expr positive :=
  fun error =>
    @of_prefancy positive Pos.succ error consts next_name _ (invert_expr.smart_App_curried (e _) x).

Section Proofs.
  Fixpoint var_pairs {t var1 var2}
    : type.for_each_lhs_of_arrow var1 t
      -> type.for_each_lhs_of_arrow var2 t
      -> list {t : Compilers.type base.type.type & (var1 t * var2 t)%type } :=
    match t as t0 return
          (type.for_each_lhs_of_arrow var1 t0
           -> type.for_each_lhs_of_arrow var2 t0 -> _) with
    | type.base _ => fun _ _ => nil
    | type.arrow s d =>
      fun x1 x2 =>
        existT _ _ (fst x1, fst x2) :: var_pairs (snd x1) (snd x2)
    end.

  Local Notation tZ := (base.type.type_base base.type.Z).
  Local Notation existZ := (existT _  (type.base tZ)).
  Local Notation existZZ := (existT _  (type.base (tZ * tZ)%etype)).

  Fixpoint make_ctx (var_list : list (positive * Z)) : positive -> Z :=
    match var_list with
    | [] => fun _ => 0
    | (n, v) :: l' => fun m => if (m =? n)%positive then v else make_ctx l' m
    end.

  Definition make_pairs :
    list (positive * Z) -> list {t : Compilers.type base.type.type & (var positive t * @type.interp base.type base.interp t)%type } := List.map (fun x => existZ x).

  Fixpoint make_consts (consts_list : list (positive * Z)) : Z -> option positive :=
    match consts_list with
    | [] => fun _ => None
    | (n, v) :: l' => fun x => if x =? v then Some n else make_consts l' x
    end.

  Local Ltac ez :=
    repeat match goal with
           | _ => progress intros
           | _ => progress subst
           | H : _ \/ _ |- _ => destruct H
           | H : _ |- _ => rewrite Pos.eqb_eq in H
           | H : _ |- _ => rewrite Pos.eqb_neq in H
           | _ => progress break_innermost_match
           | _ => progress break_match_hyps
           | _ => progress inversion_sigma
           | _ => progress inversion_option
           | _ => progress inversion_prod
           | _ => progress HProp.eliminate_hprop_eq
           | _ => progress Z.ltb_to_lt
           | _ => reflexivity
           | _ => congruence
           | _ => solve [eauto]
           end.

  Lemma make_consts_ok consts_list n v :
    make_consts consts_list v = Some n ->
    In (existZ (n, v)%zrange) (make_pairs consts_list).
  Proof.
    cbv [make_pairs]; induction consts_list as [|[ ? ? ] ?]; cbn; ez.
  Qed.

  Lemma make_pairs_ok consts_list:
    forall v1 v2,
      In (existZ (v1, v2)%zrange) (make_pairs consts_list) ->
      In (v1, v2) consts_list.
  Proof.
    cbv [make_pairs]. induction consts_list as [| [ n v ] ? ]; cbn; [ tauto | ]. ez.
  Qed.
  Lemma make_ctx_ok consts_list:
    (forall n v1 v2, In (n, v1) consts_list ->
                     In (n, v2) consts_list -> v1 = v2) ->
    forall n v,
      In (n, v) consts_list ->
    make_ctx consts_list n = v.
  Proof.
    induction consts_list as [| [ n v ] ? ]; cbn; [ tauto | ].
    repeat match goal with
           | _ => progress cbn [eq_rect fst snd] in *
           | _ => progress ez
           end.
  Qed.

  Lemma make_ctx_cases consts_list n :
    make_ctx consts_list n = 0 \/
    In (n, make_ctx consts_list n) consts_list.
  Proof. induction consts_list; cbn; ez. Qed.

  Lemma only_integers consts_list t v1 v2 :
    In (existT (fun t : type => (var positive t * type.interp base.interp t)%type) (type.base t)
               (v1, v2)%zrange) (make_pairs consts_list) ->
    t = tZ.
  Proof.
    induction consts_list; cbn; [ tauto | ].
    destruct 1; congruence || tauto.
  Qed.

  Lemma no_pairs consts_list v1 v2 :
    In (existZZ (v1, v2)%zrange) (make_pairs consts_list) -> False.
  Proof. intro H; apply only_integers in H. congruence. Qed.

  Definition make_cc last_wrote ctx carry_flag : CC.state :=
    {| CC.cc_c := carry_flag;
       CC.cc_m := cc_spec CC.M (ctx (last_wrote CC.M));
       CC.cc_l := cc_spec CC.L (ctx (last_wrote CC.L));
       CC.cc_z := cc_spec CC.Z (ctx (last_wrote CC.Z)
                                + (if (last_wrote CC.C =? last_wrote CC.Z)%positive
                                   then wordmax * Z.b2z carry_flag else 0));
    |}.

  Hint Resolve Pos.lt_trans Pos.lt_irrefl Pos.lt_succ_diag_r Pos.eqb_refl : core.
  Hint Resolve in_or_app : core.
  Hint Resolve make_consts_ok make_pairs_ok make_ctx_ok no_pairs : core.
  (* TODO : probably not all of these preconditions are necessary -- prune them sometime *)
  Lemma of_Expr_correct next_name consts_list arg_list error
        (carry_flag : bool)
        (last_wrote : CC.code -> positive) (* variables which last wrote to each flag; put RegZero if flag empty *)
        t (e : Expr t)
        (x1 : type.for_each_lhs_of_arrow (var positive) t)
        (x2 : type.for_each_lhs_of_arrow _ t) result :
    let e1 := (invert_expr.smart_App_curried (e _) x1) in
    let e2 := (invert_expr.smart_App_curried (e _) x2) in
    let ctx := make_ctx (consts_list ++ arg_list) in
    let consts := make_consts consts_list in
    let cc := make_cc last_wrote ctx carry_flag in
    let G := make_pairs consts_list ++ make_pairs arg_list in
    (forall c, last_wrote c < next_name)%positive ->
    (forall n v, In (n, v) (consts_list ++ arg_list) -> (n < next_name)%positive) ->
    (In (last_wrote CC.C, Z.b2z carry_flag) consts_list) ->
    (forall n v1 v2, In (n, v1) (consts_list ++ arg_list) ->
                     In (n, v2) (consts_list ++ arg_list) -> v1 = v2) (* no duplicate names *) ->
    (forall v1 v2, In (v1, v2) consts_list -> v2 mod 2 ^ 256 = v2) ->
    (forall v1 v2, In (v1, v2) arg_list -> v2 mod 2 ^ 256 = v2) ->
    (API.wf G e1 e2) ->
    valid_expr _ error consts _ last_wrote e1 ->
    interp_if_Z e2 = Some result ->
    interp Pos.eqb wordmax cc_spec (of_Expr next_name consts e x1 error) cc ctx = result.
  Proof.
    cbv [of_Expr]; intros.
    eapply of_prefancy_correct with (name_lt := Pos.lt)
                                    (cctx := fun n => if (n =? last_wrote CC.C)%positive
                                                      then carry_flag
                                                      else match make_consts consts_list 1 with
                                                           | Some n1 => (n =? n1)%positive
                                                           | _ => false
                                                           end);
      cbv [id]; eauto;
      try apply Pos.eqb_neq; intros;
        try solve [apply make_ctx_ok; auto; apply make_pairs_ok;
                   cbv [make_pairs]; rewrite map_app; auto ];
        repeat match goal with
               | H : _ |- _ => apply in_app_or in H; destruct H
               | H : In _ (make_pairs _) |- context [ _ = base.type.type_base _] => apply only_integers in H
               | H : In _ (make_pairs _) |- context [interp_base] =>
                 pose proof (only_integers _ _ _ _ H); subst; cbn [interp_base]
               | _ => solve [apply Pos.eqb_eq; auto]
               | _ => solve [eauto]
               | _ => solve [exfalso; eauto]
               end; [ | ].
    (* TODO : clean this up *)
    { cbv [cc_good make_cc]; repeat split; intros;
        [ rewrite Pos.eqb_refl; reflexivity | | ];
      break_innermost_match; try rewrite Pos.eqb_eq in *;
        subst; try reflexivity;
          repeat match goal with
                 | H : make_consts _ _ = Some _ |- _ =>
                   try rewrite H;
                   apply make_consts_ok, make_pairs_ok in H
                 | _ => apply Pos.eqb_neq; intro; subst
                 | _ => inversion_option; congruence
                 end;
          match goal with
          | H : In (?n, ?x) consts_list,
                H': In (?n, ?y) consts_list,
                    H'' : forall n x y, In (n,x) _ ->
                                        In (n,y) _ -> x = y |- _ =>
            assert (x = y) by (eapply H''; eauto)
          end; destruct carry_flag; cbn [Z.b2z] in *; congruence. }
    { match goal with |- context [make_ctx ?l ?n] =>
                      let H := fresh in
                      destruct (make_ctx_cases l n) as [H | H];
                        [ rewrite H | apply in_app_or in H; destruct H ]
      end; eauto. }
  Qed.
End Proofs.

(* Tactics and lemmas that help to prove two Fancy expressions are
equivalent. *)
Section Equivalence.
  Context {name} (name_eqb : name -> name -> bool) (wordmax : Z).
  Let interp := interp name_eqb wordmax cc_spec.
  Lemma interp_step i rd args cont cc ctx :
    interp (Instr i rd args cont) cc ctx =
    let result := spec i (Tuple.map ctx args) cc in
    let new_cc := CC.update (writes_conditions i) result cc_spec cc in
    let new_ctx := fun n => if name_eqb n rd then result mod wordmax else ctx n in interp cont new_cc new_ctx.
  Proof using Type. reflexivity. Qed.

  Lemma interp_state_equiv e :
    forall cc ctx cc' ctx',
      cc = cc' -> (forall r, ctx r = ctx' r) ->
      interp e cc ctx = interp e cc' ctx'.
  Proof using name_eqb wordmax.
    induction e; intros; subst; cbn; [solve[auto]|].
    apply IHe; rewrite Tuple.map_ext with (g:=ctx') by auto;
      [reflexivity|].
    intros; break_match; auto.
  Qed.

  Local Ltac prove_comm H :=
    rewrite !interp_step; cbn - [interp];
    intros; rewrite H; try reflexivity.

  Lemma mulll_comm rd x y cont cc ctx :
    interp (Instr MUL128LL rd (x, y) cont) cc ctx =
    interp (Instr MUL128LL rd (y, x) cont) cc ctx.
  Proof using name_eqb wordmax. prove_comm Z.mul_comm. Qed.

  Lemma mulhh_comm rd x y cont cc ctx :
    interp (Instr MUL128UU rd (x, y) cont) cc ctx =
    interp (Instr MUL128UU rd (y, x) cont) cc ctx.
  Proof using name_eqb wordmax. prove_comm Z.mul_comm. Qed.

  Lemma mullh_mulhl rd x y cont cc ctx :
    interp (Instr MUL128LU rd (x, y) cont) cc ctx =
    interp (Instr MUL128UL rd (y, x) cont) cc ctx.
  Proof using name_eqb wordmax. prove_comm Z.mul_comm. Qed.

  Lemma add_comm rd x y cont cc ctx :
    0 <= ctx x < 2^256 ->
    0 <= ctx y < 2^256 ->
    interp (Instr (ADD 0) rd (x, y) cont) cc ctx =
    interp (Instr (ADD 0) rd (y, x) cont) cc ctx.
  Proof using name_eqb wordmax.
    prove_comm Z.add_comm.
    rewrite !(Z.mod_small (ctx _)) by lia.
    reflexivity.
  Qed.

  Lemma addc_comm rd x y cont cc ctx :
    0 <= ctx x < 2^256 ->
    0 <= ctx y < 2^256 ->
    interp (Instr (ADDC 0) rd (x, y) cont) cc ctx =
    interp (Instr (ADDC 0) rd (y, x) cont) cc ctx.
  Proof using name_eqb wordmax.
    prove_comm (Z.add_comm (ctx x)).
    rewrite !(Z.mod_small (ctx _)) by lia.
    reflexivity.
  Qed.
End Equivalence.

(* Tactics for proving equivalence *)

(* pose the equivalence of a single instruction's result *)
Ltac results_equiv :=
  match goal with
    |- ?lhs = ?rhs =>
    match lhs with
      context [spec ?li ?largs ?lcc] =>
      match rhs with
        context [spec ?ri ?rargs ?rcc] =>
        replace (spec li largs lcc) with (spec ri rargs rcc)
      end
    end
  end.

(* simplify the updating of condition codes to prevent exponential blowup in [step] *)
Ltac simplify_cc :=
  match goal with
    |- context [CC.update ?to_write ?result ?cc_spec ?old_state] =>
    let e := eval cbv -[spec cc_spec CC.cc_l CC.cc_m CC.cc_z CC.cc_c] in
    (CC.update to_write result cc_spec old_state) in
        change (CC.update to_write result cc_spec old_state) with e
  end.

(* remember the result of an instruction (again, prevents exponential blowup in [step] *)
Ltac remember_single_result :=
  match goal with |- context [(spec ?i ?args ?cc) mod ?w] =>
                  let x := fresh "x" in
                  let y := fresh "y" in
                  remember (spec i args cc) as x;
                  remember (x mod w) as y
  end.

(* rewrite with interp_step and remember the results *)
Ltac step :=
  match goal with
    |- interp _ _ _ (Instr ?i ?rd1 ?args1 ?cont1) ?cc1 ?ctx1 =
       interp _ _ _ (Instr ?i ?rd2 ?args2 ?cont2) ?cc2 ?ctx2 =>
    rewrite (interp_step _ _ i rd1 args1 cont1);
    rewrite (interp_step _ _ i rd2 args2 cont2)
  end;
  cbn - [interp spec cc_spec];
  repeat progress
         rewrite ?Registers.reg_eqb_neq, ?Registers.reg_eqb_refl by congruence;
  results_equiv; [ remember_single_result; repeat simplify_cc | try reflexivity ].


(** Extraction of [Z] into Haskell's [Int] *)

Require Coq.extraction.Extraction.

Require Import ZArith.
Require Import ExtrHaskellZNum.

(**
 * Disclaimer: trying to obtain efficient certified programs
 * by extracting [Z] into [Int] is definitively *not* a good idea.
 * See comments in [ExtrOcamlNatInt.v].
 *)

Extract Inductive positive => "Prelude.Int" [
  "(\x -> 2 Prelude.* x Prelude.+ 1)"
  "(\x -> 2 Prelude.* x)"
  "1" ]
  "(\fI fO fH n -> if n Prelude.== 1 then fH () else
                   if Prelude.odd n
                   then fI (n `Prelude.div` 2)
                   else fO (n `Prelude.div` 2))".

Extract Inductive Z => "Prelude.Int" [ "0" "(\x -> x)" "Prelude.negate" ]
  "(\fO fP fN n -> if n Prelude.== 0 then fO () else
                   if n Prelude.> 0 then fP n else
                   fN (Prelude.negate n))".


From Coq Require Import Bool String List BinPos Compare_dec Lia Arith.
From Equations Require Import Equations.
From Translation
Require Import util Sorts SAst SLiftSubst SCommon Equality.
Import ListNotations.

Open Scope s_scope.

(*! Typing *)

Section ITyping.

Context `{Sort_notion : Sorts.notion}.

Reserved Notation " Σ ;;; Γ '|-i' t : T " (at level 50, Γ, t, T at next level).

Inductive typing (Σ : sglobal_context) : scontext -> sterm -> sterm -> Prop :=
| type_Rel Γ n :
    wf Σ Γ ->
    forall A,
      nth_error Γ n = Some A ->
      Σ ;;; Γ |-i (sRel n) : lift0 (S n) A

| type_Sort Γ s :
    wf Σ Γ ->
    Σ ;;; Γ |-i (sSort s) : sSort (Sorts.succ s)

| type_Prod Γ n t b s1 s2 :
    Σ ;;; Γ |-i t : sSort s1 ->
    Σ ;;; Γ ,, t |-i b : sSort s2 ->
    Σ ;;; Γ |-i (sProd n t b) : sSort (Sorts.prod_sort s1 s2)

| type_Lambda Γ n n' t b s1 s2 bty :
    Σ ;;; Γ |-i t : sSort s1 ->
    Σ ;;; Γ ,, t |-i bty : sSort s2 ->
    Σ ;;; Γ ,, t |-i b : bty ->
    Σ ;;; Γ |-i (sLambda n t bty b) : sProd n' t bty

| type_App Γ n s1 s2 t A B u :
    Σ ;;; Γ |-i A : sSort s1 ->
    Σ ;;; Γ ,, A |-i B : sSort s2 ->
    Σ ;;; Γ |-i t : sProd n A B ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i (sApp t A B u) : B{ 0 := u }

| type_Sum Γ n t b s1 s2 :
    Σ ;;; Γ |-i t : sSort s1 ->
    Σ ;;; Γ ,, t |-i b : sSort s2 ->
    Σ ;;; Γ |-i (sSum n t b) : sSort (Sorts.sum_sort s1 s2)

| type_Pair Γ n A B u v s1 s2 :
    Σ ;;; Γ |-i A : sSort s1 ->
    Σ ;;; Γ ,, A |-i B : sSort s2 ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i v : B{ 0 := u } ->
    Σ ;;; Γ |-i sPair A B u v : sSum n A B

| type_Pi1 Γ n A B s1 s2 p :
    Σ ;;; Γ |-i p : sSum n A B ->
    Σ ;;; Γ |-i A : sSort s1 ->
    Σ ;;; Γ ,, A |-i B : sSort s2 ->
    Σ ;;; Γ |-i sPi1 A B p : A

| type_Pi2 Γ n A B s1 s2 p :
    Σ ;;; Γ |-i p : sSum n A B ->
    Σ ;;; Γ |-i A : sSort s1 ->
    Σ ;;; Γ ,, A |-i B : sSort s2 ->
    Σ ;;; Γ |-i sPi2 A B p : B{ 0 := sPi1 A B p }

| type_Eq Γ s A u v :
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i v : A ->
    Σ ;;; Γ |-i sEq A u v : sSort (Sorts.eq_sort s)

| type_Refl Γ s A u :
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i sRefl A u : sEq A u u

| type_J Γ s1 s2 A u v P p w :
    Σ ;;; Γ |-i A : sSort s1 ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i v : A ->
    Σ ;;; Γ ,, A ,, (sEq (lift0 1 A) (lift0 1 u) (sRel 0)) |-i P : sSort s2 ->
    Σ ;;; Γ |-i p : sEq A u v ->
    Σ ;;; Γ |-i w : P{ 1 := u }{ 0 := sRefl A u } ->
    Σ ;;; Γ |-i sJ A u P w v p : P{ 1 := v }{ 0 := p }

| type_Transport Γ s T1 T2 p t :
    Σ ;;; Γ |-i T1 : sSort s ->
    Σ ;;; Γ |-i T2 : sSort s ->
    Σ ;;; Γ |-i p : sEq (sSort s) T1 T2 ->
    Σ ;;; Γ |-i t : T1 ->
    Σ ;;; Γ |-i sTransport T1 T2 p t : T2

| type_Beta Γ A B t u n s :
    Σ ;;; Γ ,, A |-i t : B ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i sBeta t u : sEq (B{ 0 := u })
                               (sApp (sLambda n A B t) A B u)
                               (t{ 0 := u })

| type_Heq Γ A a B b s :
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i B : sSort s ->
    Σ ;;; Γ |-i a : A ->
    Σ ;;; Γ |-i b : B ->
    Σ ;;; Γ |-i sHeq A a B b : sSort (heq_sort s)

| type_HeqToEq Γ A u v p s :
    Σ ;;; Γ |-i p : sHeq A u A v ->
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i v : A ->
    Σ ;;; Γ |-i sHeqToEq p : sEq A u v

| type_HeqRefl Γ A a s :
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i a : A ->
    Σ ;;; Γ |-i sHeqRefl A a : sHeq A a A a

| type_HeqSym Γ A a B b p s :
    Σ ;;; Γ |-i p : sHeq A a B b ->
    Σ ;;; Γ |-i a : A ->
    Σ ;;; Γ |-i b : B ->
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i B : sSort s ->
    Σ ;;; Γ |-i sHeqSym p : sHeq B b A a

| type_HeqTrans Γ A a B b C c p q s :
    Σ ;;; Γ |-i p : sHeq A a B b ->
    Σ ;;; Γ |-i q : sHeq B b C c ->
    Σ ;;; Γ |-i a : A ->
    Σ ;;; Γ |-i b : B ->
    Σ ;;; Γ |-i c : C ->
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i B : sSort s ->
    Σ ;;; Γ |-i C : sSort s ->
    Σ ;;; Γ |-i sHeqTrans p q : sHeq A a C c

| type_HeqTransport Γ A B p t s :
    Σ ;;; Γ |-i p : sEq (sSort s) A B ->
    Σ ;;; Γ |-i t : A ->
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i B : sSort s ->
    Σ ;;; Γ |-i sHeqTransport p t : sHeq A t B (sTransport A B p t)

| type_CongProd Γ s z nx ny A1 A2 B1 B2 pA pB :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ ,, (sPack A1 A2)
    |-i pB : sHeq (sSort z) ((lift 1 1 B1){ 0 := sProjT1 (sRel 0) })
                 (sSort z) ((lift 1 1 B2){ 0 := sProjT2 (sRel 0) }) ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ ,, A1 |-i B1 : sSort z ->
    Σ ;;; Γ ,, A2 |-i B2 : sSort z ->
    Σ ;;; Γ |-i sCongProd B1 B2 pA pB :
    sHeq (sSort (Sorts.prod_sort s z)) (sProd nx A1 B1)
         (sSort (Sorts.prod_sort s z)) (sProd ny A2 B2)

| type_CongLambda Γ s z nx ny A1 A2 B1 B2 t1 t2 pA pB pt :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ ,, (sPack A1 A2)
    |-i pB : sHeq (sSort z) ((lift 1 1 B1){ 0 := sProjT1 (sRel 0) })
                 (sSort z) ((lift 1 1 B2){ 0 := sProjT2 (sRel 0) }) ->
    Σ ;;; Γ ,, (sPack A1 A2)
    |-i pt : sHeq ((lift 1 1 B1){ 0 := sProjT1 (sRel 0) })
                 ((lift 1 1 t1){ 0 := sProjT1 (sRel 0) })
                 ((lift 1 1 B2){ 0 := sProjT2 (sRel 0) })
                 ((lift 1 1 t2){ 0 := sProjT2 (sRel 0) }) ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ ,, A1 |-i B1 : sSort z ->
    Σ ;;; Γ ,, A2 |-i B2 : sSort z ->
    Σ ;;; Γ ,, A1 |-i t1 : B1 ->
    Σ ;;; Γ ,, A2 |-i t2 : B2 ->
    Σ ;;; Γ |-i sCongLambda B1 B2 t1 t2 pA pB pt :
               sHeq (sProd nx A1 B1) (sLambda nx A1 B1 t1)
                    (sProd ny A2 B2) (sLambda ny A2 B2 t2)

| type_CongApp Γ s z nx ny A1 A2 B1 B2 u1 u2 v1 v2 pA pB pu pv :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ ,, (sPack A1 A2)
    |-i pB : sHeq (sSort z) ((lift 1 1 B1){ 0 := sProjT1 (sRel 0) })
                 (sSort z) ((lift 1 1 B2){ 0 := sProjT2 (sRel 0) }) ->
    Σ ;;; Γ |-i pu : sHeq (sProd nx A1 B1) u1 (sProd ny A2 B2) u2 ->
    Σ ;;; Γ |-i pv : sHeq A1 v1 A2 v2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ ,, A1 |-i B1 : sSort z ->
    Σ ;;; Γ ,, A2 |-i B2 : sSort z ->
    Σ ;;; Γ |-i u1 : sProd nx A1 B1 ->
    Σ ;;; Γ |-i u2 : sProd ny A2 B2 ->
    Σ ;;; Γ |-i v1 : A1 ->
    Σ ;;; Γ |-i v2 : A2 ->
    Σ ;;; Γ |-i sCongApp B1 B2 pu pA pB pv :
               sHeq (B1{0 := v1}) (sApp u1 A1 B1 v1)
                    (B2{0 := v2}) (sApp u2 A2 B2 v2)

| type_CongSum Γ s z nx ny A1 A2 B1 B2 pA pB :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ ,, (sPack A1 A2)
    |-i pB : sHeq (sSort z) ((lift 1 1 B1){ 0 := sProjT1 (sRel 0) })
                 (sSort z) ((lift 1 1 B2){ 0 := sProjT2 (sRel 0) }) ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ ,, A1 |-i B1 : sSort z ->
    Σ ;;; Γ ,, A2 |-i B2 : sSort z ->
    Σ ;;; Γ |-i sCongSum B1 B2 pA pB :
    sHeq (sSort (Sorts.sum_sort s z)) (sSum nx A1 B1)
         (sSort (Sorts.sum_sort s z)) (sSum ny A2 B2)

| type_CongPair Γ s z nx ny A1 A2 B1 B2 u1 u2 v1 v2 pA pB pu pv :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ ,, (sPack A1 A2)
    |-i pB : sHeq (sSort z) ((lift 1 1 B1){ 0 := sProjT1 (sRel 0) })
                 (sSort z) ((lift 1 1 B2){ 0 := sProjT2 (sRel 0) }) ->
    Σ ;;; Γ |-i pu : sHeq A1 u1 A2 u2 ->
    Σ ;;; Γ |-i pv : sHeq (B1{ 0 := u1 }) v1 (B2{ 0 := u2 }) v2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ ,, A1 |-i B1 : sSort z ->
    Σ ;;; Γ ,, A2 |-i B2 : sSort z ->
    Σ ;;; Γ |-i u1 : A1 ->
    Σ ;;; Γ |-i u2 : A2 ->
    Σ ;;; Γ |-i v1 : B1{ 0 := u1 } ->
    Σ ;;; Γ |-i v2 : B2{ 0 := u2 } ->
    Σ ;;; Γ |-i sCongPair B1 B2 pA pB pu pv :
    sHeq (sSum nx A1 B1) (sPair A1 B1 u1 v1)
         (sSum ny A2 B2) (sPair A2 B2 u2 v2)

| type_CongPi1 Γ s z nx ny A1 A2 B1 B2 p1 p2 pA pB pp :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ ,, (sPack A1 A2)
    |-i pB : sHeq (sSort z) ((lift 1 1 B1){ 0 := sProjT1 (sRel 0) })
                 (sSort z) ((lift 1 1 B2){ 0 := sProjT2 (sRel 0) }) ->
    Σ ;;; Γ |-i pp : sHeq (sSum nx A1 B1) p1 (sSum ny A2 B2) p2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ ,, A1 |-i B1 : sSort z ->
    Σ ;;; Γ ,, A2 |-i B2 : sSort z ->
    Σ ;;; Γ |-i p1 : sSum nx A1 B1 ->
    Σ ;;; Γ |-i p2 : sSum ny A2 B2 ->
    Σ ;;; Γ |-i sCongPi1 B1 B2 pA pB pp : sHeq A1 (sPi1 A1 B1 p1)
                                              A2 (sPi1 A2 B2 p2)

| type_CongPi2 Γ s z nx ny A1 A2 B1 B2 p1 p2 pA pB pp :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ ,, (sPack A1 A2)
    |-i pB : sHeq (sSort z) ((lift 1 1 B1){ 0 := sProjT1 (sRel 0) })
                 (sSort z) ((lift 1 1 B2){ 0 := sProjT2 (sRel 0) }) ->
    Σ ;;; Γ |-i pp : sHeq (sSum nx A1 B1) p1 (sSum ny A2 B2) p2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ ,, A1 |-i B1 : sSort z ->
    Σ ;;; Γ ,, A2 |-i B2 : sSort z ->
    Σ ;;; Γ |-i p1 : sSum nx A1 B1 ->
    Σ ;;; Γ |-i p2 : sSum ny A2 B2 ->
    Σ ;;; Γ |-i sCongPi2 B1 B2 pA pB pp :
               sHeq (B1{ 0 := sPi1 A1 B1 p1}) (sPi2 A1 B1 p1)
                    (B2{ 0 := sPi1 A2 B2 p2}) (sPi2 A2 B2 p2)

| type_CongEq Γ s A1 A2 u1 u2 v1 v2 pA pu pv :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ |-i pu : sHeq A1 u1 A2 u2 ->
    Σ ;;; Γ |-i pv : sHeq A1 v1 A2 v2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ |-i u1 : A1 ->
    Σ ;;; Γ |-i u2 : A2 ->
    Σ ;;; Γ |-i v1 : A1 ->
    Σ ;;; Γ |-i v2 : A2 ->
    Σ ;;; Γ |-i sCongEq pA pu pv :
               sHeq (sSort (Sorts.eq_sort s)) (sEq A1 u1 v1)
                    (sSort (Sorts.eq_sort s)) (sEq A2 u2 v2)

| type_CongRefl Γ s A1 A2 u1 u2 pA pu :
    Σ ;;; Γ |-i pA : sHeq (sSort s) A1 (sSort s) A2 ->
    Σ ;;; Γ |-i pu : sHeq A1 u1 A2 u2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ |-i u1 : A1 ->
    Σ ;;; Γ |-i u2 : A2 ->
    Σ ;;; Γ |-i sCongRefl pA pu :
               sHeq (sEq A1 u1 u1) (sRefl A1 u1) (sEq A2 u2 u2) (sRefl A2 u2)

| type_EqToHeq Γ A u v p s :
    Σ ;;; Γ |-i p : sEq A u v ->
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i v : A ->
    Σ ;;; Γ |-i sEqToHeq p : sHeq A u A v

| type_HeqTypeEq Γ A u B v p s :
    Σ ;;; Γ |-i p : sHeq A u B v ->
    Σ ;;; Γ |-i A : sSort s ->
    Σ ;;; Γ |-i B : sSort s ->
    Σ ;;; Γ |-i u : A ->
    Σ ;;; Γ |-i v : B ->
    Σ ;;; Γ |-i sHeqTypeEq A B p : sEq (sSort s) A B

| type_Pack Γ A1 A2 s :
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ |-i sPack A1 A2 : sSort (pack_sort s)

| type_ProjT1 Γ A1 A2 p s :
    Σ ;;; Γ |-i p : sPack A1 A2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ |-i sProjT1 p : A1

| type_ProjT2 Γ A1 A2 p s :
    Σ ;;; Γ |-i p : sPack A1 A2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ |-i sProjT2 p : A2

| type_ProjTe Γ A1 A2 p s :
    Σ ;;; Γ |-i p : sPack A1 A2 ->
    Σ ;;; Γ |-i A1 : sSort s ->
    Σ ;;; Γ |-i A2 : sSort s ->
    Σ ;;; Γ |-i sProjTe p : sHeq A1 (sProjT1 p) A2 (sProjT2 p)

| type_Ax Γ id ty :
    wf Σ Γ ->
    lookup_glob Σ id = Some ty ->
    Σ ;;; Γ |-i sAx id : ty

| type_rename Γ t A B :
    Σ ;;; Γ |-i t : A ->
    nl A = nl B ->
    Σ ;;; Γ |-i t : B

where " Σ ;;; Γ '|-i' t : T " := (@typing Σ Γ t T) : i_scope

with wf (Σ : sglobal_context) : scontext -> Prop :=
| wf_nil :
    wf Σ nil

| wf_snoc s Γ A :
    wf Σ Γ ->
    Σ ;;; Γ |-i A : sSort s ->
    wf Σ (Γ ,, A)
.

End ITyping.

Notation " Σ ;;; Γ '|-i' t : T " :=
  (@typing _ Σ Γ t T) (at level 50, Γ, t, T at next level) : i_scope.

Derive Signature for typing.
Derive Signature for wf.

Delimit Scope i_scope with i.

(* Syntactic restriction.
   We define the notion of ETT compatibility to restrict the syntax
   to the one that is allowed in ETT.
   [TODO Move?]
 *)
Section Xcomp.

Context `{Sort_notion : Sorts.notion}.

Inductive Xcomp : sterm -> Type :=
| xcomp_Rel n : Xcomp (sRel n)
| xcomp_Sort s : Xcomp (sSort s)
| xcomp_Prod na A B :
    Xcomp A ->
    Xcomp B ->
    Xcomp (sProd na A B)
| xcomp_Lambda na A B t :
    Xcomp A ->
    Xcomp B ->
    Xcomp t ->
    Xcomp (sLambda na A B t)
| xcomp_App u A B v :
    Xcomp u ->
    Xcomp A ->
    Xcomp B ->
    Xcomp v ->
    Xcomp (sApp u A B v)
| xcomp_Sum na A B :
    Xcomp A ->
    Xcomp B ->
    Xcomp (sSum na A B)
| xcomp_Pair A B u v :
    Xcomp A ->
    Xcomp B ->
    Xcomp u ->
    Xcomp v ->
    Xcomp (sPair A B u v)
| xcomp_Pi1 A B p :
    Xcomp A ->
    Xcomp B ->
    Xcomp p ->
    Xcomp (sPi1 A B p)
| xcomp_Pi2 A B p :
    Xcomp A ->
    Xcomp B ->
    Xcomp p ->
    Xcomp (sPi2 A B p)
| xcomp_Eq A u v :
    Xcomp A ->
    Xcomp u ->
    Xcomp v ->
    Xcomp (sEq A u v)
| xcomp_Refl A u :
    Xcomp A ->
    Xcomp u ->
    Xcomp (sRefl A u)
| xcomp_Ax id :
    Xcomp (sAx id)
.

End Xcomp.

Derive Signature for Xcomp.

Section Global.

Context `{Sort_notion : Sorts.notion}.

Open Scope i_scope.

Definition isType (Σ : sglobal_context) (Γ : scontext) (t : sterm) :=
  exists s, Σ ;;; Γ |-i t : sSort s.

Inductive fresh_glob (id : ident) : sglobal_context -> Prop :=
| fresh_glob_nil : fresh_glob id []
| fresh_glob_cons Σ d :
    fresh_glob id Σ ->
    (dname d) <> id ->
    fresh_glob id (d :: Σ).

Inductive type_glob : sglobal_context -> Type :=
| type_glob_nil : type_glob []
| type_glob_cons Σ d :
    type_glob Σ ->
    fresh_glob (dname d) Σ ->
    isType Σ [] (dtype d) ->
    Xcomp (dtype d) ->
    type_glob (d :: Σ).

End Global.

Derive Signature for fresh_glob.
Derive Signature for type_glob.


From Tealeaves Require Export
  Classes.Decorated.Monad
  Classes.Traversable.Monad
  Classes.DT.Functor.

Import Product.Notations.
Import Monad.Notations.
Import Comonad.Notations.

#[local] Generalizable Variables T F G W A B C.

(** * Decorated-traversable monads *)
(******************************************************************************)
Section DecoratedTraverableMonad.

  Context
  (W : Type)
  (T : Type -> Type)
  `{op : Monoid_op W}
  `{unit : Monoid_unit W}
  `{Fmap T} `{Return T} `{Join T}
  `{Decorate W T} `{Dist T}.

  Class DecoratedTraversableMonad :=
    { dtmon_decorated :> DecoratedMonad W T;
      dtmon_traversable :> TraversableMonad T;
      dtmon_functor :> DecoratedTraversableFunctor W T;
    }.

End DecoratedTraverableMonad.

(** Now we verify that the sub-classes can be inferred as well. *)
(******************************************************************************)
Section test_typeclasses.

  Context
    `{DecoratedTraversableMonad W T}.

  Goal Functor T. typeclasses eauto. Qed.
  Goal DecoratedFunctor W T. typeclasses eauto. Qed.
  (*Goal SetlikeFunctor T. typeclasses eauto. Qed.
  Goal ListableFunctor T. typeclasses eauto. Qed.
  Goal TraversableFunctor T. typeclasses eauto. Qed.
  Goal DecoratedTraversableFunctor W T. typeclasses eauto. Qed.

  Goal Monad T. typeclasses eauto. Qed.
  Goal DecoratedMonad W T. typeclasses eauto. Qed.
  Goal SetlikeMonad T. Fail typeclasses eauto. Abort.
  Goal ListableMonad T. Fail typeclasses eauto. Abort.
  Goal TraversableMonad T. typeclasses eauto. Qed. *)

End test_typeclasses.

(** * Notations *)
(******************************************************************************)
Module Notations.
  Include Monad.Notations.
  Include Comonad.Notations.
End Notations.

(** * Kleisli presentation of D-T monads *)
(******************************************************************************)

From Tealeaves Require
  Classes.Kleisli.DT.Monad.

Module ToKleisli.

  Import Classes.Kleisli.DT.Monad.

  Import Classes.Kleisli.DT.Monad.Notations.
  Import Classes.Monad.Notations.
  Import Classes.Traversable.Functor.Notations.
  Import Data.Strength.Notations.

  Section operation.

    Context
      (T : Type -> Type)
        `{Fmap T} `{Decorate W T} `{Dist T} `{Join T}.

    (* Define <<Binddt>> from dist/dec/join *)
    #[export] Instance Binddt_ddj : Binddt W T T :=
      fun G `{Fmap G} `{Pure G} `{Mult G} A B
        (f : W * A -> G (T B)) =>
        fmap G (join T) ∘ dist T G ∘ fmap T f ∘ dec T.

  End operation.

  #[local] Tactic Notation "bring" constr(x) "and" constr(y) "together" :=
    change (?f ∘ x ∘ y ∘ ?g) with (f ∘ (x ∘ y) ∘ g).

  Section adapter.

    Context
      (T : Type -> Type)
      `{Classes.DT.Monad.DecoratedTraversableMonad W T}.

    Context
      (G1 : Type -> Type) `{Fmap G1} `{Pure G1} `{Mult G1} `{! Applicative G1}
        (G2 : Type -> Type) `{Fmap G2} `{Pure G2} `{Mult G2} `{! Applicative G2}.

    Definition kcompose_dtm_alt :=
      fun `(g : W * B -> G2 (T C))
        `(f : W * A -> G1 (T B))
      => (fmap G1 ((fmap G2 (μ T))
                    ∘ δ T G2
                    ∘ fmap T (g ∘ μ (W ×))
                    ∘ σ T
                    ∘ fmap (W ×) (dec T)))
          ∘ σ G1
          ∘ cobind (W ×) f.

    Lemma equiv' :
      forall  `(g : W * B -> G2 (T C))
         `(f : W * A -> G1 (T B)),
        g ⋆dtm f =
          kcompose_dtm_alt g f.
    Proof.
      intros. unfold kcompose_dtm.
      unfold_ops @Binddt_ddj.
      ext [w a].
      unfold kcompose_dtm_alt.
      unfold compose at 6. cbn.
      unfold compose at 6.
      cbn. compose near (f (w, a)) on right.
      rewrite (fun_fmap_fmap G1).
      repeat fequal.
      ext t.
      unfold compose. cbn.
      compose near (dec T t).
      rewrite (fun_fmap_fmap T).
      unfold compose.
      repeat fequal.
      now ext [w' b].
    Qed.

  End adapter.

  Section with_monad.

    Context
      (T : Type -> Type)
      `{Classes.DT.Monad.DecoratedTraversableMonad W T}.

    Import Decorated.Monad.ToKleisli.
    Import Kleisli.Decorated.Monad.

    Theorem kdtm_binddt1_T {A} : binddt T (fun A => A) (ret T ∘ extract (prod W)) = @id (T A).
    Proof.
      introv. unfold_ops @Binddt_ddj.
      change (fmap (fun A => A) ?f) with f.
      rewrite (dist_unit T).
      change_left (bindd T (η T ∘ extract (prod W) (A := A))).
      ext t. now rewrite (kmond_bindd1 T).
    Qed.

    Theorem kdtm_binddt0_T : forall
        (A B : Type)
        (G : Type -> Type) `{Fmap G} `{Pure G} `{Mult G} `{! Applicative G}
        (f : W * A -> G (T B)),
        binddt T G f ∘ η T = f ∘ η (prod W).
    Proof.
      intros.
      unfold_ops @Binddt_ddj.
      reassociate -> on left.
      rewrite (dmon_ret W T).
      reassociate <- on left.
      reassociate -> near (η T).
      rewrite (natural (ϕ := @ret T _)).
      unfold_ops @Fmap_I.
      reassociate <- on left.
      reassociate -> near (η T).
      rewrite (trvmon_ret T).
      rewrite (fun_fmap_fmap G).
      rewrite (mon_join_ret T).
      rewrite (fun_fmap_id G).
      reflexivity.
    Qed.

    Section binddt_binddt.

      Context
        {A B C : Type}
        (G1 : Type -> Type) `{Fmap G1} `{Pure G1} `{Mult G1} `{! Applicative G1}
        (G2 : Type -> Type) `{Fmap G2} `{Pure G2} `{Mult G2} `{! Applicative G2}
        (g : W * B -> G2 (T C)) (f : W * A -> G1 (T B)).

      Lemma binddt_binddt1 :
        (fmap G1 (dec T ∘ μ T) ∘ δ T G1 : T (G1 (T B)) -> G1 (T (W * B)))
        = fmap G1 (fmap T (μ (prod W)) ∘ μ T ∘ fmap T (σ T ∘ fmap (prod W) (dec T)))
            ∘ δ T G1 ∘ fmap T (σ G1) ∘ dec T.
      Proof.
        rewrite (dmon_join W T). unfold shift.
        change (?f ∘ δ T G1 ∘ fmap T (σ G1) ∘ dec T)
          with (f ∘ (δ T G1 ∘ fmap T (σ G1) ∘ dec T)).
        rewrite (dtfun_compat W T (G := G1)).
        reassociate <- on right. fequal.
        rewrite (fun_fmap_fmap G1).
        rewrite (natural (ϕ := @join T _)).
        fequal. rewrite <- (fun_fmap_fmap T _ _ _ _ (σ T)).
        repeat reassociate <- on right.
        change (fmap T (fmap (prod W) (dec T) (A := ?A)))
          with (fmap (T ∘ (W ×)) (dec T) (A := A)).
        reassociate -> near (dec T).
        #[local] Set Keyed Unification.
        rewrite (natural (ϕ := @dec W T _) (dec T)).
        do 2 reassociate <-.
        reassociate -> near (fmap T (σ T)).
        now rewrite (fun_fmap_fmap T).
        #[local] Unset Keyed Unification.
      Qed.

      (* Note that we *cannot* immediately cancel out the right-most two <<dec T>> sub-expressions *)
      Lemma binddt_binddt2 :
        fmap G1 (fmap T g ∘ dec T ∘ μ T) ∘ δ T G1 ∘ fmap T f ∘ dec T =
          fmap G1 (μ T ∘ fmap T (fmap T (g ∘ μ (prod W)) ∘ σ T ∘ fmap (prod W) (dec T))) ∘ δ T G1
            ∘ fmap T (σ G1 ∘ cobind (prod W) f) ∘ dec T.
      Proof.
        rewrite <- (fun_fmap_fmap T _ _ _ (cobind (W ×) f)).
        do 4 reassociate -> on right.
        rewrite (cobind_dec T f).
        do 5 reassociate <- on right. fequal. fequal.
        change (fmap T (g ∘ μ (prod W)) ∘ σ T ∘ fmap (prod W) (dec T))
          with (fmap T (g ∘ μ (prod W)) ∘ (σ T ∘ fmap (prod W) (dec T))).
        rewrite <- (fun_fmap_fmap T).
        change (fmap T (fmap T (g ∘ μ (prod W)))) with (fmap (T ∘ T) (g ∘ μ (W ×))).
        reassociate <- on right.
        rewrite <- (natural (ϕ := @join T _) (g ∘ μ (W ×))).
        change (fmap T g ∘ dec T ∘ μ T) with (fmap T g ∘ (dec T ∘ μ T)).
        rewrite <- (fun_fmap_fmap G1).
        rewrite <- (fun_fmap_fmap T).
        change ((fmap T g ∘ fmap T (μ (prod W)) ∘ μ T
                   ∘ fmap T (σ T ∘ fmap (prod W) (dec T))))
          with ((fmap T g ∘ (fmap T (μ (prod W)) ∘ μ T
                               ∘ fmap T (σ T ∘ fmap (prod W) (dec T))))).
        rewrite <- (fun_fmap_fmap G1 _ _ _ _ (fmap T g)).
        reassociate -> on left.
        do 4 reassociate -> on right.
        fequal. do 3 reassociate <- on right.
        apply binddt_binddt1.
      Qed.

      Theorem kdtm_binddt2_T :
        fmap G1 (binddt T G2 g) ∘ binddt T G1 f = binddt T (G1 ∘ G2) (g ⋆dtm f).
      Proof.
        unfold binddt at 1 2 3; unfold Binddt_ddj at 1 2 3.
        repeat reassociate <-.
        rewrite (dist_linear T). reassociate <- on right.
        #[local] Set Keyed Unification.
        rewrite 2(fun_fmap_fmap G1).
        #[local] Unset Keyed Unification.
        rewrite (equiv' T G1 G2 g f).
        unfold kcompose_dtm_alt.
        reassociate -> near (cobind (W ×) f).
        rewrite <- (fun_fmap_fmap T).
        change (fmap T (fmap G1 ?f)) with (fmap (T ∘ G1) f).
        reassociate <- on right.
        change (fmap G1 (fmap G2 (μ T) ∘ δ T G2) ∘ ?dist ∘ ?op)
          with (fmap G1 (fmap G2 (μ T) ∘ δ T G2) ∘ (dist ∘ op)).
        unfold_compose_in_compose.
        rewrite <- (natural (ϕ := @dist T _ G1 _ _ _)).
        change (fmap (G1 ∘ T) ?f) with (fmap G1 (fmap T f)).
        reassociate <-.
        #[local] Set Keyed Unification.
        rewrite (fun_fmap_fmap G1).
        #[local] Unset Keyed Unification.
        change (fmap T (fmap G2 (μ T) ∘ δ T G2 ∘ fmap T (g ∘ μ (prod W)) ∘ σ T ∘ fmap (prod W) (dec T)))
          with (fmap T (fmap G2 (μ T) ∘ (δ T G2 ∘ fmap T (g ∘ μ (prod W)) ∘ σ T ∘ fmap (prod W) (dec T)))).
        rewrite <- (fun_fmap_fmap T).
        change (fmap T (fmap G2 ?f)) with (fmap (T ∘ G2) f).
        reassociate <- on right. reassociate -> near (fmap (T ∘ G2) (μ T)).
        unfold_compose_in_compose.
        rewrite <- (natural (ϕ := @dist T _ G2 _ _ _)).
        reassociate <- on right.
        #[local] Set Keyed Unification.
        rewrite (fun_fmap_fmap G2).
        rewrite <- (mon_join_join T).
        #[local] Unset Keyed Unification.
        rewrite <- (fun_fmap_fmap G2).
        reassociate -> near (δ T G2).
        change (δ T G2 ∘ fmap T (g ∘ μ (prod W)) ∘ σ T ∘ fmap (prod W) (dec T))
          with (δ T G2 ∘ (fmap T (g ∘ μ (prod W)) ∘ σ T ∘ fmap (prod W) (dec T))).
        rewrite <- (fun_fmap_fmap T).
        do 2 reassociate <- on right.
        change (fmap G2 (μ T) ∘ fmap G2 (μ T) ∘ δ T G2 ∘ fmap T (δ T G2) (A := ?A))
          with (fmap G2 (μ T) ∘ (fmap G2 (μ T) ∘ δ T G2 ∘ fmap T (δ T G2) (A := A))).
        change (fmap G2 (μ T) ∘ δ T G2 ∘ fmap T (δ T G2) (A := T (G2 ?A)))
          with (fmap G2 (μ T) ∘ δ (T ∘ T) G2 (A := A)).
        rewrite <- (trvmon_join T).
        reassociate <- on right.
        change_left (fmap G1 (fmap G2 (μ T) ∘ δ T G2 ∘ (fmap T g ∘ dec T ∘ μ T)) ∘ δ T G1 ∘ fmap T f ∘ dec T).
        change_right (fmap G1 (fmap G2 (μ T) ∘ δ T G2 ∘
                                 (μ T ∘ fmap T (fmap T (g ∘ μ (prod W)) ∘ σ T ∘ fmap (prod W) (dec T))))
                        ∘ δ T G1
                        ∘ fmap T (σ G1 ∘ cobind (prod W) f) ∘ dec T).
        (* cancel out the common sub-expressions and apply a lemma to the remaining. *)
        rewrite <- (fun_fmap_fmap G1).
        rewrite <- (fun_fmap_fmap G1 _ _ _ _ (fmap G2 (μ T) ∘ δ T G2)).
        do 3 reassociate -> on left.
        do 3 reassociate -> on right.
        fequal. repeat reassociate <-.
        apply binddt_binddt2.
      Qed.

    End binddt_binddt.

    Lemma kdtm_morph_T :
      forall (G1 G2 : Type -> Type)
        `{Fmap G1} `{Pure G1} `{Mult G1}
        `{Fmap G2} `{Pure G2} `{Mult G2}
        (ϕ : forall A : Type, G1 A -> G2 A),
        ApplicativeMorphism G1 G2 ϕ ->
        forall (A B : Type) (f : W * A -> G1 (T B)),
          ϕ (T B) ∘ binddt T G1 f = binddt T G2 (ϕ (T B) ∘ f).
    Proof.
      introv morph. intros.
      unfold_ops @Binddt_ddj.
      do 3 reassociate <- on left.
      fequal.
      inversion morph.
      unfold compose. ext t.
      rewrite appmor_natural.
      fequal.
      compose near (fmap T f t).
      rewrite <- (dist_morph T).
      unfold compose. compose near t on left.
      now rewrite (fun_fmap_fmap T).
    Qed.

    #[export] Instance: Kleisli.DT.Monad.Monad W T :=
      {| kdtm_binddt0 := @kdtm_binddt0_T;
        kdtm_binddt1 := @kdtm_binddt1_T;
        kdtm_binddt2 := @kdtm_binddt2_T;
        kdtm_morph := @kdtm_morph_T;
      |}.

  End with_monad.

End ToKleisli.


From mathcomp
Require Import ssreflect ssrbool ssrnat eqtype seq ssrfun.
From fcsl
Require Import prelude pred pcm unionmap heap.
From HTT
Require Import stmod stsep stlog stlogR.
From SSL
Require Import core.
From Hammer Require Import Hammer.
(* Configure Hammer *)
Set Hammer ATPLimit 60.
Unset Hammer Eprover.
Unset Hammer Vampire.
Add Search Blacklist "fcsl.".
Add Search Blacklist "HTT.".
Add Search Blacklist "Coq.ssr.ssrfun".
Add Search Blacklist "mathcomp.ssreflect.ssrfun".
Add Search Blacklist "mathcomp.ssreflect.bigop".
Add Search Blacklist "mathcomp.ssreflect.choice".
Add Search Blacklist "mathcomp.ssreflect.div".
Add Search Blacklist "mathcomp.ssreflect.finfun".
Add Search Blacklist "mathcomp.ssreflect.fintype".
Add Search Blacklist "mathcomp.ssreflect.path".
Add Search Blacklist "mathcomp.ssreflect.tuple".

Inductive sll (x : ptr) (s : seq nat) (h : heap) : Prop :=
| sll_1 of (x) == (null) of
  @perm_eq nat_eqType (s) (@nil nat) /\ h = empty
| sll_2 of ~~ ((x) == (null)) of
  exists (v : nat) (s1 : seq nat) (nxt : ptr),
  exists h_sll_nxts1_1,
  @perm_eq nat_eqType (s) (([:: v]) ++ (s1)) /\ h = x :-> (v) \+ x .+ 1 :-> (nxt) \+ h_sll_nxts1_1 /\ sll nxt s1 h_sll_nxts1_1.

Inductive dll (x : ptr) (z : ptr) (s : seq nat) (h : heap) : Prop :=
| dll_1 of (x) == (null) of
  @perm_eq nat_eqType (s) (@nil nat) /\ h = empty
| dll_2 of ~~ ((x) == (null)) of
  exists (v : nat) (s1 : seq nat) (w : ptr),
  exists h_dll_wxs1_0,
  @perm_eq nat_eqType (s) (([:: v]) ++ (s1)) /\ h = x :-> (v) \+ x .+ 1 :-> (w) \+ x .+ 2 :-> (z) \+ h_dll_wxs1_0 /\ dll w x s1 h_dll_wxs1_0.


Require Import Coq.Lists.List.

Import ListNotations.

(** Some useful functions on lists. *)
Module List.
  Fixpoint map_pair {A B C : Type} (f : A -> B -> C) (l : list (A * B))
    : list C :=
    match l with
    | [] => []
    | (a, b) :: l => f a b :: map_pair f l
    end.

  Fixpoint map_triple {A B C D : Type} (f : A -> B -> C -> D)
    (l : list (A * B * C)) : list D :=
    match l with
    | [] => []
    | (a, b, c) :: l => f a b c :: map_triple f l
    end.

  Fixpoint map_quad {A B C D E : Type} (f : A -> B -> C -> D -> E)
    (l : list (A * B * C * D)) : list E :=
    match l with
    | [] => []
    | (a, b, c, d) :: l => f a b c d :: map_quad f l
    end.
End List.


Require Import Bedrock.Platform.Cito.SyntaxExpr Bedrock.Platform.Cito.SemanticsExpr.
Require Import Coq.Lists.List.

Fixpoint varsIn expr:=
match expr with
  | Var s => s :: nil
  | Const w => nil
  | SyntaxExpr.Binop op a b => varsIn a ++ varsIn b
  | SyntaxExpr.TestE op a b => varsIn a ++ varsIn b
end.


Require Import Category.Lib.
Require Import Category.Instance.Lambda.Ltac.
Require Import Category.Instance.Lambda.Ty.
Require Import Category.Instance.Lambda.Exp.
Require Import Category.Instance.Lambda.Value.
Require Import Category.Instance.Lambda.Ren.
Require Import Category.Instance.Lambda.Sub.
Require Import Category.Instance.Lambda.Sem.
Require Import Category.Instance.Lambda.Step.

From Equations Require Import Equations.
Set Equations With UIP.

Section Multi.

Generalizable All Variables.

Inductive multi `(R : crelation X) : crelation X :=
  | multi_refl (x : X) : multi x x
  | multi_step (x y z : X) : R x y → multi y z → multi x z.

Derive Signature for multi.

Theorem multi_R `(R : crelation X) (x y : X) :
  R x y → multi R x y.
Proof.
  intros.
  eapply multi_step; eauto.
  now constructor.
Qed.

Theorem multi_trans `(R : crelation X) (x y z : X) :
  multi R x y →
  multi R y z →
  multi R x z.
Proof.
  intros.
  induction X0; auto.
  now eapply multi_step; eauto.
Qed.

#[export]
Program Instance multi_PreOrder `(R : crelation X) :
  PreOrder (multi R).
Next Obligation. now constructor. Qed.
Next Obligation. now eapply multi_trans; eauto. Qed.

Notation " t '--->*' t' " := (multi Step t t') (at level 40).

#[export]
Program Instance multi_respects_Step {Γ τ} :
  Proper (Step --> Step ==> arrow) (multi (Step (Γ:=Γ) (τ:=τ))).
Next Obligation.
  econstructor; eauto.
  generalize dependent y0.
  generalize dependent y.
  induction X; intros; eauto.
  - now do 4 (econstructor; eauto).
  - unfold flip in *.
    now econstructor; eauto.
Qed.

#[export]
Program Instance multi_respects_multi `(R : crelation X) :
  Proper (multi R --> multi R ==> arrow) (multi R).
Next Obligation.
  unfold flip in *.
  transitivity x; eauto.
  now transitivity x0; eauto.
Qed.

#[local] Hint Constructors ValueP Step : core.

#[local] Hint Extern 7 (_ ---> _) => repeat econstructor : core.

Corollary values_final {Γ τ} {e e' : Exp Γ τ} :
  e --->* e' → ValueP e → e = e'.
Proof.
  intros.
  apply value_is_nf in H.
  unfold normal_form in H.
  induction X; auto.
  now intuition.
Qed.

Lemma multistep_AppR {Γ dom cod} {e e' : Γ ⊢ dom} {v : Γ ⊢ (dom ⟶ cod)} :
  (e --->* e') → ValueP v → APP v e --->* APP v e'.
Proof.
  intros.
  induction X.
  - now apply multi_refl.
  - rewrite <- IHX; clear IHX; auto.
    inv H.
    eapply (AppR_LAM (e:=e)) in r.
    now apply multi_R.
Qed.

Ltac simpl_multistep :=
  intros;
  match goal with
  | [ H : _ --->* _ |- _ ] => induction H
  end;
  [ now apply multi_refl
  | eapply multi_step; eauto;
    now econstructor; eauto ].

Lemma multistep_Pair1 {Γ τ1 τ2} {e1 e1' : Γ ⊢ τ1} {e2 : Γ ⊢ τ2} :
  (e1 --->* e1') → Pair e1 e2 --->* Pair e1' e2.
Proof. now simpl_multistep. Qed.

Lemma multistep_Pair2 {Γ τ1 τ2} {e1 : Γ ⊢ τ1} {e2 e2' : Γ ⊢ τ2} :
  ValueP e1 → (e2 --->* e2') → Pair e1 e2 --->* Pair e1 e2'.
Proof. now simpl_multistep. Qed.

Lemma multistep_Pair {Γ τ1 τ2} {e1 e1' : Γ ⊢ τ1} {e2 e2' : Γ ⊢ τ2} :
  ValueP e1' →
  (e1 --->* e1') → (e2 --->* e2') → Pair e1 e2 --->* Pair e1' e2'.
Proof.
  intros.
  erewrite multistep_Pair1; eauto.
  erewrite multistep_Pair2; eauto.
  now apply multi_refl.
Qed.

Lemma multistep_Fst1 {Γ τ1 τ2} {p p' : Γ ⊢ (τ1 × τ2)} :
  (p --->* p') → Fst p --->* Fst p'.
Proof. now simpl_multistep. Qed.

Lemma multistep_Snd1 {Γ τ1 τ2} {p p' : Γ ⊢ (τ1 × τ2)} :
  (p --->* p') → Snd p --->* Snd p'.
Proof. now simpl_multistep. Qed.

End Multi.

Notation " t '--->*' t' " := (multi Step t t') (at level 40).


Require Export SystemFR.RedTactics.

Lemma open_reducible_err:
  forall Θ Γ T,
    [ Θ; Γ ⊨ tfalse ≡ ttrue ] ->
    [ Θ; Γ ⊨ notype_err : T ].
Proof.
  unfold open_reducible, open_equivalent;
    repeat step || t_instantiate_sat3;
    eauto using false_true_not_equivalent with exfalso.
Qed.


Require Import Crypto.Reflection.SmartMap.
Require Import Crypto.Reflection.Named.Syntax.
Require Import Crypto.Reflection.Named.MapCast.
Require Import Crypto.Reflection.Named.InterpretToPHOAS.
Require Import Crypto.Reflection.Named.Compile.
Require Import Crypto.Reflection.Named.PositiveContext.
Require Import Crypto.Reflection.Named.PositiveContext.Defaults.
Require Import Crypto.Reflection.Linearize.
Require Import Crypto.Reflection.Eta.
Require Import Crypto.Reflection.Syntax.

Section language.
  Context {base_type_code : Type}
          {op : flat_type base_type_code -> flat_type base_type_code -> Type}
          (base_type_code_beq : base_type_code -> base_type_code -> bool)
          (base_type_code_bl_transparent : forall x y, base_type_code_beq x y = true -> x = y)
          (failb : forall var t, @Syntax.exprf base_type_code op var (Tbase t))
          {interp_base_type_bounds : base_type_code -> Type}
          (interp_op_bounds : forall src dst, op src dst -> interp_flat_type interp_base_type_bounds src -> interp_flat_type interp_base_type_bounds dst)
          (pick_typeb : forall t, interp_base_type_bounds t -> base_type_code).
  Local Notation pick_type v := (SmartFlatTypeMap pick_typeb v).
  Context (cast_op : forall t tR (opc : op t tR) args_bs,
              op (pick_type args_bs) (pick_type (interp_op_bounds t tR opc args_bs))).

  Local Notation PContext var := (PositiveContext _ var _ base_type_code_bl_transparent).

  Section MapCast.
    Context {t} (e : Expr base_type_code op t)
            (input_bounds : interp_flat_type interp_base_type_bounds (domain t)).

    Definition MapCastPreProcess
      := ExprEta (Linearize e).
    Definition MapCastCompile (e' : Expr base_type_code op (domain t -> codomain t))
      := compile (e' _) (DefaultNamesFor e').
    Definition MapCastDoCast (e' : option (Named.expr base_type_code op BinNums.positive (domain t -> codomain t)))
      := option_map
           (fun e'' => map_cast
                         (t:=Arrow (domain t) (codomain t))
                         interp_op_bounds pick_typeb cast_op
                         (BoundsContext:=PContext _)
                         empty
                         e''
                         input_bounds)
           e'.
    Definition MapCastDoInterp
               (e' : option
                       (option
                          { output_bounds : interp_flat_type interp_base_type_bounds (codomain t) &
                                           Named.expr _ _ _ (Arrow (pick_type input_bounds) (pick_type output_bounds)) }))
      : option { output_bounds : interp_flat_type interp_base_type_bounds (codomain t)
                                 & Expr base_type_code op (Arrow (pick_type input_bounds) (pick_type output_bounds)) }
      := match e' with
         | Some (Some (existT output_bounds e''))
           => Some (existT _ output_bounds (Eta.ExprEta (InterpToPHOAS (Context:=fun var => PContext var) failb e'')))
         | Some None | None => None
         end.
    Definition MapCast
      : option { output_bounds : interp_flat_type interp_base_type_bounds (codomain t)
                                 & Expr base_type_code op (Arrow (pick_type input_bounds) (pick_type output_bounds)) }
      := MapCastDoInterp (MapCastDoCast (MapCastCompile MapCastPreProcess)).
  End MapCast.
End language.