feat: instant library_search (#3404)

Well, not quite instant, but very much faster. This works by creating an extra file in the `build` directory containing a cache of the `library_search` discrimination tree.

* If you modify a file, then open another file, `library_search` will not see the modified declarations.
* However after a `lake exe cache get` or a `lake env lean Extras/LibrarySearch.lean`, you will see the modifications.
* In particular CI will automatically do the work of building this cache.
* This will increase the size of a `lake exe cache get` download (by about 20mb).

To accommodate the cache file, we create a new `build/extra/` directory, that can store `.olean` format cache files. Our `cache` executable now manages files in this directory too.



Co-authored-by: Scott Morrison <scott.morrison@gmail.com>
Co-authored-by: Jeremy Tan Jie Rui <e0191785@u.nus.edu>
Co-authored-by: Pol_tta <65514131+Ruben-VandeVelde@users.noreply.github.com>
Co-authored-by: EmilieUthaiwat <emiliepathum@gmail.com>
Co-authored-by: Eric Wieser <wieser.eric@gmail.com>
Co-authored-by: Yaël Dillies <yael.dillies@gmail.com>
Co-authored-by: Jireh Loreaux <loreaujy@gmail.com>
Co-authored-by: Gabriel Ebner <gebner@gebner.org>
Co-authored-by: Violeta Hernández <vi.hdz.p@gmail.com>

Author: 9 people

.github/workflows/bors.yml

@@ -95,6 +95,9 @@ jobs:
       - name: build mathlib
         run: env LEAN_ABORT_ON_PANIC=1 lake build
 
+      - name: build library_search cache
+        run: lake build MathlibExtras
+
       - name: upload cache
         if: always()
         run: |

.github/workflows/build.yml

@@ -101,6 +101,9 @@ jobs:
       - name: build mathlib
         run: env LEAN_ABORT_ON_PANIC=1 lake build
 
+      - name: build library_search cache
+        run: lake build MathlibExtras
+
       - name: upload cache
         if: always()
         run: |

.github/workflows/build.yml.in

@@ -81,6 +81,9 @@ jobs:
       - name: build mathlib
         run: env LEAN_ABORT_ON_PANIC=1 lake build
 
+      - name: build library_search cache
+        run: lake build MathlibExtras
+
       - name: upload cache
         if: always()
         run: |

.github/workflows/build_fork.yml

@@ -99,6 +99,9 @@ jobs:
       - name: build mathlib
         run: env LEAN_ABORT_ON_PANIC=1 lake build
 
+      - name: build library_search cache
+        run: lake build MathlibExtras
+
       - name: upload cache
         if: always()
         run: |

Cache/Hashing.lean

@@ -102,8 +102,11 @@ partial def getFileHash (filePath : FilePath) : HashM $ Option UInt64 := do
         cache   := stt.cache.insert   filePath (some fileHash)
         depsMap := stt.depsMap.insert filePath fileImports })
 
+/-- Files to start hashing from. -/
+def roots : Array FilePath := #["Mathlib.lean", "MathlibExtras.lean"]
+
 /-- Main API to retrieve the hashes of the Lean files -/
 def getHashMemo : IO HashMemo :=
-  return (← StateT.run (getFileHash ⟨"Mathlib.lean"⟩) default).2
+  return (← StateT.run (roots.mapM getFileHash) default).2
 
 end Cache.Hashing

Cache/IO.lean

@@ -70,7 +70,8 @@ def getPackageDirs : IO PackageDirs := return .ofList [
   ("Mathlib", if ← isMathlibRoot then "." else mathlibDepPath),
   ("Aesop", LAKEPACKAGESDIR / "aesop"),
   ("Std", LAKEPACKAGESDIR / "std"),
-  ("Qq", LAKEPACKAGESDIR / "Qq")
+  ("Qq", LAKEPACKAGESDIR / "Qq"),
+  ("MathlibExtras", ".")
 ]
 
 initialize pkgDirs : PackageDirs ← getPackageDirs
@@ -146,18 +147,23 @@ end HashMap
 def mkDir (path : FilePath) : IO Unit := do
   if !(← path.pathExists) then IO.FS.createDirAll path
 
