Delay in Lean

delay.lean

-- CoInductive delay (A : Type) : Type :=
--   | now : A -> delay A
--   | later : delay A -> delay A

-- See:
-- https://www.cl.cam.ac.uk/techreports/UCAM-CL-TR-304.pdf
-- https://courses.grainger.illinois.edu/cs576/sp2010/doc/ind-defs.pdf

import Mathlib.Tactic
import Mathlib.Data.Set.Basic
import Mathlib.Data.Set.Lattice
open Set

set_option autoImplicit true
set_option linter.unusedTactic false

-- Type ⋂ using '\' + 'Inter'
def lfp (f : Set α → Set α) : Set α := ⋂₀ {x : Set α | f x ⊆ x}

def lfp_fix (Hm : Monotone f) : lfp f = f (lfp f) := by
  have H₁ : f (lfp f) ⊆ lfp f := by
    simp [lfp]; intro P HP; trans (f P) <;> try assumption
    apply Hm; apply sInter_subset_of_mem
    simp; assumption
  have H₂ : lfp f ⊆ f (lfp f) := by
    simp[lfp]; apply sInter_subset_of_mem; apply Hm; apply H₁
  apply Subset.antisymm <;> assumption
  done

def gfp (f : Set α → Set α) : Set α := ⋃₀ {x : Set α | x ⊆ f x}

def gfp_fix (Hm : Monotone f) : gfp f = f (gfp f) := by
  have H₁ : gfp f ⊆ f (gfp f) := by
    simp [gfp]; intro P HP; trans (f P) <;> try assumption
    apply Hm; apply subset_sUnion_of_subset
    · rfl
    · assumption
  have H₂ : f (gfp f) ⊆ gfp f := by
    intro a H; simp [gfp]; exists (f (gfp f)); constructor <;> try simp [*]
    apply Hm; assumption
  apply Subset.antisymm <;> assumption
  done

def gfp_greatest (H : P ⊆ f P) : P ⊆ gfp f := by
  intro a H; simp [gfp]; exists P

def coind (f : Set α → Set α) (P : Set α) (H : P ⊆ f P) :
  ∀ a, a ∈ P → a ∈ gfp f := by
  intro a H; simp [gfp]; exists P

abbrev node (α : Type) := Set (Nat ⊕ α)

instance : Preorder (node α) := by
  refine Preorder.mk ?r₁ ?r₂
  intro a; simp
  intro a b c Ha Hb; trans b <;> assumption
  done

def succInl : Nat ⊕ α → Nat ⊕ α
  | .inl n => .inl (n + 1)
  | .inr a => .inr a

def Now (a : α) : node α := λx ↦ x = Sum.inr a

def Later (l : node α) : node α := {Sum.inl 0} ∪ (succInl '' l)

def Now_neq_Later : Now a ≠ Later l := by
  intro C
  have G : Later l ⊆ Now a := by simp [*]
  rw [subset_def] at G
  specialize (G (Sum.inl 0) (by simp [Later]))
  contradiction

def DelayF : Set (node α) → Set (node α) :=
  λT ↦ { Now x | x : α } ∪ { Later l | l ∈ T }

def monotonic_DelayF : Monotone DelayF (α := Set (node α)) := by
  intro a b Hab; simp [DelayF] at *
  intro x Hx; right
  obtain ⟨l, Hl, Hx⟩ := Hx
  exists l; simp [*]; apply Hab; assumption
  done

def Delay : Set (node α) := gfp DelayF

def Delay_eta : x ∈ Delay → (∃ a, x = Now a) ∨ (∃ l, l ∈ Delay ∧ x = Later l) := by
  intro H; simp [Delay, gfp] at H
  obtain ⟨t, ⟨tM, xt⟩⟩ := H
  have tD : t ⊆ DelayF t := by assumption
  simp [DelayF, subset_def] at tM
  specialize (tM x xt)
  cases tM with
  | inl n => obtain ⟨a, Hn⟩ := n; left; exists a; simp [*]
  | inr l =>
    obtain ⟨l, Hl⟩ := l; right; exists l
    simp [*]; unfold Delay gfp; simp; exists t; simp [*]
  done

def Now_inj : Now a = Now a' → a = a' := by
  intro H; unfold Now at H; apply congrFun at H; aesop

