Introduce cmdEnv? context field

Author: Kha

src/Lean/Elab/Command.lean

@@ -377,7 +377,8 @@ private def mkInfoTree (elaborator : Name) (stx : Syntax) (trees : PersistentArr
   let scope := s.scopes.head!
   let tree := InfoTree.node (Info.ofCommandInfo { elaborator, stx }) trees
   let ctx := PartialContextInfo.commandCtx {
-    env := s.env, fileMap := ctx.fileMap, mctx := {}, currNamespace := scope.currNamespace,
+    env := s.env, cmdEnv? := some s.env, fileMap := ctx.fileMap, mctx := {},
+    currNamespace := scope.currNamespace,
     openDecls := scope.openDecls, options := scope.opts, ngen := s.ngen
   }
   return InfoTree.context ctx tree

src/Lean/Elab/InfoTree/Main.lean

@@ -52,7 +52,7 @@ def PartialContextInfo.mergeIntoOuter?
   | .autoImplicitCtx _, none =>
     panic! "Unexpected incomplete InfoTree context info."
   | .commandCtx innerInfo, some outer =>
-    some { outer with toCommandContextInfo := innerInfo }
+    some { outer with toCommandContextInfo := { innerInfo with cmdEnv? := outer.cmdEnv? <|> innerInfo.cmdEnv? } }
   | .parentDeclCtx innerParentDecl, some outer =>
     some { outer with parentDecl? := innerParentDecl }
   | .autoImplicitCtx innerAutoImplicits, some outer =>

src/Lean/Elab/InfoTree/Types.lean

@@ -18,12 +18,23 @@ public section
 namespace Lean.Elab
 
 /--
-Context after executing `liftTermElabM`.
+Context at the `CommandElabM/TermElabM` level. Created by `elabCommand` at the top level and then
+nestedly when relevant fields are affected (e.g. just before discarding the `mctx` when exiting from
+`TermElabM`).
+
 Note that the term information collected during elaboration may contain metavariables, and their
 assignments are stored at `mctx`.
 -/
 structure CommandContextInfo where
   env           : Environment
+  /--
+  Final environment at the end of `elabCommand`; empty for nested contexts. This environment can be
+  used to access information about the fully-elaborated current declaration such as declaration
+  ranges. It may not be a strict superset of `env` in case of backtracking, so `env` should be
+  preferred to access information about the elaboration context at the time this context object was
+  created.
+  -/
+  cmdEnv?       : Option Environment := none
   fileMap       : FileMap
   mctx          : MetavarContext := {}
   options       : Options        := {}

src/Lean/Server/References.lean

@@ -85,17 +85,16 @@ def addRef (i : RefInfo) (ref : Reference) : RefInfo :=
 def toLspRefInfo (i : RefInfo) : BaseIO Lsp.RefInfo := do
   let refToRefInfoLocation (ref : Reference) : BaseIO RefInfo.Location := do
     let parentDeclName? := ref.ci.parentDecl?
-    let .ok parentDeclRanges? ← EIO.toBaseIO <| ref.ci.runCoreM do
-        let some parentDeclName := parentDeclName?
-          | return none
-        -- Use `local` as it avoids unnecessary blocking, which is especially important when called
-        -- from the snapshot reporter. Specifically, if `ref` is from a tactic of an async theorem,
-        -- `parentDeclName` will not be available in the current environment and we would block only
-        -- to return `none` in the end anyway. At the end of elaboration, we rerun this function on
-        -- the full info tree with the main environment, so the access will succeed immediately.
-        return declRangeExt.find? (asyncMode := .local) (← getEnv) parentDeclName
-      -- we only use `CoreM` to get access to a `MonadEnv`, but these are currently all `IO`
-      | unreachable!
+    let parentDeclRanges? : Option DeclarationRanges := do
+      -- Use final command env that has decl ranges for current declaration as well
+      let cmdEnv ← ref.ci.cmdEnv?
+      let parentDeclName ← parentDeclName?
+      -- Use `local` as it avoids unnecessary blocking, which is especially important when called
+      -- from the snapshot reporter. Specifically, if `ref` is from a tactic of an async theorem,
+      -- `parentDeclName` will not be available in the current environment and we would block only
+      -- to return `none` in the end anyway. At the end of elaboration, we rerun this function on
+      -- the full info tree with the main environment, so the access will succeed immediately.
+      declRangeExt.find? (asyncMode := .local) cmdEnv parentDeclName
     return {
       range := ref.range
       parentDecl? := do