-/-- Given a path to a Lean file, concatenates the paths to its build files -/
-def mkBuildPaths (path : FilePath) : IO $ Array FilePath := do
+/--
+Given a path to a Lean file, concatenates the paths to its build files.
+Each build file also has a `Bool` indicating whether that file is required for caching to proceed.
+-/
+def mkBuildPaths (path : FilePath) : IO $ Array (FilePath × Bool) := do
   let packageDir ← getPackageDir path
   return #[
-    packageDir / LIBDIR / path.withExtension "olean",
-    packageDir / LIBDIR / path.withExtension "ilean",
-    packageDir / LIBDIR / path.withExtension "trace",
-    packageDir / IRDIR  / path.withExtension "c"]
-
-def allExist (paths : Array FilePath) : IO Bool := do
-  for path in paths do
-    if !(← path.pathExists) then return false
+    (packageDir / LIBDIR / path.withExtension "olean", true),
+    (packageDir / LIBDIR / path.withExtension "ilean", true),
+    (packageDir / LIBDIR / path.withExtension "trace", true),
+    (packageDir / IRDIR  / path.withExtension "c", true),
+    (packageDir / LIBDIR / path.withExtension "extra", false)]
+
+/-- Check that all required build files exist. -/
+def allExist (paths : Array (FilePath × Bool)) : IO Bool := do
+  for (path, required) in paths do
+    if required && !(← path.pathExists) then return false
   pure true
 
 /-- Compresses build files into the local cache and returns an array with the compressed files -/
@@ -172,7 +178,7 @@ def packCache (hashMap : HashMap) (overwrite : Bool) : IO $ Array String := do
     if ← allExist buildPaths then
       if (overwrite || !(← zipPath.pathExists)) then
         discard $ runCmd "tar" $ #["-I", "gzip -9", "-cf", zipPath.toString] ++
-          (buildPaths.map toString)
+          ((← buildPaths.filterM (·.1.pathExists)) |>.map (·.1.toString))
       acc := acc.push zip
   return acc
 

Mathlib.lean

@@ -2023,6 +2023,7 @@ import Mathlib.Util.Export
 import Mathlib.Util.IncludeStr
 import Mathlib.Util.LongNames
 import Mathlib.Util.MemoFix
+import Mathlib.Util.Pickle
 import Mathlib.Util.Syntax
 import Mathlib.Util.SynthesizeUsing
 import Mathlib.Util.Tactic

Mathlib/Tactic/LibrarySearch.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Gabriel Ebner, Scott Morrison
 -/
 import Mathlib.Tactic.TryThis
+import Mathlib.Util.Pickle
 import Mathlib.Lean.Expr.Basic
 import Mathlib.Lean.Meta.DiscrTree
 import Mathlib.Tactic.Cache
@@ -46,23 +47,58 @@ inductive DeclMod
 | none | symm | mp | mpr
 deriving DecidableEq
 
