From iris.program_logic Require Export language ectx_language ectxi_language.
From st.prelude Require Export autosubst generic lang_base.

(** The syntax of expressions *)
Inductive expr :=
| Var (x : var)
| LetIn (e1 : expr) (e2 : {bind 1 of expr})
| Lam (e : {bind 1 of expr})
(* | Fix (e : expr) *)
(* | Rec (e : {bind 2 of expr}) *)
| App (e1 e2 : expr)
(* Base Types *)
| Lit (l : base_lit)
| BinOp (op : bin_op) (e1 e2 : expr)
| If (e0 e1 e2 : expr)
| Seq (e1 e2 : expr)
(* Products *)
| Pair (e1 e2 : expr)
| Fst (e : expr)
| Snd (e : expr)
(* Sums *)
| InjL (e : expr)
| InjR (e : expr)
| Case (e0 : expr) (e1 : {bind expr}) (e2 : {bind expr})
(* Recursive Types *)
| Fold (e : expr)
| Unfold (e : expr)
(** Polymorphic Types *)
(* | TLam (e : expr) *)
(* | TApp (e : expr). *).

Instance Ids_expr : Ids expr. derive. Defined.
Instance Rename_expr : Rename expr. derive. Defined.
Instance Subst_expr : Subst expr. derive. Defined.
Instance SubstLemmas_expr : SubstLemmas expr. derive. Qed.

Inductive val :=
| LamV (e : {bind 1 of expr})
| LitV (v : base_lit)
| PairV (v1 v2 : val)
| InjLV (v : val)
| InjRV (v : val)
| FoldV (v : val).

Fixpoint of_val (v : val) : expr :=
 match v with
 | LamV e => Lam e
 | LitV v => Lit v
 | PairV v1 v2 => Pair (of_val v1) (of_val v2)
 | InjLV v => InjL (of_val v)
 | InjRV v => InjR (of_val v)
 | FoldV v => Fold (of_val v)
 end.

Fixpoint to_val (e : expr) : option val :=
 match e with
 | Lam e => Some (LamV e)
 | Lit e => Some (LitV e)
 | Pair e1 e2 => v1 ← to_val e1; v2 ← to_val e2; Some (PairV v1 v2)
 | InjL e => InjLV <$> to_val e
 | InjR e => InjRV <$> to_val e
 | Fold e => v ← to_val e; Some (FoldV v)
 | _ => None
 end.

Fixpoint val_subst (v : val) (σ : var → expr) : val :=
  match v with
  | LamV e => LamV (e.[up σ])
  | LitV v => LitV v
  | PairV v1 v2 => PairV (val_subst v1 σ) (val_subst v2 σ)
  | InjLV v => InjLV (val_subst v σ)
  | InjRV v => InjRV (val_subst v σ)
  | FoldV v => FoldV (val_subst v σ)
  end.

Lemma to_of_val v : to_val (of_val v) = Some v.
Proof.
 by induction v; try simplify_option_eq; repeat f_equal; try apply (proof_irrel _).
Qed.
Lemma of_to_val e v : to_val e = Some v → of_val v = e.
Proof.
 revert v; induction e; intros v ?; simplify_option_eq; auto with f_equal.
Qed.

(** Equality and other typeclass stuff *)
Instance of_val_inj : Inj (=) (=) of_val.
Proof. by intros ?? Hv; apply (inj Some); rewrite -!to_of_val Hv. Qed.

Instance base_lit_eq_dec : EqDecision base_lit.
Proof. solve_decision. Defined.
Instance bin_op_eq_dec : EqDecision bin_op.
Proof. solve_decision. Defined.
Instance expr_eq_dec : EqDecision expr.
Proof. solve_decision. Defined.
Instance val_eq_dec : EqDecision val.
Proof.
 refine (λ v v', cast_if (decide (of_val v = of_val v')));
   abstract naive_solver.
Defined.

Global Instance val_inhabited : Inhabited val := populate (LitV LitUnit).

(** Evaluation contexts *)
Inductive ectx_item :=
| LetInCtx (e2 : expr)
| AppLCtx (e2 : expr)
| AppRCtx (v1 : val)
| PairLCtx (e2 : expr)
| PairRCtx (v1 : val)
| FstCtx
| SndCtx
| InjLCtx
| InjRCtx
| CaseCtx (e1 : {bind expr}) (e2 : {bind expr})
| IfCtx (e2 : expr) (e3 : expr)
| BinOpLCtx (op : bin_op) (e2 : expr)
| BinOpRCtx (op : bin_op) (v1 : val)
| SeqCtx (e2 : expr)
| FoldCtx
| UnfoldCtx.

Definition fill_item (Ki : ectx_item) (e : expr) : expr :=
 match Ki with
 | LetInCtx e2 => LetIn e e2
 | AppLCtx e2 => App e e2
 | AppRCtx v1 => App (of_val v1) e
 | PairLCtx e2 => Pair e e2
 | PairRCtx v1 => Pair (of_val v1) e
 | FstCtx => Fst e
 | SndCtx => Snd e
 | InjLCtx => InjL e
 | InjRCtx => InjR e
 | CaseCtx e1 e2 => Case e e1 e2
 | IfCtx e1 e2 => If e e1 e2
 | BinOpLCtx op e2 => BinOp op e e2
 | BinOpRCtx op v1 => BinOp op (of_val v1) e
 | SeqCtx e2 => Seq e e2
 | FoldCtx => Fold e
 | UnfoldCtx => Unfold e
 end.

(** The stepping relation *)

Definition bin_op_eval (op : bin_op) (z1 z2 : Z) : val :=
 match op with
 | PlusOp => LitV $ LitInt (z1 + z2)%Z
 | MinusOp => LitV $ LitInt (z1 - z2)
 | LeOp => LitV $ LitBool $ bool_decide (z1 ≤ z2)%Z
 | LtOp => LitV $ LitBool $ bool_decide (z1 < z2)%Z
 | EqOp => LitV $ LitBool $ bool_decide (z1 = z2)
 end.

Definition state : Type := ().

Inductive head_step : expr → state → list Empty_set → expr → state → list expr → Prop :=
(* β *)
  LetIn_head_step e1 v1 e2 σ :
   to_val e1 = Some v1 →
   head_step (LetIn e1 e2) σ [] e2.[e1/] σ []
| App_Lam_head_step e1 e2 v2 σ :
   to_val e2 = Some v2 →
   head_step (App (Lam e1) e2) σ [] e1.[e2/] σ []
(* | App_Rec_head_step e1 e2 v2 σ : *)
   (* to_val e2 = Some v2 → *)
   (* head_step (App (Rec e1) e2) σ [] e1.[(Rec e1), e2/] σ [] *)
(* fix *)
(* | Fix_head_step e σ : *)
    (* head_step (Fix (Lam e)) σ [] e.[Fix (Lam e)/] σ [] *)
(* binary operation *)
| BinOp_head_step op e1 e2 z1 z2 σ :
   to_val e1 = Some (LitV $ LitInt z1) → to_val e2 = Some (LitV $ LitInt z2) →
   head_step (BinOp op e1 e2) σ [] (of_val (bin_op_eval op z1 z2)) σ []
(* if *)
| If_True_head_step e1 e2 σ :
   head_step (If (Lit $ LitBool true) e1 e2) σ [] e1 σ []
| If_False_head_step e1 e2 σ :
   head_step (If (Lit $ LitBool false) e1 e2) σ [] e2 σ []
(* seq *)
| Seq_Unit_head_step e1 e2 σ :
    to_val e1 = Some (LitV LitUnit) →
    head_step (Seq e1 e2) σ [] e2 σ []
(* Products *)
| Fst_Pair_head_step e1 v1 e2 v2 σ :
   to_val e1 = Some v1 → to_val e2 = Some v2 →
   head_step (Fst (Pair e1 e2)) σ [] e1 σ []
| Snd_Pair_head_step e1 v1 e2 v2 σ :
   to_val e1 = Some v1 → to_val e2 = Some v2 →
   head_step (Snd (Pair e1 e2)) σ [] e2 σ []
(* Sums *)
| Case_InjL_head_step e0 v0 e1 e2 σ :
   to_val e0 = Some v0 →
   head_step (Case (InjL e0) e1 e2) σ [] e1.[e0/] σ []
| Case_InjR_head_step e0 v0 e1 e2 σ :
   to_val e0 = Some v0 →
   head_step (Case (InjR e0) e1 e2) σ [] e2.[e0/] σ []
(* Recursive Types *)
| Unfold_Fold_head_step e v σ :
   to_val e = Some v →
   head_step (Unfold (Fold e)) σ [] e σ []
(* Polymorphic Types *)
(* | TBeta e σ : *)
   (* head_step (TApp (TLam e)) σ [] e σ []. *).

Instance fill_item_inj Ki : Inj (=) (=) (fill_item Ki).
Proof. destruct Ki; intros ???; simplify_eq/=; auto with f_equal. Qed.
Lemma fill_item_val Ki e :
 is_Some (to_val (fill_item Ki e)) → is_Some (to_val e).
Proof. intros [v ?]. destruct Ki; simplify_option_eq; eauto. Qed.
Lemma val_stuck e1 σ1 κ e2 σ2 efs :
 head_step e1 σ1 κ e2 σ2 efs → to_val e1 = None.
Proof. destruct 1; done. Qed.
Lemma head_ctx_step_val Ki e σ1 κ e2 σ2 efs :
 head_step (fill_item Ki e) σ1 κ e2 σ2 efs → is_Some (to_val e).
Proof.
 destruct Ki; inversion_clear 1; simplify_option_eq; eauto.
Qed.
Lemma fill_item_no_val_inj Ki1 Ki2 e1 e2 :
 to_val e1 = None → to_val e2 = None →
 fill_item Ki1 e1 = fill_item Ki2 e2 → Ki1 = Ki2.
Proof.
 destruct Ki1, Ki2; intros; try discriminate; simplify_eq/=;
  repeat match goal with
  | H : to_val (of_val _) = None |- _ => by rewrite to_of_val in H
  end; auto.
Qed.

Lemma st_ectxi_lang_mixin : EctxiLanguageMixin of_val to_val fill_item head_step.
Proof.
  split; eauto using to_of_val, of_to_val,
  val_stuck, fill_item_val, fill_item_no_val_inj,
 head_ctx_step_val, fill_item_inj.
Qed.

Canonical Structure STLCmu_ectxi_lang : ectxiLanguage := EctxiLanguage st_ectxi_lang_mixin.
Canonical Structure STLCmu_ectx_lang : ectxLanguage := EctxLanguageOfEctxi STLCmu_ectxi_lang.
Canonical Structure STLCmu_lang : language := LanguageOfEctx STLCmu_ectx_lang.

Lemma fill_val (e : expr) (K : list ectx_item):
  is_Some (to_val (fill K e)) -> is_Some (to_val e).
Proof.
  move=> [v h]. destruct (to_val e) eqn:eq.
    by exists v0.
    have fill_not_val: to_val (fill K e) = None. eauto using fill_not_val.
    congruence.
Qed.

(* Arguments val_stuck {_ _ _ _ _} _. *)
(* Arguments fill_val {_ _} _. *)

(* Wrapper around prim_step *)

Definition STLCmu_step (e1 e2 : expr) : Prop := prim_step e1 tt [] e2 tt [].

(* We do not use forks, nor prophecy variables. *)

Lemma head_step_no_forks e σ κ e' σ' efs : head_step e σ κ e' σ' efs → efs = [].
Proof. intros H. by inversion H. Qed.

Lemma prim_step_no_forks (e : expr) σ κ e' σ' efs : prim_step e σ κ e' σ' efs → efs = [].
Proof. intros H. inversion H. by eapply head_step_no_forks. Qed.

Lemma head_step_no_obs e σ κ e' σ' efs : head_step e σ κ e' σ' efs → κ = [].
Proof. intros H. by inversion H. Qed.

Lemma prim_step_no_obs (e : expr) σ κ e' σ' efs : prim_step e σ κ e' σ' efs → κ = [].
Proof. intros H. inversion H. by eapply head_step_no_obs. Qed.

(* Our language is deterministic *)

Lemma head_step_det e e1 σ1 κ1 σ1' efs1 e2 σ2 κ2 σ2' efs2 : head_step e σ1 κ1 e1 σ1' efs1 → head_step e σ2 κ2 e2 σ2' efs2 → e1 = e2.
Proof. intros H1 H2. inversion H1; inversion H2; ((by simplify_eq) || (try done) || simplify_eq; inversion G2). Qed.

Lemma prim_step_det (e e1 e2 : expr) σ κ : prim_step e σ κ e1 σ [] → prim_step e σ κ e2 σ [] → e1 = e2.
Proof.
  intros H1 H2.
  inversion H1. inversion H2. simplify_eq. simpl in *.
  assert (K = K0) as <-.
  { destruct (step_by_val K K0 _ _ σ κ e2'0 σ [] H4) as [Kred eq] ; try done; try by eapply val_stuck.
    assert (H4' : fill K0 e1'0 = fill K e1'); first done.
    destruct (step_by_val K0 K _ _ σ κ e2' σ [] H4') as [Kred' eq'] ; try done; try by eapply val_stuck.
    rewrite eq in eq'. simpl in *. assert (length K = length (Kred' ++ Kred ++ K)). simpl in *. by rewrite -eq'.
    do 2 rewrite app_length in H. assert (Kred = []) as ->. apply length_zero_iff_nil. lia. by rewrite eq. }
  f_equal. assert (e1' = e1'0) as ->. apply (fill_inj K _ _ H4). by eapply head_step_det.
Qed.

(* Our language is pure *)

Lemma prim_step_pure (e1 e2 : expr) σ1 σ2 κ efs : prim_step e1 σ1 κ e2 σ2 efs → pure_step e1 e2.
Proof.
  intros Hprim.
  assert (efs = []) as ->. by eapply prim_step_no_forks.
  assert (κ = []) as ->. by eapply prim_step_no_obs.
  destruct σ1, σ2.
  split.
  intros σ. destruct σ. rewrite /reducible_no_obs. by exists e2, tt, [].
  intros.
  assert (efs = []) as ->. by eapply prim_step_no_forks.
  assert (κ = []) as ->. by eapply prim_step_no_obs.
  destruct σ1, σ2. by erewrite (prim_step_det _ _ _ _ _ H).
Qed.

(* Wrappers around lemmas *)

