1 files changed, 70 insertions, 0 deletions
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