Add an AVL tree in the examples (#283)

* tests: add AVL example

Signed-off-by: Ryan Lahfa <[email protected]>

* tests: add Lean formalization for AVL

Signed-off-by: Ryan Lahfa <[email protected]>

* Fix minor extraction issues

* Start fixing the AVL tree example

* Commit some missing files

* Make a minor modification

* Update the code of the AVL tree

* Fix the code of the AVL tree

* Regenerate the code of the AVL tree

* Make progress on updating the proofs of the AVL trees

* Make progress on updating the proofs of the AVL tree

* Make minor modifications to the AVLs

* Move and update the code of the AVL tree

* Regenerate the Lean model for the AVL tree

* Generate simp lemmas for the custom field projectors

* Regenerate the types for the AVL test

* Make progress on the proof of the AVL tree

* Make progress on the proofs

* Make progress on the proofs of the AVL tree

* Make progress on the proofs of the AVL

* Cleanup a bit

* Cleanup a bit

* Do more cleanup

* Make progress on the AVL tree

* Make good progress on the avl

* Make minor modifications

* Fix pspec for mutually recursive functions

* Regenerate the test files

* Improve pspec and progress and cleanup a bit

* Add a comment

* Fix an issue with the CI and regenerate the tests

* Cleanup

* Do more cleanup

* Make progress generate warnings instead of errors if there are too many ids

* Fix a CI issue

* Fix a bug in destructure_abs

* Cleanup a bit

* Fix a minor issue in the flake.nix

---------

Signed-off-by: Ryan Lahfa <[email protected]>
Co-authored-by: Ryan Lahfa <[email protected]>

backends/lean/Base.lean

@@ -3,5 +3,6 @@ import Base.Diverge
 import Base.IList
 import Base.Primitives
 import Base.Progress
+import Base.SimpLemmas
 import Base.Utils
 import Base.Termination

backends/lean/Base/Arith/Int.lean

@@ -95,7 +95,7 @@ def intTacPreprocess (extraPreprocess :  Tactic.TacticM Unit) : Tactic.TacticM U
   -- might have proven the goal, hence the `Tactic.allGoals`
   let dsimp :=
     Tactic.allGoals do tryTac (
-      -- We set `simpOnly` at false on purpose
+      -- We set `simpOnly` at false on purpose.
       dsimpAt false {} intTacSimpRocs
         -- Declarations to unfold
         []
@@ -107,6 +107,17 @@ def intTacPreprocess (extraPreprocess :  Tactic.TacticM Unit) : Tactic.TacticM U
   Tactic.allGoals (Utils.tryTac (Utils.normCastAtAll))
   -- norm_cast does weird things with negative numbers so we reapply simp
   dsimp
+  -- Int.subNatNat is very annoying - TODO: there is probably something more general thing to do
+  Utils.tryTac (
+    Utils.simpAt true {}
+               -- Simprocs
+               []
+               -- Unfoldings
+               []
+                -- Simp lemmas
+                [``Int.subNatNat_eq_coe]
+                -- Hypotheses
+                [] .wildcard)
   -- We also need this, in case the goal is: ¬ False
   Tactic.allGoals do tryTac (
     Utils.simpAt true {}
@@ -178,4 +189,10 @@ example (a : Prop) (x : Int) (h0: (0 : Nat) < x) (h1: x < 0) : a := by
 example (x : Int) (h : x ≤ -3) : x ≤ -2 := by
   int_tac
 
+example (x y : Int) (h : x + y = 3) :
+  let z := x + y
+  z = 3 := by
+  intro z
+  omega
+
 end Arith

backends/lean/Base/Arith/Scalar.lean

@@ -89,4 +89,8 @@ example (x : I32) : -100000000000 < x.val := by
 example : (Usize.ofInt 2).val ≠ 0 := by
   scalar_tac
 
+example (x y : Nat) (z : Int) (h : Int.subNatNat x y + z = 0) : (x : Int) - (y : Int) + z = 0 := by
+  scalar_tac_preprocess
+  omega
+
 end Arith

backends/lean/Base/Diverge/Elab.lean

@@ -661,12 +661,10 @@ mutual
  -/
 partial def proveExprIsValid (k_var kk_var : Expr) (e : Expr) : MetaM Expr := do
   trace[Diverge.def.valid] "proveExprIsValid: {e}"
-  -- Normalize to eliminate the lambdas - TODO: this is slightly dangerous.
+  -- Normalize to eliminate the let-bindings
   let e ← do
     if e.isLet ∧ normalize_let_bindings then do
-      let updt_config (config : Lean.Meta.Config) :=
-        { config with transparency := .reducible }
-      let e ← withConfig updt_config (whnf e)
+      let e ← normalizeLetBindings e
       trace[Diverge.def.valid] "e (after normalization): {e}"
       pure e
     else pure e

backends/lean/Base/Primitives/Base.lean

@@ -76,6 +76,7 @@ def eval_global {α: Type u} (x: Result α) (_: ok? x := by prove_eval_global) :
   | fail _ | div => by contradiction
   | ok x => x
 
+@[simp]
 def Result.ofOption {a : Type u} (x : Option a) (e : Error) : Result a :=
   match x with
   | some x => ok x

backends/lean/Base/Primitives/Core.lean

@@ -51,17 +51,6 @@ def clone.CloneBool : clone.Clone Bool := {
   clone := fun b => ok (clone.impls.CloneBool.clone b)
 }
 
-namespace option -- core.option
-
-/- [core::option::{core::option::Option<T>}::unwrap] -/
-def Option.unwrap (T : Type) (x : Option T) : Result T :=
-  Result.ofOption x Error.panic
-
-end option -- core.option
-
-/- [core::option::Option::take] -/
-@[simp] def Option.take (T: Type) (self: Option T): Option T × Option T := (self, .none)
-
 /- [core::mem::replace]
 
    This acts like a swap effectively in a functional pure world.
@@ -73,3 +62,13 @@ end option -- core.option
 @[simp] def mem.swap (T: Type) (a b: T): T × T := (b, a)
 
 end core
+
+/- [core::option::{core::option::Option<T>}::unwrap] -/
+@[simp] def core.option.Option.unwrap (T : Type) (x : Option T) : Result T :=
+  Result.ofOption x Error.panic
+
+/- [core::option::Option::take] -/
+@[simp] def core.option.Option.take (T: Type) (self: Option T): Option T × Option T := (self, .none)
+
+/- [core::option::Option::is_none] -/
+@[simp] def core.option.Option.is_none (T: Type) (self: Option T): Bool := self.isNone

backends/lean/Base/Progress/Base.lean

@@ -81,7 +81,7 @@ section Methods
     trace[Progress] "After splitting the conjunction:\n- eq: {th}\n- post: {post}"
     -- Destruct the equality
     let (mExpr, ret) ← destEq th.consumeMData
-    trace[Progress] "After splitting the equality:\n- lhs: {th}\n- rhs: {ret}"
+    trace[Progress] "After splitting the equality:\n- lhs: {mExpr}\n- rhs: {ret}"
     -- Recursively destruct the monadic application to dive into the binds,
     -- if necessary (this is for when we use `withPSpec` inside of the `progress` tactic),
     -- and destruct the application to get the function name
@@ -154,18 +154,36 @@ initialize pspecAttr : PSpecAttr ← do
         throwError "invalid attribute 'pspec', must be global"
       -- Lookup the theorem
       let env ← getEnv
-      let thDecl := env.constants.find! thName
-      let fKey ← MetaM.run' (do
-        let fExpr ← getPSpecFunArgsExpr false thDecl.type
-        trace[Progress] "Registering spec theorem for {fExpr}"
-        -- Convert the function expression to a discrimination tree key
-        -- We use the default configuration
-        let config : WhnfCoreConfig := {}
-        DiscrTree.mkPath fExpr config)
-      let env := ext.addEntry env (fKey, thName)
-      setEnv env
-      trace[Progress] "Saved the environment"
-      pure ()
+      -- If we apply the attribute to a definition in a group of mutually recursive definitions
+      -- (say, to `foo` in the group [`foo`, `bar`]), the attribute gets applied to `foo` but also to
+      -- the recursive definition which encodes `foo` and `bar` (Lean encodes mutually recursive
+      -- definitions in one recursive definition, e.g., `foo._mutual`, before deriving the individual
+      -- definitions, e.g., `foo` and `bar`, from this one). This definition should be named `foo._mutual`
+      -- or `bar._mutual`, and we want to ignore it.
+      -- TODO: this is a hack
+      if let .str _ "_mutual" := thName then
+        -- Ignore: this is the fixed point of a mutually recursive definition -
+        -- the attribute will also be applied to the definitions revealed to the user:
+        -- we want to apply the attribute for those
+        trace[Progress] "Ignoring a mutually recursive definition: {thName}"
+      else
+        trace[Progress] "Registering spec theorem for {thName}"
+        let thDecl := env.constants.find! thName
+        let fKey ← MetaM.run' (do
+          trace[Progress] "Theorem: {thDecl.type}"
+          -- Normalize to eliminate the let-bindings
+          let ty ← normalizeLetBindings thDecl.type
+          trace[Progress] "Theorem after normalization (to eliminate the let bindings): {ty}"
+          let fExpr ← getPSpecFunArgsExpr false ty
+          trace[Progress] "Registering spec theorem for expr: {fExpr}"
+          -- Convert the function expression to a discrimination tree key
+          -- We use the default configuration
+          let config : WhnfCoreConfig := {}
+          DiscrTree.mkPath fExpr config)
+        let env := ext.addEntry env (fKey, thName)
+        setEnv env
+        trace[Progress] "Saved the environment"
+        pure ()
   }
   registerBuiltinAttribute attrImpl
   pure { attr := attrImpl, ext := ext }

