Start working on a version of Diverge.FixI more suited to higher-order functions

3 files changed, 510 insertions, 67 deletions

diff --git a/backends/lean/Base/Diverge/Base.lean b/backends/lean/Base/Diverge/Base.lean
index 6a52387d..9d986f4e 100644
--- a/backends/lean/Base/Diverge/Base.lean
+++ b/backends/lean/Base/Diverge/Base.lean
@@ -12,6 +12,78 @@ set_option profiler.threshold 100
 
 namespace Diverge
 
+/- Auxiliary lemmas -/
+namespace Lemmas
+  -- TODO: not necessary anymore
+  def for_all_fin_aux {n : Nat} (f : Fin n → Prop) (m : Nat) (h : m ≤ n) : Prop :=
+    if heq: m = n then True
+    else
+      f ⟨ m, by simp_all [Nat.lt_iff_le_and_ne] ⟩ ∧
+      for_all_fin_aux f (m + 1) (by simp_all [Arith.add_one_le_iff_le_ne])
+  termination_by for_all_fin_aux n _ m h => n - m
+  decreasing_by
+    simp_wf
+    apply Nat.sub_add_lt_sub <;> try simp
+    simp_all [Arith.add_one_le_iff_le_ne]
+
+  def for_all_fin {n : Nat} (f : Fin n → Prop) := for_all_fin_aux f 0 (by simp)
+
+  theorem for_all_fin_aux_imp_forall {n : Nat} (f : Fin n → Prop) (m : Nat) :
+    (h : m ≤ n) →
+    for_all_fin_aux f m h → ∀ i, m ≤ i.val → f i
+    := by
+    generalize h: (n - m) = k
+    revert m
+    induction k -- TODO: induction h rather?
+    case zero =>
+      simp_all
+      intro m h1 h2
+      have h: n = m := by
+        linarith
+      unfold for_all_fin_aux; simp_all
+      simp_all
+      -- There is no i s.t. m ≤ i
+      intro i h3; cases i; simp_all
+      linarith
+    case succ k hi =>
+      intro m hk hmn
+      intro hf i hmi
+      have hne: m ≠ n := by
+        have hineq := Nat.lt_of_sub_eq_succ hk
+        linarith
+      -- m = i?
+      if heq: m = i then
+        -- Yes: simply use the `for_all_fin_aux` hyp
+        unfold for_all_fin_aux at hf
+        simp_all
+      else
+        -- No: use the induction hypothesis
+        have hlt: m < i := by simp_all [Nat.lt_iff_le_and_ne]
+        have hineq: m + 1 ≤ n := by
+          have hineq := Nat.lt_of_sub_eq_succ hk
+          simp [*, Nat.add_one_le_iff]
+        have heq1: n - (m + 1) = k := by
+          -- TODO: very annoying arithmetic proof
+          simp [Nat.sub_eq_iff_eq_add hineq]
+          have hineq1: m ≤ n := by linarith
+          simp [Nat.sub_eq_iff_eq_add hineq1] at hk
+          simp_arith [hk]
+        have hi := hi (m + 1) heq1 hineq
+        apply hi <;> simp_all
+        . unfold for_all_fin_aux at hf
+          simp_all
+        . simp_all [Arith.add_one_le_iff_le_ne]
+
+  -- TODO: this is not necessary anymore
+  theorem for_all_fin_imp_forall (n : Nat) (f : Fin n → Prop) :
+    for_all_fin f → ∀ i, f i
+    := by
+    intro Hf i
+    apply for_all_fin_aux_imp_forall <;> try assumption
+    simp
+
+end Lemmas
+
 namespace Fix
 
   open Primitives
@@ -436,6 +508,10 @@ namespace FixI
   /- Indexed fixed-point: definitions with indexed types, convenient to use for mutually
      recursive definitions. We simply port the definitions and proofs from Fix to a more
      specific case.
+
+     Remark: the index designates the function in the mutually recursive group
+     (it should be a finite type). We make the return type depend on the input
+     type because we group the type parameters in the input type.
    -/
   open Primitives Fix
 
@@ -538,73 +614,6 @@ namespace FixI
         let j: Fin tys1.length := ⟨ j, jLt ⟩
         Eq.mp (by simp) (get_fun tl j)
 
