Automated prefix state sharing optimization

REPL/Frontend.lean

@@ -4,11 +4,44 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison
 -/
 import Lean.Elab.Frontend
+import Std.Data.HashMap
 
-open Lean Elab
+open Lean Elab Language
 
 namespace Lean.Elab.IO
 
+partial def IO.processCommandsIncrementally' (inputCtx : Parser.InputContext)
+    (parserState : Parser.ModuleParserState) (commandState : Command.State)
+    (old? : Option IncrementalState) :
+    BaseIO (List (IncrementalState × Option InfoTree)) := do
+  let task ← Language.Lean.processCommands inputCtx parserState commandState
+    (old?.map fun old => (old.inputCtx, old.initialSnap))
+  return go task.get task #[]
+where
+  go initialSnap t commands :=
+    let snap := t.get
+    let commands := commands.push snap
+    -- Opting into reuse also enables incremental reporting, so make sure to collect messages from
+    -- all snapshots
+    let messages := toSnapshotTree initialSnap
+      |>.getAll.map (·.diagnostics.msgLog)
+      |>.foldl (· ++ ·) {}
+    -- In contrast to messages, we should collect info trees only from the top-level command
+    -- snapshots as they subsume any info trees reported incrementally by their children.
+    let trees := commands.map (·.elabSnap.infoTreeSnap.get.infoTree?) |>.filterMap id |>.toPArray'
+    let result : (IncrementalState × Option InfoTree) :=
+      ({ commandState := { snap.elabSnap.resultSnap.get.cmdState with messages, infoState.trees := trees }
+         , parserState := snap.parserState
+         , cmdPos := snap.parserState.pos
+         , commands := commands.map (·.stx)
+         , inputCtx := inputCtx
+         , initialSnap := initialSnap
+       }, snap.elabSnap.infoTreeSnap.get.infoTree?)
+    if let some next := snap.nextCmdSnap? then
+      result :: go initialSnap next.task commands
+    else
+      [result]
+
 /--
 Wrapper for `IO.processCommands` that enables info states, and returns
 * the new command state
@@ -17,41 +50,48 @@ Wrapper for `IO.processCommands` that enables info states, and returns
 -/
 def processCommandsWithInfoTrees
     (inputCtx : Parser.InputContext) (parserState : Parser.ModuleParserState)
-    (commandState : Command.State) : IO (Command.State × List Message × List InfoTree) := do
+    (commandState : Command.State) (incrementalState? : Option IncrementalState := none)
+    : IO (List (IncrementalState × Option InfoTree) × List Message) := do
   let commandState := { commandState with infoState.enabled := true }
-  let s ← IO.processCommands inputCtx parserState commandState <&> Frontend.State.commandState
-  pure (s, s.messages.toList, s.infoState.trees.toList)
+  let incStates ← IO.processCommandsIncrementally' inputCtx parserState commandState incrementalState?
+  pure (incStates, (incStates.getLast?.map (·.1.commandState.messages.toList)).getD {})
 
 /--
 Process some text input, with or without an existing command state.
-If there is no existing environment, we parse the input for headers (e.g. import statements),
-and create a new environment.
-Otherwise, we add to the existing environment.
+Supports header caching to avoid reprocessing the same imports repeatedly.
 
 Returns:
-1. The header-only command state (only useful when cmdState? is none)
-2. The resulting command state after processing the entire input
-3. List of messages
-4. List of info trees
+1. The resulting command state after processing the entire input
+2. List of messages
+3. List of info trees along with Command.State from the incremental processing
+4. Updated header cache mapping import keys to Command.State
 -/
 def processInput (input : String) (cmdState? : Option Command.State)
+    (incrementalState? : Option IncrementalState := none)
+    (headerCache : Std.HashMap String Command.State)
     (opts : Options := {}) (fileName : Option String := none) :
