Improve the Lean backend

2 files changed, 25 insertions, 8 deletions

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 0ad16ab6..935af3f5 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -82,17 +82,21 @@ section Methods
     -- Destruct the equality
     let (mExpr, ret) ← destEq th.consumeMData
     trace[Progress] "After splitting the equality:\n- lhs: {th}\n- rhs: {ret}"
-    -- Destruct the monadic application to dive into the bind, if necessary (this
-    -- is for when we use `withPSpec` inside of the `progress` tactic), and
-    -- destruct the application to get the function name
-    mExpr.consumeMData.withApp fun mf margs => do
-    trace[Progress] "After stripping the arguments of the monad expression:\n- mf: {mf}\n- margs: {margs}"
-    let (fArgsExpr, f, args) ← do
+    -- 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
+    let rec strip_monad mExpr := do
+      mExpr.consumeMData.withApp fun mf margs => do
+      trace[Progress] "After stripping the arguments of the monad expression:\n- mf: {mf}\n- margs: {margs}"
       if mf.isConst ∧ mf.constName = ``Bind.bind then do
         -- Dive into the bind
         let fExpr := (margs.get! 4).consumeMData
-        fExpr.withApp fun f args => pure (fExpr, f, args)
-      else pure (mExpr, mf, margs)
+        -- Recursve
+        strip_monad fExpr
+      else
+        -- No bind
+        pure (mExpr, mf, margs)
+    let (fArgsExpr, f, args) ← strip_monad mExpr
     trace[Progress] "After stripping the arguments of the function call:\n- f: {f}\n- args: {args}"
     if ¬ f.isConst then throwError "Not a constant: {f}"
     -- *Sanity check* (activated if we are analyzing a theorem to register it in a DB)
diff --git a/backends/lean/Base/Progress/Progress.lean b/backends/lean/Base/Progress/Progress.lean
index a6a4e82a..0fb276aa 100644
--- a/backends/lean/Base/Progress/Progress.lean
+++ b/backends/lean/Base/Progress/Progress.lean
@@ -421,6 +421,19 @@ namespace Test
     progress
     simp [*]
 
+  /- Checking that progress can handle nested blocks -/
+  example {α : Type} (v: Vec α) (i: Usize) (x : α)
+    (hbounds : i.val < v.length) :
+    ∃ nv,
+      (do
+         (do
+            let _ ← v.update_usize i x
+            .ret ())
+         .ret ()) = ret nv
+      := by
+    progress
+    simp [*]
+
 end Test
 
 end Progress