feat: snapshot trees and language processors

src/Lean/Server/Completion.lean

@@ -81,39 +81,53 @@ open Elab
 open Meta
 open FuzzyMatching
 
+abbrev EligibleHeaderDecls := HashMap Name ConstantInfo
+
 /-- Cached header declarations for which `allowCompletion headerEnv decl` is true. -/
-builtin_initialize eligibleHeaderDeclsRef : IO.Ref (HashMap Name ConstantInfo) ← IO.mkRef {}
+builtin_initialize eligibleHeaderDeclsRef : IO.Ref (Option EligibleHeaderDecls) ←
+  IO.mkRef none
 
-/-- Caches the declarations in the header for which `allowCompletion headerEnv decl` is true. -/
-def fillEligibleHeaderDecls (headerEnv : Environment) : IO Unit := do
-  eligibleHeaderDeclsRef.modify fun startEligibleHeaderDecls =>
-    (·.2) <| StateT.run (m := Id) (s := startEligibleHeaderDecls) do
-      headerEnv.constants.forM fun declName c => do
+/--
+Returns the declarations in the header for which `allowCompletion env decl` is true, caching them
+if not already cached.
+-/
+def getEligibleHeaderDecls (env : Environment) : IO EligibleHeaderDecls := do
+  if let some eligibleHeaderDecls ← eligibleHeaderDeclsRef.get then
+    return eligibleHeaderDecls
+
+  -- enter critical section to avoid running the computation multiple times
+  -- `take` safety: we call `set` in all branches, without a chance for exceptions in between
+  if let some eligibleHeaderDecls ← unsafe eligibleHeaderDeclsRef.take then
+    -- someone else was quicker!
+    eligibleHeaderDeclsRef.set eligibleHeaderDecls
+    return eligibleHeaderDecls
+  let (_, eligibleHeaderDecls) :=
+    StateT.run (m := Id) (s := {}) do
+      -- `map₁` are the header decls
+      env.constants.map₁.forM fun declName c => do
         modify fun eligibleHeaderDecls =>
-          if allowCompletion headerEnv declName then
+          if allowCompletion env declName then
             eligibleHeaderDecls.insert declName c
           else
             eligibleHeaderDecls
+  eligibleHeaderDeclsRef.set (some eligibleHeaderDecls)
+  return eligibleHeaderDecls
 
 /-- Iterate over all declarations that are allowed in completion results. -/
 private def forEligibleDeclsM [Monad m] [MonadEnv m] [MonadLiftT (ST IO.RealWorld) m]
-    (f : Name → ConstantInfo → m PUnit) : m PUnit := do
+    [MonadLiftT IO m] (f : Name → ConstantInfo → m PUnit) : m PUnit := do
   let env ← getEnv
-  (← eligibleHeaderDeclsRef.get).forM f
+  (← getEligibleHeaderDecls (← getEnv)).forM f
   -- map₂ are exactly the local decls
   env.constants.map₂.forM fun name c => do
     if allowCompletion env name then
       f name c
 
 /-- Checks whether this declaration can appear in completion results. -/
-private def allowCompletion [Monad m] [MonadEnv m] [MonadLiftT (ST IO.RealWorld) m]
-    (declName : Name) : m Bool := do
-  let env ← getEnv
-  if (← eligibleHeaderDeclsRef.get).contains declName then
-    return true
-  if env.constants.map₂.contains declName && Lean.Meta.allowCompletion env declName then
-    return true
-  return false
+private def allowCompletion (env : Environment) (eligibleHeaderDecls : EligibleHeaderDecls)
+    (declName : Name) : Bool :=
+  eligibleHeaderDecls.contains declName ||
+    env.constants.map₂.contains declName && Lean.Meta.allowCompletion env declName
 
 /--
 Sorts `items` descendingly according to their score and ascendingly according to their label
@@ -371,16 +385,17 @@ private def idCompletionCore
       matchAtomic id (alias.replacePrefix ns Name.anonymous)
     else
       none
