ofe.v

From iris.algebra Require Export base.
Set Default Proof Using "Type".
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 : x ≡ y ↔ ∀ n, x ≡{n}≡ y;
  mixin_dist_equivalence n : Equivalence (dist n);
  mixin_dist_S n x y : 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 {A AC AiH} f {Hf} := fixpoint_aux.(unseal) A AC AiH f Hf.
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.
  Local Unset Default Proof Using.

  Context `{Cofe A, Cofe B, !Inhabited A, !Inhabited B}.
  Context (fA : A → B → A).
  Context (fB : A → B → B).
  Context `{∀ n, Proper (dist_later n ==> dist n ==> dist n) fA}.
  Context `{∀ 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. 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.
    apply ne_proper_2=> n x x' ? y y' ?. f_contractive; auto using dist_S.
  Qed.
  Instance: Proper ((≡) ==> (≡) ==> (≡)) fB.
  Proof.
    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 → oFunc