Lemma STLCmu_pure e1 e2 : STLCmu_step e1 e2 <-> pure_step e1 e2.
Proof.
  split. apply prim_step_pure. intro H. inversion H.
  destruct (pure_step_safe tt) as [e2' [σ [efs Hp]]].
  destruct σ. by destruct (pure_step_det _ _ _ _ _ Hp) as [a [b [-> ->]]].
Qed.

Lemma STLCmu_step_ctx K `{!LanguageCtx K} e1 e2 : STLCmu_step e1 e2 → STLCmu_step (K e1) (K e2).
Proof. intro. apply STLCmu_pure. apply pure_step_ctx. auto. by apply STLCmu_pure. Qed.

Lemma rtc_STLCmu_step_ctx K `{!LanguageCtx K} e1 e2 : rtc STLCmu_step e1 e2 → rtc STLCmu_step (K e1) (K e2).
Proof. eauto using rtc_congruence, STLCmu_step_ctx. Qed.

Lemma nsteps_STLCmu_step_ctx K `{!LanguageCtx K} n e1 e2 : nsteps STLCmu_step n e1 e2 → nsteps STLCmu_step n (K e1) (K e2).
Proof. eauto using nsteps_congruence, STLCmu_step_ctx. Qed.

Lemma nsteps_PureExec (e1 e2 : expr) n : nsteps STLCmu_step n e1 e2 <-> PureExec True n e1 e2.
Proof.
  split. intros s t. eapply nsteps_congruence with (f := id). by apply STLCmu_pure. auto.
  intro H. eapply nsteps_congruence with (f := id). apply STLCmu_pure. apply pure_exec. auto.
Qed.

Lemma rtc_PureExec (e1 e2 : expr) : rtc STLCmu_step e1 e2 <-> ∃ n, PureExec True n e1 e2.
Proof.
  split.
  intro H. assert (H' : rtc pure_step e1 e2).
  eapply rtc_subrel. by apply STLCmu_pure. auto. destruct (iffLR (rtc_nsteps _ _) H') as [n H'']. exists n. intros _. done.
  intro d. destruct d as [n H]. eapply rtc_nsteps. exists n. by eapply nsteps_PureExec.
Qed.

Lemma step_PureExec (e1 e2 : expr) : STLCmu_step e1 e2 → PureExec True 1 e1 e2.
Proof. intros s t. apply nsteps_once. by apply STLCmu_pure. Qed.

Definition STLCmu_halts (e : STLCmu.lang.expr) : Prop :=
  ∃ (v : STLCmu.lang.val), rtc STLCmu_step e (STLCmu.lang.of_val v).


(* File: Cons_Counter_Model.v  (last edited on 27/10/2000) (c) Klaus Weich  *)

From IPC Require Export Disjunct NDeco_Sound.

Fixpoint n2forest (n : nested_imps) : Forest atoms :=
  match n with
  | nil => Nil_Forest atoms
  | Undecorated _ :: n => n2forest n
  | Decorated _ k :: n => Cons_Forest atoms k (n2forest n)
  end.

Remark cons_counter_model_suc :
 forall (x : nimp) (k : kripke_tree) (ni : nested_imps) (a : atoms),
 In (Decorated x k) ni -> Suc k (node atoms a (n2forest ni)).
intros x k ni a in_k.
unfold Suc in |- *; apply successor_trans with k.
generalize in_k; clear in_k.
elim ni; clear ni.
intros in_k.
inversion_clear in_k.
intros y ni ih in_k.
unfold n2forest in |- *; simpl in |- *; fold n2forest in |- *.
inversion_clear in_k.
 rewrite H; clear H y.
apply in_forest_head.
generalize (ih H); clear ih H.
intros ih.
case y; clear y.
intros; assumption.
intros; apply in_forest_tail; assumption.
apply successor_refl.
Qed.

Remark in_forest_ex_a0a1b :
 forall (k : kripke_tree) (ni : nested_imps),
 In_Forest atoms k (n2forest ni) -> exists x : nimp, In (Decorated x k) ni.
intros k ni.
elim ni; clear ni.
intros in_k.
inversion_clear in_k.
unfold n2forest in |- *; simpl in |- *; fold n2forest in |- *.
intros x; case x; clear x.
intros x ni ih in_k.
elim ih; clear ih.
intros y in_y.
exists y.
right; assumption.
assumption.
intros x k0 ni ih in_k.
inversion_clear in_k.
exists x.
left; trivial.
elim ih; try assumption.
intros x0 in_k.
exists x0.
right; assumption.
Qed.

(**********************************************************************)

Remark deco_sound_in_forest_forces :
 forall (work : nf_list) (ds : disjs) (ni : nested_imps)
   (ai : atomic_imps) (a : atoms) (k : kripke_tree)
   (c : normal_form),
 deco_sound work ds ni ai a ->
 In_Forest atoms k (n2forest ni) ->
 in_ngamma work ds ni ai a c -> forces_t k (nf2form c).
intros work ds ni ai a k c complete in_k in_ngamma.
elim in_forest_ex_a0a1b with k ni; try assumption.
intros x; case x; clear x.
intros a0 a1 b in_x.
elim (complete k a0 a1 b in_x); clear complete.
intros k_is_mon k_forces_ngamma k_notforces_a0a1.
apply k_forces_ngamma.
assumption.
Qed.

(**********************************************************************)

Remark cons_counter_model_mon :
 forall (ni : nested_imps) (ai : atomic_imps) (a : atoms),
 (forall x : nimp, In (Undecorated x) ni -> False) ->
 deco_sound nf_nil DNil ni ai a ->
 Is_Monotone_kripke_tree (node atoms a (n2forest ni)).
intros ni ai a.
elim ni; clear ni.
intros all_ref complete.
unfold n2forest in |- *.
simpl in |- *.
unfold Is_Monotone_kripke_tree in |- *.
apply is_monotone_tree_intro.
apply is_monotone_forest_nil.

intros ni; case ni; clear ni.
intros x ni ih all_ref complete.
elimtype False.
apply (all_ref x).
left; trivial.
intros x; case x; clear x.
intros a0 a1 b k ni ih all_ref complete.
unfold n2forest in |- *; simpl in |- *; fold n2forest in |- *.
elim (complete k a0 a1 b).
intros k_is_mon k_forces_ngamma k_notforces_a0a1.
unfold Is_Monotone_kripke_tree in |- *.
apply is_monotone_tree_intro.
apply is_monotone_forest_cons.
intros i forces_i.
change (forces_t k (Atom i)) in |- *.
apply k_forces_ngamma with (c := NAtom i).
apply In_Atoms; assumption.
assumption.
cut (Is_Monotone_kripke_tree (node atoms a (n2forest ni))).
intros claim.
inversion_clear claim; assumption.
apply ih; clear ih.
intros x in_x.
apply all_ref with x.
right; assumption.
apply deco_sound_cons_ni_tail with (Decorated (NImp a0 a1 b) k); assumption.
left; trivial.
Qed.

(**********************************************************************)

Inductive Cons_Counter_Model_Spec (goal : Int) (ni : nested_imps)
(ai : atomic_imps) (a : atoms) : Set :=
    cons_counter_model_intro :
      forall k : kripke_tree,
      Is_Monotone_kripke_tree k ->
      forces_ngamma nf_nil DNil ni ai a k ->
      (forces_t k (Atom goal) -> False) ->
      Cons_Counter_Model_Spec goal ni ai a.

Remark all_ref_rev_app_nil :
 forall (dni : decorated_nested_imps) (x : nimp),
 In (Undecorated x) (rev_app dni NNil) -> False.
intros dni x.
elim dni; clear dni.
simpl in |- *.
intros; assumption.
intros a; case a; clear a.
intros y k dni ih.
simpl in |- *.
 rewrite (rev_app_app dni (Decorated y k :: NNil)).
intros in_x.
apply ih.
 rewrite (rev_app_app dni NNil).
apply in_or_app.
cut
 (In (Undecorated x) (rev_app dni NNil) \/
  In (Undecorated x) (Decorated y k :: NNil)).
intros claim.
elim claim; clear claim; intro claim.
left; assumption.
right; inversion_clear claim.
 discriminate H.
assumption.
apply in_app_or.
assumption.
Qed.

Remark nth_nimp__nth_nested_imp :
 forall (x : nimp) (n : nat) (ni : nested_imps),
 my_nth nimp n (nested_imps2nimps ni) x ->
 {y : nested_imp | my_nth nested_imp n ni y /\ nested_imp2nimp y = x}.
intros x n.
elim n; clear n.
intros ni; case ni; clear ni.
intros nth; elimtype False; inversion_clear nth.
simpl in |- *; intros y ni nth.
exists y.
split.
apply My_NthO.
inversion_clear nth.
trivial.
intros n ih ni.
case ni; clear ni.
intros nth; elimtype False; inversion_clear nth.
intros y ni nth.
elim (ih ni).
intros y' nth'.
exists y'.
elim nth'; clear nth'; intros nth' eq.
split.
apply My_NthS; assumption.
assumption.
inversion_clear nth; assumption.
Qed.

(***********************************************************************)

Lemma cons_counter_model :
 forall (i : Int) (dni : decorated_nested_imps) (ai : atomic_imps)
   (a : atoms),
 deco_sound nf_nil DNil (rev_app dni NNil) ai a ->
 a_ai_disj a ai ->
 a_goal_disj a i -> Cons_Counter_Model_Spec i (rev_app dni NNil) ai a.
intros i dni ai a complete a_ai_disjunct a_goal_disj.
cut (Is_Monotone_kripke_tree (node atoms a (n2forest (rev_app dni NNil)))).
intro mon.
exists (node atoms a (n2forest (rev_app dni NNil))).

apply mon.

unfold forces_ngamma in |- *.
intros c in_c.
elim in_c; clear in_c c.

intros n c nth.
inversion_clear nth.

intros n a0 a1 nth.
inversion_clear nth.

intros n x; case x; clear x.
intros a0 a1 b nth.
simpl in |- *; apply forces_t_imp.
assumption.
intros forces_a0a1.
elimtype False.
elim (nth_nimp__nth_nested_imp (NImp a0 a1 b) n (rev_app dni NNil) nth).
intros x; case x; clear x.
intros x nth_x.
apply all_ref_rev_app_nil with dni x.
apply nth_in with n.
elim nth_x; intros; assumption.
intros x k nth_x.
elim (complete k a0 a1 b).
intros k_is_mon k_forces_ngamma k_notforces_a0a1.
apply k_notforces_a0a1.
apply forces_t_mon with (node atoms a (n2forest (rev_app dni NNil))).
assumption.
assumption.
apply cons_counter_model_suc with x.
apply nth_in with n.
elim nth_x; intros; assumption.
apply nth_in with n.
elim nth_x; clear nth_x; intros nth_x eq.
 rewrite <- eq.
assumption.
simpl in |- *.
intros k' in_k'.
change (forces_t k' (nf2form (NImp_NF (NImp a0 a1 b)))) in |- *.
apply deco_sound_in_forest_forces with nf_nil DNil (rev_app dni NNil) ai a.
assumption.
assumption.
apply In_Nested_Imps with n; assumption.

intros j b n bs lookup_j nth.
simpl in |- *; apply forces_t_imp.
assumption.
intros forces_j.
elimtype False.
apply a_ai_disjunct with j bs; assumption.
simpl in |- *.
intros k' in_k'.
change (forces_t k' (nf2form (AImp j b))) in |- *.
apply deco_sound_in_forest_forces with nf_nil DNil (rev_app dni NNil) ai a.
assumption.
assumption.
apply In_Atomic_Imps with (i := j) (b := b) (n := n) (bs := bs); assumption.

intros j lookup_j.
assumption.

assumption.
apply cons_counter_model_mon with ai; try assumption.
intros x in_x.
apply all_ref_rev_app_nil with dni x; assumption.
Qed.


From MetaCoq.PCUIC Require Import PCUICAst PCUICAstUtils PCUICLiftSubst.
From MetaCoq.PCUIC Require TemplateToPCUIC PCUICToTemplate.

From MetaCoq.Template Require Import config monad_utils utils TemplateMonad.
From MetaCoq.Template Require Ast.
Import MonadNotation.

Require Import List String Relation_Operators.
Import ListNotations.

Require Import non_uniform.

Definition clift0 (n : nat) (t : context_decl) : context_decl :=
  {| decl_name := t.(decl_name);
     decl_body := match t.(decl_body) with
                   | Some q => Some (lift0 n q)
                   | None => None
                  end;
     decl_type := lift0 n t.(decl_type)
  |}.

Definition subterms_for_constructor
           (refi : inductive)
           (ref   : term) (* we need the term for exactly this inductive just for the substition *)
           (ntypes : nat) (* number of types in the mutual inductive *)
           (npars : nat) (* number of parameters in the type *)
           (nind : nat) (* number of proper indeces in the type *)
           (ct    : term) (* type of the constructor *)
           (ncons : nat) (* index of the constructor in the inductive *)
           (nargs : nat) (* number of arguments in this constructor *)
                  : list (nat * term * nat)
  := let indrel := (ntypes - inductive_ind refi - 1) in
    let '(ctx, ap) := decompose_prod_assum [] ct in
    (*    ^ now ctx is a reversed list of assumptions and definitions *)
    let len := List.length ctx in
    let params := List.skipn (len - npars) (ctx) in
    let inds := List.skipn npars (snd (decompose_app ap)) in
    (* d is a list of occurences of subterms in arguments to the constructor *)
    let d :=List.concat (
           (* so this i represents distance from return object to `t` *)
              mapi (fun i t =>
                      let '(ctx, ar) := decompose_prod_assum [] (decl_type t)
                      in let p := (indrel + (len - i - 1) + List.length (ctx))
                      in let (f, s) := decompose_app ar
                      in match f with
                         | tRel j => if Nat.eqb p j
                                    then [(i, ctx, s)]
                                    else []
                         | _ => []
                         end) ctx) in
    let '(ctx_sbst, _) := decompose_prod_assum [] (subst1 ref indrel ct) in
    let construct_cons :=
        fun (* index of a subterm in this constructor *)
          (i: nat)
          (* these are arguments for the function
             that is a parameter of the constructor
             and if applied fully returns something of the needed type *)
          (ctx': context)
          (* these are arguments of the type of the subterm *)
          (args' : list term) =>
          let len' := List.length ctx' in
          let ctxl' :=
              let pr := it_mkProd_or_LetIn ctx' (tVar "ignore me") in
              let lifted := lift0 (1+i) pr in
              let (c,_) := decompose_prod_assum [] lifted in
              c in
          it_mkProd_or_LetIn
             (ctxl' ++ ctx_sbst)
             (mkApps (tRel (len + len'))
                   ((map (lift0 (len' + len - npars))
                         (to_extended_list params)) ++
                    (map (lift (1 + i) len') (List.skipn npars args')) ++
                    (map (lift0 len') inds) ++
                    [mkApps (tRel (i + len'))
                          (to_extended_list ctxl');
                     mkApps (tConstruct refi ncons [])
                         (map (lift0 len') (to_extended_list ctx_sbst))])) in
    mapi (fun i '(n, c, a) => (i, construct_cons n c a, len + List.length c)) d.

Definition subterm_for_ind
           (refi : inductive)
           (ref   : term)
           (allparams : nat)
           (ntypes : nat) (* number of types in the mutual inductive *)
           (ind   : one_inductive_body)
                  : one_inductive_body
  := let (pai, _) := decompose_prod_assum [] ind.(ind_type) in
    let sort := (tSort (Universe.make' UnivExpr.lProp)) in
    let npars := getParamCount ind allparams in
    let pars := List.skipn (List.length pai - npars) pai in
    let inds := List.firstn (List.length pai - npars) pai in
    let ninds := List.length inds in
    let aptype1 :=
        mkApps ref ((map (lift0 (2 * ninds)) (to_extended_list pars)) ++
                    (map (lift0 ninds) (to_extended_list inds))) in
    let aptype2 :=
        mkApps ref ((map (lift0 (1 + 2 * ninds)) (to_extended_list pars)) ++
                    (map (lift0 1) (to_extended_list inds))) in
    let renamer name i := (name ++ "_subterm" ++ (string_of_nat i))%string in
    {| ind_name := (ind.(ind_name) ++ "_direct_subterm")%string;
       ind_type  := it_mkProd_or_LetIn
                      pars
                   (it_mkProd_or_LetIn
                      (inds)
                   (it_mkProd_or_LetIn
                      (map (clift0 (ninds)) inds)
                   (it_mkProd_or_LetIn
                       [mkdecl nAnon None aptype2; mkdecl nAnon None aptype1]
                       sort)));
       ind_kelim := InProp;
       ind_ctors :=List.concat
                     (mapi (fun n '(id', ct, k) => (
                       map (fun '(si, st, sk) => (renamer id' si, st, sk))
                       (subterms_for_constructor refi ref ntypes npars ninds ct n k)))
                       ind.(ind_ctors));
       ind_projs := [] |}.

Definition direct_subterm_for_mutual_ind
            (mind : mutual_inductive_body)
            (ind0 : inductive) (* internal metacoq representation of inductive, part of tInd *)
            (ref  : term) (* reference term for the inductive type, like (tInd {| inductive_mind := "Coq.Init.Datatypes.nat"; inductive_ind := 0 |} []) *)
                  : option mutual_inductive_body
  := let i0 := inductive_ind ind0 in
    let ntypes := List.length (ind_bodies mind) in
    b <- List.nth_error mind.(ind_bodies) i0 ;;
    let npars := getParamCount b (ind_npars mind) in
    ret {|
        ind_finite := BasicAst.Finite;
        ind_npars := npars;
        ind_universes := ind_universes mind;
        ind_params := List.firstn (ind_npars mind - npars) (ind_params mind);
        ind_bodies := [subterm_for_ind ind0 ref mind.(ind_npars) ntypes b];
        ind_variance := None
      |}.

Definition subterm (t : Ast.term)
  : TemplateMonad unit
  := match t with
    | Ast.tInd ind0 _ =>
      decl <- tmQuoteInductive (inductive_mind ind0);;
      tmPrint decl;;
      match (direct_subterm_for_mutual_ind
               (TemplateToPCUIC.trans_minductive_body decl)
               ind0
               (TemplateToPCUIC.trans t)) with
      | None =>
        tmPrint t;;
        @tmFail unit "Coulnd't construct a subterm"
      | Some d =>
        v <- tmEval lazy (PCUICToTemplate.trans_minductive_body d);;
        tmPrint v;;
        tmMkInductive' v
      end
    | _ =>
      tmPrint t;;
      @tmFail unit " is not an inductive"
    end.

Require Import MetaCoq.Template.All.
Import MonadNotation.

Notation "'Derive' 'Subterm' 'for' T" := (subterm <% T %>) (at level 0).


From iris.algebra Require Export base.
From iris Require Import options.
Set Primitive Projections.

(** This files defines (a shallow embedding of) the category of OFEs:
    Complete ordered families of equivalences. This is a cartesian closed
    category, and mathematically speaking, the entire development lives
    in this category. However, we will generally prefer to work with raw
    Coq functions plus some registered Proper instances for non-expansiveness.
    This makes writing such functions much easier. It turns out that it many
    cases, we do not even need non-expansiveness.
*)

(** Unbundled version *)
Class Dist A := dist : nat → relation A.
Instance: Params (@dist) 3 := {}.
Notation "x ≡{ n }≡ y" := (dist n x y)
  (at level 70, n at next level, format "x  ≡{ n }≡  y").
Notation "x ≡{ n }@{ A }≡ y" := (dist (A:=A) n x y)
  (at level 70, n at next level, only parsing).

Hint Extern 0 (_ ≡{_}≡ _) => reflexivity : core.
Hint Extern 0 (_ ≡{_}≡ _) => symmetry; assumption : core.
Notation NonExpansive f := (∀ n, Proper (dist n ==> dist n) f).
Notation NonExpansive2 f := (∀ n, Proper (dist n ==> dist n ==> dist n) f).

Tactic Notation "ofe_subst" ident(x) :=
  repeat match goal with
  | _ => progress simplify_eq/=
  | H:@dist ?A ?d ?n x _ |- _ => setoid_subst_aux (@dist A d n) x
  | H:@dist ?A ?d ?n _ x |- _ => symmetry in H;setoid_subst_aux (@dist A d n) x
  end.
Tactic Notation "ofe_subst" :=
  repeat match goal with
  | _ => progress simplify_eq/=
  | H:@dist ?A ?d ?n ?x _ |- _ => setoid_subst_aux (@dist A d n) x
  | H:@dist ?A ?d ?n _ ?x |- _ => symmetry in H;setoid_subst_aux (@dist A d n) x
  end.

Record OfeMixin A `{Equiv A, Dist A} := {
  mixin_equiv_dist (x y : A) : x ≡ y ↔ ∀ n, x ≡{n}≡ y;
  mixin_dist_equivalence n : Equivalence (@dist A _ n);
  mixin_dist_S n (x y : A) : x ≡{S n}≡ y → x ≡{n}≡ y
}.

(** Bundled version *)
Structure ofeT := OfeT {
  ofe_car :> Type;
  ofe_equiv : Equiv ofe_car;
  ofe_dist : Dist ofe_car;
  ofe_mixin : OfeMixin ofe_car
}.
Arguments OfeT _ {_ _} _.
Add Printing Constructor ofeT.
Hint Extern 0 (Equiv _) => eapply (@ofe_equiv _) : typeclass_instances.
Hint Extern 0 (Dist _) => eapply (@ofe_dist _) : typeclass_instances.
Arguments ofe_car : simpl never.
Arguments ofe_equiv : simpl never.
Arguments ofe_dist : simpl never.
Arguments ofe_mixin : simpl never.

(** When declaring instances of subclasses of OFE (like CMRAs and unital CMRAs)
we need Coq to *infer* the canonical OFE instance of a given type and take the
mixin out of it. This makes sure we do not use two different OFE instances in
different places (see for example the constructors [CmraT] and [UcmraT] in the
file [cmra.v].)

In order to infer the OFE instance, we use the definition [ofe_mixin_of'] which
is inspired by the [clone] trick in ssreflect. It works as follows, when type
checking [@ofe_mixin_of' A ?Ac id] Coq faces a unification problem:

  ofe_car ?Ac  ~  A

which will resolve [?Ac] to the canonical OFE instance corresponding to [A]. The
definition [@ofe_mixin_of' A ?Ac id] will then provide the corresponding mixin.
Note that type checking of [ofe_mixin_of' A id] will fail when [A] does not have
a canonical OFE instance.

The notation [ofe_mixin_of A] that we define on top of [ofe_mixin_of' A id]
hides the [id] and normalizes the mixin to head normal form. The latter is to
ensure that we do not end up with redundant canonical projections to the mixin,
i.e. them all being of the shape [ofe_mixin_of' A id]. *)
Definition ofe_mixin_of' A {Ac : ofeT} (f : Ac → A) : OfeMixin Ac := ofe_mixin Ac.
Notation ofe_mixin_of A :=
  ltac:(let H := eval hnf in (ofe_mixin_of' A id) in exact H) (only parsing).

(** Lifting properties from the mixin *)
Section ofe_mixin.
  Context {A : ofeT}.
  Implicit Types x y : A.
  Lemma equiv_dist x y : x ≡ y ↔ ∀ n, x ≡{n}≡ y.
  Proof. apply (mixin_equiv_dist _ (ofe_mixin A)). Qed.
  Global Instance dist_equivalence n : Equivalence (@dist A _ n).
  Proof. apply (mixin_dist_equivalence _ (ofe_mixin A)). Qed.
  Lemma dist_S n x y : x ≡{S n}≡ y → x ≡{n}≡ y.
  Proof. apply (mixin_dist_S _ (ofe_mixin A)). Qed.
End ofe_mixin.

Hint Extern 1 (_ ≡{_}≡ _) => apply equiv_dist; assumption : core.

(** Discrete OFEs and discrete OFE elements *)
Class Discrete {A : ofeT} (x : A) := discrete y : x ≡{0}≡ y → x ≡ y.
Arguments discrete {_} _ {_} _ _.
Hint Mode Discrete + ! : typeclass_instances.
Instance: Params (@Discrete) 1 := {}.

Class OfeDiscrete (A : ofeT) := ofe_discrete_discrete (x : A) :> Discrete x.

(** OFEs with a completion *)
Record chain (A : ofeT) := {
  chain_car :> nat → A;
  chain_cauchy n i : n ≤ i → chain_car i ≡{n}≡ chain_car n
}.
Arguments chain_car {_} _ _.
Arguments chain_cauchy {_} _ _ _ _.

Program Definition chain_map {A B : ofeT} (f : A → B)
    `{!NonExpansive f} (c : chain A) : chain B :=
  {| chain_car n := f (c n) |}.
Next Obligation. by intros A B f Hf c n i ?; apply Hf, chain_cauchy. Qed.

Notation Compl A := (chain A%type → A).
Class Cofe (A : ofeT) := {
  compl : Compl A;
  conv_compl n c : compl c ≡{n}≡ c n;
}.
Arguments compl : simpl never.
Hint Mode Cofe ! : typeclass_instances.

Lemma compl_chain_map `{Cofe A, Cofe B} (f : A → B) c `(NonExpansive f) :
  compl (chain_map f c) ≡ f (compl c).
Proof. apply equiv_dist=>n. by rewrite !conv_compl. Qed.

Program Definition chain_const {A : ofeT} (a : A) : chain A :=
  {| chain_car n := a |}.
Next Obligation. by intros A a n i _. Qed.

Lemma compl_chain_const {A : ofeT} `{!Cofe A} (a : A) :
  compl (chain_const a) ≡ a.
Proof. apply equiv_dist=>n. by rewrite conv_compl. Qed.

(** General properties *)
Section ofe.
  Context {A : ofeT}.
  Implicit Types x y : A.
  Global Instance ofe_equivalence : Equivalence ((≡) : relation A).
  Proof.
    split.
    - by intros x; rewrite equiv_dist.
    - by intros x y; rewrite !equiv_dist.
    - by intros x y z; rewrite !equiv_dist; intros; trans y.
  Qed.
  Global Instance dist_ne n : Proper (dist n ==> dist n ==> iff) (@dist A _ n).
  Proof.
    intros x1 x2 ? y1 y2 ?; split; intros.
    - by trans x1; [|trans y1].
    - by trans x2; [|trans y2].
  Qed.
  Global Instance dist_proper n : Proper ((≡) ==> (≡) ==> iff) (@dist A _ n).
  Proof.
    by move => x1 x2 /equiv_dist Hx y1 y2 /equiv_dist Hy; rewrite (Hx n) (Hy n).
  Qed.
  Global Instance dist_proper_2 n x : Proper ((≡) ==> iff) (dist n x).
  Proof. by apply dist_proper. Qed.
  Global Instance Discrete_proper : Proper ((≡) ==> iff) (@Discrete A).
  Proof. intros x y Hxy. rewrite /Discrete. by setoid_rewrite Hxy. Qed.

  Lemma dist_le n n' x y : x ≡{n}≡ y → n' ≤ n → x ≡{n'}≡ y.
  Proof. induction 2; eauto using dist_S. Qed.
  Lemma dist_le' n n' x y : n' ≤ n → x ≡{n}≡ y → x ≡{n'}≡ y.
  Proof. intros; eauto using dist_le. Qed.
  (** [ne_proper] and [ne_proper_2] are not instances to improve efficiency of
  type class search during setoid rewriting.
  Instances of [NonExpansive{,2}] are hence accompanied by instances of
  [Proper] built using these lemmas. *)
  Lemma ne_proper {B : ofeT} (f : A → B) `{!NonExpansive f} :
    Proper ((≡) ==> (≡)) f.
  Proof. by intros x1 x2; rewrite !equiv_dist; intros Hx n; rewrite (Hx n). Qed.
  Lemma ne_proper_2 {B C : ofeT} (f : A → B → C) `{!NonExpansive2 f} :
    Proper ((≡) ==> (≡) ==> (≡)) f.
  Proof.
     unfold Proper, respectful; setoid_rewrite equiv_dist.
     by intros x1 x2 Hx y1 y2 Hy n; rewrite (Hx n) (Hy n).
  Qed.

  Lemma conv_compl' `{Cofe A} n (c : chain A) : compl c ≡{n}≡ c (S n).
  Proof.
    transitivity (c n); first by apply conv_compl. symmetry.
    apply chain_cauchy. lia.
  Qed.

  Lemma discrete_iff n (x : A) `{!Discrete x} y : x ≡ y ↔ x ≡{n}≡ y.
  Proof.
    split; intros; auto. apply (discrete _), dist_le with n; auto with lia.
  Qed.
  Lemma discrete_iff_0 n (x : A) `{!Discrete x} y : x ≡{0}≡ y ↔ x ≡{n}≡ y.
  Proof. by rewrite -!discrete_iff. Qed.
End ofe.

(** Contractive functions *)
Definition dist_later `{Dist A} (n : nat) (x y : A) : Prop :=
  match n with 0 => True | S n => x ≡{n}≡ y end.
Arguments dist_later _ _ !_ _ _ /.

Global Instance dist_later_equivalence (A : ofeT) n : Equivalence (@dist_later A _ n).
Proof. destruct n as [|n]. by split. apply dist_equivalence. Qed.

Lemma dist_dist_later {A : ofeT} n (x y : A) : dist n x y → dist_later n x y.
Proof. intros Heq. destruct n; first done. exact: dist_S. Qed.

Lemma dist_later_dist {A : ofeT} n (x y : A) : dist_later (S n) x y → dist n x y.
Proof. done. Qed.

(* We don't actually need this lemma (as our tactics deal with this through
   other means), but technically speaking, this is the reason why
   pre-composing a non-expansive function to a contractive function
   preserves contractivity. *)
Lemma ne_dist_later {A B : ofeT} (f : A → B) :
  NonExpansive f → ∀ n, Proper (dist_later n ==> dist_later n) f.
Proof. intros Hf [|n]; last exact: Hf. hnf. by intros. Qed.

Notation Contractive f := (∀ n, Proper (dist_later n ==> dist n) f).

Instance const_contractive {A B : ofeT} (x : A) : Contractive (@const A B x).
Proof. by intros n y1 y2. Qed.

Section contractive.
  Local Set Default Proof Using "Type*".
  Context {A B : ofeT} (f : A → B) `{!Contractive f}.
  Implicit Types x y : A.

  Lemma contractive_0 x y : f x ≡{0}≡ f y.
  Proof. by apply (_ : Contractive f). Qed.
  Lemma contractive_S n x y : x ≡{n}≡ y → f x ≡{S n}≡ f y.
  Proof. intros. by apply (_ : Contractive f). Qed.

  Global Instance contractive_ne : NonExpansive f | 100.
  Proof. by intros n x y ?; apply dist_S, contractive_S. Qed.
  Global Instance contractive_proper : Proper ((≡) ==> (≡)) f | 100.
  Proof. apply (ne_proper _). Qed.
End contractive.

Ltac f_contractive :=
  match goal with
  | |- ?f _ ≡{_}≡ ?f _ => simple apply (_ : Proper (dist_later _ ==> _) f)
  | |- ?f _ _ ≡{_}≡ ?f _ _ => simple apply (_ : Proper (dist_later _ ==> _ ==> _) f)
  | |- ?f _ _ ≡{_}≡ ?f _ _ => simple apply (_ : Proper (_ ==> dist_later _ ==> _) f)
  end;
  try match goal with
  | |- @dist_later ?A _ ?n ?x ?y =>
         destruct n as [|n]; [exact I|change (@dist A _ n x y)]
  end;
  try simple apply reflexivity.

Ltac solve_contractive :=
  solve_proper_core ltac:(fun _ => first [f_contractive | f_equiv]).

(** Limit preserving predicates *)
Class LimitPreserving `{!Cofe A} (P : A → Prop) : Prop :=
  limit_preserving (c : chain A) : (∀ n, P (c n)) → P (compl c).
Hint Mode LimitPreserving + + ! : typeclass_instances.

Section limit_preserving.
  Context `{Cofe A}.
  (* These are not instances as they will never fire automatically...
     but they can still be helpful in proving things to be limit preserving. *)

  Lemma limit_preserving_ext (P Q : A → Prop) :
    (∀ x, P x ↔ Q x) → LimitPreserving P → LimitPreserving Q.
  Proof. intros HP Hlimit c ?. apply HP, Hlimit=> n; by apply HP. Qed.

  Global Instance limit_preserving_const (P : Prop) : LimitPreserving (λ _ : A, P).
  Proof. intros c HP. apply (HP 0). Qed.

  Lemma limit_preserving_discrete (P : A → Prop) :
    Proper (dist 0 ==> impl) P → LimitPreserving P.
  Proof. intros PH c Hc. by rewrite (conv_compl 0). Qed.

  Lemma limit_preserving_and (P1 P2 : A → Prop) :
    LimitPreserving P1 → LimitPreserving P2 →
    LimitPreserving (λ x, P1 x ∧ P2 x).
  Proof. intros Hlim1 Hlim2 c Hc. split. apply Hlim1, Hc. apply Hlim2, Hc. Qed.

  Lemma limit_preserving_impl (P1 P2 : A → Prop) :
    Proper (dist 0 ==> impl) P1 → LimitPreserving P2 →
    LimitPreserving (λ x, P1 x → P2 x).
  Proof.
    intros Hlim1 Hlim2 c Hc HP1. apply Hlim2=> n; apply Hc.
    eapply Hlim1, HP1. apply dist_le with n; last lia. apply (conv_compl n).
  Qed.

  Lemma limit_preserving_forall {B} (P : B → A → Prop) :
    (∀ y, LimitPreserving (P y)) →
    LimitPreserving (λ x, ∀ y, P y x).
  Proof. intros Hlim c Hc y. by apply Hlim. Qed.

  Lemma limit_preserving_equiv `{!Cofe B} (f g : A → B) :
    NonExpansive f → NonExpansive g → LimitPreserving (λ x, f x ≡ g x).
  Proof.
    intros Hf Hg c Hc. apply equiv_dist=> n.
    by rewrite -!compl_chain_map !conv_compl /= Hc.
  Qed.
End limit_preserving.

(** Fixpoint *)
Program Definition fixpoint_chain {A : ofeT} `{Inhabited A} (f : A → A)
  `{!Contractive f} : chain A := {| chain_car i := Nat.iter (S i) f inhabitant |}.
Next Obligation.
  intros A ? f ? n.
  induction n as [|n IH]=> -[|i] //= ?; try lia.
  - apply (contractive_0 f).
  - apply (contractive_S f), IH; auto with lia.
Qed.

Program Definition fixpoint_def `{Cofe A, Inhabited A} (f : A → A)
  `{!Contractive f} : A := compl (fixpoint_chain f).
Definition fixpoint_aux : seal (@fixpoint_def). Proof. by eexists. Qed.
Definition fixpoint := fixpoint_aux.(unseal).
Arguments fixpoint {A _ _} f {_}.
Definition fixpoint_eq : @fixpoint = @fixpoint_def := fixpoint_aux.(seal_eq).

Section fixpoint.
  Context `{Cofe A, Inhabited A} (f : A → A) `{!Contractive f}.

  Lemma fixpoint_unfold : fixpoint f ≡ f (fixpoint f).
  Proof.
    apply equiv_dist=>n.
    rewrite fixpoint_eq /fixpoint_def (conv_compl n (fixpoint_chain f)) //.
    induction n as [|n IH]; simpl; eauto using contractive_0, contractive_S.
  Qed.

  Lemma fixpoint_unique (x : A) : x ≡ f x → x ≡ fixpoint f.
  Proof.
    rewrite !equiv_dist=> Hx n. induction n as [|n IH]; simpl in *.
    - rewrite Hx fixpoint_unfold; eauto using contractive_0.
    - rewrite Hx fixpoint_unfold. apply (contractive_S _), IH.
  Qed.

  Lemma fixpoint_ne (g : A → A) `{!Contractive g} n :
    (∀ z, f z ≡{n}≡ g z) → fixpoint f ≡{n}≡ fixpoint g.
  Proof.
    intros Hfg. rewrite fixpoint_eq /fixpoint_def
      (conv_compl n (fixpoint_chain f)) (conv_compl n (fixpoint_chain g)) /=.
    induction n as [|n IH]; simpl in *; [by rewrite !Hfg|].
    rewrite Hfg; apply contractive_S, IH; auto using dist_S.
  Qed.
  Lemma fixpoint_proper (g : A → A) `{!Contractive g} :
    (∀ x, f x ≡ g x) → fixpoint f ≡ fixpoint g.
  Proof. setoid_rewrite equiv_dist; naive_solver eauto using fixpoint_ne. Qed.

  Lemma fixpoint_ind (P : A → Prop) :
    Proper ((≡) ==> impl) P →
    (∃ x, P x) → (∀ x, P x → P (f x)) →
    LimitPreserving P →
    P (fixpoint f).
  Proof.
    intros ? [x Hx] Hincr Hlim. set (chcar i := Nat.iter (S i) f x).
    assert (Hcauch : ∀ n i : nat, n ≤ i → chcar i ≡{n}≡ chcar n).
    { intros n. rewrite /chcar. induction n as [|n IH]=> -[|i] //=;
        eauto using contractive_0, contractive_S with lia. }
    set (fp2 := compl {| chain_cauchy := Hcauch |}).
    assert (f fp2 ≡ fp2).
    { apply equiv_dist=>n. rewrite /fp2 (conv_compl n) /= /chcar.
      induction n as [|n IH]; simpl; eauto using contractive_0, contractive_S. }
    rewrite -(fixpoint_unique fp2) //.
    apply Hlim=> n /=. by apply Nat_iter_ind.
  Qed.
End fixpoint.

(** Fixpoint of f when f^k is contractive. **)
Definition fixpointK `{Cofe A, Inhabited A} k (f : A → A)
  `{!Contractive (Nat.iter k f)} := fixpoint (Nat.iter k f).

Section fixpointK.
  Local Set Default Proof Using "Type*".
  Context `{Cofe A, Inhabited A} (f : A → A) (k : nat).
  Context {f_contractive : Contractive (Nat.iter k f)} {f_ne : NonExpansive f}.
  (* Note than f_ne is crucial here:  there are functions f such that f^2 is contractive,
     but f is not non-expansive.
     Consider for example f: SPred → SPred (where SPred is "downclosed sets of natural numbers").
     Define f (using informative excluded middle) as follows:
     f(N) = N  (where N is the set of all natural numbers)
     f({0, ..., n}) = {0, ... n-1}  if n is even (so n-1 is at least -1, in which case we return the empty set)
     f({0, ..., n}) = {0, ..., n+2} if n is odd
     In other words, if we consider elements of SPred as ordinals, then we decreaste odd finite
     ordinals by 1 and increase even finite ordinals by 2.
     f is not non-expansive:  Consider f({0}) = ∅ and f({0,1}) = f({0,1,2,3}).
     The arguments are clearly 0-equal, but the results are not.

     Now consider g := f^2. We have
     g(N) = N
     g({0, ..., n}) = {0, ... n+1}  if n is even
     g({0, ..., n}) = {0, ..., n+4} if n is odd
     g is contractive.  All outputs contain 0, so they are all 0-equal.
     Now consider two n-equal inputs. We have to show that the outputs are n+1-equal.
     Either they both do not contain n in which case they have to be fully equal and
     hence so are the results.  Or else they both contain n, so the results will
     both contain n+1, so the results are n+1-equal.
   *)

  Let f_proper : Proper ((≡) ==> (≡)) f := ne_proper f.
  Local Existing Instance f_proper.

  Lemma fixpointK_unfold : fixpointK k f ≡ f (fixpointK k f).
  Proof.
    symmetry. rewrite /fixpointK. apply fixpoint_unique.
    by rewrite -Nat_iter_S_r Nat_iter_S -fixpoint_unfold.
  Qed.

  Lemma fixpointK_unique (x : A) : x ≡ f x → x ≡ fixpointK k f.
  Proof.
    intros Hf. apply fixpoint_unique. clear f_contractive.
    induction k as [|k' IH]=> //=. by rewrite -IH.
  Qed.

  Section fixpointK_ne.
    Context (g : A → A) `{g_contractive : !Contractive (Nat.iter k g)}.
    Context {g_ne : NonExpansive g}.

    Lemma fixpointK_ne n : (∀ z, f z ≡{n}≡ g z) → fixpointK k f ≡{n}≡ fixpointK k g.
    Proof.
      rewrite /fixpointK=> Hfg /=. apply fixpoint_ne=> z.
      clear f_contractive g_contractive.
      induction k as [|k' IH]=> //=. by rewrite IH Hfg.
    Qed.

    Lemma fixpointK_proper : (∀ z, f z ≡ g z) → fixpointK k f ≡ fixpointK k g.
    Proof. setoid_rewrite equiv_dist; naive_solver eauto using fixpointK_ne. Qed.
  End fixpointK_ne.

  Lemma fixpointK_ind (P : A → Prop) :
    Proper ((≡) ==> impl) P →
    (∃ x, P x) → (∀ x, P x → P (f x)) →
    LimitPreserving P →
    P (fixpointK k f).
  Proof.
    intros. rewrite /fixpointK. apply fixpoint_ind; eauto.
    intros; apply Nat_iter_ind; auto.
  Qed.
End fixpointK.

(** Mutual fixpoints *)
Section fixpointAB.
  Context `{Cofe A, Cofe B, !Inhabited A, !Inhabited B}.
  Context (fA : A → B → A).
  Context (fB : A → B → B).
  Context {fA_contractive : ∀ n, Proper (dist_later n ==> dist n ==> dist n) fA}.
  Context {fB_contractive : ∀ n, Proper (dist_later n ==> dist_later n ==> dist n) fB}.

  Local Definition fixpoint_AB (x : A) : B := fixpoint (fB x).
  Local Instance fixpoint_AB_contractive : Contractive fixpoint_AB.
  Proof.
    intros n x x' Hx; rewrite /fixpoint_AB.
    apply fixpoint_ne=> y. by f_contractive.
  Qed.

  Local Definition fixpoint_AA (x : A) : A := fA x (fixpoint_AB x).
  Local Instance fixpoint_AA_contractive : Contractive fixpoint_AA.
  Proof using fA_contractive. solve_contractive. Qed.

  Definition fixpoint_A : A := fixpoint fixpoint_AA.
  Definition fixpoint_B : B := fixpoint_AB fixpoint_A.

  Lemma fixpoint_A_unfold : fA fixpoint_A fixpoint_B ≡ fixpoint_A.
  Proof. by rewrite {2}/fixpoint_A (fixpoint_unfold _). Qed.
  Lemma fixpoint_B_unfold : fB fixpoint_A fixpoint_B ≡ fixpoint_B.
  Proof. by rewrite {2}/fixpoint_B /fixpoint_AB (fixpoint_unfold _). Qed.

  Instance: Proper ((≡) ==> (≡) ==> (≡)) fA.
  Proof using fA_contractive.
    apply ne_proper_2=> n x x' ? y y' ?. f_contractive; auto using dist_S.
  Qed.
  Instance: Proper ((≡) ==> (≡) ==> (≡)) fB.
  Proof using fB_contractive.
    apply ne_proper_2=> n x x' ? y y' ?. f_contractive; auto using dist_S.
  Qed.

  Lemma fixpoint_A_unique p q : fA p q ≡ p → fB p q ≡ q → p ≡ fixpoint_A.
  Proof.
    intros HfA HfB. rewrite -HfA. apply fixpoint_unique. rewrite /fixpoint_AA.
    f_equiv=> //. apply fixpoint_unique. by rewrite HfA HfB.
  Qed.
  Lemma fixpoint_B_unique p q : fA p q ≡ p → fB p q ≡ q → q ≡ fixpoint_B.
  Proof. intros. apply fixpoint_unique. by rewrite -fixpoint_A_unique. Qed.
End fixpointAB.

Section fixpointAB_ne.
  Context `{Cofe A, Cofe B, !Inhabited A, !Inhabited B}.
  Context (fA fA' : A → B → A).
  Context (fB fB' : A → B → B).
  Context `{∀ n, Proper (dist_later n ==> dist n ==> dist n) fA}.
  Context `{∀ n, Proper (dist_later n ==> dist n ==> dist n) fA'}.
  Context `{∀ n, Proper (dist_later n ==> dist_later n ==> dist n) fB}.
  Context `{∀ n, Proper (dist_later n ==> dist_later n ==> dist n) fB'}.

  Lemma fixpoint_A_ne n :
    (∀ x y, fA x y ≡{n}≡ fA' x y) → (∀ x y, fB x y ≡{n}≡ fB' x y) →
    fixpoint_A fA fB ≡{n}≡ fixpoint_A fA' fB'.
  Proof.
    intros HfA HfB. apply fixpoint_ne=> z.
    rewrite /fixpoint_AA /fixpoint_AB HfA. f_equiv. by apply fixpoint_ne.
  Qed.
  Lemma fixpoint_B_ne n :
    (∀ x y, fA x y ≡{n}≡ fA' x y) → (∀ x y, fB x y ≡{n}≡ fB' x y) →
    fixpoint_B fA fB ≡{n}≡ fixpoint_B fA' fB'.
  Proof.
    intros HfA HfB. apply fixpoint_ne=> z. rewrite HfB. f_contractive.
    apply fixpoint_A_ne; auto using dist_S.
  Qed.

  Lemma fixpoint_A_proper :
    (∀ x y, fA x y ≡ fA' x y) → (∀ x y, fB x y ≡ fB' x y) →
    fixpoint_A fA fB ≡ fixpoint_A fA' fB'.
  Proof. setoid_rewrite equiv_dist; naive_solver eauto using fixpoint_A_ne. Qed.
  Lemma fixpoint_B_proper :
    (∀ x y, fA x y ≡ fA' x y) → (∀ x y, fB x y ≡ fB' x y) →
    fixpoint_B fA fB ≡ fixpoint_B fA' fB'.
  Proof. setoid_rewrite equiv_dist; naive_solver eauto using fixpoint_B_ne. Qed.
End fixpointAB_ne.

(** Non-expansive function space *)
Record ofe_mor (A B : ofeT) : Type := OfeMor {
  ofe_mor_car :> A → B;
  ofe_mor_ne : NonExpansive ofe_mor_car
}.
Arguments OfeMor {_ _} _ {_}.
Add Printing Constructor ofe_mor.
Existing Instance ofe_mor_ne.

Notation "'λne' x .. y , t" :=
  (@OfeMor _ _ (λ x, .. (@OfeMor _ _ (λ y, t) _) ..) _)
  (at level 200, x binder, y binder, right associativity).

Section ofe_mor.
  Context {A B : ofeT}.
  Global Instance ofe_mor_proper (f : ofe_mor A B) : Proper ((≡) ==> (≡)) f.
  Proof. apply ne_proper, ofe_mor_ne. Qed.
  Instance ofe_mor_equiv : Equiv (ofe_mor A B) := λ f g, ∀ x, f x ≡ g x.
  Instance ofe_mor_dist : Dist (ofe_mor A B) := λ n f g, ∀ x, f x ≡{n}≡ g x.
  Definition ofe_mor_ofe_mixin : OfeMixin (ofe_mor A B).
  Proof.
    split.
    - intros f g; split; [intros Hfg n k; apply equiv_dist, Hfg|].
      intros Hfg k; apply equiv_dist=> n; apply Hfg.
    - intros n; split.
      + by intros f x.
      + by intros f g ? x.
      + by intros f g h ?? x; trans (g x).
    - by intros n f g ? x; apply dist_S.
  Qed.
  Canonical Structure ofe_morO := OfeT (ofe_mor A B) ofe_mor_ofe_mixin.

  Program Definition ofe_mor_chain (c : chain ofe_morO)
    (x : A) : chain B := {| chain_car n := c n x |}.
  Next Obligation. intros c x n i ?. by apply (chain_cauchy c). Qed.
  Program Definition ofe_mor_compl `{Cofe B} : Compl ofe_morO := λ c,
    {| ofe_mor_car x := compl (ofe_mor_chain c x) |}.
  Next Obligation.
    intros ? c n x y Hx. by rewrite (conv_compl n (ofe_mor_chain c x))
      (conv_compl n (ofe_mor_chain c y)) /= Hx.
  Qed.
  Global Program Instance ofe_mor_cofe `{Cofe B} : Cofe ofe_morO :=
    {| compl := ofe_mor_compl |}.
  Next Obligation.
    intros ? n c x; simpl.
    by rewrite (conv_compl n (ofe_mor_chain c x)) /=.
  Qed.

  Global Instance ofe_mor_car_ne :
    NonExpansive2 (@ofe_mor_car A B).
  Proof. intros n f g Hfg x y Hx; rewrite Hx; apply Hfg. Qed.
  Global Instance ofe_mor_car_proper :
    Proper ((≡) ==> (≡) ==> (≡)) (@ofe_mor_car A B) := ne_proper_2 _.
  Lemma ofe_mor_ext (f g : ofe_mor A B) : f ≡ g ↔ ∀ x, f x ≡ g x.
  Proof. done. Qed.
End ofe_mor.

Arguments ofe_morO : clear implicits.
Notation "A -n> B" :=
  (ofe_morO A B) (at level 99, B at level 200, right associativity).
Instance ofe_mor_inhabited {A B : ofeT} `{Inhabited B} :
  Inhabited (A -n> B) := populate (λne _, inhabitant).

(** Identity and composition and constant function *)
Definition cid {A} : A -n> A := OfeMor id.
Instance: Params (@cid) 1 := {}.
Definition cconst {A B : ofeT} (x : B) : A -n> B := OfeMor (const x).
Instance: Params (@cconst) 2 := {}.

Definition ccompose {A B C}
  (f : B -n> C) (g : A -n> B) : A -n> C := OfeMor (f ∘ g).
Instance: Params (@ccompose) 3 := {}.
Infix "◎" := ccompose (at level 40, left associativity).
Global Instance ccompose_ne {A B C} :
  NonExpansive2 (@ccompose A B C).
Proof. intros n ?? Hf g1 g2 Hg x. rewrite /= (Hg x) (Hf (g2 x)) //. Qed.

(* Function space maps *)
Definition ofe_mor_map {A A' B B'} (f : A' -n> A) (g : B -n> B')
  (h : A -n> B) : A' -n> B' := g ◎ h ◎ f.
Instance ofe_mor_map_ne {A A' B B'} n :
  Proper (dist n ==> dist n ==> dist n ==> dist n) (@ofe_mor_map A A' B B').
Proof. intros ??? ??? ???. by repeat apply ccompose_ne. Qed.

Definition ofe_morO_map {A A' B B'} (f : A' -n> A) (g : B -n> B') :
  (A -n> B) -n> (A' -n>  B') := OfeMor (ofe_mor_map f g).
Instance ofe_morO_map_ne {A A' B B'} :
  NonExpansive2 (@ofe_morO_map A A' B B').
Proof.
  intros n f f' Hf g g' Hg ?. rewrite /= /ofe_mor_map.
  by repeat apply ccompose_ne.
Qed.

(** * Unit type *)
Section unit.
  Instance unit_dist : Dist unit := λ _ _ _, True.
  Definition unit_ofe_mixin : OfeMixin unit.
  Proof. by repeat split; try exists 0. Qed.
  Canonical Structure unitO : ofeT := OfeT unit unit_ofe_mixin.

  Global Program Instance unit_cofe : Cofe unitO := { compl x := () }.
  Next Obligation. by repeat split; try exists 0. Qed.

  Global Instance unit_ofe_discrete : OfeDiscrete unitO.
  Proof. done. Qed.
End unit.

(** * Empty type *)
Section empty.
  Instance Empty_set_dist : Dist Empty_set := λ _ _ _, True.
  Definition Empty_set_ofe_mixin : OfeMixin Empty_set.
  Proof. by repeat split; try exists 0. Qed.
  Canonical Structure Empty_setO : ofeT := OfeT Empty_set Empty_set_ofe_mixin.

  Global Program Instance Empty_set_cofe : Cofe Empty_setO := { compl x := x 0 }.
  Next Obligation. by repeat split; try exists 0. Qed.

  Global Instance Empty_set_ofe_discrete : OfeDiscrete Empty_setO.
  Proof. done. Qed.
End empty.

(** * Product type *)
Section product.
  Context {A B : ofeT}.

  Instance prod_dist : Dist (A * B) := λ n, prod_relation (dist n) (dist n).
  Global Instance pair_ne :
    NonExpansive2 (@pair A B) := _.
  Global Instance fst_ne : NonExpansive (@fst A B) := _.
  Global Instance snd_ne : NonExpansive (@snd A B) := _.
  Definition prod_ofe_mixin : OfeMixin (A * B).
  Proof.
    split.
    - intros x y; unfold dist, prod_dist, equiv, prod_equiv, prod_relation.
      rewrite !equiv_dist; naive_solver.
    - apply _.
    - by intros n [x1 y1] [x2 y2] [??]; split; apply dist_S.
  Qed.
  Canonical Structure prodO : ofeT := OfeT (A * B) prod_ofe_mixin.

  Global Program Instance prod_cofe `{Cofe A, Cofe B} : Cofe prodO :=
    { compl c := (compl (chain_map fst c), compl (chain_map snd c)) }.
  Next Obligation.
    intros ?? n c; split. apply (conv_compl n (chain_map fst c)).
    apply (conv_compl n (chain_map snd c)).
  Qed.

  Global Instance prod_discrete (x : A * B) :
    Discrete (x.1) → Discrete (x.2) → Discrete x.
  Proof. by intros ???[??]; split; apply (discrete _). Qed.
  Global Instance prod_ofe_discrete :
    OfeDiscrete A → OfeDiscrete B → OfeDiscrete prodO.
  Proof. intros ?? [??]; apply _. Qed.
End product.

Arguments prodO : clear implicits.
Typeclasses Opaque prod_dist.

Instance prod_map_ne {A A' B B' : ofeT} n :
  Proper ((dist n ==> dist n) ==> (dist n ==> dist n) ==>
           dist n ==> dist n) (@prod_map A A' B B').
Proof. by intros f f' Hf g g' Hg ?? [??]; split; [apply Hf|apply Hg]. Qed.
Definition prodO_map {A A' B B'} (f : A -n> A') (g : B -n> B') :
  prodO A B -n> prodO A' B' := OfeMor (prod_map f g).
Instance prodO_map_ne {A A' B B'} :
  NonExpansive2 (@prodO_map A A' B B').
Proof. intros n f f' Hf g g' Hg [??]; split; [apply Hf|apply Hg]. Qed.

(** * COFE → OFE Functors *)
Record oFunctor := OFunctor {
  oFunctor_car : ∀ A `{!Cofe A} B `{!Cofe B}, ofeT;
  oFunctor_map `{!Cofe A1, !Cofe A2, !Cofe B1, !Cofe B2} :
    ((A2 -n> A1) * (B1 -n> B2)) → oFunctor_car A1 B1 -n> oFunctor_car A2 B2;
  oFunctor_map_ne `{!Cofe A1, !Cofe A2, !Cofe B1, !Cofe B2} :
    NonExpansive (@oFunctor_map A1 _ A2 _ B1 _ B2 _);
  oFunctor_map_id `{!Cofe A, !Cofe B} (x : oFunctor_car A B) :
    oFunctor_map (cid,cid) x ≡ x;
  oFunctor_map_compose `{!Cofe A1, !Cofe A2, !Cofe A3, !Cofe B1, !Cofe B2, !Cofe B3}
      (f : A2 -n> A1) (g : A3 -n> A2) (f' : B1 -n> B2) (g' : B2 -n> B3) x :
    oFunctor_map (f◎g, g'◎f') x ≡ oFunctor_map (g,g') (oFunctor_map (f,f') x)
}.
Existing Instance oFunctor_map_ne.
Instance: Params (@oFunctor_map) 9 := {}.

Declare Scope oFunctor_scope.
Delimit Scope oFunctor_scope with OF.
Bind Scope oFunctor_scope with oFunctor.

Class oFunctorContractive (F : oFunctor) :=
  oFunctor_map_contractive `{!Cofe A1, !Cofe A2, !Cofe B1, !Cofe B2} :>
    Contractive (@oFunctor_map F A1 _ A2 _ B1 _ B2 _).
Hint Mode oFunctorContractive ! : typeclass_instances.

(** Not a coercion due to the [Cofe] type class argument, and to avoid
ambiguous coercion paths, see https://gitlab.mpi-sws.org/iris/iris/issues/240. *)
Definition oFunctor_apply (F: oFunctor) (A: ofeT) `{!Cofe A} : ofeT :=
  oFunctor_car F A A.

Program Definition oFunctor_oFunctor_compose (F1 F2 : oFunctor)
  `{!∀ `{Cofe A, Cofe B}, Cofe (oFunctor_car F2 A B)} : oFunctor := {|
  oFunctor_car A _ B _ := oFunctor_car F1 (oFunctor_car F2 B A) (oFunctor_car F2 A B);
  oFunctor_map A1 _ A2 _ B1 _ B2 _ 'fg :=
    oFunctor_map F1 (oFunctor_map F2 (fg.2,fg.1),oFunctor_map F2 fg)
|}.
Next Obligation.
  intros F1 F2 ? A1 ? A2 ? B1 ? B2 ? n [f1 g1] [f2 g2] [??]; simpl in *.
  apply oFunctor_map_ne; split; apply oFunctor_map_ne; by split.
Qed.
Next Obligation.
  intros F1 F2 ? A ? B ? x; simpl in *. rewrite -{2}(oFunctor_map_id F1 x).
  apply equiv_dist=> n. apply oFunctor_map_ne.
  split=> y /=; by rewrite !oFunctor_map_id.
Qed.
Next Obligation.
  intros F1 F2 ? A1 ? A2 ? A3 ? B1 ? B2 ? B3 ? f g f' g' x; simpl in *.
  rewrite -oFunctor_map_compose. apply equiv_dist=> n. apply oFunctor_map_ne.
  split=> y /=; by rewrite !oFunctor_map_compose.
Qed.
Instance oFunctor_oFunctor_compose_contractive_1 (F1 F2 : oFunctor)
    `{!∀ `{Cofe A, Cofe B}, Cofe (oFunctor_car F2 A B)} :
  oFunctorContractive F1 → oFunctorContractive (oFunctor_oFunctor_compose F1 F2).
Proof.
  intros ? A1 ? A2 ? B1 ? B2 ? n [f1 g1] [f2 g2] Hfg; simpl in *.
  f_contractive; destruct Hfg; split; simpl in *; apply oFunctor_map_ne; by split.
Qed.
Instance oFunctor_oFunctor_compose_contractive_2 (F1 F2 : oFunctor)
    `{!∀ `{Cofe A, Cofe B}, Cofe (oFunctor_car F2 A B)} :
  oFunctorContractive F2 → oFunctorContractive (oFunctor_oFunctor_compose F1 F2).
Proof.
  intros ? A1 ? A2 ? B1 ? B2 ? n [f1 g1] [f2 g2] Hfg; simpl in *.
  f_equiv; split; simpl in *; f_contractive; destruct Hfg; by split.
Qed.

Program Definition constOF (B : ofeT) : oFunctor :=
  {| oFunctor_car A1 A2 _ _ := B; oFunctor_map A1 _ A2 _ B1 _ B2 _ f := cid |}.
Solve Obligations with done.
Coercion constOF : ofeT >-> oFunctor.

Instance constOF_contractive B : oFunctorContractive (constOF B).
Proof. rewrite /oFunctorContractive; apply _. Qed.

Program Definition idOF : oFunctor :=
  {| oFunctor_car A1 _ A2 _ := A2; oFunctor_map A1 _ A2 _ B1 _ B2 _ f := f.2 |}.
Solve Obligations with done.
Notation "∙" := idOF : oFunctor_scope.

Program Definition prodOF (F1 F2 : oFunctor) : oFunctor := {|
  oFunctor_car A _ B _ := prodO (oFunctor_car F1 A B) (oFunctor_car F2 A B);
  oFunctor_map A1 _ A2 _ B1 _ B2 _ fg :=
    prodO_map (oFunctor_map F1 fg) (oFunctor_map F2 fg)
|}.
Next Obligation.
  intros ?? A1 ? A2 ? B1 ? B2 ? n ???; by apply prodO_map_ne; apply oFunctor_map_ne.
Qed.
Next Obligation. by intros F1 F2 A ? B ? [??]; rewrite /= !oFunctor_map_id. Qed.
Next Obligation.
  intros F1 F2 A1 ? A2 ? A3 ? B1 ? B2 ? B3 ? f g f' g' [??]; simpl.
  by rewrite !oFunctor_map_compose.
Qed.
Notation "F1 * F2" := (prodOF F1%OF F2%OF) : oFunctor_scope.

Instance prodOF_contractive F1 F2 :
  oFunctorContractive F1 → oFunctorContractive F2 →
  oFunctorContractive (prodOF F1 F2).
Proof.
  intros ?? A1 ? A2 ? B1 ? B2 ? n ???;
    by apply prodO_map_ne; apply oFunctor_map_contractive.
Qed.

Program Definition ofe_morOF (F1 F2 : oFunctor) : oFunctor := {|
  oFunctor_car A _ B _ := oFunctor_car F1 B A -n> oFunctor_car F2 A B;
  oFunctor_map A1 _ A2 _ B1 _ B2 _ fg :=
    ofe_morO_map (oFunctor_map F1 (fg.2, fg.1)) (oFunctor_map F2 fg)
|}.
Next Obligation.
  intros F1 F2 A1 ? A2 ? B1 ? B2 ? n [f g] [f' g'] Hfg; simpl in *.
  apply ofe_morO_map_ne; apply oFunctor_map_ne; split; by apply Hfg.
Qed.
Next Obligation.
  intros F1 F2 A ? B ? [f ?] ?; simpl. rewrite /= !oFunctor_map_id.
  apply (ne_proper f). apply oFunctor_map_id.
Qed.
Next Obligation.
  intros F1 F2 A1 ? A2 ? A3 ? B1 ? B2 ? B3 ? f g f' g' [h ?] ?; simpl in *.
  rewrite -!oFunctor_map_compose. do 2 apply (ne_proper _). apply oFunctor_map_compose.
Qed.
Notation "F1 -n> F2" := (ofe_morOF F1%OF F2%OF) : oFunctor_scope.

Instance ofe_morOF_contractive F1 F2 :
  oFunctorContractive F1 → oFunctorContractive F2 →
  oFunctorContractive (ofe_morOF F1 F2).
Proof.
  intros ?? A1 ? A2 ? B1 ? B2 ? n [f g] [f' g'] Hfg; simpl in *.
  apply ofe_morO_map_ne; apply oFunctor_map_contractive; destruct n, Hfg; by split.
Qed.

(** * Sum type *)
Section sum.
  Context {A B : ofeT}.

  Instance sum_dist : Dist (A + B) := λ n, sum_relation (dist n) (dist n).
  Global Instance inl_ne : NonExpansive (@inl A B) := _.
  Global Instance inr_ne : NonExpansive (@inr A B) := _.
  Global Instance inl_ne_inj {n} : Inj (dist n) (dist n) (@inl A B) := _.
  Global Instance inr_ne_inj {n} : Inj (dist n) (dist n) (@inr A B) := _.

  Definition sum_ofe_mixin : OfeMixin (A + B).
  Proof.
    split.
    - intros x y; split=> Hx.
      + destruct Hx=> n; constructor; by apply equiv_dist.
      + destruct (Hx 0); constructor; apply equiv_dist=> n; by apply (inj _).
    - apply _.
    - destruct 1; constructor; by apply dist_S.
  Qed.
  Canonical Structure sumO : ofeT := OfeT (A + B) sum_ofe_mixin.

  Program Definition inl_chain (c : chain sumO) (a : A) : chain A :=
    {| chain_car n := match c n return _ with inl a' => a' | _ => a end |}.
  Next Obligation. intros c a n i ?; simpl. by destruct (chain_cauchy c n i). Qed.
  Program Definition inr_chain (c : chain sumO) (b : B) : chain B :=
    {| chain_car n := match c n return _ with inr b' => b' | _ => b end |}.
  Next Obligation. intros c b n i ?; simpl. by destruct (chain_cauchy c n i). Qed.

  Definition sum_compl `{Cofe A, Cofe B} : Compl sumO := λ c,
    match c 0 with
    | inl a => inl (compl (inl_chain c a))
    | inr b => inr (compl (inr_chain c b))
    end.
  Global Program Instance sum_cofe `{Cofe A, Cofe B} : Cofe sumO :=
    { compl := sum_compl }.
  Next Obligation.
    intros ?? n c; rewrite /compl /sum_compl.
    feed inversion (chain_cauchy c 0 n); first by auto with lia; constructor.
    - rewrite (conv_compl n (inl_chain c _)) /=. destruct (c n); naive_solver.
    - rewrite (conv_compl n (inr_chain c _)) /=. destruct (c n); naive_solver.
  Qed.

  Global Instance inl_discrete (x : A) : Discrete x → Discrete (inl x).
  Proof. inversion_clear 2; constructor; by apply (discrete _). Qed.
  Global Instance inr_discrete (y : B) : Discrete y → Discrete (inr y).
  Proof. inversion_clear 2; constructor; by apply (discrete _). Qed.
  Global Instance sum_ofe_discrete :
    OfeDiscrete A → OfeDiscrete B → OfeDiscrete sumO.
  Proof. intros ?? [?|?]; apply _. Qed.
End sum.

Arguments sumO : clear implicits.
Typeclasses Opaque sum_dist.

Instance sum_map_ne {A A' B B' : ofeT} n :
  Proper ((dist n ==> dist n) ==> (dist n ==> dist n) ==>
           dist n ==> dist n) (@sum_map A A' B B').
Proof.
  intros f f' Hf g g' Hg ??; destruct 1; constructor; [by apply Hf|by apply Hg].
Qed.
Definition sumO_map {A A' B B'} (f : A -n> A') (g : B -n> B') :
  sumO A B -n> sumO A' B' := OfeMor (sum_map f g).
Instance sumO_map_ne {A A' B B'} :
  NonExpansive2 (@sumO_map A A' B B').
Proof. intros n f f' Hf g g' Hg [?|?]; constructor; [apply Hf|apply Hg]. Qed.

Program Definition sumOF (F1 F2 : oFunctor) : oFunctor := {|
  oFunctor_car A _ B _ := sumO (oFunctor_car F1 A B) (oFunctor_car F2 A B);
  oFunctor_map A1 _ A2 _ B1 _ B2 _ fg :=
    sumO_map (oFunctor_map F1 fg) (oFunctor_map F2 fg)
|}.
Next Obligation.
  intros ?? A1 ? A2 ? B1 ? B2 ? n ???; by apply sumO_map_ne; apply oFunctor_map_ne.
Qed.
Next Obligation. by intros F1 F2 A ? B ? [?|?]; rewrite /= !oFunctor_map_id. Qed.
Next Obligation.
  intros F1 F2 A1 ? A2 ? A3 ? B1 ? B2 ? B3 ? f g f' g' [?|?]; simpl;
    by rewrite !oFunctor_map_compose.
Qed.
Notation "F1 + F2" := (sumOF F1%OF F2%OF) : oFunctor_scope.

Instance sumOF_contractive F1 F2 :
  oFunctorContractive F1 → oFunctorContractive F2 →
  oFunctorContractive (sumOF F1 F2).
Proof.
  intros ?? A1 ? A2 ? B1 ? B2 ? n ???;
    by apply sumO_map_ne; apply oFunctor_map_contractive.
Qed.

(** * Discrete OFEs *)
Section discrete_ofe.
  Context `{Equiv A} (Heq : @Equivalence A (≡)).

  Instance discrete_dist : Dist A := λ n x y, x ≡ y.
  Definition discrete_ofe_mixin : OfeMixin A.
  Proof using Type*.
    split.
    - intros x y; split; [done|intros Hn; apply (Hn 0)].
    - done.
    - done.
  Qed.

  Global Instance discrete_ofe_discrete : OfeDiscrete (OfeT A discrete_ofe_mixin).
  Proof. by intros x y. Qed.

  Global Program Instance discrete_cofe : Cofe (OfeT A discrete_ofe_mixin) :=
    { compl c := c 0 }.
  Next Obligation.
    intros n c. rewrite /compl /=;
    symmetry; apply (chain_cauchy c 0 n). lia.
  Qed.
End discrete_ofe.

Notation discreteO A := (OfeT A (discrete_ofe_mixin _)).
(** Force the [Equivalence] proof to be [eq_equivalence] so that it does not
find another one, like [ofe_equivalence], in the case of aliases. See also
https://gitlab.mpi-sws.org/iris/iris/issues/299 *)
Notation leibnizO A := (OfeT A (@discrete_ofe_mixin _ equivL eq_equivalence)).

(** In order to define a discrete CMRA with carrier [A] (in the file [cmra.v])
we need to determine the [Equivalence A] proof that was used to construct the
OFE instance of [A] (note that this proof is not the same as the one we obtain
via [ofe_equivalence]).

We obtain the proof of [Equivalence A] by inferring the canonical OFE mixin
using [ofe_mixin_of A], and then check whether it is indeed a discrete OFE. This
will fail if no OFE, or an OFE other than the discrete OFE, was registered. *)
Notation discrete_ofe_equivalence_of A := ltac:(
  match constr:(ofe_mixin_of A) with
  | discrete_ofe_mixin ?H => exact H
  end) (only parsing).

Instance leibnizO_leibniz A : LeibnizEquiv (leibnizO A).
Proof. by intros x y. Qed.

(** * Basic Coq types *)
Canonical Structure boolO := leibnizO bool.
Canonical Structure natO := leibnizO nat.
Canonical Structure positiveO := leibnizO positive.
Canonical Structure NO := leibnizO N.
Canonical Structure ZO := leibnizO Z.

Section prop.
  Instance Prop_equiv : Equiv Prop := iff.
  Instance Prop_equivalence : Equivalence (≡@{Prop}) := _.
  Canonical Structure PropO := discreteO Prop.
End prop.

(** * Option type *)
Section option.
  Context {A : ofeT}.

  Instance option_dist : Dist (option A) := λ n, option_Forall2 (dist n).
  Lemma dist_option_Forall2 n mx my : mx ≡{n}≡ my ↔ option_Forall2 (dist n) mx my.
  Proof. done. Qed.

  Definition option_ofe_mixin : OfeMixin (option A).
  Proof.
    split.
    - intros mx my; split; [by destruct 1; constructor; apply equiv_dist|].
      intros Hxy; destruct (Hxy 0); constructor; apply equiv_dist.
      by intros n; feed inversion (Hxy n).
    - apply _.
    - destruct 1; constructor; by apply dist_S.
  Qed.
  Canonical Structure optionO := OfeT (option A) option_ofe_mixin.

  Program Definition option_chain (c : chain optionO) (x : A) : chain A :=
    {| chain_car n := default x (c n) |}.
  Next Obligation. intros c x n i ?; simpl. by destruct (chain_cauchy c n i). Qed.
  Definition option_compl `{Cofe A} : Compl optionO := λ c,
    match c 0 with Some x => Some (compl (option_chain c x)) | None => None end.
  Global Program Instance option_cofe `{Cofe A} : Cofe optionO :=
    { compl := option_compl }.
  Next Obligation.
    intros ? n c; rewrite /compl /option_compl.
    feed inversion (chain_cauchy c 0 n); auto with lia; [].
    constructor. rewrite (conv_compl n (option_chain c _)) /=.
    destruct (c n); naive_solver.
  Qed.

  Global Instance option_ofe_discrete : OfeDiscrete A → OfeDiscrete optionO.
  Proof. destruct 2; constructor; by apply (discrete _). Qed.

  Global Instance Some_ne : NonExpansive (@Some A).
  Proof. by constructor. Qed.
  Global Instance is_Some_ne n : Proper (dist n ==> iff) (@is_Some A).
  Proof. destruct 1; split; eauto. Qed.
  Global Instance Some_dist_inj {n} : Inj (dist n) (dist n) (@Some A).
  Proof. by inversion_clear 1. Qed.
  Global Instance from_option_ne {B} (R : relation B) (f : A → B) n :
    Proper (dist n ==> R) f → Proper (R ==> dist n ==> R) (from_option f).
  Proof. destruct 3; simpl; auto. Qed.

  Global Instance None_discrete : Discrete (@None A).
  Proof. inversion_clear 1; constructor. Qed.
  Global Instance Some_discrete x : Discrete x → Discrete (Some x).
  Proof. by intros ?; inversion_clear 1; constructor; apply discrete. Qed.

  Lemma dist_None n mx : mx ≡{n}≡ None ↔ mx = None.
  Proof. split; [by inversion_clear 1|by intros ->]. Qed.
  Lemma dist_Some_inv_l n mx my x :
    mx ≡{n}≡ my → mx = Some x → ∃ y, my = Some y ∧ x ≡{n}≡ y.
  Proof. destruct 1; naive_solver. Qed.
  Lemma dist_Some_inv_r n mx my y :
    mx ≡{n}≡ my → my = Some y → ∃ x, mx = Some x ∧ x ≡{n}≡ y.
  Proof. destruct 1; naive_solver. Qed.
  Lemma dist_Some_inv_l' n my x : Some x ≡{n}≡ my → ∃ x', Some x' = my ∧ x ≡{n}≡ x'.
  Proof. intros ?%(dist_Some_inv_l _ _ _ x); naive_solver. Qed.
  Lemma dist_Some_inv_r' n mx y : mx ≡{n}≡ Some y → ∃ y', mx = Some y' ∧ y ≡{n}≡ y'.
  Proof. intros ?%(dist_Some_inv_r _ _ _ y); naive_solver. Qed.
End option.

Typeclasses Opaque option_dist.
Arguments optionO : clear implicits.

Instance option_fmap_ne {A B : ofeT} n:
  Proper ((dist n ==> dist n) ==> dist n ==> dist n) (@fmap option _ A B).
Proof. intros f f' Hf ?? []; constructor; auto. Qed.
Instance option_mbind_ne {A B : ofeT} n:
  Proper ((dist n ==> dist n) ==> dist n ==> dist n) (@mbind option _ A B).
Proof. destruct 2; simpl; auto. Qed.
Instance option_mjoin_ne {A : ofeT} n:
  Proper (dist n ==> dist n) (@mjoin option _ A).
Proof. destruct 1 as [?? []|]; simpl; by constructor. Qed.

Lemma fmap_Some_dist {A B : ofeT} (f : A → B) (mx : option A) (y : B) n :
  f <$> mx ≡{n}≡ Some y ↔ ∃ x : A, mx = Some x ∧ y ≡{n}≡ f x.
Proof.
  split; [|by intros (x&->&->)].
  intros (?&?%fmap_Some&?)%dist_Some_inv_r'; naive_solver.
Qed.

Definition optionO_map {A B} (f : A -n> B) : optionO A -n> optionO B :=
  OfeMor (fmap f : optionO A → optionO B).
Instance optionO_map_ne A B : NonExpansive (@optionO_map A B).
Proof. by intros n f f' Hf []; constructor; apply Hf. Qed.

Program Definition optionOF (F : oFunctor) : oFunctor := {|
  oFunctor_car A _ B _ := optionO (oFunctor_car F A B);
  oFunctor_map A1 _ A2 _ B1 _ B2 _ fg := optionO_map (oFunctor_map F fg)
|}.
Next Obligation.
  by intros F A1 ? A2 ? B1 ? B2 ? n f g Hfg; apply optionO_map_ne, oFunctor_map_ne.
Qed.
Next Obligation.
  intros F A ? B ? x. rewrite /= -{2}(option_fmap_id x).
  apply option_fmap_equiv_ext=>y; apply oFunctor_map_id.
Qed.
Next Obligation.
  intros F A1 ? A2 ? A3 ? B1 ? B2 ? B3 ? f g f' g' x. rewrite /= -option_fmap_compose.
  apply option_fmap_equiv_ext=>y; apply oFunctor_map_compose.
Qed.

Instance optionOF_contractive F :
  oFunctorContractive F → oFunctorContractive (optionOF F).
Proof.
  by intros ? A1 ? A2 ? B1 ? B2 ? n f g Hfg;
    apply optionO_map_ne, oFunctor_map_contractive.
Qed.

(** * Later type *)
(** Note that the projection [later_car] is not non-expansive (see also the
lemma [later_car_anti_contractive] below), so it cannot be used in the logic.
If you need to get a witness out, you should use the lemma [Next_uninj]
instead. *)
Record later (A : Type) : Type := Next { later_car : A }.
Add Printing Constructor later.
Arguments Next {_} _.
Arguments later_car {_} _.
Instance: Params (@Next) 1 := {}.

Section later.
  Context {A : ofeT}.
  Instance later_equiv : Equiv (later A) := λ x y, later_car x ≡ later_car y.
  Instance later_dist : Dist (later A) := λ n x y,
    dist_later n (later_car x) (later_car y).
  Definition later_ofe_mixin : OfeMixin (later A).
  Proof.
    split.
    - intros x y; unfold equiv, later_equiv; rewrite !equiv_dist.
      split. intros Hxy [|n]; [done|apply Hxy]. intros Hxy n; apply (Hxy (S n)).
    - split; rewrite /dist /later_dist.
      + by intros [x].
      + by intros [x] [y].
      + by intros [x] [y] [z] ??; trans y.
    - intros [|n] [x] [y] ?; [done|]; rewrite /dist /later_dist; by apply dist_S.
  Qed.
  Canonical Structure laterO : ofeT := OfeT (later A) later_ofe_mixin.

  Program Definition later_chain (c : chain laterO) : chain A :=
    {| chain_car n := later_car (c (S n)) |}.
  Next Obligation. intros c n i ?; apply (chain_cauchy c (S n)); lia. Qed.
  Global Program Instance later_cofe `{Cofe A} : Cofe laterO :=
    { compl c := Next (compl (later_chain c)) }.
  Next Obligation.
    intros ? [|n] c; [done|by apply (conv_compl n (later_chain c))].
  Qed.

  Global Instance Next_contractive : Contractive (@Next A).
  Proof. by intros [|n] x y. Qed.
  Global Instance Later_inj n : Inj (dist n) (dist (S n)) (@Next A).
  Proof. by intros x y. Qed.

  Lemma Next_uninj x : ∃ a, x ≡ Next a.
  Proof. by exists (later_car x). Qed.
  Instance later_car_anti_contractive n :
    Proper (dist n ==> dist_later n) later_car.
  Proof. move=> [x] [y] /= Hxy. done. Qed.

  (** [f] is contractive iff it can factor into [Next] and a non-expansive
  function. *)
  Lemma contractive_alt {B : ofeT} (f : A → B) :
    Contractive f ↔ ∃ g : later A → B, NonExpansive g ∧ ∀ x, f x ≡ g (Next x).
  Proof.
    split.
    - intros Hf. exists (f ∘ later_car); split=> // n x y ?. by f_equiv.
    - intros (g&Hg&Hf) n x y Hxy. rewrite !Hf. by apply Hg.
  Qed.
End later.

Arguments laterO : clear implicits.

Definition later_map {A B} (f : A → B) (x : later A) : later B :=
  Next (f (later_car x)).
Instance later_map_ne {A B : ofeT} (f : A → B) n :
  Proper (dist (pred n) ==> dist (pred n)) f →
  Proper (dist n ==> dist n) (later_map f) | 0.
Proof. destruct n as [|n]; intros Hf [x] [y] ?; do 2 red; simpl; auto. Qed.
Instance later_map_proper {A B : ofeT} (f : A → B) :
  Proper ((≡) ==> (≡)) f →
  Proper ((≡) ==> (≡)) (later_map f).
Proof. solve_proper. Qed.
Lemma later_map_id {A} (x : later A) : later_map id x = x.
Proof. by destruct x. Qed.
Lemma later_map_compose {A B C} (f : A → B) (g : B → C) (x : later A) :
  later_map (g ∘ f) x = later_map g (later_map f x).
Proof. by destruct x. Qed.
Lemma later_map_ext {A B : ofeT} (f g : A → B) x :
  (∀ x, f x ≡ g x) → later_map f x ≡ later_map g x.
Proof. destruct x; intros Hf; apply Hf. Qed.
Definition laterO_map {A B} (f : A -n> B) : laterO A -n> laterO B :=
  OfeMor (later_map f).
Instance laterO_map_contractive (A B : ofeT) : Contractive (@laterO_map A B).
Proof. intros [|n] f g Hf n'; [done|]; apply Hf; lia. Qed.

Program Definition laterOF (F : oFunctor) : oFunctor := {|
  oFunctor_car A _ B _ := laterO (oFunctor_car F A B);
  oFunctor_map A1 _ A2 _ B1 _ B2 _ fg := laterO_map (oFunctor_map F fg)
|}.
Next Obligation.
  intros F A1 ? A2 ? B1 ? B2 ? n fg fg' ?.
  by apply (contractive_ne laterO_map), oFunctor_map_ne.
Qed.
Next Obligation.
  intros F A ? B ? x; simpl. rewrite -{2}(later_map_id x).
  apply later_map_ext=>y. by rewrite oFunctor_map_id.
Qed.
Next Obligation.
  intros F A1 ? A2 ? A3 ? B1 ? B2 ? B3 ? f g f' g' x; simpl. rewrite -later_map_compose.
  apply later_map_ext=>y; apply oFunctor_map_compose.
Qed.
Notation "▶ F"  := (laterOF F%OF) (at level 20, right associativity) : oFunctor_scope.

Instance laterOF_contractive F : oFunctorContractive (laterOF F).
Proof.
  intros A1 ? A2 ? B1 ? B2 ? n fg fg' Hfg. apply laterO_map_contractive.
  destruct n as [|n]; simpl in *; first done. apply oFunctor_map_ne, Hfg.
Qed.

(** * Dependently-typed functions over a discrete domain *)
(** This separate notion is useful whenever we need dependent functions, and
whenever we want to avoid the hassle of the bundled non-expansive function type.

Note that non-dependent functions over a discrete domain, [A -d> B] (following
the notation we introduce below) are non-expansive if they are
[Proper ((≡) ==> (≡))]. In other words, since the domain is discrete,
non-expansiveness and respecting [(≡)] are the same. If the domain is moreover
Leibniz ([LeibnizEquiv A]), we get both for free.

We make [discrete_fun] a definition so that we can register it as a canonical
structure.  We do not bundle the [Proper] proof to keep [discrete_fun] easier to
use. It turns out all the desired OFE and functorial properties do not rely on
this [Proper] instance. *)
Definition discrete_fun {A} (B : A → ofeT) := ∀ x : A, B x.

Section discrete_fun.
  Context {A : Type} {B : A → ofeT}.
  Implicit Types f g : discrete_fun B.

  Instance discrete_fun_equiv : Equiv (discrete_fun B) := λ f g, ∀ x, f x ≡ g x.
  Instance discrete_fun_dist : Dist (discrete_fun B) := λ n f g, ∀ x, f x ≡{n}≡ g x.
  Definition discrete_fun_ofe_mixin : OfeMixin (discrete_fun B).
  Proof.
    split.
    - intros f g; split; [intros Hfg n k; apply equiv_dist, Hfg|].
      intros Hfg k; apply equiv_dist=> n; apply Hfg.
    - intros n; split.
      + by intros f x.
      + by intros f g ? x.
      + by intros f g h ?? x; trans (g x).
    - by intros n f g ? x; apply dist_S.
  Qed.
  Canonical Structure discrete_funO := OfeT (discrete_fun B) discrete_fun_ofe_mixin.

  Program Definition discrete_fun_chain `(c : chain discrete_funO)
    (x : A) : chain (B x) := {| chain_car n := c n x |}.
  Next Obligation. intros c x n i ?. by apply (chain_cauchy c). Qed.
  Global Program Instance discrete_fun_cofe `{∀ x, Cofe (B x)} : Cofe discrete_funO :=
    { compl c x := compl (discrete_fun_chain c x) }.
  Next Obligation. intros ? n c x. apply (conv_compl n (discrete_fun_chain c x)). Qed.

  Global Instance discrete_fun_inhabited `{∀ x, Inhabited (B x)} : Inhabited discrete_funO :=
    populate (λ _, inhabitant).
  Global Instance discrete_fun_lookup_discrete `{EqDecision A} f x :
    Discrete f → Discrete (f x).
  Proof.
    intros Hf y ?.
    set (g x' := if decide (x = x') is left H then eq_rect _ B y _ H else f x').
    trans (g x).
    { apply Hf=> x'. unfold g. by destruct (decide _) as [[]|]. }
    unfold g. destruct (decide _) as [Hx|]; last done.
    by rewrite (proof_irrel Hx eq_refl).
  Qed.
End discrete_fun.

Arguments discrete_funO {_} _.
Notation "A -d> B" :=
  (@discrete_funO A (λ _, B)) (at level 99, B at level 200, right associativity).

Definition discrete_fun_map {A} {B1 B2 : A → ofeT} (f : ∀ x, B1 x → B2 x)
  (g : discrete_fun B1) : discrete_fun B2 := λ x, f _ (g x).

Lemma discrete_fun_map_ext {A} {B1 B2 : A → ofeT} (f1 f2 : ∀ x, B1 x → B2 x)
  (g : discrete_fun B1) :
  (∀ x, f1 x (g x) ≡ f2 x (g x)) → discrete_fun_map f1 g ≡ discrete_fun_map f2 g.
Proof. done. Qed.
Lemma discrete_fun_map_id {A} {B : A → ofeT} (g : discrete_fun B) :
  discrete_fun_map (λ _, id) g = g.
Proof. done. Qed.
Lemma discrete_fun_map_compose {A} {B1 B2 B3 : A → ofeT}
    (f1 : ∀ x, B1 x → B2 x) (f2 : ∀ x, B2 x → B3 x) (g : discrete_fun B1) :
  discrete_fun_map (λ x, f2 x ∘ f1 x) g = discrete_fun_map f2 (discrete_fun_map f1 g).
Proof. done. Qed.

Instance discrete_fun_map_ne {A} {B1 B2 : A → ofeT} (f : ∀ x, B1 x → B2 x) n :
  (∀ x, Proper (dist n ==> dist n) (f x)) →
  Proper (dist n ==> dist n) (discrete_fun_map f).
Proof. by intros ? y1 y2 Hy x; rewrite /discrete_fun_map (Hy x). Qed.

Definition discrete_funO_map {A} {B1 B2 : A → ofeT} (f : discrete_fun (λ x, B1 x -n> B2 x)) :
  discrete_funO B1 -n> discrete_funO B2 := OfeMor (discrete_fun_map f).
Instance discrete_funO_map_ne {A} {B1 B2 : A → ofeT} :
  NonExpansive (@discrete_funO_map A B1 B2).
Proof. intros n f1 f2 Hf g x; apply Hf. Qed.

Program Definition discrete_funOF {C} (F : C → oFunctor) : oFunctor := {|
  oFunctor_car A _ B _ := discrete_funO (λ c, oFunctor_car (F c) A B);
  oFunctor_map A1 _ A2 _ B1 _ B2 _ fg := discrete_funO_map (λ c, oFunctor_map (F c) fg)
|}.
Next Obligation.
  intros C F A1 ? A2 ? B1 ? B2 ? n ?? g.
  by apply discrete_funO_map_ne=>?; apply oFunctor_map_ne.
Qed.
Next Obligation.
  intros C F A ? B ? g; simpl. rewrite -{2}(discrete_fun_map_id g).
  apply discrete_fun_map_ext=> y; apply oFunctor_map_id.
Qed.
Next Obligation.
  intros C F A1 ? A2 ? A3 ? B1 ? B2 ? B3 ? f1 f2 f1' f2' g.
  rewrite /= -discrete_fun_map_compose.
  apply discrete_fun_map_ext=>y; apply oFunctor_map_compose.
Qed.

Notation "T -d> F" := (@discrete_funOF T%type (λ _, F%OF)) : oFunctor_scope.

Instance discrete_funOF_contractive {C} (F : C → oFunctor) :
  (∀ c, oFunctorContractive (F c)) → oFunctorContractive (discrete_funOF F).
Proof.
  intros ? A1 ? A2 ? B1 ? B2 ? n ?? g.
  by apply discrete_funO_map_ne=>c; apply oFunctor_map_contractive.
Qed.

(** * Constructing isomorphic OFEs *)
Lemma iso_ofe_mixin {A : ofeT} {B : Type} `{!Equiv B, !Dist B} (g : B → A)
  (g_equiv : ∀ y1 y2, y1 ≡ y2 ↔ g y1 ≡ g y2)
  (g_dist : ∀ n y1 y2, y1 ≡{n}≡ y2 ↔ g y1 ≡{n}≡ g y2) : OfeMixin B.
Proof.
  split.
  - intros y1 y2. rewrite g_equiv. setoid_rewrite g_dist. apply equiv_dist.
  - split.
    + intros y. by apply g_dist.
    + intros y1 y2. by rewrite !g_dist.
    + intros y1 y2 y3. rewrite !g_dist. intros ??; etrans; eauto.
  - intros n y1 y2. rewrite !g_dist. apply dist_S.
Qed.

Section iso_cofe_subtype.
  Context {A B : ofeT} `{Cofe A} (P : A → Prop) (f : ∀ x, P x → B) (g : B → A).
  Context (g_dist : ∀ n y1 y2, y1 ≡{n}≡ y2 ↔ g y1 ≡{n}≡ g y2).
  Let Hgne : NonExpansive g.
  Proof. intros n y1 y2. apply g_dist. Qed.
  Existing Instance Hgne.
  Context (gf : ∀ x Hx, g (f x Hx) ≡ x).
  Context (Hlimit : ∀ c : chain B, P (compl (chain_map g c))).
  Program Definition iso_cofe_subtype : Cofe B :=
    {| compl c := f (compl (chain_map g c)) _ |}.
  Next Obligation. apply Hlimit. Qed.
  Next Obligation.
    intros n c; simpl. apply g_dist. by rewrite gf conv_compl.
  Qed.
End iso_cofe_subtype.

Lemma iso_cofe_subtype' {A B : ofeT} `{Cofe A}
  (P : A → Prop) (f : ∀ x, P x → B) (g : B → A)
  (Pg : ∀ y, P (g y))
  (g_dist : ∀ n y1 y2, y1 ≡{n}≡ y2 ↔ g y1 ≡{n}≡ g y2)
  (gf : ∀ x Hx, g (f x Hx) ≡ x)
  (Hlimit : LimitPreserving P) : Cofe B.
Proof. apply: (iso_cofe_subtype P f g)=> // c. apply Hlimit=> ?; apply Pg. Qed.

Definition iso_cofe {A B : ofeT} `{Cofe A} (f : A → B) (g : B → A)
  (g_dist : ∀ n y1 y2, y1 ≡{n}≡ y2 ↔ g y1 ≡{n}≡ g y2)
  (gf : ∀ x, g (f x) ≡ x) : Cofe B.
Proof. by apply (iso_cofe_subtype (λ _, True) (λ x _, f x) g). Qed.

(** * Sigma type *)
Section sigma.
  Context {A : ofeT} {P : A → Prop}.
  Implicit Types x : sig P.

  (* TODO: Find a better place for this Equiv instance. It also
     should not depend on A being an OFE. *)
  Instance sig_equiv : Equiv (sig P) := λ x1 x2, `x1 ≡ `x2.
  Instance sig_dist : Dist (sig P) := λ n x1 x2, `x1 ≡{n}≡ `x2.

  Definition sig_equiv_alt x y : x ≡ y ↔ `x ≡ `y := reflexivity _.
  Definition sig_dist_alt n x y : x ≡{n}≡ y ↔ `x ≡{n}≡ `y := reflexivity _.

  Lemma exist_ne n a1 a2 (H1 : P a1) (H2 : P a2) :
    a1 ≡{n}≡ a2 → a1 ↾ H1 ≡{n}≡ a2 ↾ H2.
  Proof. done. Qed.

  Global Instance proj1_sig_ne : NonExpansive (@proj1_sig _ P).
  Proof. by intros n [a Ha] [b Hb] ?. Qed.
  Definition sig_ofe_mixin : OfeMixin (sig P).
  Proof. by apply (iso_ofe_mixin proj1_sig). Qed.
  Canonical Structure sigO : ofeT := OfeT (sig P) sig_ofe_mixin.

  Global Instance sig_cofe `{Cofe A, !LimitPreserving P} : Cofe sigO.
  Proof. apply (iso_cofe_subtype' P (exist P) proj1_sig)=> //. by intros []. Qed.

  Global Instance sig_discrete (x : sig P) :  Discrete (`x) → Discrete x.
  Proof. intros ? y. rewrite sig_dist_alt sig_equiv_alt. apply (discrete _). Qed.
  Global Instance sig_ofe_discrete : OfeDiscrete A → OfeDiscrete sigO.
  Proof. intros ??. apply _. Qed.
End sigma.

Arguments sigO {_} _.

(** * SigmaT type *)
(** Ofe for [sigT]. The first component must be discrete and use Leibniz
equality, while the second component might be any OFE. *)
Section sigT.
  Import EqNotations.

  Context {A : Type} {P : A → ofeT}.
  Implicit Types x : sigT P.

  (**
    The distance for [{ a : A & P }] uses Leibniz equality on [A] to
    transport the second components to the same type,
    and then step-indexed distance on the second component.
    Unlike in the topos of trees, with (C)OFEs we cannot use step-indexed equality
    on the first component.
  *)
  Instance sigT_dist : Dist (sigT P) := λ n x1 x2,
    ∃ Heq : projT1 x1 = projT1 x2, rew Heq in projT2 x1 ≡{n}≡ projT2 x2.

  (**
    Usually we'd give a direct definition, and show it equivalent to
    [∀ n, x1 ≡{n}≡ x2] when proving the [equiv_dist] OFE axiom.
    But here the equivalence requires UIP — see [sigT_equiv_eq_alt].
    By defining [equiv] in terms of [dist], we can define an OFE
    without assuming UIP, at the cost of complex reasoning on [equiv].
  *)
  Instance sigT_equiv : Equiv (sigT P) := λ x1 x2,
    ∀ n, x1 ≡{n}≡ x2.

  (** Unfolding lemmas.
      Written with [↔] not [=] to avoid https://github.com/coq/coq/issues/3814. *)
  Definition sigT_equiv_eq x1 x2 : (x1 ≡ x2) ↔ ∀ n, x1 ≡{n}≡ x2 :=
      reflexivity _.

  Definition sigT_dist_eq x1 x2 n : (x1 ≡{n}≡ x2) ↔
    ∃ Heq : projT1 x1 = projT1 x2, (rew Heq in projT2 x1) ≡{n}≡ projT2 x2 :=
      reflexivity _.

  Definition sigT_dist_proj1 n {x y} : x ≡{n}≡ y → projT1 x = projT1 y := proj1_ex.
  Definition sigT_equiv_proj1 {x y} : x ≡ y → projT1 x = projT1 y := λ H, proj1_ex (H 0).

  Definition sigT_ofe_mixin : OfeMixin (sigT P).
  Proof.
    split => // n.
    - split; hnf; setoid_rewrite sigT_dist_eq.
      + intros. by exists eq_refl.
      + move => [xa x] [ya y] /=. destruct 1 as [-> Heq].
        by exists eq_refl.
      + move => [xa x] [ya y] [za z] /=.
        destruct 1 as [-> Heq1].
        destruct 1 as [-> Heq2]. exists eq_refl => /=. by trans y.
    - setoid_rewrite sigT_dist_eq.
      move => [xa x] [ya y] /=. destruct 1 as [-> Heq].
      exists eq_refl. exact: dist_S.
  Qed.

  Canonical Structure sigTO : ofeT := OfeT (sigT P) sigT_ofe_mixin.

  Lemma sigT_equiv_eq_alt `{!∀ a b : A, ProofIrrel (a = b)} x1 x2 :
    x1 ≡ x2 ↔
    ∃ Heq : projT1 x1 = projT1 x2, rew Heq in projT2 x1 ≡ projT2 x2.
  Proof.
    setoid_rewrite equiv_dist; setoid_rewrite sigT_dist_eq; split => Heq.
    - move: (Heq 0) => [H0eq1 _].
      exists H0eq1 => n. move: (Heq n) => [] Hneq1.
      by rewrite (proof_irrel H0eq1 Hneq1).
    - move: Heq => [Heq1 Heqn2] n. by exists Heq1.
  Qed.

  (** [projT1] is non-expansive and proper. *)
  Global Instance projT1_ne : NonExpansive (projT1 : sigTO → leibnizO A).
  Proof. solve_proper. Qed.

  Global Instance projT1_proper : Proper ((≡) ==> (≡)) (projT1 : sigTO → leibnizO A).
  Proof. apply ne_proper, projT1_ne. Qed.

  (** [projT2] is "non-expansive"; the properness lemma [projT2_ne] requires UIP. *)
  Lemma projT2_ne n (x1 x2 : sigTO) (Heq : x1 ≡{n}≡ x2) :
    rew (sigT_dist_proj1 n Heq) in projT2 x1 ≡{n}≡ projT2 x2.
  Proof. by destruct Heq. Qed.

  Lemma projT2_proper `{!∀ a b : A, ProofIrrel (a = b)} (x1 x2 : sigTO) (Heqs : x1 ≡ x2):
    rew (sigT_equiv_proj1 Heqs) in projT2 x1 ≡ projT2 x2.
  Proof.
    move: x1 x2 Heqs => [a1 x1] [a2 x2] Heqs.
    case: (proj1 (sigT_equiv_eq_alt _ _) Heqs) => /=. intros ->.
    rewrite (proof_irrel (sigT_equiv_proj1 Heqs) eq_refl) /=. done.
  Qed.

  (** [existT] is "non-expansive" — general, dependently-typed statement. *)
  Lemma existT_ne n {i1 i2} {v1 : P i1} {v2 : P i2} :
    ∀ (Heq : i1 = i2), (rew f_equal P Heq in v1 ≡{n}≡ v2) →
      existT i1 v1 ≡{n}≡ existT i2 v2.
  Proof. intros ->; simpl. exists eq_refl => /=. done. Qed.

  Lemma existT_proper {i1 i2} {v1 : P i1} {v2 : P i2} :
    ∀ (Heq : i1 = i2), (rew f_equal P Heq in v1 ≡ v2) →
      existT i1 v1 ≡ existT i2 v2.
  Proof. intros Heq Heqv n. apply (existT_ne n Heq), equiv_dist, Heqv. Qed.

  (** [existT] is "non-expansive" — non-dependently-typed version. *)
  Global Instance existT_ne_2 a : NonExpansive (@existT A P a).
  Proof. move => ??? Heq. apply (existT_ne _ eq_refl Heq). Qed.

  Global Instance existT_proper_2 a : Proper ((≡) ==> (≡)) (@existT A P a).
  Proof. apply ne_proper, _. Qed.

  Implicit Types (c : chain sigTO).

  Global Instance sigT_discrete x : Discrete (projT2 x) → Discrete x.
  Proof.
    move: x => [xa x] ? [ya y] [] /=; intros -> => /= Hxy n.
    exists eq_refl => /=. apply equiv_dist, (discrete _), Hxy.
  Qed.

  Global Instance sigT_ofe_discrete : (∀ a, OfeDiscrete (P a)) → OfeDiscrete sigTO.
  Proof. intros ??. apply _. Qed.

  Lemma sigT_chain_const_proj1 c n : projT1 (c n) = projT1 (c 0).
  Proof. refine (sigT_dist_proj1 _ (chain_cauchy c 0 n _)). lia. Qed.

  (* For this COFE construction we need UIP (Uniqueness of Identity Proofs)
    on [A] (i.e. [∀ x y : A, ProofIrrel (x = y)]. UIP is most commonly obtained
    from decidable equality (by Hedberg’s theorem, see
    [stdpp.proof_irrel.eq_pi]). *)
  Section cofe.
    Context `{!∀ a b : A, ProofIrrel (a = b)} `{!∀ a, Cofe (P a)}.

    Program Definition chain_map_snd c : chain (P (projT1 (c 0))) :=
      {| chain_car n := rew (sigT_chain_const_proj1 c n) in projT2 (c n) |}.
    Next Obligation.
      move => c n i Hle /=.
      (* [Hgoal] is our thesis, up to casts: *)
      case: (chain_cauchy c n i Hle) => [Heqin Hgoal] /=.
      (* Pretty delicate. We have two casts to [projT1 (c 0)].
        We replace those by one cast. *)
      move: (sigT_chain_const_proj1 c i) (sigT_chain_const_proj1 c n)
        => Heqi0 Heqn0.
      (* Rewrite [projT1 (c 0)] to [projT1 (c n)] in goal and [Heqi0]: *)
      destruct Heqn0.
      by rewrite /= (proof_irrel Heqi0 Heqin).
    Qed.

    Definition sigT_compl : Compl sigTO :=
      λ c, existT (projT1 (chain_car c 0)) (compl (chain_map_snd c)).

    Global Program Instance sigT_cofe : Cofe sigTO := { compl := sigT_compl }.
    Next Obligation.
      intros n c. rewrite /sigT_compl sigT_dist_eq /=.
      exists (symmetry (sigT_chain_const_proj1 c n)).
      (* Our thesis, up to casts: *)
      pose proof (conv_compl n (chain_map_snd c)) as Hgoal.
      move: (compl (chain_map_snd c)) Hgoal => pc0 /=.
      destruct (sigT_chain_const_proj1 c n); simpl. done.
    Qed.
  End cofe.
End sigT.

Arguments sigTO {_} _.

Section sigTOF.
  Context {A : Type}.

  Program Definition sigT_map {P1 P2 : A → ofeT} :
    discrete_funO (λ a, P1 a -n> P2 a) -n>
    sigTO P1 -n> sigTO P2 :=
    λne f xpx, existT _ (f _ (projT2 xpx)).
  Next Obligation.
    move => ?? f n [x px] [y py] [/= Heq]. destruct Heq; simpl.
    exists eq_refl => /=. by f_equiv.
  Qed.
  Next Obligation.
    move => ?? n f g Heq [x px] /=. exists eq_refl => /=. apply Heq.
  Qed.

  Program Definition sigTOF (F : A → oFunctor) : oFunctor := {|
    oFunctor_car A CA B CB := sigTO (λ a, oFunctor_car (F a) A B);
    oFunctor_map A1 _ A2 _ B1 _ B2 _ fg := sigT_map (λ a, oFunctor_map (F a) fg)
  |}.
  Next Obligation.
    repeat intro. exists eq_refl => /=. solve_proper.
  Qed.
  Next Obligation.
    simpl; intros. apply (existT_proper eq_refl), oFunctor_map_id.
  Qed.
  Next Obligation.
    simpl; intros. apply (existT_proper eq_refl), oFunctor_map_compose.
  Qed.

  Global Instance sigTOF_contractive {F} :
    (∀ a, oFunctorContractive (F a)) → oFunctorContractive (sigTOF F).
  Proof.
    repeat intro. apply sigT_map => a. exact: oFunctor_map_contractive.
  Qed.
End sigTOF.
Arguments sigTOF {_} _%OF.

Notation "{ x  &  P }" := (sigTOF (λ x, P%OF)) : oFunctor_scope.
Notation "{ x : A &  P }" := (@sigTOF A%type (λ x, P%OF)) : oFunctor_scope.

(** * Isomorphisms between OFEs *)
Record ofe_iso (A B : ofeT) := OfeIso {
  ofe_iso_1 : A -n> B;
  ofe_iso_2 : B -n> A;
  ofe_iso_12 y : ofe_iso_1 (ofe_iso_2 y) ≡ y;
  ofe_iso_21 x : ofe_iso_2 (ofe_iso_1 x) ≡ x;
}.
Arguments OfeIso {_ _} _ _ _ _.
Arguments ofe_iso_1 {_ _} _.
Arguments ofe_iso_2 {_ _} _.
Arguments ofe_iso_12 {_ _} _ _.
Arguments ofe_iso_21 {_ _} _ _.

Section ofe_iso.
  Context {A B : ofeT}.

  Instance ofe_iso_equiv : Equiv (ofe_iso A B) := λ I1 I2,
    ofe_iso_1 I1 ≡ ofe_iso_1 I2 ∧ ofe_iso_2 I1 ≡ ofe_iso_2 I2.

  Instance ofe_iso_dist : Dist (ofe_iso A B) := λ n I1 I2,
    ofe_iso_1 I1 ≡{n}≡ ofe_iso_1 I2 ∧ ofe_iso_2 I1 ≡{n}≡ ofe_iso_2 I2.

  Global Instance ofe_iso_1_ne : NonExpansive (ofe_iso_1 (A:=A) (B:=B)).
  Proof. by destruct 1. Qed.
  Global Instance ofe_iso_2_ne : NonExpansive (ofe_iso_2 (A:=A) (B:=B)).
  Proof. by destruct 1. Qed.

  Lemma ofe_iso_ofe_mixin : OfeMixin (ofe_iso A B).
  Proof. by apply (iso_ofe_mixin (λ I, (ofe_iso_1 I, ofe_iso_2 I))). Qed.
  Canonical Structure ofe_isoO : ofeT := OfeT (ofe_iso A B) ofe_iso_ofe_mixin.

  Global Instance ofe_iso_cofe `{!Cofe A, !Cofe B} : Cofe ofe_isoO.
  Proof.
    apply (iso_cofe_subtype'
      (λ I : prodO (A -n> B) (B -n> A),
        (∀ y, I.1 (I.2 y) ≡ y) ∧ (∀ x, I.2 (I.1 x) ≡ x))
      (λ I HI, OfeIso (I.1) (I.2) (proj1 HI) (proj2 HI))
      (λ I, (ofe_iso_1 I, ofe_iso_2 I))); [by intros []|done|done|].
    apply limit_preserving_and; apply limit_preserving_forall=> ?;
      apply limit_preserving_equiv; first [intros ???; done|solve_proper].
  Qed.
End ofe_iso.

Arguments ofe_isoO : clear implicits.

Program Definition iso_ofe_refl {A} : ofe_iso A A := OfeIso cid cid _ _.
Solve Obligations with done.

Definition iso_ofe_sym {A B : ofeT} (I : ofe_iso A B) : ofe_iso B A :=
  OfeIso (ofe_iso_2 I) (ofe_iso_1 I) (ofe_iso_21 I) (ofe_iso_12 I).
Instance iso_ofe_sym_ne {A B} : NonExpansive (iso_ofe_sym (A:=A) (B:=B)).
Proof. intros n I1 I2 []; split; simpl; by f_equiv. Qed.

Program Definition iso_ofe_trans {A B C}
    (I : ofe_iso A B) (J : ofe_iso B C) : ofe_iso A C :=
  OfeIso (ofe_iso_1 J ◎ ofe_iso_1 I) (ofe_iso_2 I ◎ ofe_iso_2 J) _ _.
Next Obligation. intros A B C I J z; simpl. by rewrite !ofe_iso_12. Qed.
Next Obligation. intros A B C I J z; simpl. by rewrite !ofe_iso_21. Qed.
Instance iso_ofe_trans_ne {A B C} : NonExpansive2 (iso_ofe_trans (A:=A) (B:=B) (C:=C)).
Proof. intros n I1 I2 [] J1 J2 []; split; simpl; by f_equiv. Qed.

Program Definition iso_ofe_cong (F : oFunctor) `{!Cofe A, !Cofe B}
    (I : ofe_iso A B) : ofe_iso (oFunctor_apply F A) (oFunctor_apply F B) :=
  OfeIso (oFunctor_map F (ofe_iso_2 I, ofe_iso_1 I))
    (oFunctor_map F (ofe_iso_1 I, ofe_iso_2 I)) _ _.
Next Obligation.
  intros F A ? B ? I x. rewrite -oFunctor_map_compose -{2}(oFunctor_map_id F x).
  apply equiv_dist=> n.
  apply oFunctor_map_ne; split=> ? /=; by rewrite ?ofe_iso_12 ?ofe_iso_21.
Qed.
Next Obligation.
  intros F A ? B ? I y. rewrite -oFunctor_map_compose -{2}(oFunctor_map_id F y).
  apply equiv_dist=> n.
  apply oFunctor_map_ne; split=> ? /=; by rewrite ?ofe_iso_12 ?ofe_iso_21.
Qed.
Instance iso_ofe_cong_ne (F : oFunctor) `{!Cofe A, !Cofe B} :
  NonExpansive (iso_ofe_cong F (A:=A) (B:=B)).
Proof. intros n I1 I2 []; split; simpl; by f_equiv. Qed.
Instance iso_ofe_cong_contractive (F : oFunctor) `{!Cofe A, !Cofe B} :
  oFunctorContractive F → Contractive (iso_ofe_cong F (A:=A) (B:=B)).
Proof. intros ? n I1 I2 HI; split; simpl; f_contractive; by destruct HI. Qed.