W-type example for lean4

W-type-test.lean

inductive W (A : Type) (B : A → Type) where
  | mk (a : A) (f : B a → W A B) : W A B

def wInd {A : Type} {B : A → Type} {p : W A B → Type}
  (step : (a : A) → (f : B a → W A B) → (∀ (b : B a), p (f b)) → p (W.mk a f))
    (t : W A B) : p t :=
  match t with
  | W.mk a f => step a f (fun b => wInd step (f b))

-- N

inductive NatShape where
  | zero_s | succ_s

def ℕ := W NatShape (λ s => match s with | .zero_s => Empty | .succ_s => Unit)

def ℕ.zero : ℕ := W.mk NatShape.zero_s Empty.elim

def ℕ.succ : ℕ → ℕ := λ n =>
  W.mk NatShape.succ_s (λ _ => n)

#check Nat.rec

def ℕ.recEasy {T} (zero : T) (succ : ℕ → T → T) (n : ℕ) : T :=
  match n with
  | ⟨.zero_s, _⟩ => zero
  | ⟨.succ_s, f⟩ => let pred := f ()
    succ pred $ ℕ.recEasy zero succ pred

def ℕ.rec {motive : ℕ → Sort u} (zero : motive ℕ.zero) (succ : (n : ℕ) → motive n → motive n.succ)
  (t : ℕ) : motive t :=
    match t with
    | ⟨.zero_s, f⟩ => have : f = Empty.elim := funext (λ x => x.elim)
        this ▸ zero
    | ⟨.succ_s, f⟩ => let pred := f ()
    succ pred $ ℕ.rec zero succ pred

#print axioms ℕ.rec

#check ℕ.succ $ ℕ.succ $ ℕ.zero

section
def ℕ.rec' {motive : ℕ → Sort u}
  -- zero 现在接受任何 f，只要它的形状是 zero_s
  (zero : ∀ (f : Empty → ℕ), motive (W.mk .zero_s f))
  -- succ 接受任何 f，以及递归结果
  (succ : ∀ (f : Unit → ℕ), motive (f ()) → motive (W.mk .succ_s f))
  (t : ℕ) : motive t :=
    match t with
    | ⟨.zero_s, f⟩ => zero f
    | ⟨.succ_s, f⟩ => succ f (ℕ.rec' zero succ (f ()))
-- 改变 zero 和 succ 的类型签名，使它们接受任何 f，而不是特定的 f。绕过等价性证明，但是更麻烦。

def ℕ.rec'' {motive : ℕ → Type}
  (zero : ∀ f, motive (W.mk .zero_s f))
  (succ : ∀ f, motive (f ()) → motive (W.mk .succ_s f))
  (t : ℕ) : motive t :=
  wInd (λ a f ih =>
    match a with
    | .zero_s => zero f
    | .succ_s => succ f (ih ())
  ) t -- use wInd
end

def myAdd (n : ℕ) (m : ℕ) : ℕ :=
  ℕ.rec (motive := λ _ => ℕ)
    m
    (λ n' ih => ℕ.succ ih)
    n

def myMul (n : ℕ) (m : ℕ) : ℕ :=
  ℕ.rec (motive := λ _ => ℕ)
    ℕ.zero
    (λ n' ih => myAdd m ih)
    n

def factorial (n : ℕ) : ℕ :=
  ℕ.rec (motive := λ _ => ℕ)
    (ℕ.succ ℕ.zero) -- 0! = 1
    (λ n' ih => myMul (ℕ.succ n') ih)
    n

def toNat (n : ℕ) : Nat :=
  ℕ.rec 0 (λ _ ih => ih + 1) n

#eval toNat $ myAdd (ℕ.succ $ ℕ.succ $ ℕ.zero) (ℕ.succ $ ℕ.zero)
#eval toNat $ myMul ℕ.zero.succ.succ ℕ.zero.succ.succ.succ
#eval toNat $ factorial ℕ.zero.succ.succ.succ.succ  -- 4! = 24

-- list α

inductive lstS (α:Type) : Type where
  | nil_s | cons_s : α → lstS α

def Lst α := W (lstS α) (λ s => match s with | .nil_s => Empty | .cons_s _ => α)

def Lst.nil : Lst α := W.mk lstS.nil_s Empty.elim

def Lst.cons : α → Lst α → Lst α := λ a n =>
  W.mk (lstS.cons_s a) (λ _ => n)

#check Lst.cons 5 $ Lst.cons 7 $ Lst.nil

-- tree

inductive TreeS where
  | nil_s | node_s

def Tree := W TreeS (λ s => match s with | .nil_s => Empty | .node_s => Bool)

def Tree.nil : Tree := W.mk TreeS.nil_s Empty.elim

def Tree.node : Tree → Tree → Tree := λ l r =>
  W.mk TreeS.node_s (λ b => if b = true then l else r)

#check Tree.node Tree.nil $ Tree.node (Tree.node Tree.nil Tree.nil) Tree.nil

/-
inductive Tree' where
  | nil | node (left right : Tree') : Tree'
#check Tree'.rec
-/

def Tree.rec {motive : Tree → Sort u} (nil : motive Tree.nil)
  (node : (l r : Tree) → motive l → motive r → motive (Tree.node l r))
  (t : Tree) : motive t :=
  match t with
  | ⟨.nil_s, f⟩ => have : f = Empty.elim := funext (λ x => x.elim)
      this ▸ nil
  | ⟨.node_s, f⟩ =>
    let l := f true
    let r := f false
    have : f = (λ b => if b = true then l else r) :=
      funext (λ b => by cases b <;> rfl)
    this ▸ node l r (Tree.rec nil node l) (Tree.rec nil node r)

#print axioms Tree.rec