feat: look for .ilean.mmap files on watchdog load

src/Lean/Server/References.lean

@@ -216,6 +216,55 @@ structure Ilean where
   decls         : Lsp.Decls
   deriving FromJson, ToJson
 
+/-- Extra metadata data stored with a loaded `.ilean` and written to disk for mmaped load -/
+structure Milean extends Ilean where
+  /-- The `.ilean.hash` value for the stored `.ilean`, if it existed when the `Milean` was created. -/
+  hash : Option UInt64
+
+/--
+Given `<path>.hash`, attempt to read the 64-bit hex value it contains.
+Duplicates, and needs to be kept functionally in sync with, `Lake.Hash.load?`.
+-/
+private def loadHash (hashFile : System.FilePath) : IO UInt64 := do
+  let contents ← FS.readBinFile hashFile
+  if contents.size != 16 then throw (.userError s!"hash file does not have expected form")
+
+  let mut result: UInt64 := 0
+  for byte in contents do
+    result := result.shiftLeft 4
+    if '0'.toUInt8 ≤ byte && byte ≤ '9'.toUInt8 then result := result + (byte - '0'.toUInt8).toUInt64
+    else if 'a'.toUInt8 ≤ byte && byte ≤ 'f'.toUInt8 then result := result + (byte - 'a'.toUInt8 + 10).toUInt64
+    else if 'A'.toUInt8 ≤ byte && byte ≤ 'F'.toUInt8 then result := result + (byte - 'A'.toUInt8 + 10).toUInt64
+    else throw (.userError s!"unexpected byte {byte} in hash")
+  return result
+
+/--
+The extension for a memory-mapped ilean file
+-/
+def Milean.ext := "ilean.mmap"
+
+/--
+Given a path `<path>.ilean`, attempt to read `<path>.ilean.mmap`. Returns `.none` if the file does
+not exist, `.some ilean` if the file hashes match, indicating that `<path>.ilean.mmap` is a
+mmapped version of `<path>.ilean`, throws an error otherwise.
+-/
+def Milean.load (ileanPath : System.FilePath) : IO (Option Ilean) := do
+  if ileanPath.extension ≠ .some "ilean" then throw (.userError s!"Milean.load: '{ileanPath}' not an .ilean")
+  let mileanPath := ileanPath.withExtension Milean.ext
+  if not (← mileanPath.pathExists) then return .none
+
+  let hashPath := ileanPath.addExtension "hash"
+
+  -- nb: ignoring the region means we can't unmap this ilean later
+  let (milean, _region) ← unsafe CompactedRegion.read (α := Milean) mileanPath #[]
+  if let .some memHash := milean.hash then
+    if not (← hashPath.pathExists) then throw (.userError s!"Milean.load: '{mileanPath}' expected '{hashPath}' to exist, but that file does not exist")
+    let fsHash ← loadHash hashPath
+    if memHash ≠ fsHash then throw (.userError s!"Milean.load: '{mileanPath}' expects a .ilean.hash value {String.ofList <| Nat.toDigits 16 memHash.toNat}, but .ilean.hash contains {String.ofList <| Nat.toDigits 16 fsHash.toNat}")
+  else
+    if (← hashPath.pathExists) then throw (.userError s!"Milean.load: '{mileanPath}' expected no '{hashPath}' to exist, but that file does exist")
+  return milean.toIlean
+
 namespace Ilean
 
 open Lean.IO

src/Lean/Server/Watchdog.lean

@@ -15,6 +15,7 @@ public import Lean.Server.Completion.CompletionUtils
 public import Init.Data.List.Sort
 public import Std.Sync.Channel
 public import Lean.Server.Logging
+import all Lean.Elab.ErrorUtils
 
 public section
 
@@ -1678,9 +1679,30 @@ results in requests that need references.
 def startLoadingReferences (referenceData : Std.Mutex ReferenceData) : IO Unit := do
   let task ← ServerTask.IO.asTask do
     let oleanSearchPath ← Lean.searchPathRef.get
+    let mut mileanCount := 0
+    let mut mileanErrors := 0
+    let mut visitedPaths : Std.HashSet String := {}
     for path in ← oleanSearchPath.findAllWithExt "ilean" do
+      -- Avoid loading an ilean twice
+      let realPathString := (← FS.realPath path).toString
+      if visitedPaths.contains realPathString then continue
+      visitedPaths := visitedPaths.insert realPathString
+
+      let milean : Option Ilean ← try
+          let .some milean ← Milean.load path | pure .none
+          mileanCount := mileanCount + 1
+          pure milean
+        catch _ =>
+          -- milean errors should not be fatal, but we *should* log them
+          -- individually to help debug the global error message we give on error
+          mileanErrors := mileanErrors + 1
+          mileanCount := mileanCount + 1 -- errors mean there *was* a milean but something went wrong
+          pure .none
+
       try
-        let ilean ← Ilean.load path
+        let ilean ← match milean with
+          | .none => Ilean.load path
+          | .some il => pure il
         referenceData.atomically do
           let rd ← get
           let rd ← rd.modifyReferencesM (·.addIlean path ilean)
@@ -1690,6 +1712,8 @@ def startLoadingReferences (referenceData : Std.Mutex ReferenceData) : IO Unit :
         -- ilean load errors should not be fatal, but we *should* log them
         -- when we add logging to the server
         pure ()
+    if mileanErrors > 0 then
+      (← getStderr).putStrLn s!"Attempted to load {mileanCount} .{Milean.ext} {mileanCount.plural "file"}, but failed to load {mileanErrors} of them"
   referenceData.atomically <| modify fun rd =>
     { rd with loadingTask := task.mapCheap fun _ => () }