feat: partial context info (#3159)

This PR facilitates augmenting the context of an `InfoTree` with
*partial* contexts while elaborating a command. Using partial contexts,
this PR also adds support for tracking the parent declaration name of a
term in the `InfoTree`. The parent declaration name is needed to compute
the call hierarchy in #3082.

Specifically, the `Lean.Elab.InfoTree.context` constructor is refactored
to take a value of the new type `Lean.Elab.PartialContextInfo` instead
of a `Lean.Elab.ContextInfo`, which now refers to a full `InfoTree`
context. The `PartialContextInfo` is then merged into a `ContextInfo`
while traversing the tree using
`Lean.Elab.PartialContextInfo.mergeIntoOuter?`. The partial context
after executing `liftTermElabM` is stored in values of a new type
`Lean.Elab.CommandContextInfo`.

As a result of this, `Lean.Elab.ContextInfo.save` moves to
`Lean.Elab.CommandContextInfo.save`.

For obtaining the parent declaration for a term, a new typeclass
`MonadParentDecl` is introduced to save the parent declaration in
`Lean.Elab.withSaveParentDeclInfoContext`. `Lean.Elab.Term.withDeclName
x` now calls `withSaveParentDeclInfoContext x` to save the declaration
name.

### Migration

**The changes to the `InfoTree.context` constructor break backwards
compatibility with all downstream users that traverse the `InfoTree`
manually instead of going through the functions in `InfoUtils.lean`.**
To fix this, you can merge the outer `ContextInfo` in a traversal with
the `PartialContextInfo` of an `InfoTree.context` node using
`PartialContextInfo.mergeIntoOuter?`. See e.g.
`Lean.Elab.InfoTree.foldInfo` for an example:
```lean
partial def InfoTree.foldInfo (f : ContextInfo → Info → α → α) (init : α) : InfoTree → α :=
  go none init
where go ctx? a
  | context ctx t => go (ctx.mergeIntoOuter? ctx?) a t
  | node i ts =>
    let a := match ctx? with
      | none => a
      | some ctx => f ctx i a
    ts.foldl (init := a) (go <| i.updateContext? ctx?)
  | _ => a
```

Downstream users that manually save `InfoTree`s may need to adjust calls
to `ContextInfo.save` to use `CommandContextInfo.save` instead and
potentially wrap their `CommandContextInfo` in a
`PartialContextInfo.commandCtx` constructor when storing it in an
`InfoTree` or `ContextInfo.mk` when creating a full context.

### Motivation

As of now, `ContextInfo`s are always *full* contexts, constructed as if
they were always created in `liftTermElabM` after running the
`TermElabM` action. This is not strictly true; we already create
`ContextInfo`s in several places other than `liftTermElabM` and work
around the limitation that `ContextInfo`s are always full contexts in
certain places (e.g. `Info.updateContext?` is a crux that we need
because we can't always create partial contexts at the term-level), but
it has mostly worked out so far. Note that one must be very careful when
saving a `ContextInfo` in places other than `liftTermElabM` because the
context may not be as complete as we would like (e.g. it may lack
meta-variable assignments, potentially leading to a language server
panic).

Unfortunately, the parent declaration of a term is another example of a
context that cannot be provided in `liftTermElabM`: The parent
declaration is usually set via `withDeclName`, which itself lives in
`TermElabM`. So by the time we are trying to save the full
`ContextInfo`, the declaration name is already gone. There is no easy
fix for this like in the other cases where we would really just like to
augment the context with an extra field.

The refactor that we decided on to resolve the issue is to refactor the
`InfoTree` to take a `PartialContextInfo` instead of a `ContextInfo` and
have code that traverses the `InfoTree` merge inner contexts with outer
contexts to produce a full `ContextInfo` value.

### Bumps for downstream projects

- `lean-pr-testing-3159` branch at Std, not yet opened as a PR
- `lean-pr-testing-3159` branch at Mathlib, not yet opened as a PR
- leanprover/LeanInk#57
- hargoniX/LeanInk#1
- tydeu/lean4-alloy#7
- leanprover-community/repl#29

---------

Co-authored-by: Sebastian Ullrich <[email protected]>

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

@@ -327,33 +327,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`.
 -/
@@ -400,9 +373,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 :=
@@ -433,6 +409,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)