2 files changed, 68 insertions, 3 deletions
diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 6f820a84..76a92795 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -167,7 +167,8 @@ structure PSpecClassExprAttr where
   deriving Inhabited
 
 -- TODO: the original function doesn't define correctly the `addImportedFn`. Do a PR?
-def mkMapDeclarationExtension [Inhabited α] (name : Name := by exact decl_name%) : IO (MapDeclarationExtension α) :=
+def mkMapDeclarationExtension [Inhabited α] (name : Name := by exact decl_name%) :
+  IO (MapDeclarationExtension α) :=
   registerSimplePersistentEnvExtension {
     name          := name,
     addImportedFn := fun a => a.foldl (fun s a => a.foldl (fun s (k, v) => s.insert k v) s) RBMap.empty,
@@ -175,6 +176,54 @@ def mkMapDeclarationExtension [Inhabited α] (name : Name := by exact decl_name%
     toArrayFn     := fun es => es.toArray.qsort (fun a b => Name.quickLt a.1 b.1)
   }
 
+-- Declare an extension of maps of maps (using [RBMap]).
+-- The important point is that we need to merge the bound values (which are maps).
+def mkMapMapDeclarationExtension [Inhabited β] (ord : α → α → Ordering)
+  (name : Name := by exact decl_name%) :
+  IO (MapDeclarationExtension (RBMap α β ord)) :=
+  registerSimplePersistentEnvExtension {
+    name          := name,
+    addImportedFn := fun a =>
+      a.foldl (fun s a => a.foldl (
+        -- We need to merge the maps
+        fun s (k0, k1_to_v) =>
+        match s.find? k0 with
+        | none =>
+          -- No binding: insert one
+          s.insert k0 k1_to_v
+        | some m =>
+          -- There is already a binding: merge
+          let m := RBMap.fold (fun m k v => m.insert k v) m k1_to_v
+          s.insert k0 m)
+          s) RBMap.empty,
+    addEntryFn    := fun s n => s.insert n.1 n.2 ,
+    toArrayFn     := fun es => es.toArray.qsort (fun a b => Name.quickLt a.1 b.1)
+  }
+
+-- Declare an extension of maps of maps (using [HashMap]).
+-- The important point is that we need to merge the bound values (which are maps).
+def mkMapHashMapDeclarationExtension [BEq α] [Hashable α] [Inhabited β]
+  (name : Name := by exact decl_name%) :
+  IO (MapDeclarationExtension (HashMap α β)) :=
+  registerSimplePersistentEnvExtension {
+    name          := name,
+    addImportedFn := fun a =>
+      a.foldl (fun s a => a.foldl (
+        -- We need to merge the maps
+        fun s (k0, k1_to_v) =>
+        match s.find? k0 with
+        | none =>
+          -- No binding: insert one
+          s.insert k0 k1_to_v
+        | some m =>
+          -- There is already a binding: merge
+          let m := HashMap.fold (fun m k v => m.insert k v) m k1_to_v
+          s.insert k0 m)
+          s) RBMap.empty,
+    addEntryFn    := fun s n => s.insert n.1 n.2 ,
+    toArrayFn     := fun es => es.toArray.qsort (fun a b => Name.quickLt a.1 b.1)
+  }
+
 /- The persistent map from function to pspec theorems. -/
 initialize pspecAttr : PSpecAttr ← do
   let ext ← mkMapDeclarationExtension `pspecMap
@@ -200,7 +249,8 @@ initialize pspecAttr : PSpecAttr ← do
 
 /- The persistent map from type classes to pspec theorems -/
 initialize pspecClassAttr : PSpecClassAttr ← do
-  let ext : MapDeclarationExtension (NameMap Name) ← mkMapDeclarationExtension `pspecClassMap
+  let ext : MapDeclarationExtension (NameMap Name) ←
+    mkMapMapDeclarationExtension Name.quickCmp `pspecClassMap
   let attrImpl : AttributeImpl  := {
     name := `cpspec
     descr := "Marks theorems to use for type classes with the `progress` tactic"
@@ -231,7 +281,8 @@ initialize pspecClassAttr : PSpecClassAttr ← do
 
 /- The 2nd persistent map from type classes to pspec theorems -/
 initialize pspecClassExprAttr : PSpecClassExprAttr ← do
-  let ext : MapDeclarationExtension (HashMap Expr Name) ← mkMapDeclarationExtension `pspecClassExprMap
+  let ext : MapDeclarationExtension (HashMap Expr Name) ←
+    mkMapHashMapDeclarationExtension `pspecClassExprMap
   let attrImpl : AttributeImpl  := {
     name := `cepspec
     descr := "Marks theorems to use for type classes with the `progress` tactic"
diff --git a/backends/lean/Base/Progress/Progress.lean b/backends/lean/Base/Progress/Progress.lean
index 4fd88e36..8b0759c5 100644
--- a/backends/lean/Base/Progress/Progress.lean
+++ b/backends/lean/Base/Progress/Progress.lean
@@ -243,21 +243,26 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
     tryLookupApply keep ids splitPost asmTac fExpr "pspec theorem" pspec do
     -- It failed: try to lookup a *class* expr spec theorem (those are more
     -- specific than class spec theorems)
+    trace[Progress] "Failed using a pspec theorem: trying to lookup a pspec class expr theorem"
     let pspecClassExpr ← do
       match getFirstArg args with
       | none => pure none
       | some arg => do
+        trace[Progress] "Using: f:{fName}, arg: {arg}"
         let thName ← pspecClassExprAttr.find? fName arg
         pure (thName.map fun th => .Theorem th)
     tryLookupApply keep ids splitPost asmTac fExpr "pspec class expr theorem" pspecClassExpr do
     -- It failed: try to lookup a *class* spec theorem
+    trace[Progress] "Failed using a pspec class expr theorem: trying to lookup a pspec class theorem"
     let pspecClass ← do
       match ← getFirstArgAppName args with
       | none => pure none
       | some argName => do
+        trace[Progress] "Using: f: {fName}, arg: {argName}"
         let thName ← pspecClassAttr.find? fName argName
         pure (thName.map fun th => .Theorem th)
     tryLookupApply keep ids splitPost asmTac fExpr "pspec class theorem" pspecClass do
+    trace[Progress] "Failed using a pspec class theorem: trying to use a recursive assumption"
     -- Try a recursive call - we try the assumptions of kind "auxDecl"
     let ctx ← Lean.MonadLCtx.getLCtx
     let decls ← ctx.getAllDecls
@@ -346,11 +351,14 @@ elab "progress" args:progressArgs : tactic =>
 namespace Test
   open Primitives Result
 
+  -- Show the traces
   -- set_option trace.Progress true
   -- set_option pp.rawOnError true
 
+  -- The following commands display the databases of theorems
   -- #eval showStoredPSpec
   -- #eval showStoredPSpecClass
+  -- #eval showStoredPSpecExprClass
 
   example {ty} {x y : Scalar ty}
     (hmin : Scalar.min ty ≤ x.val + y.val)
@@ -366,6 +374,12 @@ namespace Test
     progress keep h with Scalar.add_spec as ⟨ z ⟩
     simp [*, h]
 
+  example {x y : U32}
+    (hmax : x.val + y.val ≤ U32.max) :
+    ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
+    progress keep _ as ⟨ z, h1 .. ⟩
+    simp [*, h1]
+
   /- Checking that universe instantiation works: the original spec uses
      `α : Type u` where u is quantified, while here we use `α : Type 0` -/
   example {α : Type} (v: Vec α) (i: Usize) (x : α)