+  let eligibleHeaderDecls ← getEligibleHeaderDecls env
   -- Auxiliary function for `alias`
   let addAlias (alias : Name) (declNames : List Name) (score : Float) : M Unit := do
     declNames.forM fun declName => do
-      if ← allowCompletion declName then
+      if allowCompletion env eligibleHeaderDecls declName then
         addUnresolvedCompletionItemForDecl alias.getString! declName score
   -- search explicitly open `ids`
   for openDecl in ctx.openDecls do
     match openDecl with
     | OpenDecl.explicit openedId resolvedId =>
-      if ← allowCompletion resolvedId then
+      if allowCompletion env eligibleHeaderDecls resolvedId then
         if let some score := matchAtomic id openedId then
           addUnresolvedCompletionItemForDecl openedId.getString! resolvedId score
     | OpenDecl.simple ns _      =>

src/Lean/Server/FileWorker.lean

@@ -238,8 +238,6 @@ def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
     return .error { diagnostics := .empty, success? := none }
 
   let imports := Elab.headerToImports stx
-  -- Prepare header-based caches that requests may use
-  --runHeaderCachingHandlers headerEnv
   let fileSetupResult ← setupFile meta imports fun stderrLine => do
     let progressDiagnostic := {
       range      := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -640,7 +640,6 @@ builtin_initialize
     CompletionParams
     CompletionList
     handleCompletion
-    Completion.fillEligibleHeaderDecls
   registerLspRequestHandler
     "completionItem/resolve"
     CompletionItem

src/Lean/Server/Requests.lean

@@ -181,13 +181,8 @@ section HandlerTable
 open Lsp
 
 structure RequestHandler where
-  fileSource          : Json → Except RequestError Lsp.DocumentUri
-  handle              : Json → RequestM (RequestTask Json)
-  /--
-  Handler that is called by the file worker after processing the header with the header environment.
-  Enables request handlers to cache data related to imports.
-  -/
-  handleHeaderCaching : Environment → IO Unit
+  fileSource : Json → Except RequestError Lsp.DocumentUri
+  handle : Json → RequestM (RequestTask Json)
 
 builtin_initialize requestHandlers : IO.Ref (PersistentHashMap String RequestHandler) ←
   IO.mkRef {}
@@ -207,9 +202,7 @@ as LSP error responses. -/
 def registerLspRequestHandler (method : String)
     paramType [FromJson paramType] [FileSource paramType]
     respType [ToJson respType]
-    (handler : paramType → RequestM (RequestTask respType))
-    (headerCachingHandler : Environment → IO Unit := fun _ => pure ())
-    : IO Unit := do
+    (handler : paramType → RequestM (RequestTask respType)) : IO Unit := do
   if !(← Lean.initializing) then
     throw <| IO.userError s!"Failed to register LSP request handler for '{method}': only possible during initialization"
   if (← requestHandlers.get).contains method then
@@ -220,8 +213,8 @@ def registerLspRequestHandler (method : String)
     let params ← liftExcept <| parseRequestParams paramType j
     let t ← handler params
     pure <| t.map <| Except.map ToJson.toJson
-  let handleHeaderCaching := headerCachingHandler
-  requestHandlers.modify fun rhs => rhs.insert method { fileSource, handle, handleHeaderCaching }
+
+  requestHandlers.modify fun rhs => rhs.insert method { fileSource, handle }
 
 def lookupLspRequestHandler (method : String) : IO (Option RequestHandler) :=
   return (← requestHandlers.get).find? method
@@ -252,14 +245,6 @@ def chainLspRequestHandler (method : String)
   else
     throw <| IO.userError s!"Failed to chain LSP request handler for '{method}': no initial handler registered"
 
-/--
-Runs the header caching handler for every single registered request handler using the given
-`headerEnv`.
--/
-def runHeaderCachingHandlers (headerEnv : Environment) : IO Unit := do
-  (← requestHandlers.get).forM fun _ handler =>
-    handler.handleHeaderCaching headerEnv
-
 def routeLspRequest (method : String) (params : Json) : IO (Except RequestError DocumentUri) := do
   match (← lookupLspRequestHandler method) with
   | none => return Except.error <| RequestError.methodNotFound method