-initialize librarySearchLemmas : DeclCache (DiscrTree (Name × DeclMod) true) ←
-  DeclCache.mk "librarySearch: init cache" {} fun name constInfo lemmas => do
-    if constInfo.isUnsafe then return lemmas
-    if ← name.isBlackListed then return lemmas
-    withNewMCtxDepth do withReducible do
-      let (_, _, type) ← forallMetaTelescopeReducing constInfo.type
-      let keys ← DiscrTree.mkPath type
-      let lemmas := lemmas.insertIfSpecific keys (name, .none)
-      match type.getAppFnArgs with
-      | (``Eq, #[_, lhs, rhs]) => do
-        let keys_symm ← DiscrTree.mkPath (← mkEq rhs lhs)
-        pure (lemmas.insertIfSpecific keys_symm (name, .symm))
-      | (``Iff, #[lhs, rhs]) => do
-        let keys_mp ← DiscrTree.mkPath rhs
-        let keys_mpr ← DiscrTree.mkPath lhs
-        pure <| (lemmas.insertIfSpecific keys_mp (name, .mp)).insertIfSpecific keys_mpr (name, .mpr)
-      | _ => pure lemmas
+/-- Insert a lemma into the discrimination tree. -/
+def addLemma (name : Name) (constInfo : ConstantInfo)
+    (lemmas : DiscrTree (Name × DeclMod) true) : MetaM (DiscrTree (Name × DeclMod) true) := do
+  if constInfo.isUnsafe then return lemmas
+  if ← name.isBlackListed then return lemmas
+  withNewMCtxDepth do withReducible do
+    let (_, _, type) ← forallMetaTelescopeReducing constInfo.type
+    let keys ← DiscrTree.mkPath type
+    let lemmas := lemmas.insertIfSpecific keys (name, .none)
+    match type.getAppFnArgs with
+    | (``Eq, #[_, lhs, rhs]) => do
+      let keys_symm ← DiscrTree.mkPath (← mkEq rhs lhs)
+      pure (lemmas.insertIfSpecific keys_symm (name, .symm))
+    | (``Iff, #[lhs, rhs]) => do
+      let keys_mp ← DiscrTree.mkPath rhs
+      let keys_mpr ← DiscrTree.mkPath lhs
+      pure <| (lemmas.insertIfSpecific keys_mp (name, .mp)).insertIfSpecific keys_mpr (name, .mpr)
+    | _ => pure lemmas
+
+/-- Construct the discrimination tree of all lemmas. -/
+def buildDiscrTree : IO (DeclCache (DiscrTree (Name × DeclMod) true)) :=
+  DeclCache.mk "librarySearch: init cache" {} addLemma
+
+open System (FilePath)
+
+def cachePath : IO FilePath :=
+  try
+    return (← findOLean `MathlibExtras.LibrarySearch).withExtension "extra"
+  catch _ =>
+    return "build" / "lib" / "MathlibExtras" / "LibrarySearch.extra"
+
+/--
+A structure that holds the cached discrimination tree,
+and possibly a pointer to a memory region, if we unpickled the tree from disk.
+-/
+structure CachedData where
+  pointer? : Option CompactedRegion
+  cache : DeclCache (DiscrTree (Name × DeclMod) true)
+deriving Nonempty
+
+initialize cachedData : CachedData ← unsafe do
+  let path ← cachePath
+  if (← path.pathExists) then
+    let (d, r) ← unpickle (DiscrTree (Name × DeclMod) true) path
+    return ⟨r, (← Cache.mk (pure d), addLemma)⟩
+  else
+    return ⟨none, ← buildDiscrTree⟩
+
+/--
+Retrieve the current current of lemmas.
+-/
+def librarySearchLemmas : DeclCache (DiscrTree (Name × DeclMod) true) := cachedData.cache
 
 /-- Shortcut for calling `solveByElim`. -/
 def solveByElim (goals : List MVarId) (required : List Expr) (depth) := do

Mathlib/Util/Pickle.lean

@@ -0,0 +1,44 @@
+/-
+Copyright (c) 2023 Mario Carneiro. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Mario Carneiro
+-/
+import Std.Util.TermUnsafe
+
+/-!
+# Pickling and unpickling objects
+
+By abusing `saveModuleData` and `readModuleData` we can pickle and unpickle objects to disk.
+-/
+
+open Lean System
+
+/--
+Save an object to disk.
+If you need to write multiple objects from within a single declaration,
+you will need to provide a unique `key` for each.
+-/
+def pickle {α : Type} (path : FilePath) (x : α) (key : Name := by exact decl_name%) : IO Unit :=
+  saveModuleData path key (unsafe unsafeCast x)
+
+/--
+Load an object from disk.
+Note: The returned `CompactedRegion` can be used to free the memory behind the value
+of type `α`, using `CompactedRegion.free` (which is only safe once all references to the `α` are
+released). Ignoring the `CompactedRegion` results in the data being leaked.
+Use `withUnpickle` to call `CompactedRegion.free` automatically.
+
+This function is unsafe because the data being loaded may not actually have type `α`, and this
+may cause crashes or other bad behavior.
+-/
+unsafe def unpickle (α : Type) (path : FilePath) : IO (α × CompactedRegion) := do
+  let (x, region) ← readModuleData path
+  pure (unsafeCast x, region)
+
+/-- Load an object from disk and run some continuation on it, freeing memory afterwards. -/
+unsafe def withUnpickle [Monad m] [MonadLiftT IO m] {α β : Type}
+    (path : FilePath) (f : α → m β) : m β := do
+  let (x, region) ← unpickle α path
+  let r ← f x
+  region.free
+  pure r

MathlibExtras.lean

@@ -0,0 +1 @@
+import MathlibExtras.LibrarySearch