Require Import Classical Sets ClassicalEpsilon FunctionalExtensionality.

Set Implicit Arguments.

Module Type S.
  Parameter fmap : Type -> Type -> Type.

  Parameter empty : forall A B, fmap A B.
  Parameter lookup : forall A B, fmap A B -> A -> option B.
  Parameter add : forall A B, fmap A B -> A -> B -> fmap A B.

  Parameter remove : forall A B, fmap A B -> A -> fmap A B.
  Parameter join : forall A B, fmap A B -> fmap A B -> fmap A B.
  Parameter merge : forall A B, (option B -> option B -> option B) -> fmap A B -> fmap A B -> fmap A B.
  Parameter restrict : forall A B, (A -> Prop) -> fmap A B -> fmap A B.
  Parameter includes : forall A B, fmap A B -> fmap A B -> Prop.

  Notation "$0" := (empty _ _).
  Notation "m $+ ( k , v )" := (add m k v) (at level 50, left associativity).
  Infix "$-" := remove (at level 50, left associativity).
  Infix "$++" := join (at level 50, left associativity).
  Infix "$?" := lookup (at level 50, no associativity).
  Infix "$<=" := includes (at level 75).

  Parameter dom : forall A B, fmap A B -> set A.

  Axiom fmap_ext : forall A B (m1 m2 : fmap A B),
    (forall k, m1 $? k = m2 $? k)
    -> m1 = m2.

  Axiom lookup_empty : forall A B k, empty A B $? k = None.

  Axiom includes_lookup : forall A B (m m' : fmap A B) k v,
    m $? k = Some v
    -> m $<= m'
    -> lookup m' k = Some v.

  Axiom includes_add : forall A B (m m' : fmap A B) k v,
    m $<= m'
    -> add m k v $<= add m' k v.

  Axiom lookup_add_eq : forall A B (m : fmap A B) k1 k2 v,
    k1 = k2
    -> add m k1 v $? k2 = Some v.

  Axiom lookup_add_ne : forall A B (m : fmap A B) k k' v,
    k' <> k
    ->  add m k v $? k' = m $? k'.

  Axiom lookup_remove_eq : forall A B (m : fmap A B) k1 k2,
    k1 = k2
    -> remove m k1 $? k2 = None.

  Axiom lookup_remove_ne : forall A B (m : fmap A B) k k',
    k' <> k
    ->  remove m k $? k' = m $? k'.

  Axiom lookup_join1 : forall A B (m1 m2 : fmap A B) k,
    k \in dom m1
    -> (m1 $++ m2) $? k = m1 $? k.

  Axiom lookup_join2 : forall A B (m1 m2 : fmap A B) k,
    ~k \in dom m1
    -> (m1 $++ m2) $? k = m2 $? k.

  Axiom join_comm : forall A B (m1 m2 : fmap A B),
    dom m1 \cap dom m2 = constant nil
    -> m1 $++ m2 = m2 $++ m1.

  Axiom join_assoc : forall A B (m1 m2 m3 : fmap A B),
    (m1 $++ m2) $++ m3 = m1 $++ (m2 $++ m3).

  Axiom lookup_merge : forall A B f (m1 m2 : fmap A B) k,
    merge f m1 m2 $? k = f (m1 $? k) (m2 $? k).

  Axiom merge_empty1 : forall A B f (m : fmap A B),
    (forall x, f None x = x)
    -> merge f (@empty _ _) m = m.

  Axiom merge_empty2 : forall A B f (m : fmap A B),
    (forall x, f x None = x)
    -> merge f m (@empty _ _) = m.

  Axiom merge_empty1_alt : forall A B f (m : fmap A B),
    (forall x, f None x = None)
    -> merge f (@empty _ _) m = @empty _ _.

  Axiom merge_empty2_alt : forall A B f (m : fmap A B),
    (forall x, f x None = None)
    -> merge f m (@empty _ _) = @empty _ _.

  Axiom merge_add1 : forall A B f (m1 m2 : fmap A B) k v,
    (forall x y, f (Some x) y = None -> False)
    -> ~k \in dom m1
    -> merge f (add m1 k v) m2 = match f (Some v) (lookup m2 k) with
                                 | None => merge f m1 m2
                                 | Some v => add (merge f m1 m2) k v
                                 end.

  Axiom merge_add2 : forall A B f (m1 m2 : fmap A B) k v,
    (forall x y, f x (Some y) = None -> False)
    -> ~k \in dom m2
    -> merge f m1 (add m2 k v) = match f (lookup m1 k) (Some v) with
                                 | None => merge f m1 m2
                                 | Some v => add (merge f m1 m2) k v
                                 end.

  Axiom merge_add1_alt : forall A B f (m1 m2 : fmap A B) k v,
    (forall x y, f (Some x) (Some y) = None -> False)
    -> ~k \in dom m1
    -> k \in dom m2
    -> merge f (add m1 k v) m2 = match f (Some v) (lookup m2 k) with
                                 | None => merge f m1 m2
                                 | Some v => add (merge f m1 m2) k v
                                 end.

  Axiom empty_includes : forall A B (m : fmap A B), empty A B $<= m.

  Axiom dom_empty : forall A B, dom (empty A B) = constant nil.

  Axiom dom_add : forall A B (m : fmap A B) (k : A) (v : B),
    dom (add m k v) = constant (k :: nil) \cup dom m.

  Axiom lookup_restrict_true : forall A B (P : A -> Prop) (m : fmap A B) k,
    P k
    -> lookup (restrict P m) k = lookup m k.

  Axiom lookup_restrict_false : forall A B (P : A -> Prop) (m : fmap A B) k,
    ~P k
    -> lookup (restrict P m) k = None.

  Axiom lookup_restrict_true_fwd : forall A B (P : A -> Prop) (m : fmap A B) k v,
    lookup (restrict P m) k = Some v
    -> P k.

  Global Hint Extern 1 => match goal with
                     | [ H : lookup (empty _ _) _ = Some _ |- _ ] =>
                       rewrite lookup_empty in H; discriminate
                   end : core.

  Global Hint Resolve includes_lookup includes_add empty_includes : core.

  Global Hint Rewrite lookup_empty lookup_add_eq lookup_add_ne lookup_remove_eq lookup_remove_ne
       lookup_merge lookup_restrict_true lookup_restrict_false using congruence.

  Global Hint Rewrite dom_empty dom_add.

  Ltac maps_equal :=
    apply fmap_ext; intros;
      repeat (subst; autorewrite with core; try reflexivity;
        match goal with
          | [ |- context[lookup (add _ ?k _) ?k' ] ] => destruct (classic (k = k')); subst
        end).

  Global Hint Extern 3 (_ = _) => maps_equal : core.

  Axiom lookup_split : forall A B (m : fmap A B) k v k' v',
    (m $+ (k, v)) $? k' = Some v'
    -> (k' <> k /\ m $? k' = Some v') \/ (k' = k /\ v' = v).

  Global Hint Rewrite merge_empty1 merge_empty2 using solve [ eauto 1 ].
  Global Hint Rewrite merge_empty1_alt merge_empty2_alt using congruence.

  Global Hint Rewrite merge_add1 using solve [ eauto | unfold Sets.In; autorewrite with core in *; simpl in *; try (normalize_set; simpl); intuition congruence ].
  Global Hint Rewrite merge_add1_alt using solve [ congruence | unfold Sets.In; autorewrite with core in *; simpl in *; try (normalize_set; simpl); intuition congruence ].

  Axiom includes_intro : forall K V (m1 m2 : fmap K V),
      (forall k v, m1 $? k = Some v -> m2 $? k = Some v)
      -> m1 $<= m2.

  Axiom lookup_Some_dom : forall K V (m : fmap K V) k v,
      m $? k = Some v
      -> k \in dom m.

  Axiom lookup_None_dom : forall K V (m : fmap K V) k,
      m $? k = None
      -> ~ k \in dom m.

  (* Bits meant for separation logic *)
  Section splitting.
    Variables K V : Type.

    Definition disjoint (h1 h2 : fmap K V) : Prop :=
      forall a, h1 $? a <> None
                -> h2 $? a <> None
                -> False.

    Definition split (h h1 h2 : fmap K V) : Prop :=
      h = h1 $++ h2.

    Axiom split_empty_fwd : forall h h1,
        split h h1 $0
        -> h = h1.

    Axiom split_empty_fwd' : forall h h1,
      split h $0 h1
      -> h = h1.

    Axiom split_empty_bwd : forall h,
      split h h $0.

    Axiom split_empty_bwd' : forall h,
      split h $0 h.

    Axiom disjoint_hemp : forall h,
      disjoint h $0.

    Axiom disjoint_hemp' : forall h,
      disjoint $0 h.

    Axiom disjoint_comm : forall h1 h2,
      disjoint h1 h2
      -> disjoint h2 h1.

    Axiom split_comm : forall h h1 h2,
      disjoint h1 h2
      -> split h h1 h2
      -> split h h2 h1.

    Axiom split_assoc1 : forall h h1 h' h2 h3,
      split h h1 h'
      -> split h' h2 h3
      -> split h (join h1 h2) h3.

    Axiom split_assoc2' : forall h h1 h' h2 h3,
      split h h1 h'
      -> split h' h2 h3
      -> disjoint h1 h'
      -> disjoint h2 h3
      -> split h h2 (join h3 h1).

    Axiom split_assoc2 : forall h h1 h' h2 h3,
      split h h' h1
      -> split h' h2 h3
      -> disjoint h' h1
      -> disjoint h2 h3
      -> split h h2 (join h3 h1).

    Axiom disjoint_assoc1 : forall h h1 h' h2 h3,
      split h h1 h'
      -> split h' h2 h3
      -> disjoint h1 h'
      -> disjoint h2 h3
      -> disjoint (join h1 h2) h3.

    Axiom disjoint_assoc2 : forall h h1 h' h2 h3,
      split h h' h1
      -> split h' h2 h3
      -> disjoint h' h1
      -> disjoint h2 h3
      -> disjoint h2 (join h3 h1).

    Axiom split_join : forall h1 h2,
      split (join h1 h2) h1 h2.

    Axiom split_disjoint : forall h h1 h2 h' h3,
      split h h1 h'
      -> split h' h2 h3
      -> disjoint h1 h'
      -> disjoint h2 h3
      -> disjoint h1 h2.

    Axiom disjoint_assoc3 : forall h h1 h2 h3,
      disjoint h h2
      -> split h h1 h3
      -> disjoint h1 h3
      -> disjoint h3 h2.
  End splitting.

  Global Hint Immediate disjoint_comm split_comm : core.
  Global Hint Immediate split_empty_bwd disjoint_hemp disjoint_hemp' split_assoc1 split_assoc2 : core.
  Global Hint Immediate disjoint_assoc1 disjoint_assoc2 split_join split_disjoint disjoint_assoc3 : core.
End S.

Module M : S.
  Definition fmap (A B : Type) := A -> option B.

  Definition empty A B : fmap A B := fun _ => None.

  Section decide.
    Variable P : Prop.

    Lemma decided : inhabited (sum P (~P)).
    Proof.
      destruct (classic P).
      constructor; exact (inl _ H).
      constructor; exact (inr _ H).
    Qed.

    Definition decide : sum P (~P) :=
      epsilon decided (fun _ => True).
  End decide.

  Definition add A B (m : fmap A B) (k : A) (v : B) : fmap A B :=
    fun k' => if decide (k' = k) then Some v else m k'.
  Definition remove A B (m : fmap A B) (k : A) : fmap A B :=
    fun k' => if decide (k' = k) then None else m k'.
  Definition join A B (m1 m2 : fmap A B) : fmap A B :=
    fun k => match m1 k with
               | None => m2 k
               | x => x
             end.
  Definition merge A B f (m1 m2 : fmap A B) : fmap A B :=
    fun k => f (m1 k) (m2 k).
  Definition lookup A B (m : fmap A B) (k : A) := m k.
  Definition restrict A B (P : A -> Prop) (m : fmap A B) : fmap A B :=
    fun k => if decide (P k) then m k else None.
  Definition includes A B (m1 m2 : fmap A B) :=
    forall k v, m1 k = Some v -> m2 k = Some v.

  Definition dom A B (m : fmap A B) : set A := fun x => m x <> None.

  Theorem fmap_ext : forall A B (m1 m2 : fmap A B),
    (forall k, lookup m1 k = lookup m2 k)
    -> m1 = m2.
  Proof.
    intros; extensionality k; auto.
  Qed.

  Theorem lookup_empty : forall A B (k : A), lookup (empty B) k = None.
  Proof.
    auto.
  Qed.

  Theorem includes_lookup : forall A B (m m' : fmap A B) k v,
    lookup m k = Some v
    -> includes m m'
    -> lookup m' k = Some v.
  Proof.
    auto.
  Qed.

  Theorem includes_add : forall A B (m m' : fmap A B) k v,
    includes m m'
    -> includes (add m k v) (add m' k v).
  Proof.
    unfold includes, add; intuition.
    destruct (decide (k0 = k)); auto.
  Qed.

  Theorem lookup_add_eq : forall A B (m : fmap A B) k1 k2 v,
    k1 = k2
    -> lookup (add m k1 v) k2 = Some v.
  Proof.
    unfold lookup, add; intuition.
    destruct (decide (k2 = k1)); try tauto.
    congruence.
  Qed.

  Theorem lookup_add_ne : forall A B (m : fmap A B) k k' v,
    k' <> k
    -> lookup (add m k v) k' = lookup m k'.
  Proof.
    unfold lookup, add; intuition.
    destruct (decide (k' = k)); intuition.
  Qed.

  Theorem lookup_remove_eq : forall A B (m : fmap A B) k1 k2,
    k1 = k2
    -> lookup (remove m k1) k2 = None.
  Proof.
    unfold lookup, remove; intuition.
    destruct (decide (k2 = k1)); try tauto.
    congruence.
  Qed.

  Theorem lookup_remove_ne : forall A B (m : fmap A B) k k',
    k' <> k
    -> lookup (remove m k) k' = lookup m k'.
  Proof.
    unfold lookup, remove; intuition.
    destruct (decide (k' = k)); try tauto.
  Qed.

  Theorem lookup_join1 : forall A B (m1 m2 : fmap A B) k,
    k \in dom m1
    -> lookup (join m1 m2) k = lookup m1 k.
  Proof.
    unfold lookup, join, dom, In; intros.
    destruct (m1 k); congruence.
  Qed.

  Theorem lookup_join2 : forall A B (m1 m2 : fmap A B) k,
    ~k \in dom m1
    -> lookup (join m1 m2) k = lookup m2 k.
  Proof.
    unfold lookup, join, dom, In; intros.
    destruct (m1 k); try congruence.
    exfalso; apply H; congruence.
  Qed.

  Theorem join_comm : forall A B (m1 m2 : fmap A B),
    dom m1 \cap dom m2 = constant nil
    -> join m1 m2 = join m2 m1.
  Proof.
    intros; apply fmap_ext; unfold join, lookup; intros.
    apply (f_equal (fun f => f k)) in H.
    unfold dom, intersection, constant in H; simpl in H.
    destruct (m1 k), (m2 k); auto.
    exfalso; rewrite <- H.
    intuition congruence.
  Qed.

  Theorem join_assoc : forall A B (m1 m2 m3 : fmap A B),
    join (join m1 m2) m3 = join m1 (join m2 m3).
  Proof.
    intros; apply fmap_ext; unfold join, lookup; intros.
    destruct (m1 k); auto.
  Qed.

  Theorem lookup_merge : forall A B f (m1 m2 : fmap A B) k,
      lookup (merge f m1 m2) k = f (m1 k) (m2 k).
  Proof.
    auto.
  Qed.

  Theorem merge_empty1 : forall A B f (m : fmap A B),
    (forall x, f None x = x)
    -> merge f (@empty _ _) m = m.
  Proof.
    intros; apply fmap_ext; unfold lookup, merge; auto.
  Qed.

  Theorem merge_empty2 : forall A B f (m : fmap A B),
    (forall x, f x None = x)
    -> merge f m (@empty _ _) = m.
  Proof.
    intros; apply fmap_ext; unfold lookup, merge; auto.
  Qed.

  Theorem merge_empty1_alt : forall A B f (m : fmap A B),
    (forall x, f None x = None)
    -> merge f (@empty _ _) m = @empty _ _.
  Proof.
    intros; apply fmap_ext; unfold lookup, merge; auto.
  Qed.

  Theorem merge_empty2_alt : forall A B f (m : fmap A B),
    (forall x, f x None = None)
    -> merge f m (@empty _ _) = @empty _ _.
  Proof.
    intros; apply fmap_ext; unfold lookup, merge; auto.
  Qed.

  Theorem merge_add1 : forall A B f (m1 m2 : fmap A B) k v,
    (forall x y, f (Some x) y = None -> False)
    -> ~k \in dom m1
    -> merge f (add m1 k v) m2 = match f (Some v) (lookup m2 k) with
                                 | None => merge f m1 m2
                                 | Some v => add (merge f m1 m2) k v
                                 end.
  Proof.
    intros; apply fmap_ext; unfold lookup, merge, add; intros.
    destruct (decide (k0 = k)); auto; subst.
    case_eq (f (Some v) (m2 k)); intros.
    case_eq (decide (k = k)); congruence.
    exfalso; eauto.

    case_eq (f (Some v) (m2 k)); intros.
    destruct (decide (k0 = k)); congruence.
    auto.
  Qed.

  Theorem merge_add2 : forall A B f (m1 m2 : fmap A B) k v,
    (forall x y, f x (Some y) = None -> False)
    -> ~k \in dom m2
    -> merge f m1 (add m2 k v) = match f (lookup m1 k) (Some v) with
                                 | None => merge f m1 m2
                                 | Some v => add (merge f m1 m2) k v
                                 end.
  Proof.
    intros; apply fmap_ext; unfold lookup, merge, add; intros.
    destruct (decide (k0 = k)); auto; subst.
    case_eq (f (m1 k) (Some v)); intros.
    case_eq (decide (k = k)); congruence.
    exfalso; eauto.

    case_eq (f (m1 k) (Some v)); intros.
    destruct (decide (k0 = k)); congruence.
    auto.
  Qed.

  Theorem merge_add1_alt : forall A B f (m1 m2 : fmap A B) k v,
    (forall x y, f (Some x) (Some y) = None -> False)
    -> ~k \in dom m1
    -> k \in dom m2
    -> merge f (add m1 k v) m2 = match f (Some v) (lookup m2 k) with
                                 | None => merge f m1 m2
                                 | Some v => add (merge f m1 m2) k v
                                 end.
  Proof.
    intros; apply fmap_ext; unfold lookup, merge, add; intros.
    destruct (decide (k0 = k)); auto; subst.
    case_eq (f (Some v) (m2 k)); intros.
    case_eq (decide (k = k)); congruence.
    case_eq (m2 k); intros.
    rewrite H3 in H2.
    exfalso; eauto.
    congruence.

    case_eq (f (Some v) (m2 k)); intros.
    destruct (decide (k0 = k)); congruence.
    auto.
  Qed.

  Theorem empty_includes : forall A B (m : fmap A B), includes (empty (A := A) B) m.
  Proof.
    unfold includes, empty; intuition congruence.
  Qed.

  Theorem dom_empty : forall A B, dom (empty (A := A) B) = constant nil.
  Proof.
    unfold dom, empty; intros; sets idtac.
  Qed.

  Theorem dom_add : forall A B (m : fmap A B) (k : A) (v : B),
    dom (add m k v) = constant (k :: nil) \cup dom m.
  Proof.
    unfold dom, add; simpl; intros.
    sets ltac:(simpl in *; try match goal with
                                 | [ _ : context[if ?E then _ else _] |- _ ] => destruct E
                               end; intuition congruence).
  Qed.

  Lemma lookup_split : forall A B (m : fmap A B) k v k' v',
    lookup (add m k v) k' = Some v'
    -> (k' <> k /\ lookup m k' = Some v') \/ (k' = k /\ v' = v).
  Proof.
    unfold lookup, add; simpl; intros.
    destruct (decide (k' = k)); intuition congruence.
  Qed.

  Theorem lookup_restrict_true : forall A B (P : A -> Prop) (m : fmap A B) k,
    P k
    -> lookup (restrict P m) k = lookup m k.
  Proof.
    unfold lookup, restrict; intros.
    destruct (decide (P k)); tauto.
  Qed.

  Theorem lookup_restrict_false : forall A B (P : A -> Prop) (m : fmap A B) k,
    ~P k
    -> lookup (restrict P m) k = None.
  Proof.
    unfold lookup, restrict; intros.
    destruct (decide (P k)); tauto.
  Qed.

  Theorem lookup_restrict_true_fwd : forall A B (P : A -> Prop) (m : fmap A B) k v,
    lookup (restrict P m) k = Some v
    -> P k.
  Proof.
    unfold lookup, restrict; intros.
    destruct (decide (P k)); intuition congruence.
  Qed.

  Lemma includes_intro : forall K V (m1 m2 : fmap K V),
      (forall k v, lookup m1 k = Some v -> lookup m2 k = Some v)
      -> includes m1 m2.
  Proof.
    auto.
  Qed.

  Lemma lookup_Some_dom : forall K V (m : fmap K V) k v,
      lookup m k = Some v
      -> k \in dom m.
  Proof.
    unfold lookup, dom, In; congruence.
  Qed.

  Lemma lookup_None_dom : forall K V (m : fmap K V) k,
      lookup m k = None
      -> ~ k \in dom m.
  Proof.
    unfold lookup, dom, In; congruence.
  Qed.

  Section splitting.
    Variables K V : Type.

    Notation "$0" := (@empty K V).
    Notation "m $+ ( k , v )" := (add m k v) (at level 50, left associativity).
    Infix "$-" := remove (at level 50, left associativity).
    Infix "$++" := join (at level 50, left associativity).
    Infix "$?" := lookup (at level 50, no associativity).
    Infix "$<=" := includes (at level 90).

    Definition disjoint (h1 h2 : fmap K V) : Prop :=
      forall a, h1 $? a <> None
                -> h2 $? a <> None
                -> False.

    Definition split (h h1 h2 : fmap K V) : Prop :=
      h = h1 $++ h2.

    Hint Extern 2 (_ <> _) => congruence : core.

    Ltac splt := unfold disjoint, split, join, lookup in *; intros; subst;
                 try match goal with
                     | [ |- @eq (fmap K V) _ _ ] => let a := fresh "a" in extensionality a; simpl
                     end;
                 repeat match goal with
                        | [ a : K, H : forall a : K, _ |- _ ] => specialize (H a)
                        end;
                 repeat match goal with
                        | [ H : _ |- _ ] => rewrite H
                        | [ |- context[match ?E with Some _ => _ | None => _ end] ] => destruct E
                        | [ _ : context[match ?E with Some _ => _ | None => _ end] |- _  ] => destruct E
                        end; eauto; try solve [ exfalso; eauto ].

    Lemma split_empty_fwd : forall h h1,
        split h h1 $0
        -> h = h1.
    Proof.
      splt.
    Qed.

    Lemma split_empty_fwd' : forall h h1,
      split h $0 h1
      -> h = h1.
    Proof.
      splt.
    Qed.

    Lemma split_empty_bwd : forall h,
      split h h $0.
    Proof.
      splt.
    Qed.

    Lemma split_empty_bwd' : forall h,
      split h $0 h.
    Proof.
      splt.
    Qed.

    Lemma disjoint_hemp : forall h,
      disjoint h $0.
    Proof.
      splt.
    Qed.

    Lemma disjoint_hemp' : forall h,
      disjoint $0 h.
    Proof.
      splt.
    Qed.

    Lemma disjoint_comm : forall h1 h2,
      disjoint h1 h2
      -> disjoint h2 h1.
    Proof.
      splt.
    Qed.

    Lemma split_comm : forall h h1 h2,
      disjoint h1 h2
      -> split h h1 h2
      -> split h h2 h1.
    Proof.
      splt.
    Qed.

    Hint Immediate disjoint_comm split_comm : core.

    Lemma split_assoc1 : forall h h1 h' h2 h3,
      split h h1 h'
      -> split h' h2 h3
      -> split h (join h1 h2) h3.
    Proof.
      splt.
    Qed.

    Lemma split_assoc2' : forall h h1 h' h2 h3,
      split h h1 h'
      -> split h' h2 h3
      -> disjoint h1 h'
      -> disjoint h2 h3
      -> split h h2 (join h3 h1).
    Proof.
      splt.
    Qed.

    Lemma split_assoc2 : forall h h1 h' h2 h3,
      split h h' h1
      -> split h' h2 h3
      -> disjoint h' h1
      -> disjoint h2 h3
      -> split h h2 (join h3 h1).
    Proof.
      intros; eapply split_assoc2'; eauto.
    Qed.

    Lemma disjoint_assoc1 : forall h h1 h' h2 h3,
      split h h1 h'
      -> split h' h2 h3
      -> disjoint h1 h'
      -> disjoint h2 h3
      -> disjoint (join h1 h2) h3.
    Proof.
      splt.
    Qed.

    Lemma disjoint_assoc2 : forall h h1 h' h2 h3,
      split h h' h1
      -> split h' h2 h3
      -> disjoint h' h1
      -> disjoint h2 h3
      -> disjoint h2 (join h3 h1).
    Proof.
      splt.
    Qed.

    Lemma split_join : forall h1 h2,
      split (join h1 h2) h1 h2.
    Proof.
      splt.
    Qed.

    Lemma split_disjoint : forall h h1 h2 h' h3,
      split h h1 h'
      -> split h' h2 h3
      -> disjoint h1 h'
      -> disjoint h2 h3
      -> disjoint h1 h2.
    Proof.
      splt.
    Qed.

    Lemma disjoint_assoc3 : forall h h1 h2 h3,
      disjoint h h2
      -> split h h1 h3
      -> disjoint h1 h3
      -> disjoint h3 h2.
    Proof.
      splt.
    Qed.
  End splitting.
End M.

Export M.


From mm Require Import util abt abtutil.

Module op.
  Inductive t' :=
  | arrow
  | all
  | bool

  | abs
  | app
  | tyabs
  | tyapp
  | tt
  | ff
  | If
  .
  Definition t := t'.

  Definition arity (op : t) : arity.t :=
    match op with
    | arrow => [0; 0]
    | all => [1]
    | bool => []

    | abs => [0; 1]
    | app => [0; 0]
    | tyabs => [1]
    | tyapp => [0; 0]
    | tt => []
    | ff => []
    | If => [0; 0; 0]
    end.

  Definition eq_dec : forall x y : t, {x = y} + {x <> y}.
    decide equality.
  Defined.
End op.

Module abt := abt.abt op.

Module ast.
  Inductive t :=
  | var (x : nat) : t
  | arrow : t -> t -> t
  | all : t -> t
  | bool : t
  | abs : t -> t -> t
  | app : t -> t -> t
  | tyabs : t -> t
  | tyapp : t -> t -> t
  | tt : t
  | ff : t
  | If : t -> t -> t -> t
  .
End ast.

Module basis.
  Module A := abt.

  Import ast.
  Definition t := t.

  Fixpoint to_abt (term : t) : A.t :=
    match term with
    | var n => A.var n
    | arrow ty1 ty2 => A.op op.arrow [A.bind 0 (to_abt ty1); A.bind 0 (to_abt ty2)]
    | all ty => A.op op.all [A.bind 1 (to_abt ty)]
    | bool => A.op op.bool []
    | abs ty e => A.op op.abs [A.bind 0 (to_abt ty); A.bind 1 (to_abt e)]
    | app e1 e2 => A.op op.app [A.bind 0 (to_abt e1); A.bind 0 (to_abt e2)]
    | tyabs e => A.op op.tyabs [A.bind 1 (to_abt e)]
    | tyapp e ty => A.op op.tyapp [A.bind 0 (to_abt e); A.bind 0 (to_abt ty)]
    | tt => A.op op.tt []
    | ff => A.op op.ff []
    | If e1 e2 e3 => A.op op.If [A.bind 0 (to_abt e1);
                               A.bind 0 (to_abt e2);
                               A.bind 0 (to_abt e3)]
    end.

  Fixpoint of_abt (term : A.t) : t :=
    match term with
    | A.var n => var n
    | A.op op.arrow [A.bind 0 a1; A.bind 0 a2] => arrow (of_abt a1) (of_abt a2)
    | A.op op.all [A.bind 1 a'] => all (of_abt a')
    | A.op op.bool [] => bool
    | A.op op.abs [A.bind 0 a_ty; A.bind 1 a_e] => abs (of_abt a_ty) (of_abt a_e)
    | A.op op.app [A.bind 0 a1; A.bind 0 a2] => app (of_abt a1) (of_abt a2)
    | A.op op.tyabs [A.bind 1 a'] => tyabs (of_abt a')
    | A.op op.tyapp [A.bind 0 a_e; A.bind 0 a_ty] => tyapp (of_abt a_e) (of_abt a_ty)
    | A.op op.tt [] => tt
    | A.op op.ff [] => ff
    | A.op op.If [A.bind 0 a1; A.bind 0 a2; A.bind 0 a3] =>
      If (of_abt a1) (of_abt a2) (of_abt a3)
    | _ => var 0 (* bogus *)
    end.

  Fixpoint t_map ov c (term : t) : t :=
    match term with
    | var x => ov c x
    | arrow ty1 ty2 => arrow (t_map ov c ty1) (t_map ov c ty2)
    | all ty => all (t_map ov (S c) ty)
    | bool => bool
    | abs ty e => abs (t_map ov c ty) (t_map ov (S c) e)
    | app e1 e2 => app (t_map ov c e1) (t_map ov c e2)
    | tyabs e => tyabs (t_map ov (S c) e)
    | tyapp e ty => tyapp (t_map ov c e) (t_map ov c ty)
    | tt => tt
    | ff => ff
    | If e1 e2 e3 => If (t_map ov c e1) (t_map ov c e2) (t_map ov c e3)
    end.

  Fixpoint shift c d (term : t) : t :=
    match term with
    | var alpha => var (if alpha <? c then alpha else alpha + d)
    | arrow ty1 ty2 => arrow (shift c d ty1) (shift c d ty2)
    | all ty => all (shift (S c) d ty)
    | bool => bool
    | abs ty e => abs (shift c d ty) (shift (S c) d e)
    | app e1 e2 => app (shift c d e1) (shift c d e2)
    | tyabs e => tyabs (shift (S c) d e)
    | tyapp e ty => tyapp (shift c d e) (shift c d ty)
    | tt => tt
    | ff => ff
    | If e1 e2 e3 => If (shift c d e1) (shift c d e2) (shift c d e3)
    end.

  Fixpoint subst rho term :=
    match term with
    | var alpha => match nth_error rho alpha with
                  | Some ty' => ty'
                  | None => term
                  end
    | arrow ty1 ty2 => arrow (subst rho ty1) (subst rho ty2)
    | all ty => all (subst (var 0 :: map (shift 0 1) rho) ty)
    | bool => bool
    | abs ty e => abs (subst rho ty) (subst (var 0 :: map (shift 0 1) rho) e)
    | app e1 e2 => app (subst rho e1) (subst rho e2)
    | tyabs e => tyabs (subst (var 0 :: map (shift 0 1) rho) e)
    | tyapp e ty => tyapp (subst rho e) (subst rho ty)
    | tt => tt
    | ff => ff
    | If e1 e2 e3 => If (subst rho e1) (subst rho e2) (subst rho e3)
    end.

  Fixpoint wf n term :=
    match term with
    | var alpha => alpha < n
    | arrow ty1 ty2 => wf n ty1 /\ wf n ty2
    | all ty => wf (S n) ty
    | bool => True
    | abs ty e => wf n ty /\ wf (S n) e
    | app e1 e2 => wf n e1 /\ wf n e2
    | tyabs e => wf (S n) e
    | tyapp e ty => wf n e /\ wf n ty
    | tt => True
    | ff => True
    | If e1 e2 e3 => wf n e1 /\ wf n e2 /\ wf n e3
    end.

  Fixpoint identity_subst (n : nat) : list t :=
    match n with
    | 0 => []
    | S n => var 0 :: map (shift 0 1) (identity_subst n)
    end.

  Lemma ws_to_abt : forall e, A.ws (to_abt e).
  Proof. A.basis_util.prove_ws_to_abt. Qed.

  Lemma of_to_abt : forall e, of_abt (to_abt e) = e.
  Proof. A.basis_util.prove_of_to_abt. Qed.

  Lemma to_of_abt : forall a, A.ws a -> to_abt (of_abt a) = a.
  Proof. A.basis_util.prove_to_of_abt to_abt of_abt. Qed.

  Lemma t_map_to_abt_comm : forall ov e c,
      to_abt (t_map ov c e) = A.t_map (fun c x => to_abt (ov c x)) c (to_abt e).
  Proof. A.basis_util.prove_t_map_to_abt_comm. Qed.

  Lemma shift_to_abt_comm : forall e c d, to_abt (shift c d e) = A.shift c d (to_abt e).
  Proof. A.basis_util.prove_shift_to_abt_comm. Qed.

  Lemma map_shift_to_abt_comm :
    forall c d rho, map to_abt (map (shift c d) rho) = map (A.shift c d) (map to_abt rho).
  Proof. A.basis_util.prove_map_shift_to_abt_comm shift_to_abt_comm. Qed.

  Lemma subst_to_abt_comm : forall e rho,
      to_abt (subst rho e) = A.subst (map to_abt rho) (to_abt e).
  Proof. A.basis_util.prove_subst_to_abt_comm t map_shift_to_abt_comm. Qed.

  Lemma wf_to_abt : forall e n, wf n e <-> A.wf n (to_abt e).
  Proof. A.basis_util.prove_wf_to_abt. Qed.

  Lemma identity_subst_to_abt_comm :
    forall n, List.map to_abt (identity_subst n) = A.identity_subst n.
  Proof. A.basis_util.prove_identity_subst_to_abt_comm map_shift_to_abt_comm. Qed.

  Definition var := var.
  Arguments var /.
  Lemma var_to_abt : forall n, to_abt (var n) = A.var n.
  Proof. reflexivity. Qed.
End basis.

Module term.
  Include abt_util basis.
  Notation arrow := ast.arrow.
  Notation all := ast.all.
  Notation bool := ast.bool.
  Notation abs := ast.abs.
  Notation app := ast.app.
  Notation tyabs := ast.tyabs.
  Notation tyapp := ast.tyapp.
  Notation tt := ast.tt.
  Notation ff := ast.ff.
  Notation If := ast.If.
End term.

Module cat.
  Inductive t := TYPE | EXPR.
End cat.

Module wc.
  Inductive t : list cat.t -> term.t -> cat.t -> Prop :=
  | var : forall C n c,
      nth_error C n = Some c ->
      t C (term.var n) c
  | arrow : forall C t1 t2,
      t C t1 cat.TYPE ->
      t C t2 cat.TYPE ->
      t C (term.arrow t1 t2) cat.TYPE
  | all : forall C ty,
      t (cat.TYPE :: C) ty cat.TYPE ->
      t C (term.all ty) cat.TYPE
  | bool : forall C,
      t C ast.bool cat.TYPE
  | abs : forall C ty e,
      t C ty cat.TYPE ->
      t (cat.EXPR :: C) e cat.EXPR ->
      t C (term.abs ty e) cat.EXPR
  | app : forall C e1 e2,
      t C e1 cat.EXPR ->
      t C e2 cat.EXPR ->
      t C (term.app e1 e2) cat.EXPR
  | tyabs : forall C e,
      t (cat.TYPE :: C) e cat.EXPR ->
      t C (term.tyabs e) cat.EXPR
  | tyapp : forall C e ty,
      t C e cat.EXPR ->
      t C ty cat.TYPE ->
      t C (term.tyapp e ty) cat.EXPR
  | tt : forall C,
      t C term.tt cat.EXPR
  | ff : forall C,
      t C term.ff cat.EXPR
  | If : forall C e1 e2 e3,
      t C e1 cat.EXPR ->
      t C e2 cat.EXPR ->
      t C e3 cat.EXPR ->
      t C (term.If e1 e2 e3) cat.EXPR
  .
  Global Hint Constructors t : core.

  Lemma t_wf :
    forall C e c,
      t C e c ->
      term.wf (length C) e.
  Proof.
    induction 1; cbn; intuition.
    apply nth_error_Some.
    congruence.
  Qed.

  Lemma shift :
    forall C1 C2 C3 e c,
      t (C1 ++ C3) e c ->
      t (C1 ++ C2 ++ C3) (term.shift (length C1) (length C2) e) c.
  Proof.
    intros C1 C2 C3 e c WC.
    revert C2.
    remember (C1 ++ C3) as C eqn:EC in *.
    revert C1 C3 EC.
    induction WC; intros C1 C3 EC C2; subst C; cbn [term.shift]; constructor; auto.
    - now rewrite nth_error_shift.
    - now apply IHWC with (C2 := _ :: _).
    - now apply IHWC2 with (C2 := _ :: _).
    - now apply IHWC with (C2 := _ :: _).
  Qed.

  Lemma shift' :
    forall C e c C',
      t C e c ->
      t (C' ++ C) (term.shift 0 (List.length C') e) c.
  Proof.
    intros.
    now apply shift with (C1 := []).
  Qed.

  Lemma shift_cons :
  forall C e c c',
    t C e c ->
    t (c' :: C) (term.shift 0 1 e) c.
  Proof.
    intros C e c c' HT.
    now apply shift' with (C' := [c']).
  Qed.

  Lemma t_skipn_shift :
    forall C e c n,
      n <= length C ->
      t (skipn n C) e c ->
      t C (term.shift 0 n e) c.
  Proof.
    intros C e c n L T.
    rewrite <- firstn_skipn with (n := n) (l := C).
    apply shift' with (C' := firstn n C) in T.
    now rewrite firstn_length_le in * by assumption.
  Qed.

  Lemma shift_inv :
  forall e C1 C2 C3 c,
    t (C1 ++ C2 ++ C3) (term.shift (length C1) (length C2) e) c ->
    t (C1 ++ C3) e c.
  Proof.
    induction e; intros C1 C2 C3 c HT; invc HT; constructor; eauto.
    - now rewrite nth_error_shift in *.
    - eapply IHe with (C1 := _ :: _); eauto.
    - eapply IHe2 with (C1 := _ :: _); eauto.
    - eapply IHe with (C1 := _ :: _); eauto.
  Qed.

  Lemma shift_inv' :
    forall C e c C',
      t (C' ++ C) (term.shift 0 (List.length C') e) c ->
      t C e c.
  Proof.
    intros.
    now eapply shift_inv with (C1 := []); eauto.
  Qed.

  Lemma shift_cons_inv :
  forall C e c c',
    t (c' :: C) (term.shift 0 1 e) c ->
    t C e c.
  Proof.
    intros C e c c' HT.
    now apply shift_inv' with (C' := [c']).
  Qed.

  Lemma subst :
    forall C e c,
      t C e c ->
      forall C' rho,
        List.Forall2 (t C') rho C ->
        t C' (term.subst rho e) c.
  Proof.
    induction 1; intros C' rho F; cbn [term.subst]; try constructor;
      eauto 10 using shift_cons, Forall2_impl, Forall2_map_l.
    cbn.
    destruct (Forall2_nth_error2 F H) as [z [Hz Ht]].
    unfold term.t in Hz.
    now rewrite Hz.
  Qed.

  Lemma Forall2_identity_subst :
    forall C n,
      n = length C ->
      Forall2 (t C) (term.identity_subst n) C.
  Proof.
    intros C n ?. subst n.
    induction C; cbn; constructor.
    - auto.
    - apply Forall2_map_l.
      eauto using Forall2_impl, shift_cons.
  Qed.

  Lemma Forall2_t_Forall_wf :
    forall C es tys,
      Forall2 (wc.t C) es tys ->
      Forall (term.wf (length C)) es.
  Proof.
    induction 1; eauto using t_wf.
  Qed.
End wc.

Module ctx_ent.
  Inductive t :=
  | TYVAR
  | EXVAR : term.t -> t
  .

  Definition map (f : term.t -> term.t) (ce : t) : t :=
    match ce with
    | TYVAR => TYVAR
    | EXVAR t => EXVAR (f t)
    end.

  Definition prop (P : term.t -> Prop) (ce : t) : Prop :=
    match ce with
    | TYVAR => True
    | EXVAR t => P t
    end.

  Definition to_cat (ce : t) : cat.t :=
    match ce with
    | TYVAR => cat.TYPE
    | EXVAR _ => cat.EXPR
    end.

  Module exports.
    Notation TYVAR := TYVAR.
    Notation EXVAR := EXVAR.
  End exports.
End ctx_ent.

Module tele.
  Import ctx_ent.exports.

  Definition t := list ctx_ent.t.

  Notation to_cat_ctx T := (map ctx_ent.to_cat T).
  Notation map f := (map (ctx_ent.map f)).

  Fixpoint wf (T : t) : Prop :=
    match T with
    | [] => True
    | ce :: T =>
      wf T /\ ctx_ent.prop (fun ty => wc.t (to_cat_ctx T) ty cat.TYPE) ce
    end.

  Fixpoint shift (c d : nat) (T : t) : t :=
    match T with
    | [] => []
    | ce :: T =>
      match c with
      | 0 => ctx_ent.map (term.shift c d) ce :: shift c d T
      | S c => ctx_ent.map (term.shift c d) ce :: shift c d T
      end
    end.

  Lemma shift_length :
    forall T c d,
      length (shift c d T) = length T.
  Proof.
    induction T; intros [|] d; cbn; auto.
  Qed.

  Definition is_tyvar (T : t) (n : nat) : Prop :=
    nth_error T n = Some TYVAR.

  Definition lookup_exvar (T : t) (n : nat) (expect : term.t) : Prop :=
    exists expect',
      nth_error T n = Some (EXVAR expect') /\
      expect = term.shift 0 (S n) expect'.

  Fixpoint globalize (i : nat) (T : t) : list ctx_ent.t :=
    match T with
    | [] => []
    | ce :: T =>
      ctx_ent.map (term.shift 0 (S i)) ce :: globalize (S i) T
    end.

  Lemma nth_error_globalize :
    forall T i n,
      nth_error (globalize i T) n =
      match nth_error T n with
      | None => None
      | Some e => Some (ctx_ent.map (term.shift 0 (n + S i)) e)
      end.
  Proof.
    induction T; intros i [|n]; cbn; auto.
    rewrite IHT.
    break_match; auto.
    repeat f_equal.
    lia.
  Qed.

  Lemma lookup_exvar_globalize :
    forall T n ty,
      lookup_exvar T n ty ->
      nth_error (globalize 0 T) n = Some (EXVAR ty).
  Proof.
    intros T n ty [ty' [NE ?]].
    subst ty.
    rewrite nth_error_globalize, NE.
    cbn.
    repeat f_equal.
    lia.
  Qed.

  Lemma to_cat_shift :
    forall T c d,
      to_cat_ctx (shift c d T) = to_cat_ctx T.
  Proof.
    unfold ctx_ent.to_cat.
    induction T; cbn [List.map shift]; intros c d.
    - reflexivity.
    - destruct c, a; cbn; f_equal; auto.
  Qed.

  Lemma nth_error_tele_shift_exvar1 :
    forall T c d n ty,
      nth_error (shift c d T) n = Some (EXVAR ty) ->
      exists ty',
        term.shift (c - S n) d ty' = ty /\
        nth_error T n = Some (EXVAR ty').
  Proof.
    induction T; intros [|c] d [|n] ty NE; cbn in *; try discriminate.
    - unfold ctx_ent.map in *. destruct a; try congruence.
      invc NE. eauto.
    - apply IHT in NE. assumption.
    - destruct a; cbn in *; try congruence.
      invc NE.
      rewrite Nat.sub_0_r.
      eauto.
    - auto.
  Qed.

  Lemma nth_error_tele_shift_exvar2 :
    forall T c d n ty,
      nth_error T n = Some (EXVAR ty) ->
      nth_error (shift c d T) n = Some (EXVAR (term.shift (c - S n) d ty)).
  Proof.
    induction T; intros [|c] d [|n] ty NE; cbn in *; try discriminate.
    - invc NE; reflexivity.
    - now erewrite IHT by eassumption.
    - rewrite Nat.sub_0_r.
      invc NE; reflexivity.
    - eauto.
  Qed.

  Lemma lookup_exvar_shift :
    forall T1 T2 T3 n expect,
     lookup_exvar (shift (length T1) (length T2) T1 ++ T2 ++ T3)
                  (if n <? length T1 then n else n + length T2)
                  (term.shift (length T1) (length T2) expect) <->
     lookup_exvar (T1 ++ T3) n expect.
  Proof.
    intros T1 T2 T3 n expect.
    unfold lookup_exvar.
    rewrite nth_error_shift' in * by now rewrite shift_length.
    destruct (Nat.ltb_spec n (length T1)).
    - rewrite !nth_error_app1 by now rewrite ?shift_length.
      split; intros [expect' [NE ?]]; subst.
      + apply nth_error_tele_shift_exvar1 in NE.
        destruct NE as [ty' [Ety' NE]].
        subst. eexists. split; [eassumption|].
        rewrite term.shift_shift in * by lia.
        replace (length T1 - S n + S n) with (length T1) in * by lia.
        eauto using term.shift_inj.
      + erewrite nth_error_tele_shift_exvar2 by eassumption.
        eexists. split; [reflexivity|].
        rewrite term.shift_shift with (c2 := 0) by lia.
        f_equal. lia.
    - rewrite !nth_error_app2 by (rewrite ?shift_length; lia).
      rewrite shift_length.
      split; intros [expect' [NE ?]]; subst;
        (eexists; split; [eassumption|]).
      + apply term.shift_inj with (c := length T1) (d := length T2).
        now rewrite term.shift_merge by lia.
      + now rewrite term.shift_merge by lia.
  Qed.

  Lemma lookup_exvar_nth_error :
    forall T n ty,
      lookup_exvar T n ty ->
      exists ty',
        nth_error T n = Some (EXVAR ty').
  Proof.
    intros T n ty LE.
    unfold lookup_exvar in *.
    destruct LE as [expect' [NE ?]]. subst.
    eauto.
  Qed.

  Lemma nth_error_wc :
    forall T n ty,
      nth_error T n = Some (EXVAR ty) ->
      wf T ->
      wc.t (skipn (S n) (to_cat_ctx T)) ty cat.TYPE.
  Proof.
    induction T; cbn; intros n ty NE WF; destruct n; try discriminate.
    - invc NE. intuition.
    - apply IHT; intuition.
  Qed.

  Lemma lookup_exvar_wc :
    forall T n ty,
      lookup_exvar T n ty ->
      wf T ->
      wc.t (to_cat_ctx T) ty cat.TYPE.
  Proof.
    intros T n ty [expect' [NE ?]] WFT. subst.
    apply wc.t_skipn_shift.
    - apply lt_le_S. apply nth_error_Some.
      rewrite nth_error_map, NE.
      discriminate.
    - auto using nth_error_wc.
  Qed.

  Lemma to_cat_ctx_cons :
    forall ce T,
      to_cat_ctx (cons ce T) = ctx_ent.to_cat ce :: to_cat_ctx T.
  Proof. reflexivity. Qed.

  Lemma wf_cons_exvar :
    forall ty T,
      wc.t (to_cat_ctx T) ty cat.TYPE ->
      wf T ->
      wf (cons (EXVAR ty) T).
  Proof. cbn. auto. Qed.

  Lemma wf_cons_tyvar :
    forall T,
      wf T ->
      wf (cons TYVAR T).
  Proof. cbn. auto. Qed.

  Lemma to_cat_ctx_globalize :
    forall (G : t) i,
      to_cat_ctx (globalize i G) =
      to_cat_ctx G.
  Proof.
    induction G; cbn; intros i.
    - reflexivity.
    - f_equal; [|now auto].
      destruct a; reflexivity.
  Qed.

  Lemma to_cat_ctx_subst :
    forall rho G,
      to_cat_ctx (List.map (ctx_ent.map (term.subst rho)) G) =
      to_cat_ctx G.
  Proof.
    induction G; cbn.
    - reflexivity.
    - f_equal; [|now auto].
      destruct a; reflexivity.
  Qed.

  Lemma globalize_length :
    forall G i,
      length (tele.globalize i G) = length G.
  Proof.
    induction G; cbn; intros i; auto.
  Qed.

  Lemma globalize_0 :
    forall G i,
      globalize i G = List.map (ctx_ent.map (term.shift 0 i)) (globalize 0 G).
  Proof.
    induction G; intros i; cbn; f_equal.
    - destruct a; cbn; f_equal; auto.
      now rewrite term.shift_merge by lia.
    - rewrite IHG with (i := S i).
      rewrite IHG with (i := 1).
      rewrite map_map.
      apply map_ext.
      clear.
      intros [|ty]; cbn; f_equal; auto.
      now rewrite term.shift_merge by lia.
  Qed.

  Lemma wf_globalize :
    forall G,
      tele.wf G ->
      Forall (ctx_ent.prop (term.wf (length G))) (tele.globalize 0 G).
  Proof.
    intros G WF.
    induction G; cbn in *; constructor; destruct WF as [WF HA].
    - destruct a as [|ty]; cbn in *; auto.
      apply wc.t_wf in HA.
      apply term.wf_shift'.
      now rewrite map_length in HA.
    - forward IHG.
      rewrite globalize_0.
      rewrite Forall_map.
      eapply Forall_impl; [|eassumption].
      clear.
      intros [|ty]; cbn; intros WF; auto using term.wf_shift'.
  Qed.
End tele.

Module has_type.
  Import ctx_ent.exports.

  Inductive t : tele.t -> term.t -> term.t -> Prop :=
  | var : forall G x ty,
      tele.lookup_exvar G x ty ->
      t G (term.var x) ty
  | abs : forall G ty1 ty2 ty2' e,
      wc.t (tele.to_cat_ctx G) ty1 cat.TYPE ->
      t (EXVAR ty1 :: G) e ty2 ->
      ty2 = term.shift 0 1 ty2' ->
      t G (term.abs ty1 e) (term.arrow ty1 ty2')
  | app : forall G ty1 ty2 e1 e2,
      t G e1 (term.arrow ty1 ty2) ->
      t G e2 ty1 ->
      t G (term.app e1 e2) ty2
  | tyabs : forall G e ty,
      t (TYVAR :: G) e ty ->
      t G (term.tyabs e) (term.all ty)
  | tyapp : forall G e ty_body ty,
      wc.t (tele.to_cat_ctx G) ty cat.TYPE ->
      t G e (term.all ty_body) ->
      t G (term.tyapp e ty) (term.subst (ty :: term.identity_subst (length G)) ty_body)
  | tt : forall G,
      t G term.tt term.bool
  | ff : forall G,
      t G term.ff term.bool
  | If : forall G e1 e2 e3 ty,
      t G e1 term.bool ->
      t G e2 ty ->
      t G e3 ty ->
      t G (term.If e1 e2 e3) ty
  .
  Global Hint Constructors t : core.

  Lemma t_wc_e :
    forall G e ty,
      t G e ty ->
      tele.wf G ->
      wc.t (tele.to_cat_ctx G) e cat.EXPR.
  Proof.
    induction 1; cbn in *; intros TWF; auto.
    constructor.
    destruct (tele.lookup_exvar_nth_error _ _ _ H) as [ty' NE].
    now erewrite map_nth_error by eassumption.
  Qed.

  Lemma t_wc_ty :
    forall G e ty,
      t G e ty ->
      tele.wf G ->
      wc.t (map ctx_ent.to_cat G) ty cat.TYPE.
  Proof.
    induction 1; cbn in *; intros TWF; auto.
    - destruct (tele.lookup_exvar_nth_error _ _ _ H) as [ty' NE].
      eauto using tele.lookup_exvar_wc.
    - constructor; auto.
      forward IHt.
      subst ty2.
      eauto using wc.shift_cons_inv.
    - forward IHt1.
      now invc IHt1.
    - forward IHt.
      invc IHt.
      eapply wc.subst; eauto.
      constructor; [assumption|].
      apply wc.Forall2_identity_subst.
      now rewrite map_length.
  Qed.

  Lemma shift :
  forall G1 G2 G3 e ty,
    tele.wf (G1 ++ G3) ->
    has_type.t (G1 ++ G3) e ty ->
    has_type.t (tele.shift (length G1) (length G2) G1 ++ G2 ++ G3)
               (term.shift (length G1) (length G2) e)
               (term.shift (length G1) (length G2) ty).
  Proof.
    intros G1 G2 G3 e ty WF HT.
    revert G2.
    remember (G1 ++ G3) as G eqn:EG in *.
    revert WF G1 G3 EG.
    induction HT; intros WF G1 G3 EG G2; subst G; cbn [term.shift]; eauto.
    - constructor.
      now rewrite tele.lookup_exvar_shift.
    - econstructor; [| |reflexivity].
      + rewrite !map_app in *.
        apply wc.shift with (C2 := map ctx_ent.to_cat G2) in H.
        eqapply H.
        * now rewrite !map_length.
        * now rewrite tele.to_cat_shift.
      + rewrite term.shift_shift with (c2 := 0) by lia.
        forward IHHT.
        specialize (IHHT (_ :: _) _ eq_refl G2).
        subst ty2.
        eqapply IHHT.
        f_equal. cbn. lia.
    - econstructor.
      forward IHHT.
      eapply (IHHT (_ :: _) _ eq_refl).
    - specialize (IHHT ltac:(auto) _ _ eq_refl G2).
      assert (term.wf (length (ty :: term.identity_subst (length (G1 ++ G3)))) ty_body).
      {
        cbn. rewrite term.identity_subst_length.
        apply t_wc_ty in HT; [|assumption].
        apply wc.t_wf in HT.
        rewrite map_length in HT.
        apply HT.
      }
      rewrite term.shift_subst by assumption.
      cbn [term.shift] in IHHT.
      eapply tyapp in IHHT.
      eqapply IHHT.
      + rewrite !app_length, tele.shift_length.
        rewrite term.identity_subst_app.
        rewrite term.SIS_app.
        rewrite app_comm_cons.
        rewrite term.subst_shift.
        * cbn. f_equal. f_equal.
          rewrite term.identity_subst_app, map_app.
          rewrite term.map_shift_identity_subst'.
          f_equal.
          rewrite term.SIS_merge by lia.
          f_equal.
          lia.
        * rewrite ?app_length, ?term.SIS_length in *.
          cbn in *.
          now rewrite term.identity_subst_length in *.
        * cbn.
          now rewrite term.identity_subst_length in *.
        * now rewrite term.SIS_length.
      + rewrite !map_app in *.
        rewrite tele.to_cat_shift.
        eapply wc.shift in H.
        eqapply H.
        now rewrite !map_length.
  Qed.

  Lemma shift' :
    forall G e ty G',
      tele.wf G ->
      t G e ty ->
      t (G' ++ G) (term.shift 0 (List.length G') e) (term.shift 0 (List.length G') ty).
  Proof.
    intros.
    now apply shift with (G1 := []).
  Qed.

  Lemma shift_cons :
  forall G e ty ty',
    tele.wf G ->
    has_type.t G e ty ->
    has_type.t (ty' :: G) (term.shift 0 1 e) (term.shift 0 1 ty).
  Proof.
    intros G e ty ty' WF HT.
    now apply shift' with (G' := [ty']).
  Qed.

  Definition t_ctx_ent G e ce :=
    match ce with
    | TYVAR => wc.t (tele.to_cat_ctx G) e cat.TYPE
    | EXVAR ty => t G e ty
    end.

  Lemma Forall2_t_ctx_ent_wc :
    forall G es tys,
      tele.wf G ->
      Forall2 (t_ctx_ent G) es tys ->
      Forall2 (wc.t (tele.to_cat_ctx G)) es (map ctx_ent.to_cat tys).
  Proof.
    intros G es tys WF F.
    apply Forall2_map_r.
    eapply Forall2_impl; [|eassumption].
    intros.
    destruct b; cbn in *.
    - assumption.
    - eapply t_wc_e; eauto.
  Qed.

  Lemma subst :
    forall G e ty,
      t G e ty ->
      tele.wf G ->
      forall G' rho,
        tele.wf G' ->
        List.Forall2 (t_ctx_ent G') rho
                     (map (ctx_ent.map (term.subst rho)) (tele.globalize 0 G)) ->
        t G' (term.subst rho e) (term.subst rho ty).
  Proof.
    induction 1; intros WFG G' rho WFG' F; cbn [term.subst]; eauto;
      assert (length G = length rho) as L
        by (apply Forall2_length in F;
            rewrite map_length, tele.globalize_length in F;
            congruence).
    - pose proof (tele.lookup_exvar_globalize _ _ _ ltac:(eassumption)).
      pose proof (map_nth_error (ctx_ent.map (term.subst rho)) _ _ H0).
      destruct (Forall2_nth_error2 F ltac:(eassumption)) as [z [Hz Ht]].
      cbn.
      unfold term.t in *. (* ugh *)
      now rewrite Hz.
    - apply abs with (ty2 := term.subst (term.descend 1 rho) ty2).
      + eapply wc.subst.
        eassumption.
        apply Forall2_t_ctx_ent_wc in F; [|assumption].
        eqapply F.
        now rewrite tele.to_cat_ctx_subst, tele.to_cat_ctx_globalize.
      + forward IHt.
        apply IHt.
        {
          cbn; split; [assumption|].
          eapply wc.subst; [eassumption|].
          apply Forall2_t_ctx_ent_wc in F; [|assumption].
          now rewrite tele.to_cat_ctx_subst, tele.to_cat_ctx_globalize in F.
        }
        {
          cbn.
          constructor.
          - unfold t_ctx_ent.
            constructor.
            eexists.
            cbn.
            split;[reflexivity|].
            assert (term.wf (length rho) ty1).
            {
              apply wc.t_wf in H.
              rewrite map_length in H.
              congruence.
            }
            rewrite term.shift_subst by assumption.
            now rewrite term.subst_shift_cons by now rewrite map_length.
          - assert (Forall2
                      (t_ctx_ent (EXVAR (term.subst rho ty1) :: G'))
                      (List.map (term.shift 0 1) rho)
                      (List.map
                         (ctx_ent.map (term.shift 0 1))
                         (List.map (ctx_ent.map (term.subst rho)) (tele.globalize 0 G))))
              as F'.
            {
              (* basically apply shift_cons in F. *)
              rewrite map_map.
              rewrite Forall2_map.
              rewrite Forall2_map_r_iff in F.
              eapply Forall2_impl; [|eassumption].
              cbn.
              (* this should be a lemma *)
              intros a [|ty]; cbn; intros Hty.
              - auto using wc.shift_cons.
              - auto using shift_cons.
            }

            eqapply F'.
            rewrite tele.globalize_0 with (i := 1).
            rewrite !map_map.
            apply tele.wf_globalize in WFG.
            apply map_ext_Forall.
            eapply Forall_impl; [|eassumption].
            clear -L.
            intros [|ty] WF; cbn in *; [reflexivity|f_equal].
            rewrite term.shift_subst by congruence.
            now rewrite term.subst_shift_cons by (rewrite map_length; congruence).
        }
      + subst ty2.
        apply term.subst_descend_shift_shift_subst.
        apply t_wc_ty in H0.
        apply wc.t_wf in H0.
        rewrite map_length in *.
        cbn [length] in H0.
        apply term.wf_shift_inv' in H0.
        congruence.
        cbn. split; auto.
    - econstructor.
      apply IHt; cbn; auto.
      constructor; [now cbn; auto|].
      rewrite Forall2_map_r_iff in F.
      rewrite tele.globalize_0.
      rewrite map_map.
      apply Forall2_map.
      eapply Forall_Forall2_and_r in F; [|apply tele.wf_globalize; assumption].
      eapply Forall2_impl; [|eassumption].
      cbn.
      intros a [|tyb] [HT WF]; cbn in *.
      + auto using wc.shift_cons.
      + rewrite term.subst_cons_shift_shift_subst_1 by congruence.
        now apply shift_cons; auto.
    - forward IHt.
      specialize (IHt G' rho ltac:(assumption) ltac:(assumption)).
      cbn in IHt.
      apply tyapp with (ty := term.subst rho ty) in IHt.
      2: {
        eapply wc.subst; eauto.
        apply Forall2_t_ctx_ent_wc in F; [|assumption].
        eqapply F.
        now rewrite tele.to_cat_ctx_subst, tele.to_cat_ctx_globalize.
      }
      eqapply IHt.
      rewrite !term.subst_subst.
      + f_equal. cbn. f_equal.
        rewrite map_map.
        rewrite L.
        rewrite term.map_subst_identity_subst.
        erewrite map_ext_Forall; [now apply map_id|].
        apply Forall2_t_ctx_ent_wc in F; [|assumption].
        apply wc.Forall2_t_Forall_wf in F.
        rewrite map_length in F.
        eapply Forall_impl; [|eassumption].
        intros a WF.
        rewrite term.subst_shift_cons by now rewrite term.identity_subst_length.
        now rewrite term.subst_identity.
      + cbn. rewrite term.identity_subst_length.
        (* this should be a lemma *)
        apply t_wc_ty in H0; [|assumption].
        apply wc.t_wf in H0.
        rewrite map_length in H0.
        apply H0.
      + constructor; auto.
        apply wc.t_wf in H.
        rewrite map_length in H.
        congruence.
        rewrite L.
        apply term.wf_identity_subst.
      + cbn.
        apply t_wc_ty in H0; [|assumption].
        apply wc.t_wf in H0.
        rewrite !map_length in *.
        rewrite L in H0.
        apply H0.
      + constructor; auto.
        * cbn. lia.
        * apply Forall_map.
          apply Forall2_t_ctx_ent_wc in F; [|assumption].
          apply wc.Forall2_t_Forall_wf in F.
          eapply Forall_impl; [|eassumption].
          intros a.
          cbn.
          rewrite map_length.
          rewrite term.identity_subst_length.
          auto using term.wf_shift'.
  Qed.
End has_type.
Print Assumptions has_type.subst.


(** * Push-Button Synthesis of Saturated Solinas *)
Require Import Coq.Strings.String.
Require Import Coq.micromega.Lia.
Require Import Coq.ZArith.ZArith.
Require Import Coq.MSets.MSetPositive.
Require Import Coq.Lists.List.
Require Import Coq.QArith.QArith_base Coq.QArith.Qround.
Require Import Coq.derive.Derive.
Require Import Crypto.Util.ErrorT.
Require Import Crypto.Util.ListUtil.
Require Import Crypto.Util.ListUtil.FoldBool.
Require Import Crypto.Util.Strings.Decimal.
Require Import Crypto.Util.Strings.Show.
Require Import Crypto.Util.ZRange.
Require Import Crypto.Util.ZUtil.Definitions.
Require Import Crypto.Util.ZUtil.Zselect.
Require Import Crypto.Util.ZUtil.Tactics.LtbToLt.
Require Import Crypto.Util.Tactics.HasBody.
Require Import Crypto.Util.Tactics.Head.
Require Import Crypto.Util.Tactics.SpecializeBy.
Require Import Rewriter.Language.Wf.
Require Import Rewriter.Language.Language.
Require Import Crypto.Language.API.
Require Import Crypto.AbstractInterpretation.AbstractInterpretation.
Require Import Crypto.Stringification.Language.
Require Import Crypto.Arithmetic.Core.
Require Import Crypto.Arithmetic.ModOps.
Require Import Crypto.Arithmetic.Saturated.
Require Import Crypto.BoundsPipeline.
Require Import Crypto.COperationSpecifications.
Require Import Crypto.PushButtonSynthesis.ReificationCache.
Require Import Crypto.PushButtonSynthesis.Primitives.
Require Import Crypto.PushButtonSynthesis.SaturatedSolinasReificationCache.
Require Import Crypto.Assembly.Equivalence.
Import ListNotations.
Local Open Scope string_scope. Local Open Scope Z_scope. Local Open Scope list_scope. Local Open Scope bool_scope.

Import
  Language.Wf.Compilers
  Language.Compilers
  AbstractInterpretation.Compilers
  Stringification.Language.Compilers.
Import Compilers.API.

Import COperationSpecifications.Primitives.
Import COperationSpecifications.Solinas.
Import COperationSpecifications.SaturatedSolinas.

Import Associational Positional.

Local Coercion Z.of_nat : nat >-> Z.
Local Coercion QArith_base.inject_Z : Z >-> Q.
Local Coercion Z.pos : positive >-> Z.

Local Set Keyed Unification. (* needed for making [autorewrite] fast, c.f. COQBUG(https://github.com/coq/coq/issues/9283) *)

Local Opaque reified_mul_gen. (* needed for making [autorewrite] not take a very long time *)
(* needed for making [autorewrite] with [Set Keyed Unification] fast *)
Local Opaque expr.Interp.

Section __.
  Context {output_language_api : ToString.OutputLanguageAPI}
          {pipeline_opts : PipelineOptions}
          {pipeline_to_string_opts : PipelineToStringOptions}
          {synthesis_opts : SynthesisOptions}
          (s : Z)
          (c : list (Z * Z))
          (machine_wordsize : machine_wordsize_opt).

  Local Instance override_pipeline_opts : PipelineOptions
    := {| widen_bytes := true (* true, because we don't allow byte-sized things anyway, so we should not expect carries to be widened to byte-size when emitting C code *)
       |}.

  (* We include [0], so that even after bounds relaxation, we can
       notice where the constant 0s are, and remove them. *)
  Definition possible_values_of_machine_wordsize
    := prefix_with_carry [machine_wordsize].

  Definition n : nat := Z.to_nat (Qceiling (Z.log2_up s / machine_wordsize)).
  Definition m := s - Associational.eval c.
  (* Number of reductions is calculated as follows :
         Let i be the highest limb index of c. Then, each reduction
         decreases the number of extra limbs by (n-i-1). (The -1 comes
         from possibly having an extra high partial product at the end
         of a reduction.) So, to go from the n extra limbs we have
         post-multiplication down to 0, we need ceil (n / (n - i - 1))
         reductions.  In some cases. however, [n - i <= 1], and in
         this case, we do [n] reductions (is this enough?). *)
  Definition nreductions : nat :=
    let i := fold_right Z.max 0 (List.map (fun t => Z.log2 (fst t) / machine_wordsize) c) in
    if Z.of_nat n - i <=? 1
    then n
    else Z.to_nat (Qceiling (Z.of_nat n / (Z.of_nat n - i - 1))).
  Local Notation possible_values := possible_values_of_machine_wordsize.
  Local Notation boundsn := (saturated_bounds n machine_wordsize).

  Local Existing Instance default_translate_to_fancy.
  Local Instance no_select_size : no_select_size_opt := no_select_size_of_no_select machine_wordsize.
  Local Instance split_mul_to : split_mul_to_opt := split_mul_to_of_should_split_mul machine_wordsize possible_values.
  Local Instance split_multiret_to : split_multiret_to_opt := split_multiret_to_of_should_split_multiret machine_wordsize possible_values.

  (** Note: If you change the name or type signature of this
        function, you will need to update the code in CLI.v *)
  Definition check_args {T} (requests : list string) (res : Pipeline.ErrorT T)
    : Pipeline.ErrorT T
    := check_args_of_list
         (List.map
            (fun v => (true, v))
            [((0 <? s - Associational.eval c)%Z, Pipeline.Value_not_ltZ "s - Associational.eval c ≤ 0" 0 (s - Associational.eval c))
             ; (negb (s =? 0)%Z, Pipeline.Values_not_provably_distinctZ "s ≠ 0" s 0)
             ; (negb (n =? 0)%nat, Pipeline.Values_not_provably_distinctZ "n ≠ 0" n 0)
             ; (0 <? machine_wordsize, Pipeline.Value_not_ltZ "0 < machine_wordsize" 0 machine_wordsize)
         ])
         res.

  Local Ltac use_curve_good_t :=
    repeat first [ use_requests_to_prove_curve_good_t_step
                 | assumption
                 | lia
                 | progress autorewrite with distr_length
                 | progress distr_length ].

  Context (requests : list string)
          (curve_good : check_args requests (Success tt) = Success tt).

  Lemma use_curve_good
    : 0 < s - Associational.eval c
      /\ s - Associational.eval c <> 0
      /\ s <> 0
      /\ 0 < machine_wordsize
      /\ n <> 0%nat.
  Proof using curve_good.
    prepare_use_curve_good ().
    { use_curve_good_t. }
  Qed.

  Local Notation weightf := (weight machine_wordsize 1).
  Local Notation evalf := (eval weightf n).
  Local Notation notations_for_docstring
    := (CorrectnessStringification.dyn_context.cons
          weightf "weight"
          (CorrectnessStringification.dyn_context.cons
             evalf "eval"
             CorrectnessStringification.dyn_context.nil))%string.
  Local Notation "'docstring_with_summary_from_lemma!' summary correctness"
    := (docstring_with_summary_from_lemma_with_ctx!
          notations_for_docstring
          summary
          correctness)
         (only parsing, at level 10, summary at next level, correctness at next level).

  Definition mul
    := Pipeline.BoundsPipeline
         false (* subst01 *)
         possible_values
         (reified_mul_gen
            @ GallinaReify.Reify s @ GallinaReify.Reify c @ GallinaReify.Reify (machine_wordsize:Z) @ GallinaReify.Reify n @ GallinaReify.Reify nreductions)
         (Some boundsn, (Some boundsn, tt))
         (Some boundsn, None (* Should be: Some r[0~>0]%zrange, but bounds analysis is not good enough *) ).

  Definition smul (prefix : string)
    : string * (Pipeline.M (Pipeline.ExtendedSynthesisResult _))
    := Eval cbv beta in
        FromPipelineToString!
          machine_wordsize prefix "mul" mul
          (docstring_with_summary_from_lemma!
             (fun fname : string => [text_before_function_name ++ fname ++ " multiplies two field elements."]%string)
             (mul_correct weightf n m boundsn)).

  Local Ltac solve_extra_bounds_side_conditions :=
    cbn [lower upper fst snd] in *; Bool.split_andb; Z.ltb_to_lt; lia.

  Hint Rewrite
       (fun pf => @Rows.eval_mulmod (weight machine_wordsize 1) (@wprops _ _ pf))
       using solve [ auto with zarith | congruence | solve_extra_bounds_side_conditions ] : push_eval.
  Hint Unfold mulmod : push_eval.

  Local Ltac prove_correctness _ := Primitives.prove_correctness use_curve_good.

  Lemma mul_correct res
        (Hres : mul = Success res)
    : mul_correct (weight machine_wordsize 1) n m boundsn (Interp res).
  Proof using curve_good. prove_correctness (). Qed.

  Lemma Wf_mul res (Hres : mul = Success res) : Wf res.
  Proof using Type. prove_pipeline_wf (). Qed.

  Section for_stringification.
    Local Open Scope string_scope.
    Local Open Scope list_scope.

    Definition known_functions
      := [("mul", wrap_s smul)].

    Definition valid_names : string := Eval compute in String.concat ", " (List.map (@fst _ _) known_functions).

    (** Note: If you change the name or type signature of this
          function, you will need to update the code in CLI.v *)
    Definition Synthesize (comment_header : list string) (function_name_prefix : string) (requests : list string)
      : list (synthesis_output_kind * string * Pipeline.M (list string))
      := Primitives.Synthesize
           machine_wordsize valid_names known_functions (fun _ => nil) all_typedefs!
           check_args
           ((ToString.comment_file_header_block
               (comment_header
                  ++ ["";
                     "Computed values:";
                     "";
                     "  # reductions = " ++ show nreductions]%string)))
           function_name_prefix requests.
  End for_stringification.
End __.

Module Export Hints.
#[global]
  Hint Opaque
       mul
  : wf_op_cache.
#[global]
  Hint Immediate
       Wf_mul
  : wf_op_cache.
End Hints.


Require Import Raft.
Require Import RaftRefinementInterface.
Require Import RefinementCommonDefinitions.

Section CandidateEntriesInterface.
  Context {orig_base_params : BaseParams}.
  Context {one_node_params : OneNodeParams orig_base_params}.
  Context {raft_params : RaftParams orig_base_params}.

  Definition allEntries_candidateEntries net :=
    (forall h t e,
       In (t, e) (allEntries (fst (nwState net h))) ->
       candidateEntries e (nwState net)).

  Class allEntries_candidate_entries_interface : Prop :=
    {
      allEntries_candidateEntries_invariant :
        forall (net : network),
          refined_raft_intermediate_reachable net ->
          allEntries_candidateEntries net
    }.

End CandidateEntriesInterface.