backends/lean/Base/Progress/Progress.lean

@@ -73,6 +73,9 @@ def progressWith (fExpr : Expr) (th : TheoremOrLocal)
       inferType th
     | .Local asmDecl => pure asmDecl.type
   trace[Progress] "Looked up theorem/assumption type: {thTy}"
+  -- Normalize to inline the let-bindings
+  let thTy ← normalizeLetBindings thTy
+  trace[Progress] "After normalizing the let-bindings: {thTy}"
   -- TODO: the tactic fails if we uncomment withNewMCtxDepth
   -- withNewMCtxDepth do
   let (mvars, binders, thExBody) ← forallMetaTelescope thTy
@@ -104,6 +107,13 @@ def progressWith (fExpr : Expr) (th : TheoremOrLocal)
     | .Local decl => mkAppOptM' (mkFVar decl.fvarId) (mvars.map some)
   let asmName ← do match keep with | none => mkFreshAnonPropUserName | some n => do pure n
   let thTy ← inferType th
+  trace[Progress] "thTy (after application): {thTy}"
+  -- Normalize the let-bindings (note that we already inlined the let bindings once above when analizing
+  -- the theorem, now we do it again on the instantiated theorem - there is probably a smarter way to do,
+  -- but it doesn't really matter).
+  -- TODO: actually we might want to let the user insert them in the context
+  let thTy ← normalizeLetBindings thTy
+  trace[Progress] "thTy (after normalizing let-bindings): {thTy}"
   let thAsm ← Utils.addDeclTac asmName th thTy (asLet := false)
   withMainContext do -- The context changed - TODO: remove once addDeclTac is updated
   let ngoal ← getMainGoal
@@ -138,7 +148,11 @@ def progressWith (fExpr : Expr) (th : TheoremOrLocal)
     let _ ←
       tryTac
         (simpAt true {} [] []
-               [``Primitives.bind_tc_ok, ``Primitives.bind_tc_fail, ``Primitives.bind_tc_div]
+               [``Primitives.bind_tc_ok, ``Primitives.bind_tc_fail, ``Primitives.bind_tc_div,
+                -- This last one is quite useful. In particular, it is necessary to rewrite the
+                -- conjunctions for Lean to automatically instantiate the existential quantifiers
+                -- (I don't know why).
+                ``and_assoc]
                [hEq.fvarId!] (.targets #[] true))
     -- It may happen that at this point the goal is already solved (though this is rare)
     -- TODO: not sure this is the best way of checking it
@@ -161,8 +175,8 @@ def progressWith (fExpr : Expr) (th : TheoremOrLocal)
        | none =>
          -- Sanity check
          if ¬ ids.isEmpty then
-           return (.Error m!"Too many ids provided ({ids}): there is no postcondition to split")
-         else return .Ok
+          logWarning m!"Too many ids provided ({ids}): there is no postcondition to split"
+         return .Ok
        | some hPost => do
          let rec splitPostWithIds (prevId : Name) (hPost : Expr) (ids0 : List (Option Name)) : TacticM ProgressError := do
            match ids0 with
@@ -183,7 +197,9 @@ def progressWith (fExpr : Expr) (th : TheoremOrLocal)
              trace[Progress] "\n- prevId: {prevId}\n- nid: {nid}\n- remaining ids: {ids}"
              if ← isConj (← inferType hPost) then
                splitConjTac hPost (some (prevId, nid)) (λ _ nhPost => splitPostWithIds nid nhPost ids)
-             else return (.Error m!"Too many ids provided ({ids0}) not enough conjuncts to split in the postcondition")
+             else
+              logWarning m!"Too many ids provided ({ids0}) not enough conjuncts to split in the postcondition"
+              pure .Ok
          let curPostId := (← hPost.fvarId!.getDecl).userName
          splitPostWithIds curPostId hPost ids
   match res with
@@ -459,7 +475,6 @@ namespace Test
     progress
     simp [*]
 
-
   example {α : Type} (v: Vec α) (i: Usize) (x : α)
     (hbounds : i.val < v.length) :
     ∃ nv, v.update_usize i x = ok nv ∧
@@ -505,6 +520,73 @@ namespace Test
     progress keep _ as ⟨ z, h1 .. ⟩
     simp [*, h1]
 
+  -- Testing with mutually recursive definitions
+  mutual
+    inductive Tree
+    | mk : Trees → Tree
+
+    inductive Trees
+    | nil
+    | cons : Tree → Trees → Trees
+  end
+
+  mutual
+    def Tree.size (t : Tree) : Result Int :=
+      match t with
+      | .mk trees => trees.size
+
+    def Trees.size (t : Trees) : Result Int :=
+      match t with
+      | .nil => ok 0
+      | .cons t t' => do
+        let s ← t.size
+        let s' ← t'.size
+        ok (s + s')
+  end
+
+  mutual
+    @[pspec]
+    theorem Tree.size_spec (t : Tree) :
+      ∃ i, t.size = ok i ∧ i ≥ 0 := by
+      cases t
+      simp [Tree.size]
+      progress
+      simp
+      omega
+
+    @[pspec]
+    theorem Trees.size_spec (t : Trees) :
+      ∃ i, t.size = ok i ∧ i ≥ 0 := by
+      cases t <;> simp [Trees.size]
+      progress
+      progress
+      simp
+      omega
+  end
+
+  -- Testing progress on theorems containing local let-bindings
+  def add (x y : U32) : Result U32 := x + y
+
+  @[pspec] -- TODO: give the possibility of using pspec as a local attribute
+  theorem add_spec x y (h : x.val + y.val ≤ U32.max) :
+    let tot := x.val + y.val
+    ∃ z, add x y = ok z ∧ z.val = tot := by
+    rw [add]
+    intro tot
+    progress
+    simp [*]
+
+  def add1 (x y : U32) : Result U32 := do
+    let z ← add x y
+    add z z
+
+  example (x y : U32) (h : 2 * x.val + 2 * y.val ≤ U32.max) :
+    ∃ z, add1 x y = ok z := by
+    rw [add1]
+    progress as ⟨ z1, h ⟩
+    progress as ⟨ z2, h ⟩
+    simp
+
 end Test
 
 end Progress

backends/lean/Base/SimpLemmas.lean

@@ -0,0 +1,8 @@
+import Lean
+
+-- This simplification lemma is very useful especially for the kind of theorems we prove,
+-- which are of the shape: `∃ x y ..., ... ∧ ... ∧ ...`.
+-- Using this simp lemma often triggers simplifications like the instantiation of the
+-- existential quantifiers when there is an equality somewhere:
+-- `∃ x, x = y ∧ P x` gets rewritten to `P y`
+attribute [simp] and_assoc

backends/lean/Base/Termination.lean

@@ -91,4 +91,7 @@ macro_rules
       -- Finish
       simp_all <;> scalar_tac)
 
+-- We always activate this simplification lemma because it is useful for the proofs of termination
+attribute [simp] Prod.lex_iff
+
 end Utils

backends/lean/Base/Utils.lean

@@ -807,4 +807,8 @@ def evalAesopSaturate (options : Aesop.Options') (ruleSets : Array Name) : Tacti
     |> Aesop.ElabM.runForwardElab (← getMainGoal)
   tryLiftMetaTactic1 (Aesop.saturate rs · |>.run { options }) "Aesop.saturate failed"
 
+/-- Normalize the let-bindings by inlining them -/
+def normalizeLetBindings (e : Expr) : MetaM Expr :=
+  zetaReduce e
+
 end Utils

compiler/ExtractBuiltin.ml

@@ -497,17 +497,14 @@ let mk_builtin_funs () : (pattern * bool list option * builtin_fun_info) list =
   (* Lean-only definitions *)
   @ mk_lean_only
       [
-        (* `backend_choice` first parameter is for non-Lean backends
-            By construction, we cannot write down that parameter in the output
-            in this list
-        *)
         mk_fun "core::mem::swap" ~can_fail:false ();
         mk_fun "core::option::{core::option::Option<@T>}::take"
-          ~extract_name:(Some (backend_choice "" "Option::take"))
+          ~extract_name:(Some (backend_choice "" "core::option::Option::take"))
           ~can_fail:false ();
         mk_fun "core::option::{core::option::Option<@T>}::is_none"
-          ~extract_name:(Some (backend_choice "" "Option::isNone"))
-          ~filter:(Some [ false ]) ~can_fail:false ();
+          ~extract_name:
+            (Some (backend_choice "" "core::option::Option::is_none"))
+          ~can_fail:false ();
       ]
 
 let builtin_funs : unit -> (pattern * bool list option * builtin_fun_info) list