Reorganize a bit

3 files changed, 94 insertions, 65 deletions

diff --git a/backends/lean/Base/Diverge/Elab.lean b/backends/lean/Base/Diverge/Elab.lean
index c6628486..0088fd16 100644
--- a/backends/lean/Base/Diverge/Elab.lean
+++ b/backends/lean/Base/Diverge/Elab.lean
@@ -418,7 +418,7 @@ partial def proveExprIsValid (k_var kk_var : Expr) (e : Expr) : MetaM Expr := do
   | .lit _
   | .mvar _
   | .sort _ => throwError "Unreachable"
-  | .lam .. => throwError "Unimplemented"
+  | .lam .. => throwError "Unimplemented" -- TODO
   | .forallE .. => throwError "Unreachable" -- Shouldn't get there
   | .letE .. => do
     -- Telescope all the let-bindings (remark: this also telescopes the lambdas)
@@ -585,6 +585,7 @@ partial def proveExprIsValid (k_var kk_var : Expr) (e : Expr) : MetaM Expr := do
       else do
         -- Remaining case: normal application.
         -- It shouldn't use the continuation.
+        -- TODO: actually, it can
         proveNoKExprIsValid k_var e
 
 -- Prove that a match expression is valid.
@@ -1087,17 +1088,33 @@ namespace Tests
     . intro i hpos h; simp at h; linarith
     . rename_i hd tl ih
       intro i hpos h
-      -- We can directly use `rw [list_nth]`!
+      -- We can directly use `rw [list_nth]`
       rw [list_nth]; simp
       split <;> try simp [*]
       . tauto
-      . -- TODO: we shouldn't have to do that
+      . -- We don't have to do this if we use scalar_tac
         have hneq : 0 < i := by cases i <;> rename_i a _ <;> simp_all; cases a <;> simp_all
         simp at h
         have ⟨ x, ih ⟩ := ih (i - 1) (by linarith) (by linarith)
         simp [ih]
         tauto
 
+  -- Return a continuation
+  divergent def list_nth_with_back {a: Type} (ls : List a) (i : Int) :
+    Result (a × (a → Result (List a))) :=
+    match ls with
+    | [] => .fail .panic
+    | x :: ls =>
+      if i = 0 then return (x, (λ ret => return (ret :: ls)))
+      else do
+        let (x, back) ← list_nth_with_back ls (i - 1)
+        return (x,
+          (λ ret => do
+           let ls ← back ret
+           return (x :: ls)))
+
+  #check list_nth_with_back.unfold
+
   mutual
     divergent def is_even (i : Int) : Result Bool :=
       if i = 0 then return true else return (← is_odd (i - 1))
@@ -1121,7 +1138,6 @@ namespace Tests
   #check bar.unfold
 
   -- Testing dependent branching and let-bindings
-  -- TODO: why the linter warning?
   divergent def isNonZero (i : Int) : Result Bool :=
     if _h:i = 0 then return false
     else
diff --git a/backends/lean/Base/Lookup/Base.lean b/backends/lean/Base/Lookup/Base.lean
new file mode 100644
index 00000000..54d242ea
--- /dev/null
+++ b/backends/lean/Base/Lookup/Base.lean
@@ -0,0 +1,70 @@
+/- Utilities to have databases of theorems -/
+import Lean
+import Std.Lean.HashSet
+import Base.Utils
+import Base.Primitives.Base
+
+namespace Lookup
+
+open Lean Elab Term Meta
+open Utils
+
+-- TODO: the original function doesn't define correctly the `addImportedFn`. Do a PR?
+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,
+    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 [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)
+  }
+
+end Lookup
diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 76a92795..cf2a3b70 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -2,6 +2,7 @@ import Lean
 import Std.Lean.HashSet
 import Base.Utils
 import Base.Primitives.Base
+import Base.Lookup.Base
 
 namespace Progress
 
@@ -166,67 +167,9 @@ structure PSpecClassExprAttr where
   ext  : MapDeclarationExtension (HashMap Expr Name)
   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 α) :=
-  registerSimplePersistentEnvExtension {
-    name          := name,
-    addImportedFn := fun a => a.foldl (fun s a => a.foldl (fun s (k, v) => s.insert k v) 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 [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
+  let ext ← Lookup.mkMapDeclarationExtension `pspecMap
   let attrImpl : AttributeImpl := {
     name := `pspec
     descr := "Marks theorems to use with the `progress` tactic"
@@ -250,7 +193,7 @@ initialize pspecAttr : PSpecAttr ← do
 /- The persistent map from type classes to pspec theorems -/
 initialize pspecClassAttr : PSpecClassAttr ← do
   let ext : MapDeclarationExtension (NameMap Name) ←
-    mkMapMapDeclarationExtension Name.quickCmp `pspecClassMap
+    Lookup.mkMapMapDeclarationExtension Name.quickCmp `pspecClassMap
   let attrImpl : AttributeImpl  := {
     name := `cpspec
     descr := "Marks theorems to use for type classes with the `progress` tactic"
@@ -282,7 +225,7 @@ initialize pspecClassAttr : PSpecClassAttr ← do
 /- The 2nd persistent map from type classes to pspec theorems -/
 initialize pspecClassExprAttr : PSpecClassExprAttr ← do
   let ext : MapDeclarationExtension (HashMap Expr Name) ←
-    mkMapHashMapDeclarationExtension `pspecClassExprMap
+    Lookup.mkMapHashMapDeclarationExtension `pspecClassExprMap
   let attrImpl : AttributeImpl  := {
     name := `cepspec
     descr := "Marks theorems to use for type classes with the `progress` tactic"