def succInl_inj : succInl x = succInl y → x = y := by
  intro H
  cases x with
  | inl x => cases y <;> simp [succInl] at *; simp [*]
  | inr x => cases y <;> simp [succInl] at *; simp [*]
  done

def succInl_image_inj : succInl '' l = succInl '' r → l = r := by
  rename_i α
  intro H
  have I : Function.Injective (β := Set (ℕ ⊕ α)) (Set.image succInl) := by
    rw [image_injective]
    apply succInl_inj
  apply I H
  done

def Later_inj : Later l = Later l' → l = l' := by
  rename_i α
  intros H; unfold Later at H
  have N : ∀ (l : node α), Sum.inl 0 ∉ succInl '' l := by
    intro l; simp; constructor <;> intros x a H <;> simp [succInl] at *
  have : succInl '' l = succInl '' l' := by
    have : ({Sum.inl 0} ∪ succInl '' l) \ {Sum.inl 0} = succInl '' l := by
      apply Set.union_diff_cancel_left; simp [*]
    rw [← this]
    have : ({Sum.inl 0} ∪ succInl '' l') \ {Sum.inl 0} = succInl '' l' := by
      apply Set.union_diff_cancel_left; simp [*]
    rw [← this]
    simp [*]
  apply succInl_image_inj at this
  assumption
  done

def DelayT (α : Type) : Type := { x : node α // x ∈ Delay }

def LaterDelay (l : DelayT α) : Later l.val ∈ Delay := by
  apply coind (P := {Later l.val} ∪ Delay) <;> simp
  intro x xmem; simp at xmem
  cases xmem with
  | inl xL =>
    right; simp; right
    exists l.val; simp [l.property, xL]
  | inr xD =>
    apply Delay_eta at xD
    obtain ⟨a, xN⟩ | ⟨l, ⟨lD, xL⟩⟩ := xD
    · left; simp; exists a; simp [*]
    · right; simp; right; exists l; simp [*]
  done

def NOW (a : α) : DelayT α := Subtype.mk (Now a)
  (by exists {Now a}; constructor <;> simp [DelayF])

def LATER (l : DelayT α) : DelayT α := Subtype.mk (Later l.val) (LaterDelay l)

def NOW_neq_LATER : NOW a ≠ LATER l := by
  intro C; injection C; apply Now_neq_Later; assumption

def NOW_inj : NOW a = NOW a' → a = a' := by
  intro H; apply Now_inj; injection H

def LATER_inj : LATER l = LATER l' → l = l' := by
  intro H; cases l; cases l'; congr; apply Later_inj; injection H

def delay_weak_coind : x ∈ P → P ⊆ DelayF P → x ∈ Delay := by
  intro Px Pc; apply coind <;> assumption

def delay_coind : x ∈ P → P ⊆ (DelayF P ∪ Delay) → x ∈ Delay := by
  -- Proved by Claude Opus 4.5
  intro Px Pc
  apply coind (P := P ∪ Delay)
  · intro y Hy
    cases Hy with
    | inl yP =>
      have := Pc yP
      cases this with
      | inl h =>
        -- y ∈ DelayF P, so y ∈ DelayF (P ∪ Delay) by monotonicity
        apply monotonic_DelayF (subset_union_left) h
      | inr h =>
        -- y ∈ Delay = gfp DelayF, so y ∈ DelayF Delay ⊆ DelayF (P ∪ Delay)
        simp only [Delay] at h
        rw [gfp_fix monotonic_DelayF] at h
        apply monotonic_DelayF (subset_union_right) h
    | inr yD =>
      -- y ∈ Delay, same reasoning
      simp only [Delay] at yD
      rw [gfp_fix monotonic_DelayF] at yD
      apply monotonic_DelayF (subset_union_right) yD
  · left; exact Px

def dcorf (a : β) (f : β → α ⊕ β) : Nat → node α
  | 0 => ∅
  | n + 1 =>
    match f a with
    | .inl u => Now u
    | .inr b => Later (dcorf b f n)

def delay_corec (a : β) (f : β → α ⊕ β) := ⋃ k, dcorf a f k

def delay_corec_def : ∀ (a : β) (f : β → α ⊕ β),
  delay_corec a f = match f a with
    | Sum.inl u => Now u
    | Sum.inr b => Later (delay_corec b f) := by
  -- Proved by Claude Opus 4.5
  intro a f
  apply Subset.antisymm
  · -- (⊆) delay_corec a f ⊆ RHS
    intro x hx
    unfold delay_corec at hx
    simp only [mem_iUnion] at hx
    obtain ⟨i, hx⟩ := hx
    cases i with
    | zero => simp [dcorf] at hx
    | succ n =>
      simp only [dcorf] at hx
      cases E : f a with
      | inl u => simp [E] at hx; simp [E, hx]
      | inr b =>
        simp only [E, Later, mem_union, mem_singleton_iff, mem_image] at hx ⊢
        rcases hx with h | ⟨y, hy, rfl⟩
        · left; exact h
        · right; refine ⟨y, ?_, rfl⟩; unfold delay_corec; simp only [mem_iUnion]; exact ⟨n, hy⟩
  · -- (⊇) RHS ⊆ delay_corec a f
    intro x hx
    cases E : f a with
    | inl u =>
      simp [E, Now] at hx
      unfold delay_corec; simp only [mem_iUnion]
      exact ⟨1, by simp [dcorf, E, hx]⟩
    | inr b =>
      simp only [E, Later, mem_union, mem_singleton_iff, mem_image] at hx
      unfold delay_corec; simp only [mem_iUnion]
      rcases hx with rfl | ⟨y, hy, rfl⟩
      · exact ⟨1, by simp [dcorf, E, Later]⟩
      · unfold delay_corec at hy; simp only [mem_iUnion] at hy
        obtain ⟨k, hk⟩ := hy
        refine ⟨k + 1, ?_⟩
        simp only [dcorf, E, Later, mem_union, mem_singleton_iff, mem_image]
        right; exact ⟨y, hk, rfl⟩

-- def corec {E A β : Type} (f : β → ITree.Base E A β) (b : β) : ITree E A
--   := MvQPF.Cofix.corec (n := 2) (α := (Vec.reverse (Vec.nil.append1 A ::: E))) (F := TypeFun.ofCurried (n := 3) (ITree.Base)) f b

def delay_corec_Delay : ∀ (a : β) (f : β → α ⊕ β), delay_corec a f ∈ Delay := by
  intro a f
  apply delay_coind (P := { delay_corec b f | b : β })
  · exact ⟨a, rfl⟩
  · intro x hx
    simp only [mem_setOf_eq] at hx
    obtain ⟨b, rfl⟩ := hx
    rw [delay_corec_def]
    cases E : f b with
    | inl u =>
      left; simp only [DelayF, mem_union, mem_setOf_eq]; left; exact ⟨u, rfl⟩
    | inr b' =>
      left; simp only [DelayF, mem_union, mem_setOf_eq]; right
      exact ⟨delay_corec b' f, ⟨b', rfl⟩, rfl⟩


def omega (α : Type) : node α := delay_corec () (λ_ ↦ Sum.inr ())

def omega_Delay : omega α ∈ Delay := delay_corec_Delay () _

def omega_def {α : Type} : omega α = Later (omega α) := by
  unfold omega
  conv => lhs; rw [delay_corec_def]; simp

#check NOW
#check LATER

def OMEGA (α : Type) : DelayT α := ⟨omega α, omega_Delay⟩

def OMEGA_def {α : Type} : OMEGA α = LATER (OMEGA α) := Subtype.ext omega_def

-- CoInductive diverges {A : Type} : Delay A → Prop :=
--   | diverges_later : ∀ x, diverges x → diverges (Later x)

#check gfp (?F : Set (node ?a) → Set (node ?a))

def divergesF : Set (node α) → Set (node α) :=
  λ (X : Set (node α)) ↦ {x | ∃ x₀, x = Later x₀ ∧ x₀ ∈ X}

def diverges : Set (node α) := gfp divergesF

-- TODO: Want to be able to define divergesF over DelayT, not Set (node α)

inductive itreeF (ε : Type → Type) (α : Type) (itree : Type) where
  | Ret (r : α)
  | Tau (t : itree)
  | Vis {β : Type} (e : ε β) (k : β → itree)

-- CoInductive itree : Type := go
-- { _observe : itreeF itree }.