-    IO (Command.State × Command.State × List Message × List InfoTree) := unsafe do
+    IO (Command.State × List (IncrementalState × Option InfoTree) × List Message × (Std.HashMap String Command.State)) := unsafe do
   Lean.initSearchPath (← Lean.findSysroot)
   enableInitializersExecution
   let fileName   := fileName.getD "<input>"
   let inputCtx   := Parser.mkInputContext input fileName
-
   match cmdState? with
   | none => do
-    -- Split the processing into two phases to prevent self-reference in proofs in tactic mode
     let (header, parserState, messages) ← Parser.parseHeader inputCtx
-    let (env, messages) ← processHeader header opts messages inputCtx
-    let headerOnlyState := Command.mkState env messages opts
-    let (cmdState, messages, trees) ← processCommandsWithInfoTrees inputCtx parserState headerOnlyState
-    return (headerOnlyState, cmdState, messages, trees)
-
+    let importKey := toString header.raw -- do not contain comments and whitespace
+    match headerCache.get? importKey with
+    | some cachedHeaderState => do
+      -- Header is cached, use it as the base command state
+      let (incStates, messages) ← processCommandsWithInfoTrees inputCtx parserState cachedHeaderState incrementalState?
+      return (cachedHeaderState, incStates, messages, headerCache)
+    | none => do
+      -- Header not cached, process it and cache the result
+      let (env, messages) ← processHeader header opts messages inputCtx
+      let headerOnlyState := Command.mkState env messages opts
+      let headerCache := headerCache.insert importKey headerOnlyState
+      let (incStates, messages) ← processCommandsWithInfoTrees inputCtx parserState headerOnlyState incrementalState?
+      return (headerOnlyState, incStates, messages, headerCache)
   | some cmdStateBefore => do
     let parserState : Parser.ModuleParserState := {}
-    let (cmdStateAfter, messages, trees) ← processCommandsWithInfoTrees inputCtx parserState cmdStateBefore
-    return (cmdStateBefore, cmdStateAfter, messages, trees)
+    let (incStates, messages) ← processCommandsWithInfoTrees inputCtx parserState cmdStateBefore incrementalState?
+    return (cmdStateBefore, incStates, messages, headerCache)

REPL/JSON.lean

@@ -19,6 +19,7 @@ structure CommandOptions where
   Anything else is ignored.
   -/
   infotree : Option String
+  incrementality : Option Bool := none -- whether to use incremental mode optimization
 
 /-- Run Lean commands.
 If `env = none`, starts a new session (in which you can use `import`).
@@ -47,7 +48,7 @@ deriving ToJson, FromJson
 structure Pos where
   line : Nat
   column : Nat
-deriving ToJson, FromJson
+deriving ToJson, FromJson, BEq, Hashable
 
 /-- Severity of a message. -/
 inductive Severity

REPL/Main.lean

@@ -11,6 +11,9 @@ import REPL.Lean.Environment
 import REPL.Lean.InfoTree
 import REPL.Lean.InfoTree.ToJson
 import REPL.Snapshots
+import REPL.Util.Trie
+import Lean.Data.Trie
+import Std.Data.HashMap
 
 /-!
 # A REPL for Lean.
@@ -72,6 +75,16 @@ structure State where
   and report the numerical index for the recorded state at each sorry.
   -/
   proofStates : Array ProofSnapshot := #[]
+  /--
+  Trie-based storage for fast prefix matching, organized by environment ID.
+  Map from environment ID (None for fresh env) to trie of command prefixes with incremental states.
+  -/
+  envTries : Std.HashMap (Option Nat) (Lean.Data.Trie IncrementalState) := Std.HashMap.emptyWithCapacity 8
+  /--
+  Cache for processed headers (import statements) to avoid reprocessing the same imports repeatedly.
+  Maps import raw string to the processed command state.
+  -/
+  headerCache : Std.HashMap String Command.State := Std.HashMap.emptyWithCapacity 8
 
 /--
 The Lean REPL monad.