-  def for_all_fin_aux {n : Nat} (f : Fin n → Prop) (m : Nat) (h : m ≤ n) : Prop :=
-    if heq: m = n then True
-    else
-      f ⟨ m, by simp_all [Nat.lt_iff_le_and_ne] ⟩ ∧
-      for_all_fin_aux f (m + 1) (by simp_all [Arith.add_one_le_iff_le_ne])
-  termination_by for_all_fin_aux n _ m h => n - m
-  decreasing_by
-    simp_wf
-    apply Nat.sub_add_lt_sub <;> try simp
-    simp_all [Arith.add_one_le_iff_le_ne]
-
-  def for_all_fin {n : Nat} (f : Fin n → Prop) := for_all_fin_aux f 0 (by simp)
-
-  theorem for_all_fin_aux_imp_forall {n : Nat} (f : Fin n → Prop) (m : Nat) :
-    (h : m ≤ n) →
-    for_all_fin_aux f m h → ∀ i, m ≤ i.val → f i
-    := by
-    generalize h: (n - m) = k
-    revert m
-    induction k -- TODO: induction h rather?
-    case zero =>
-      simp_all
-      intro m h1 h2
-      have h: n = m := by
-        linarith
-      unfold for_all_fin_aux; simp_all
-      simp_all
-      -- There is no i s.t. m ≤ i
-      intro i h3; cases i; simp_all
-      linarith
-    case succ k hi =>
-      intro m hk hmn
-      intro hf i hmi
-      have hne: m ≠ n := by
-        have hineq := Nat.lt_of_sub_eq_succ hk
-        linarith
-      -- m = i?
-      if heq: m = i then
-        -- Yes: simply use the `for_all_fin_aux` hyp
-        unfold for_all_fin_aux at hf
-        simp_all
-      else
-        -- No: use the induction hypothesis
-        have hlt: m < i := by simp_all [Nat.lt_iff_le_and_ne]
-        have hineq: m + 1 ≤ n := by
-          have hineq := Nat.lt_of_sub_eq_succ hk
-          simp [*, Nat.add_one_le_iff]
-        have heq1: n - (m + 1) = k := by
-          -- TODO: very annoying arithmetic proof
-          simp [Nat.sub_eq_iff_eq_add hineq]
-          have hineq1: m ≤ n := by linarith
-          simp [Nat.sub_eq_iff_eq_add hineq1] at hk
-          simp_arith [hk]
-        have hi := hi (m + 1) heq1 hineq
-        apply hi <;> simp_all
-        . unfold for_all_fin_aux at hf
-          simp_all
-        . simp_all [Arith.add_one_le_iff_le_ne]
-
-  -- TODO: this is not necessary anymore
-  theorem for_all_fin_imp_forall (n : Nat) (f : Fin n → Prop) :
-    for_all_fin f → ∀ i, f i
-    := by
-    intro Hf i
-    apply for_all_fin_aux_imp_forall <;> try assumption
-    simp
-
   /- Automating the proofs -/
   @[simp] theorem is_valid_p_same
     (k : ((i:id) → (x:a i) → Result (b i x)) → (i:id) → (x:a i) → Result (b i x)) (x : Result c) :
@@ -707,6 +716,222 @@ namespace FixI
 
 end FixI
 
