feat: partial context info

doc/flake.nix

@@ -12,7 +12,7 @@
     flake = false;
   };
   inputs.leanInk = {
-    url = "github:leanprover/LeanInk";
+    url = "github:leanprover/LeanInk/refs/pull/57/merge";
     flake = false;
   };
 

src/Lean/Elab/Command.lean

@@ -238,10 +238,11 @@ private def mkInfoTree (elaborator : Name) (stx : Syntax) (trees : PersistentArr
   let s ← get
   let scope := s.scopes.head!
   let tree := InfoTree.node (Info.ofCommandInfo { elaborator, stx }) trees
-  return InfoTree.context {
+  let ctx := PartialContextInfo.commandCtx {
     env := s.env, fileMap := ctx.fileMap, mctx := {}, currNamespace := scope.currNamespace,
     openDecls := scope.openDecls, options := scope.opts, ngen := s.ngen
-  } tree
+  }
+  return InfoTree.context ctx tree
 
 private def elabCommandUsing (s : State) (stx : Syntax) : List (KeyedDeclsAttribute.AttributeEntry CommandElab) → CommandElabM Unit
   | []                => withInfoTreeContext (mkInfoTree := mkInfoTree `no_elab stx) <| throwError "unexpected syntax{indentD stx}"

src/Lean/Elab/InfoTree/Main.lean

@@ -6,11 +6,11 @@ Authors: Wojciech Nawrocki, Leonardo de Moura, Sebastian Ullrich
 -/
 import Lean.Meta.PPGoal
 
-namespace Lean.Elab.ContextInfo
+namespace Lean.Elab.CommandContextInfo
 
 variable [Monad m] [MonadEnv m] [MonadMCtx m] [MonadOptions m] [MonadResolveName m] [MonadNameGenerator m]
 
-def saveNoFileMap : m ContextInfo := return {
+def saveNoFileMap : m CommandContextInfo := return {
     env           := (← getEnv)
     fileMap       := default
     mctx          := (← getMCtx)
@@ -20,11 +20,32 @@ def saveNoFileMap : m ContextInfo := return {
     ngen          := (← getNGen)
   }
 
-def save [MonadFileMap m] : m ContextInfo := do
+def save [MonadFileMap m] : m CommandContextInfo := do
   let ctx ← saveNoFileMap
   return { ctx with fileMap := (← getFileMap) }
 
-end ContextInfo
+end CommandContextInfo
+
+/--
+Merges the `inner` partial context into the `outer` context s.t. fields of the `inner` context
+overwrite fields of the `outer` context. Panics if the invariant described in the documentation
+for `PartialContextInfo` is violated.
+
+When traversing an `InfoTree`, this function should be used to combine the context of outer
+nodes with the partial context of their subtrees. This ensures that the traversal has the context
+from the inner node to the root node of the `InfoTree` available, with partial contexts of
+inner nodes taking priority over contexts of outer nodes.
+-/
+def PartialContextInfo.mergeIntoOuter?
+    : (inner : PartialContextInfo) → (outer? : Option ContextInfo) → Option ContextInfo
+  | .commandCtx info, none =>
+    some { info with }
+  | .parentDeclCtx _, none =>
+    panic! "Unexpected incomplete InfoTree context info."
+  | .commandCtx innerInfo, some outer =>
+    some { outer with toCommandContextInfo := innerInfo }
+  | .parentDeclCtx innerParentDecl, some outer =>
+    some { outer with parentDecl? := innerParentDecl }
 
 def CompletionInfo.stx : CompletionInfo → Syntax
   | dot i .. => i.stx
@@ -197,7 +218,7 @@ def Info.updateContext? : Option ContextInfo → Info → Option ContextInfo
 partial def InfoTree.format (tree : InfoTree) (ctx? : Option ContextInfo := none) : IO Format := do
   match tree with
   | hole id     => return .nestD f!"• ?{toString id.name}"
-  | context i t => format t i
+  | context i t => format t <| i.mergeIntoOuter? ctx?
   | node i cs   => match ctx? with
     | none => return "• <context-not-available>"
     | some ctx =>
@@ -308,20 +329,52 @@ def withInfoTreeContext [MonadFinally m] (x : m α) (mkInfoTree : PersistentArra
 @[inline] def withInfoContext [MonadFinally m] (x : m α) (mkInfo : m Info) : m α := do
   withInfoTreeContext x (fun trees => do return InfoTree.node (← mkInfo) trees)
 
-/-- Resets the trees state `t₀`, runs `x` to produce a new trees
-state `t₁` and sets the state to be `t₀ ++ (InfoTree.context Γ <$> t₁)`
-where `Γ` is the context derived from the monad state. -/
-def withSaveInfoContext [MonadNameGenerator m] [MonadFinally m] [MonadEnv m] [MonadOptions m] [MonadMCtx m] [MonadResolveName m] [MonadFileMap m] (x : m α) : m α := do
-  if (← getInfoState).enabled then
-    let treesSaved ← getResetInfoTrees
-    Prod.fst <$> MonadFinally.tryFinally' x fun _ => do
-      let st    ← getInfoState
-      let trees ← st.trees.mapM fun tree => do
-        let tree := tree.substitute st.assignment
-        pure <| InfoTree.context (← ContextInfo.save) tree
-      modifyInfoTrees fun _ => treesSaved ++ trees
-  else
-    x
+private def withSavedPartialInfoContext [MonadFinally m]
+    (x : m α)
+    (ctx? : m (Option PartialContextInfo))
+    : m α := do
+  if !(← getInfoState).enabled then
+    return ← x
+  let treesSaved ← getResetInfoTrees
+  Prod.fst <$> MonadFinally.tryFinally' x fun _ => do
+    let st    ← getInfoState
+    let trees ← st.trees.mapM fun tree => do
+      let tree := tree.substitute st.assignment
+      match (← ctx?) with
+      | none =>
+        pure tree
+      | some ctx =>
+        pure <| InfoTree.context ctx tree
+    modifyInfoTrees fun _ => treesSaved ++ trees
+
+/--
+Resets the trees state `t₀`, runs `x` to produce a new trees state `t₁` and sets the state to be
+`t₀ ++ (InfoTree.context (PartialContextInfo.commandCtx Γ) <$> t₁)` where `Γ` is the context derived
+from the monad state.
+-/
+def withSaveInfoContext
+    [MonadNameGenerator m]
+    [MonadFinally m]
+    [MonadEnv m]
+    [MonadOptions m]
+    [MonadMCtx m]
+    [MonadResolveName m]
+    [MonadFileMap m]
+    (x : m α)
+    : m α := do
+  withSavedPartialInfoContext x do
+    return some <| .commandCtx (← CommandContextInfo.save)
+
+/--
+Resets the trees state `t₀`, runs `x` to produce a new trees state `t₁` and sets the state to be
+`t₀ ++ (InfoTree.context (PartialContextInfo.parentDeclCtx Γ) <$> t₁)` where `Γ` is the parent decl
+name provided by `MonadParentDecl m`.
+-/
+def withSaveParentDeclInfoContext [MonadFinally m] [MonadParentDecl m] (x : m α) : m α := do
+  withSavedPartialInfoContext x do
+    let some declName ← getParentDeclName?
+      | return none
+    return some <| .parentDeclCtx declName
 
 def getInfoHoleIdAssignment? (mvarId : MVarId) : m (Option InfoTree) :=
   return (← getInfoState).assignment[mvarId]

src/Lean/Elab/InfoTree/Types.lean

@@ -14,10 +14,12 @@ import Lean.Widget.Types
 
 namespace Lean.Elab
 
-/-- Context after executing `liftTermElabM`.
-   Note that the term information collected during elaboration may contain metavariables, and their
-   assignments are stored at `mctx`. -/
-structure ContextInfo where
+/--
+Context after executing `liftTermElabM`.
+Note that the term information collected during elaboration may contain metavariables, and their
+assignments are stored at `mctx`.
+-/
+structure CommandContextInfo where
   env           : Environment
   fileMap       : FileMap
   mctx          : MetavarContext := {}
@@ -26,6 +28,31 @@ structure ContextInfo where
   openDecls     : List OpenDecl  := []
   ngen          : NameGenerator -- We must save the name generator to implement `ContextInfo.runMetaM` and making we not create `MVarId`s used in `mctx`.
 
+/--
+Context from the root of the `InfoTree` up to this node.
+Note that the term information collected during elaboration may contain metavariables, and their
+assignments are stored at `mctx`.
+-/
+structure ContextInfo extends CommandContextInfo where
+  parentDecl? : Option Name := none
+
+/--
+Context for a sub-`InfoTree`.
+
+Within `InfoTree`, this must fulfill the invariant that every non-`commandCtx` `PartialContextInfo`
+node is always contained within a `commandCtx` node.
+-/
+inductive PartialContextInfo where
+  | commandCtx (info : CommandContextInfo)
+  /--
+  Context for the name of the declaration that surrounds nodes contained within this `context` node.
+  For example, this makes the name of the surrounding declaration available in `InfoTree` nodes
+  corresponding to the terms within the declaration.
+  -/
+  | parentDeclCtx (parentDecl : Name)
+  -- TODO: More constructors for the different kinds of scopes `commandCtx` is currently
+  -- used for (e.g. eliminating `Info.updateContext?` would be nice!).
+
 /-- Base structure for `TermInfo`, `CommandInfo` and `TacticInfo`. -/
 structure ElabInfo where
   /-- The name of the elaborator that created this info. -/
@@ -164,8 +191,8 @@ inductive Info where
     `hole`s which are filled in later in the same way that unassigned metavariables are.
 -/
 inductive InfoTree where
-  /-- The context object is created by `liftTermElabM` at `Command.lean` -/
-  | context (i : ContextInfo) (t : InfoTree)
+  /-- The context object is created at appropriate points during elaboration -/
+  | context (i : PartialContextInfo) (t : InfoTree)
   /-- The children contain information for nested term elaboration and tactic evaluation -/
   | node (i : Info) (children : PersistentArray InfoTree)
   /-- The elaborator creates holes (aka metavariables) for tactics and postponed terms -/
@@ -191,7 +218,7 @@ structure InfoState where
   trees      : PersistentArray InfoTree := {}
   deriving Inhabited
 
-class MonadInfoTree (m : Type → Type)  where
+class MonadInfoTree (m : Type → Type) where
   getInfoState    : m InfoState
   modifyInfoState : (InfoState → InfoState) → m Unit
 
@@ -204,4 +231,9 @@ instance [MonadLift m n] [MonadInfoTree m] : MonadInfoTree n where
 def setInfoState [MonadInfoTree m] (s : InfoState) : m Unit :=
   modifyInfoState fun _ => s
 
+class MonadParentDecl (m : Type → Type) where
+  getParentDeclName? : m (Option Name)
+
+export MonadParentDecl (getParentDeclName?)
+
 end Lean.Elab

src/Lean/Elab/Term.lean

@@ -325,33 +325,6 @@ instance : AddErrorMessageContext TermElabM where
     let msg ← addMacroStack msg ctx.macroStack
     pure (ref, msg)
 
-/--
-  Execute `x` but discard changes performed at `Term.State` and `Meta.State`.
-  Recall that the `Environment` and `InfoState` are at `Core.State`. Thus, any updates to it will
-  be preserved. This method is useful for performing computations where all
-  metavariable must be resolved or discarded.
-  The `InfoTree`s are not discarded, however, and wrapped in `InfoTree.Context`
-  to store their metavariable context. -/
-def withoutModifyingElabMetaStateWithInfo (x : TermElabM α) : TermElabM α := do
-  let s ← get
-  let sMeta ← getThe Meta.State
-  try
-    withSaveInfoContext x
-  finally
-    set s
-    set sMeta
-
-/--
-  Execute `x` but discard changes performed to the state.
-  However, the info trees and messages are not discarded. -/
-private def withoutModifyingStateWithInfoAndMessagesImpl (x : TermElabM α) : TermElabM α := do
-  let saved ← saveState
-  try
-    withSaveInfoContext x
-  finally
-    let saved := { saved with meta.core.infoState := (← getInfoState), meta.core.messages := (← getThe Core.State).messages }
-    restoreState saved
-
 /--
   Execute `x` without storing `Syntax` for recursive applications. See `saveRecAppSyntax` field at `Context`.
 -/
@@ -398,9 +371,12 @@ def getLetRecsToLift : TermElabM (List LetRecToLift) := return (← get).letRecs
 /-- Return the declaration of the given metavariable -/
 def getMVarDecl (mvarId : MVarId) : TermElabM MetavarDecl := return (← getMCtx).getDecl mvarId
 
-/-- Execute `x` with `declName? := name`. See `getDeclName?`. -/
+instance : MonadParentDecl TermElabM where
+  getParentDeclName? := getDeclName?
+
+/-- Execute `withSaveParentDeclInfoContext x` with `declName? := name`. See `getDeclName?`. -/
 def withDeclName (name : Name) (x : TermElabM α) : TermElabM α :=
-  withReader (fun ctx => { ctx with declName? := name }) x
+  withReader (fun ctx => { ctx with declName? := name }) <| withSaveParentDeclInfoContext x
 
 /-- Update the universe level parameter names. -/
 def setLevelNames (levelNames : List Name) : TermElabM Unit :=
@@ -431,6 +407,33 @@ def withoutErrToSorryImp (x : TermElabM α) : TermElabM α :=
 def withoutErrToSorry [MonadFunctorT TermElabM m] : m α → m α :=
   monadMap (m := TermElabM) withoutErrToSorryImp
 
+/--
+  Execute `x` but discard changes performed at `Term.State` and `Meta.State`.
+  Recall that the `Environment` and `InfoState` are at `Core.State`. Thus, any updates to it will
+  be preserved. This method is useful for performing computations where all
+  metavariable must be resolved or discarded.
+  The `InfoTree`s are not discarded, however, and wrapped in `InfoTree.Context`
+  to store their metavariable context. -/
+def withoutModifyingElabMetaStateWithInfo (x : TermElabM α) : TermElabM α := do
+  let s ← get
+  let sMeta ← getThe Meta.State
+  try
+    withSaveInfoContext x
+  finally
+    set s
+    set sMeta
+
+/--
+  Execute `x` but discard changes performed to the state.
+  However, the info trees and messages are not discarded. -/
+private def withoutModifyingStateWithInfoAndMessagesImpl (x : TermElabM α) : TermElabM α := do
+  let saved ← saveState
+  try
+    withSaveInfoContext x
+  finally
+    let saved := { saved with meta.core.infoState := (← getInfoState), meta.core.messages := (← getThe Core.State).messages }
+    restoreState saved
+
 /-- For testing `TermElabM` methods. The #eval command will sign the error. -/
 def throwErrorIfErrors : TermElabM Unit := do
   if (← MonadLog.hasErrors) then

src/Lean/Server/FileWorker.lean

@@ -196,7 +196,7 @@ section Initialization
       (headerMsgLog : MessageLog)
       (opts : Options)
       : Elab.Command.State :=
-    let headerContextInfo : Elab.ContextInfo := {
+    let headerContextInfo : Elab.CommandContextInfo := {
       env     := headerEnv
       fileMap := m.text
       ngen    := { namePrefix := `_worker }
@@ -210,7 +210,7 @@ section Initialization
     let headerInfoTree := Elab.InfoTree.node headerInfo headerInfoNodes.toPArray'
     let headerInfoState := {
       enabled := true
-      trees := #[Elab.InfoTree.context headerContextInfo headerInfoTree].toPArray'
+      trees := #[Elab.InfoTree.context (.commandCtx headerContextInfo) headerInfoTree].toPArray'
     }
     { Elab.Command.mkState headerEnv headerMsgLog opts with infoState := headerInfoState }
 

src/Lean/Server/InfoUtils.lean

@@ -35,14 +35,15 @@ structure InfoWithCtx where
   info : Elab.Info
   children : PersistentArray InfoTree
 
-/-- Visit nodes, passing in a surrounding context (the innermost one) and accumulating results on the way back up. -/
+/-- Visit nodes, passing in a surrounding context (the innermost one combined with all outer ones)
+and accumulating results on the way back up. -/
 partial def InfoTree.visitM [Monad m]
     (preNode  : ContextInfo → Info → (children : PersistentArray InfoTree) → m Unit := fun _ _ _ => pure ())
     (postNode : ContextInfo → Info → (children : PersistentArray InfoTree) → List (Option α) → m α)
     : InfoTree → m (Option α) :=
   go none
 where go
-  | _, context ctx t => go ctx t
+  | ctx?, context ctx t => go (ctx.mergeIntoOuter? ctx?) t
   | some ctx, node i cs => do
     preNode ctx i cs
     let as ← cs.toList.mapM (go <| i.updateContext? ctx)
@@ -77,7 +78,7 @@ partial def InfoTree.deepestNodes (p : ContextInfo → Info → PersistentArray
 partial def InfoTree.foldInfo (f : ContextInfo → Info → α → α) (init : α) : InfoTree → α :=
   go none init
 where go ctx? a
-  | context ctx t => go ctx a t
+  | context ctx t => go (ctx.mergeIntoOuter? ctx?) a t
   | node i ts =>
     let a := match ctx? with
       | none => a
@@ -367,7 +368,7 @@ partial def InfoTree.termGoalAt? (t : InfoTree) (hoverPos : String.Pos) : Option
 partial def InfoTree.hasSorry : InfoTree → IO Bool :=
   go none
 where go ci?
-  | .context ci t => go ci t
+  | .context ci t => go (ci.mergeIntoOuter? ci?) t
   | .node i cs =>
     if let (some ci, .ofTermInfo ti) := (ci?, i) then do
       let expr ← ti.runMetaM ci (instantiateMVars ti.expr)