@@ -88,12 +101,43 @@ def recordCommandSnapshot (state : CommandSnapshot) : M m Nat := do
   modify fun s => { s with cmdStates := s.cmdStates.push state }
   return id
 
+/-- Add all incremental states of a command to the trie at their corresponding prefix positions -/
+def addCommandToTrie (cmdText : String)
+    (incStates : List (IncrementalState × Option InfoTree)) (envId? : Option Nat) : M m Unit := do
+  let state ← get
+  let currentTrie := state.envTries.get? envId? |>.getD Lean.Data.Trie.empty
+
+  let mut newTrie := currentTrie
+  for (incState, _) in incStates do
+    let prefixPos := incState.cmdPos.byteIdx
+    let cmdPrefix : String := (cmdText.take prefixPos).trim
+    newTrie := REPL.Util.Trie.insert newTrie cmdPrefix incState
+
+  modify fun s => { s with envTries := s.envTries.insert envId? newTrie }
+
+/-- Record command text and incremental states for prefix matching reuse. -/
+def recordCommandIncrementals (cmdText : String)
+    (incStates : List (IncrementalState × Option InfoTree)) (envId? : Option Nat) : M m Unit := do
+  addCommandToTrie cmdText incStates envId?
+
 /-- Record a `ProofSnapshot` into the REPL state, returning its index for future use. -/
 def recordProofSnapshot (proofState : ProofSnapshot) : M m Nat := do
   let id := (← get).proofStates.size
   modify fun s => { s with proofStates := s.proofStates.push proofState }
   return id
 
