feat: OFunctor defined on COFEs

[Kaptch]

Description

OFunctor now takes COFE inputs but only promises an OFE output (OFE.lean).
OFunctorPre is re-indexed with [COFE a] [COFE b] (was OFE), and the old cofe field becomes ofe.
The requirement that the result is a COFE is still in COFESolver.lean file.

Fixes the problem encountered in
https://leanprover.zulipchat.com/#narrow/channel/490604-iris-lean/topic/OFunctor.20definition/with/607339808

The problematic snippet then becomes:

import Iris.Algebra.COFESolver

open Iris Iris.COFE OFE

namespace TraceMWE

abbrev TraceF (V : Type _) [OFE V] : OFunctorPre :=
  SumOF (constOF V) (LaterOF IdOF)

instance (V : Type _) [OFE V] : Inhabited (TraceF V (ULift Unit) (ULift Unit)) :=
  ⟨.inr (Later.next default)⟩

abbrev Trace (V : Type _) [COFE V] : Type _ := OFunctor.Fix (TraceF V)

example (V : Type _) [COFE V] : Type _ := Trace V

section Functoriality
open OFunctor

@[simp] theorem laterMap_apply {A B : Type _} [OFE A] [OFE B] (t : A -n> B) (x : Later A) :
    laterMap t x = Later.next (t x.car) := rfl

theorem laterMap_id {A : Type _} [OFE A] : laterMap (Hom.id : A -n> A) = Hom.id := rfl

theorem laterMap_comp {A B C : Type _} [OFE A] [OFE B] [OFE C] (g : B -n> C) (f : A -n> B) :
    laterMap (g.comp f) = (laterMap g).comp (laterMap f) := rfl

theorem laterMap_ne {A B : Type _} [OFE A] [OFE B] {n} {t t' : A -n> B} (h : t ≡{n}≡ t') :
    laterMap t ≡{n}≡ laterMap t' := fun x _ hm => (h x.car).lt hm

theorem laterMap_contractive {A B : Type _} [OFE A] [OFE B] {n} {t t' : A -n> B}
    (h : DistLater n t t') : laterMap t ≡{n}≡ laterMap t' := fun x m hm => h m hm x.car

variable {V₁ V₂ V₃ : Type _} [COFE V₁] [COFE V₂] [COFE V₃]

def traceMapBody (h : V₁ -n> V₂) (t : Trace V₁ -n> Trace V₂) : Trace V₁ -n> Trace V₂ :=
  Fix.fold.comp (Hom.comp ⟨Sum.map h (laterMap t), inferInstance⟩ Fix.unfold)

theorem traceMapBody_contractive (h : V₁ -n> V₂) : Contractive (traceMapBody h) where
  distLater_dist {n t t'} H := by
    have hlt : laterMap t ≡{n}≡ laterMap t' := laterMap_contractive H
    intro x
    refine Fix.fold.ne.1 ?_
    show Sum.map h (laterMap t) (Fix.unfold x) ≡{n}≡ Sum.map h (laterMap t') (Fix.unfold x)
    rcases Fix.unfold x with a | b
    · exact .rfl
    · exact hlt b

def traceMap (h : V₁ -n> V₂) : Trace V₁ -n> Trace V₂ :=
  haveI : Contractive (traceMapBody h) := traceMapBody_contractive h
  ContractiveHom.fixpoint (Function.toContractiveHom (traceMapBody h))

theorem traceMap_unfold (h : V₁ -n> V₂) : traceMap h ≡ traceMapBody h (traceMap h) := by
  haveI : Contractive (traceMapBody h) := traceMapBody_contractive h
  exact fixpoint_unfold (Function.toContractiveHom (traceMapBody h))

theorem traceMap_id : traceMap (Hom.id : V₁ -n> V₁) ≡ Hom.id := by
  haveI : Contractive (traceMapBody (Hom.id : V₁ -n> V₁)) := traceMapBody_contractive _
  refine (fixpoint_unique (f := Function.toContractiveHom (traceMapBody Hom.id)) ?_).symm
  intro x
  refine (Fix.fold_unfold x).symm.trans (NonExpansive.eqv ?_)
  show Fix.unfold x ≡ Sum.map Hom.id (laterMap Hom.id) (Fix.unfold x)
  rcases Fix.unfold x with a | b <;> exact .rfl

theorem traceMap_comp (g : V₂ -n> V₃) (f : V₁ -n> V₂) :
    traceMap (g.comp f) ≡ (traceMap g).comp (traceMap f) := by
  haveI : Contractive (traceMapBody (g.comp f)) := traceMapBody_contractive _
  refine (fixpoint_unique (f := Function.toContractiveHom (traceMapBody (g.comp f))) ?_).symm
  intro x
  have e₂ : Fix.unfold (traceMap f x) ≡ Sum.map f (laterMap (traceMap f)) (Fix.unfold x) :=
    (NonExpansive.eqv (traceMap_unfold f x)).trans (Fix.unfold_fold _)
  refine (traceMap_unfold g (traceMap f x)).trans (NonExpansive.eqv (f := Fix.fold) ?_)
  refine (NonExpansive.eqv (f := Sum.map g (laterMap (traceMap g))) e₂).trans ?_
  show Sum.map g (laterMap (traceMap g)) (Sum.map f (laterMap (traceMap f)) (Fix.unfold x))
      ≡ Sum.map (g.comp f) (laterMap ((traceMap g).comp (traceMap f))) (Fix.unfold x)
  rcases Fix.unfold x with a | b <;> exact .rfl

theorem traceMap_ne {n} {h h' : V₁ -n> V₂} (e : h ≡{n}≡ h') : traceMap h ≡{n}≡ traceMap h' := by
  haveI : Contractive (traceMapBody h) := traceMapBody_contractive h
  haveI : Contractive (traceMapBody h') := traceMapBody_contractive h'
  refine NonExpansive.ne (f := ContractiveHom.fixpoint) ?_
  intro T x
  refine Fix.fold.ne.1 ?_
  show Sum.map h (laterMap T) (Fix.unfold x) ≡{n}≡ Sum.map h' (laterMap T) (Fix.unfold x)
  rcases Fix.unfold x with a | b
  · exact e a
  · exact .rfl

end Functoriality

open OFunctor in
abbrev TraceLaterOF : OFunctorPre := fun _ β _ _ => Trace (Later β)

open OFunctor in
instance : OFunctor TraceLaterOF where
  ofe := inferInstance
  map _ g := traceMap (laterMap g)
  map_ne := ⟨fun _ _ _ _ _ _ hg => traceMap_ne (laterMap_ne hg)⟩
  map_id x := traceMap_id x
  map_comp _ _ f' g' x := traceMap_comp (laterMap g') (laterMap f') x

open OFunctor OFunctorContractive in
instance : OFunctorContractive TraceLaterOF where
  map_contractive.1 {_} _ _ h :=
    traceMap_ne (laterMap_contractive fun m hm => (h m hm).2)

open OFunctor in
abbrev DSig : OFunctorPre := HomOF (LaterOF IdOF) TraceLaterOF

instance : OFunctorContractive DSig := inferInstance

abbrev D : Type _ := OFunctor.Fix DSig

open OFunctor in
def D.iso : OFE.Iso D (Later D -n> Trace (Later D)) := by
  refine ⟨Fix.unfold (F := DSig), Fix.fold (F := DSig), ?_, ?_⟩
  · intro x; exact Fix.unfold_fold (F := DSig) x
  · intro x; exact Fix.fold_unfold (F := DSig) x

end TraceMWE