+namespace FixII
+  /- Indexed fixed-point: definitions with indexed types, convenient to use for mutually
+     recursive definitions. We simply port the definitions and proofs from Fix to a more
+     specific case.
+
+     Here however, we group the types into a parameter distinct from the inputs.
+   -/
+  open Primitives Fix
+
+  -- The index type
+  variable {id : Type u}
+
+  -- The input/output types
+  variable {ty : id → Type v} {a : (i:id) → ty i → Type w} {b : (i:id) → ty i → Type x}
+
+  -- Monotonicity relation over monadic arrows (i.e., Kleisli arrows)
+  def karrow_rel (k1 k2 : (i:id) → (t:ty i) → (a i t) → Result (b i t)) : Prop :=
+    ∀ i t x, result_rel (k1 i t x) (k2 i t x)
+
+  def kk_to_gen (k : (i:id) → (t:ty i) → (x:a i t) → Result (b i t)) :
+    (x: (i:id) × (t:ty i) × (a i t)) → Result (b x.fst x.snd.fst) :=
+    λ ⟨ i, t, x ⟩ => k i t x
+
+  def kk_of_gen (k : (x: (i:id) × (t:ty i) × (a i t)) → Result (b x.fst x.snd.fst)) :
+    (i:id) → (t:ty i) → a i t → Result (b i t) :=
+    λ i t x => k ⟨ i, t, x ⟩
+
+  def k_to_gen (k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t)) :
+    ((x: (i:id) × (t:ty i) × (a i t)) → Result (b x.fst x.snd.fst)) → (x: (i:id) × (t:ty i) × (a i t)) → Result (b x.fst x.snd.fst) :=
+    λ kk => kk_to_gen (k (kk_of_gen kk))
+
+  def k_of_gen (k : ((x: (i:id) × (t:ty i) × (a i t)) → Result (b x.fst x.snd.fst)) → (x: (i:id) × (t:ty i) × (a i t)) → Result (b x.fst x.snd.fst)) :
+    ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t) :=
+    λ kk => kk_of_gen (k (kk_to_gen kk))
+
+  def e_to_gen (e : ((i:id) → (t:ty i) → a i t → Result (b i t)) → Result c) :
+    ((x: (i:id) × (t:ty i) × (a i t)) → Result (b x.fst x.snd.fst)) → Result c :=
+    λ k => e (kk_of_gen k)
+
+  def is_valid_p (k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t))
+    (e : ((i:id) → (t:ty i) → a i t → Result (b i t)) → Result c) : Prop :=
+    Fix.is_valid_p (k_to_gen k) (e_to_gen e)
+
+  def is_valid (f : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t)) : Prop :=
+    ∀ k i t x, is_valid_p k (λ k => f k i t x)
+
+  def fix
+    (f : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t)) :
+    (i:id) → (t:ty i) → a i t → Result (b i t) :=
+    kk_of_gen (Fix.fix (k_to_gen f))
+
+  theorem is_valid_fix_fixed_eq
+    {{f : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t)}}
+    (Hvalid : is_valid f) :
+    fix f = f (fix f) := by
+    have Hvalid' : Fix.is_valid (k_to_gen f) := by
+      intro k x
+      simp only [is_valid, is_valid_p] at Hvalid
+      let ⟨ i, t, x ⟩ := x
+      have Hvalid := Hvalid (k_of_gen k) i t x
+      simp only [k_to_gen, k_of_gen, kk_to_gen, kk_of_gen] at Hvalid
+      refine Hvalid
+    have Heq := Fix.is_valid_fix_fixed_eq Hvalid'
+    simp [fix]
+    conv => lhs; rw [Heq]
+
+  /- Some utilities to define the mutually recursive functions -/
+
+  -- TODO: use more
+  abbrev kk_ty (id : Type u) (ty : id → Type v) (a : (i:id) → ty i → Type w) (b : (i:id) → ty i → Type x) :=
+    (i:id) → (t:ty i) → a i t → Result (b i t)
+  abbrev k_ty (id : Type u) (ty : id → Type v) (a : (i:id) → ty i → Type w) (b : (i:id) → ty i → Type x) :=
+    kk_ty id ty a b → kk_ty id ty a b
+
+  abbrev in_out_ty : Type (imax (u + 1) (imax (v + 1) (w + 1))) :=
+    (ty : Type u) × (ty → Type v) × (ty → Type w)
+  -- TODO: remove?
+  abbrev mk_in_out_ty (ty : Type u) (in_ty : ty → Type v) (out_ty : ty → Type w) :
+    in_out_ty :=
+    Sigma.mk ty (Prod.mk in_ty out_ty)
+
+  -- Initially, we had left out the parameters id, a and b.
+  -- However, by parameterizing Funs with those parameters, we can state
+  -- and prove lemmas like Funs.is_valid_p_is_valid_p
+  inductive Funs (id : Type u) (ty : id → Type v)
+    (a : (i:id) → ty i → Type w) (b : (i:id) → ty i → Type x) :
+    List in_out_ty.{v, w, x} → Type (max (u + 1) (max (v + 1) (max (w + 1) (x + 1)))) :=
+    | Nil : Funs id ty a b []
+    | Cons {it: Type v} {ity : it → Type w} {oty : it → Type x} {tys : List in_out_ty}
+           (f : kk_ty id ty a b → (i:it) → (x:ity i) → Result (oty i)) (tl : Funs id ty a b tys) :
+           Funs id ty a b (⟨ it, ity, oty ⟩ :: tys)
+
+  def get_fun {tys : List in_out_ty} (fl : Funs id ty a b tys) :
+    (i : Fin tys.length) → kk_ty id ty a b → (t : (tys.get i).fst) →
+    ((tys.get i).snd.fst t) → Result ((tys.get i).snd.snd t) :=
+    match fl with
+    | .Nil => λ i => by have h:= i.isLt; simp at h
+    | @Funs.Cons id ty a b it ity oty tys1 f tl =>
+      λ ⟨ i, iLt ⟩ =>
+      match i with
+      | 0 =>
+        Eq.mp (by simp [List.get]) f
+      | .succ j =>
+        have jLt: j < tys1.length := by
+          simp at iLt
+          revert iLt
+          simp_arith
+        let j: Fin tys1.length := ⟨ j, jLt ⟩
+        Eq.mp (by simp) (get_fun tl j)
+
+  /- Automating the proofs -/
+  @[simp] theorem is_valid_p_same
+    (k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t)) (x : Result c) :
+    is_valid_p k (λ _ => x) := by
+    simp [is_valid_p, k_to_gen, e_to_gen]
+
+  @[simp] theorem is_valid_p_rec
+     (k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t)) (i : id) (t : ty i) (x : a i t) :
+    is_valid_p k (λ k => k i t x) := by
+    simp [is_valid_p, k_to_gen, e_to_gen, kk_to_gen, kk_of_gen]
+
+  theorem is_valid_p_ite
+    (k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t))
+    (cond : Prop) [h : Decidable cond]
+    {e1 e2 : ((i:id) → (t:ty i) → a i t → Result (b i t)) → Result c}
+    (he1: is_valid_p k e1) (he2 : is_valid_p k e2) :
+    is_valid_p k (λ k => ite cond (e1 k) (e2 k)) := by
+    split <;> assumption
+
+  theorem is_valid_p_dite
+    (k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t))
+    (cond : Prop) [h : Decidable cond]
+    {e1 : ((i:id) → (t:ty i) → a i t → Result (b i t)) → cond → Result c}
+    {e2 : ((i:id) → (t:ty i) → a i t → Result (b i t)) → Not cond → Result c}
+    (he1: ∀ x, is_valid_p k (λ k => e1 k x))
+    (he2 : ∀ x, is_valid_p k (λ k => e2 k x)) :
+    is_valid_p k (λ k => dite cond (e1 k) (e2 k)) := by
+    split <;> simp [*]
+
+  theorem is_valid_p_bind
+    {{k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t)}}
+    {{g : ((i:id) → (t:ty i) → a i t → Result (b i t)) → Result c}}
+    {{h : c → ((i:id) → (t:ty i) → a i t → Result (b i t)) → Result d}}
+    (Hgvalid : is_valid_p k g)
+    (Hhvalid : ∀ y, is_valid_p k (h y)) :
+    is_valid_p k (λ k => do let y ← g k; h y k) := by
+    apply Fix.is_valid_p_bind
+    . apply Hgvalid
+    . apply Hhvalid
+
+  def Funs.is_valid_p
+    (k : k_ty id ty a b)
+    (fl : Funs id ty a b tys) :
+    Prop :=
+    match fl with
+    | .Nil => True
+    | .Cons f fl =>
+      (∀ i x, FixII.is_valid_p k (λ k => f k i x)) ∧ fl.is_valid_p k
+
+  theorem Funs.is_valid_p_Nil (k : k_ty id ty a b) :
+    Funs.is_valid_p k Funs.Nil := by simp [Funs.is_valid_p]
+
+  def Funs.is_valid_p_is_valid_p_aux
+    {k : k_ty id ty a b}
+    {tys : List in_out_ty}
+    (fl : Funs id ty a b tys) (Hvalid : is_valid_p k fl) :
+    ∀ (i : Fin tys.length) (t : (tys.get i).fst) (x : (tys.get i).snd.fst t),
+    FixII.is_valid_p k (fun k => get_fun fl i k t x) := by
+    -- Prepare the induction
+    have ⟨ n, Hn ⟩ : { n : Nat // tys.length = n } := ⟨ tys.length, by rfl ⟩
+    revert tys fl Hvalid
+    induction n
+    --
+    case zero =>
+      intro tys fl Hvalid Hlen;
+      have Heq: tys = [] := by cases tys <;> simp_all
+      intro i x
+      simp_all
+      have Hi := i.isLt
+      simp_all
+    case succ n Hn =>
+      intro tys fl Hvalid Hlen i x;
+      cases fl <;> simp at Hlen i x Hvalid
+      rename_i ity oty tys f fl
+      have ⟨ Hvf, Hvalid ⟩ := Hvalid
+      have Hvf1: is_valid_p k fl := by
+        simp [Hvalid, Funs.is_valid_p]
+      have Hn := @Hn tys fl Hvf1 (by simp [*])
+      -- Case disjunction on i
+      match i with
+      | ⟨ 0, _ ⟩ =>
+        simp at x
+        simp [get_fun]
+        apply (Hvf x)
+      | ⟨ .succ j, HiLt ⟩ =>
+        simp_arith at HiLt
+        simp at x
+        let j : Fin (List.length tys) := ⟨ j, by simp_arith [HiLt] ⟩
+        have Hn := Hn j x
+        apply Hn
+
+  def Funs.is_valid_p_is_valid_p
+    (tys : List in_out_ty)
+    (k : k_ty (Fin (List.length tys)) (λ i => (tys.get i).fst)
+              (fun i t => (List.get tys i).snd.fst t) (fun i t => (List.get tys i).snd.snd t))
+    (fl : Funs (Fin tys.length) (λ i => (tys.get i).fst)
+               (λ i t => (tys.get i).snd.fst t) (λ i t => (tys.get i).snd.snd t) tys) :
+    fl.is_valid_p k  →
+    ∀ (i : Fin tys.length) (t : (tys.get i).fst) (x : (tys.get i).snd.fst t),
+      FixII.is_valid_p k (fun k => get_fun fl i k t x)
+    := by
+    intro Hvalid
+    apply is_valid_p_is_valid_p_aux; simp [*]
+
+end FixII
+
 namespace Ex1
   /- An example of use of the fixed-point -/
   open Primitives Fix
@@ -1133,3 +1358,214 @@ namespace Ex6
     Heqix
 
 end Ex6
+
+namespace Ex7
+  /- `list_nth` again, but this time we use FixII -/
+  open Primitives FixII
+
+  @[simp] def tys.{u} : List in_out_ty :=
+    [mk_in_out_ty (Type u) (λ a => List a × Int) (λ a => a)]
+
+  @[simp] def ty (i : Fin 1) := (tys.get i).fst
+  @[simp] def input_ty (i : Fin 1) (t : ty i) : Type u := (tys.get i).snd.fst t
+  @[simp] def output_ty (i : Fin 1) (t : ty i) : Type u := (tys.get i).snd.snd t
+
+  def list_nth_body.{u} (k : (i:Fin 1) → (t:ty i) →  input_ty i t → Result (output_ty i t))
+    (a : Type u) (x : List a × Int) : Result a :=
+    let ⟨ ls, i ⟩ := x
+    match ls with
+    | [] => .fail .panic
+    | hd :: tl =>
+      if i = 0 then .ret hd
+      else k 0 a ⟨ tl, i - 1 ⟩
+
+  @[simp] def bodies :
+    Funs (Fin 1) ty input_ty output_ty tys :=
+    Funs.Cons list_nth_body Funs.Nil
+
+  def body (k : (i : Fin 1) → (t : ty i) → (x : input_ty i t) → Result (output_ty i t)) (i: Fin 1) :
+    (t : ty i) → (x : input_ty i t) → Result (output_ty i t) := get_fun bodies i k
+
+  theorem body_is_valid: is_valid body := by
+    -- Split the proof into proofs of validity of the individual bodies
+    rw [is_valid]
+    simp only [body]
+    intro k
+    apply (Funs.is_valid_p_is_valid_p tys)
+    simp [Funs.is_valid_p]
+    (repeat (apply And.intro)); intro x; try simp at x
+    simp only [list_nth_body]
+    -- Prove the validity of the individual bodies
+    intro k x
+    split <;> try simp
+    split <;> simp
+
+  -- Writing the proof terms explicitly
+  theorem list_nth_body_is_valid' (k : k_ty (Fin 1) ty input_ty output_ty)
+    (a : Type u) (x : List a × Int) : is_valid_p k (fun k => list_nth_body k a x) :=
+    let ⟨ ls, i ⟩ := x
+    match ls with
+    | [] => is_valid_p_same k (.fail .panic)
+    | hd :: tl =>
+      is_valid_p_ite k (Eq i 0) (is_valid_p_same k (.ret hd)) (is_valid_p_rec k 0 a ⟨tl, i-1⟩)
+
+  theorem body_is_valid' : is_valid body :=
+    fun k =>
+    Funs.is_valid_p_is_valid_p tys k bodies
+      (And.intro (list_nth_body_is_valid' k) (Funs.is_valid_p_Nil k))
+
+  def list_nth {a: Type u} (ls : List a) (i : Int) : Result a :=
+    fix body 0 a ⟨ ls , i ⟩
+
+  -- The unfolding equation - diverges if `i < 0`
+  theorem list_nth_eq (ls : List a) (i : Int) :
+    list_nth ls i =
+      match ls with
+      | [] => .fail .panic
+      | hd :: tl =>
+        if i = 0 then .ret hd
+        else list_nth tl (i - 1)
+    := by
+    have Heq := is_valid_fix_fixed_eq body_is_valid
+    simp [list_nth]
+    conv => lhs; rw [Heq]
+
+  -- Write the proof term explicitly: the generation of the proof term (without tactics)
+  -- is automatable, and the proof term is actually a lot simpler and smaller when we
+  -- don't use tactics.
+  theorem list_nth_eq'.{u} {a : Type u} (ls : List a) (i : Int) :
+    list_nth ls i =
+      match ls with
+      | [] => .fail .panic
+      | hd :: tl =>
+        if i = 0 then .ret hd
+        else list_nth tl (i - 1)
+    :=
+    -- Use the fixed-point equation
+    have Heq := is_valid_fix_fixed_eq body_is_valid.{u}
+    -- Add the index
+    have Heqi := congr_fun Heq 0
+    -- Add the type parameter
+    have Heqia := congr_fun Heqi a
+    -- Add the input
+    have Heqix := congr_fun Heqia (ls, i)
+    -- Done
+    Heqix
+
+end Ex7
+
+namespace Ex8
+  /- Higher-order example, with FixII -/
+  open Primitives FixII
+
+  variable {id : Type u} {ty : id → Type v}
+  variable {a : (i:id) → ty i → Type w} {b : (i:id) → ty i → Type x}
+
+  /- An auxiliary function, which doesn't require the fixed-point -/
+  def map {a : Type y} {b : Type z} (f : a → Result b) (ls : List a) : Result (List b) :=
+    match ls with
+    | [] => .ret []
+    | hd :: tl =>
+      do
+        let hd ← f hd
+        let tl ← map f tl
+        .ret (hd :: tl)
+
+  /- The validity theorem for `map`, generic in `f` -/
+  theorem map_is_valid
+    (i : id) (t : ty i)
+    {{f : (a i t → Result (b i t)) → (a i t) → Result c}}
+    (Hfvalid : ∀ k x, is_valid_p k (λ k => f (k i t) x))
+    (k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t))
+    (ls : List (a i t)) :
+    is_valid_p k (λ k => map (f (k i t)) ls) := by
+    induction ls <;> simp [map]
+    apply is_valid_p_bind <;> try simp_all
+    intros
+    apply is_valid_p_bind <;> try simp_all
+
+  theorem map_is_valid'
+    (i : id) (t : ty i)
+    (k : ((i:id) → (t:ty i) → a i t → Result (b i t)) → (i:id) → (t:ty i) → a i t → Result (b i t))
+    (ls : List (a i t)) :
+    is_valid_p k (λ k => map (k i t) ls) := by
+    induction ls <;> simp [map]
+    apply is_valid_p_bind <;> try simp_all
+    intros
+    apply is_valid_p_bind <;> try simp_all
+
+end Ex8
+
+namespace Ex9
+  /- An example which uses map -/
+  open Primitives FixII Ex8
+
+  inductive Tree (a : Type u) :=
+  | leaf (x : a)
+  | node (tl : List (Tree a))
+
+  @[simp] def tys.{u} : List in_out_ty :=
+    [mk_in_out_ty (Type u) (λ a => Tree a) (λ a => Tree a)]
+
+  @[simp] def ty (i : Fin 1) := (tys.get i).fst
+  @[simp] def input_ty (i : Fin 1) (t : ty i) : Type u := (tys.get i).snd.fst t
+  @[simp] def output_ty (i : Fin 1) (t : ty i) : Type u := (tys.get i).snd.snd t
+
+  def id_body.{u} (k : (i:Fin 1) → (t:ty i) →  input_ty i t → Result (output_ty i t))
+    (a : Type u) (t : Tree a) : Result (Tree a) :=
+    match t with
+    | .leaf x => .ret (.leaf x)
+    | .node tl =>
+      do
+        let tl ← map (k 0 a) tl
+        .ret (.node tl)
+
+  @[simp] def bodies :
+    Funs (Fin 1) ty input_ty output_ty tys :=
+    Funs.Cons id_body Funs.Nil
+
+  theorem id_body_is_valid :
+    ∀ (k : ((i : Fin 1) → (t : ty i) → input_ty i t → Result (output_ty i t)) → (i : Fin 1) → (t : ty i) → input_ty i t → Result (output_ty i t))
+    (a : Type u) (x : Tree a),
+    @is_valid_p (Fin 1) ty input_ty output_ty (output_ty 0 a) k (λ k => id_body k a x) := by
+    intro k a x
+    simp only [id_body]
+    split <;> try simp
+    apply is_valid_p_bind <;> try simp [*]
+    -- We have to show that `map k tl` is valid
+    -- Remark: `map_is_valid` doesn't work here, we need the specialized version
+    apply map_is_valid'
+
+  def body (k : (i : Fin 1) → (t : ty i) → (x : input_ty i t) → Result (output_ty i t)) (i: Fin 1) :
+    (t : ty i) → (x : input_ty i t) → Result (output_ty i t) := get_fun bodies i k
+
+  theorem body_is_valid : is_valid body :=
+    fun k =>
+    Funs.is_valid_p_is_valid_p tys k bodies
+      (And.intro (id_body_is_valid k) (Funs.is_valid_p_Nil k))
+
+  def id {a: Type u} (t : Tree a) : Result (Tree a) :=
+    fix body 0 a t
+
+  -- Writing the proof term explicitly
+  theorem id_eq' {a : Type u} (t : Tree a) :
+    id t =
+      (match t with
+       | .leaf x => .ret (.leaf x)
+       | .node tl =>
+         do
+           let tl ← map id tl
+           .ret (.node tl))
+    :=
+    -- The unfolding equation
+    have Heq := is_valid_fix_fixed_eq body_is_valid.{u}
+    -- Add the index
+    have Heqi := congr_fun Heq 0
+    -- Add the type parameter
+    have Heqia := congr_fun Heqi a
+    -- Add the input
+    have Heqix := congr_fun Heqia t
+    -- Done
+    Heqix
+
+end Ex9
diff --git a/backends/lean/Base/Diverge/ElabBase.lean b/backends/lean/Base/Diverge/ElabBase.lean
index fedb1c74..1a0676e2 100644
--- a/backends/lean/Base/Diverge/ElabBase.lean
+++ b/backends/lean/Base/Diverge/ElabBase.lean
@@ -1,8 +1,11 @@
 import Lean
+import Base.Utils
+import Base.Primitives.Base
 
 namespace Diverge
 
 open Lean Elab Term Meta
+open Utils
 
 -- We can't define and use trace classes in the same file
 initialize registerTraceClass `Diverge.elab
@@ -12,4 +15,7 @@ initialize registerTraceClass `Diverge.def.genBody
 initialize registerTraceClass `Diverge.def.valid
 initialize registerTraceClass `Diverge.def.unfold
 
+-- For the attribute (for higher-order functions)
+initialize registerTraceClass `Diverge.attr
+
 end Diverge
diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index cf2a3b70..573f0cc5 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -154,6 +154,7 @@ structure PSpecAttr where
    specs of the scalar operations, which is what I really need, but I'm not sure it
    applies well to other situations. A better way would probably to use type classes, but
    I couldn't get them to work on those cases. It is worth retrying.
+   UPDATE: use discrimination trees (`DiscrTree`) from core Lean
 -/
 structure PSpecClassAttr where
   attr : AttributeImpl