+/-- Find the best incremental state using trie-based prefix matching -/
+def findBestIncrementalState (newCmd : String) (envId? : Option Nat) : M m (Option IncrementalState) := do
+  let state ← get
+  let trimmedCmd := newCmd.trim
+  let trie? := state.envTries.get? envId?
+  match trie? with
+  | none => return none
+  | some trie =>
+    match REPL.Util.Trie.matchPrefix trie trimmedCmd 0 with
+    | some incState => return some incState
+    | none => return none
+
 def sorries (trees : List InfoTree) (env? : Option Environment) (rootGoals? : Option (List MVarId))
 : M m (List Sorry) :=
   trees.flatMap InfoTree.sorries |>.filter (fun t => match t.2.1 with
@@ -102,11 +146,7 @@ def sorries (trees : List InfoTree) (env? : Option Environment) (rootGoals? : Op
       fun ⟨ctx, g, pos, endPos⟩ => do
         let (goal, proofState) ← match g with
         | .tactic g => do
-          let lctx ← ctx.runMetaM {} do
-              match ctx.mctx.findDecl? g with
-              | some decl => return decl.lctx
-              | none => throwError "unknown metavariable '{g}'"
-          let s ← ProofSnapshot.create ctx lctx env? [g] rootGoals?
+          let s ← ProofSnapshot.create ctx none env? [g] rootGoals?
           pure ("\n".intercalate <| (← s.ppGoals).map fun s => s!"{s}", some s)
         | .term lctx (some t) => do
           let s ← ProofSnapshot.create ctx lctx env? [] rootGoals? [t]
@@ -115,6 +155,16 @@ def sorries (trees : List InfoTree) (env? : Option Environment) (rootGoals? : Op
         let proofStateId ← proofState.mapM recordProofSnapshot
         return Sorry.of goal pos endPos proofStateId
 
+def sorriesCmd (treeList : List (IncrementalState × Option InfoTree)) (prevEnv : Environment)
+(rootGoals? : Option (List MVarId))
+: M m (List Sorry) :=
+  match treeList with
+  | [] => pure []
+  | (state, infoTree?) :: rest => do
+    let s ← sorries infoTree?.toList prevEnv rootGoals?
+    let restSorries ← sorriesCmd rest state.commandState.env rootGoals?
+    return s ++ restSorries
+
 def ppTactic (ctx : ContextInfo) (stx : Syntax) : IO Format :=
   ctx.runMetaM {} try
     Lean.PrettyPrinter.ppTactic ⟨stx⟩
@@ -130,6 +180,14 @@ def tactics (trees : List InfoTree) (env? : Option Environment) : M m (List Tact
       let proofStateId ← proofState.mapM recordProofSnapshot
       return Tactic.of goals tactic pos endPos proofStateId ns
 
+def tacticsCmd (treeList : List (IncrementalState × Option InfoTree)) (prevEnv : Environment) : M m (List Tactic) := do
+  match treeList with
+  | [] => pure []
+  | (state, infoTree?) :: rest => do
+    let ts ← tactics infoTree?.toList prevEnv
+    let restTactics ← tacticsCmd rest state.commandState.env
+    return ts ++ restTactics
+
 def collectRootGoalsAsSorries (trees : List InfoTree) (env? : Option Environment) : M m (List Sorry) := do
   trees.flatMap InfoTree.rootGoals |>.mapM
     fun ⟨ctx, goals, pos⟩ => do
@@ -138,6 +196,14 @@ def collectRootGoalsAsSorries (trees : List InfoTree) (env? : Option Environment
       let proofStateId ← proofState.mapM recordProofSnapshot
       return Sorry.of goals pos pos proofStateId
 
+def collectRootGoalsAsSorriesCmd (treeList : List (IncrementalState × Option InfoTree)) (prevEnv : Environment) :
+    M m (List Sorry) := do
+  match treeList with
+  | [] => pure []
+  | (state, infoTree?) :: rest => do
+    let sorries ← collectRootGoalsAsSorries infoTree?.toList prevEnv
+    let restSorries ← collectRootGoalsAsSorriesCmd rest state.commandState.env
+    return sorries ++ restSorries
 
 private def collectFVarsAux : Expr → NameSet
   | .fvar fvarId => NameSet.empty.insert fvarId.name
@@ -197,7 +263,7 @@ def getProofStatus (proofState : ProofSnapshot) : M m String := do
             unless (← Meta.isDefEq pft expectedType) do
               return s!"Error: proof has type {pft} but root goal has type {expectedType}"
 
-            let pf ← abstractAllLambdaFVars pf
+            let pf ← abstractAllLambdaFVars pf >>= instantiateMVars
             let pft ← Meta.inferType pf >>= instantiateMVars
 
             if pf.hasExprMVar then
@@ -302,19 +368,29 @@ def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
   if notFound then
     return .inr ⟨"Unknown environment."⟩
   let initialCmdState? := cmdSnapshot?.map fun c => c.cmdState
-  let (initialCmdState, cmdState, messages, trees) ← try
-    IO.processInput s.cmd initialCmdState?
+
+  let (initialCmdState, incStates, messages) ← try
+    if s.incrementality.getD false then
+      let bestIncrementalState? ← findBestIncrementalState s.cmd s.env
+      let (initialCmdState, incStates, messages, headerCache) ← IO.processInput s.cmd initialCmdState? bestIncrementalState? (← get).headerCache
+      -- Store the command text and incremental states for future reuse
+      modify fun st => { st with headerCache := headerCache }
+      recordCommandIncrementals s.cmd incStates s.env
+      pure (initialCmdState, incStates, messages)
+    else
+      let (initialCmdState, incStates, messages, _) ← IO.processInput s.cmd initialCmdState? none {}
+      pure (initialCmdState, incStates, messages)
   catch ex =>
     return .inr ⟨ex.toString⟩
+
+  let cmdState := (incStates.getLast?.map (fun c => c.1.commandState)).getD initialCmdState
   let messages ← messages.mapM fun m => Message.of m
-  -- For debugging purposes, sometimes we print out the trees here:
-  -- trees.forM fun t => do IO.println (← t.format)
-  let sorries ← sorries trees initialCmdState.env none
+  let sorries ← sorriesCmd incStates initialCmdState.env none
   let sorries ← match s.rootGoals with
-  | some true => pure (sorries ++ (← collectRootGoalsAsSorries trees initialCmdState.env))
+  | some true => pure (sorries ++ (← collectRootGoalsAsSorriesCmd incStates initialCmdState.env))
   | _ => pure sorries
   let tactics ← match s.allTactics with
-  | some true => tactics trees initialCmdState.env
+  | some true => tacticsCmd incStates initialCmdState.env
   | _ => pure []
   let cmdSnapshot :=
   { cmdState
@@ -324,6 +400,9 @@ def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
         snap? := none,
         cancelTk? := none } }
   let env ← recordCommandSnapshot cmdSnapshot
+  let trees := cmdState.infoState.trees.toList
+  -- For debugging purposes, sometimes we print out the trees here:
+  -- trees.forM fun t => do IO.println (← t.format)
   let jsonTrees := match s.infotree with
   | some "full" => trees
   | some "tactics" => trees.flatMap InfoTree.retainTacticInfo

REPL/Snapshots.lean

@@ -185,33 +185,44 @@ def runString (p : ProofSnapshot) (t : String) : IO ProofSnapshot :=
 /-- Pretty print the current goals in the `ProofSnapshot`. -/
 def ppGoals (p : ProofSnapshot) : IO (List Format) :=
   Prod.fst <$> p.runMetaM do p.tacticState.goals.mapM (Meta.ppGoal ·)
+
 /--
 Construct a `ProofSnapshot` from a `ContextInfo` and optional `LocalContext`, and a list of goals.
-
 For convenience, we also allow a list of `Expr`s, and these are appended to the goals
 as fresh metavariables with the given types.
 -/
-def create (ctx : ContextInfo) (lctx? : Option LocalContext) (env? : Option Environment)
+def create (ctx : ContextInfo) (lctx? : Option LocalContext) (prevEnv? : Option Environment)
     (goals : List MVarId) (rootGoals? : Option (List MVarId)) (types : List Expr := [])
     : IO ProofSnapshot := do
-  ctx.runMetaM (lctx?.getD {}) do
-    let goals := goals ++ (← types.mapM fun t => Expr.mvarId! <$> Meta.mkFreshExprMVar (some t))
-    let s ← getThe Core.State
-    let s := match env? with
-    | none => s
-    | some env => { s with env }
-    pure <|
-    { coreState := s
-      coreContext := ← readThe Core.Context
-      metaState := ← getThe Meta.State
-      metaContext := ← readThe Meta.Context
-      termState := {}
-      termContext := {}
-      tacticState := { goals }
-      tacticContext := { elaborator := .anonymous }
-      rootGoals := match rootGoals? with
-        | none => goals
-        | some gs => gs }
+  -- Get the local context from the goals if not provided.
+  let lctx ← match lctx? with
+  | none => match goals with
+    | g :: _ => ctx.runMetaM {} do pure (← g.getDecl).lctx
+    | [] => match rootGoals? with
+      | some (g :: _) => ctx.runMetaM {} do pure (← g.getDecl).lctx
+      | _ => pure {}
+  | some lctx => pure lctx
+
+  ctx.runMetaM lctx do
+    -- update local instances, which is necessary when `types` is non-empty
+    Meta.withLocalInstances (lctx.decls.toList.filterMap id) do
+      let goals := goals ++ (← types.mapM fun t => Expr.mvarId! <$> Meta.mkFreshExprMVar (some t))
+      let s ← getThe Core.State
+      let s := match prevEnv? with
+      | none => s
+      | some env => { s with env }
+      pure <|
+      { coreState := s
+        coreContext := ← readThe Core.Context
+        metaState := ← getThe Meta.State
+        metaContext := ← readThe Meta.Context
+        termState := {}
+        termContext := {}
+        tacticState := { goals }
+        tacticContext := { elaborator := .anonymous }
+        rootGoals := match rootGoals? with
+          | none => goals
+          | some gs => gs }
 
 open Lean.Core in
 /-- A copy of `Core.State` with the `Environment`, caches, and logging omitted. -/