feat: unknown identifier code action and the module system (#11164)

This PR ensures that the code action provided on unknown identifiers
correctly inserts `public` and/or `meta` in `module`s

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/Elab/MutualDef.lean

@@ -1200,6 +1200,7 @@ where
       -- encoded name in e.g. kernel errors where it's hard to replace it)
       views.any (fun view => view.kind != .example || view.modifiers.isPublic) &&
       expandedDeclIds.any (!isPrivateName ·.declName)) do
+    withSaveInfoContext do  -- save adjusted env in info tree
     let headers ← elabHeaders views expandedDeclIds bodyPromises tacPromises
     let headers ← levelMVarToParamHeaders views headers
     if let (#[view], #[declId]) := (views, expandedDeclIds) then
@@ -1324,6 +1325,7 @@ where
     -- Never export private decls from theorem bodies to make sure they stay irrelevant for rebuilds
     withOptions (fun opts =>
       if headers.any (·.kind.isTheorem) then ResolveName.backward.privateInPublic.set opts false else opts) do
+    withSaveInfoContext do  -- save adjusted env in info tree
     let headers := headers.map fun header =>
       { header with modifiers.attrs := header.modifiers.attrs.filter (!·.name ∈ [`expose, `no_expose]) }
     let values ← try

src/Lean/Server/CodeActions/UnknownIdentifier.lean

@@ -92,7 +92,7 @@ structure Insertion where
   edit     : TextEdit
 
 structure Query extends LeanModuleQuery where
-  env                : Environment
+  ctx                : Elab.ContextInfo
   determineInsertion : Name → Insertion
 
 partial def collectOpenNamespaces (currentNamespace : Name) (openDecls : List OpenDecl)
@@ -121,7 +121,7 @@ def computeIdQuery?
   return {
     identifier := id.toString
     openNamespaces := collectOpenNamespaces ctx.currNamespace ctx.openDecls
-    env := ctx.env
+    ctx
     determineInsertion decl :=
       let minimizedId := minimizeGlobalIdentifierInContext ctx.currNamespace ctx.openDecls decl
       {
@@ -155,7 +155,7 @@ def computeDotQuery?
   return some {
     identifier := String.Pos.Raw.extract text.source pos tailPos
     openNamespaces := typeNames.map (.allExcept · #[])
-    env := ctx.env
+    ctx
     determineInsertion decl :=
       {
         fullName := decl
@@ -186,7 +186,7 @@ def computeDotIdQuery?
   return some {
     identifier := id.toString
     openNamespaces := typeNames.map (.allExcept · #[])
-    env := ctx.env
+    ctx
     determineInsertion decl :=
       {
         fullName := decl
@@ -237,6 +237,22 @@ def importAllUnknownIdentifiersCodeAction (params : CodeActionParams) (kind : St
   }
 }
 
+private def mkImportText (ctx : Elab.ContextInfo) (mod : Name) :
+    String := Id.run do
+  let mut text := s!"import {mod}\n"
+  if let some parentDecl := ctx.parentDecl? then
+    if isMarkedMeta ctx.env parentDecl then
+      text := s!"meta {text}"
+      if !isPrivateName parentDecl then
+        -- As `meta` declarations go through a second, stricter visibility check in the compiler,
+        -- we should add `public` anywhere in a public definition (technically even private defs
+        -- could require public imports but that is not something we can check for here).
+        text := s!"public {text}"
+    else if ctx.env.isExporting then
+      -- Outside `meta`, add `public` only from public scope
+      text := s!"public {text}"
+  text
+
 def handleUnknownIdentifierCodeAction
     (id             : JsonRpc.RequestID)
     (params         : CodeActionParams)
@@ -276,8 +292,8 @@ def handleUnknownIdentifierCodeAction
     | return #[]
   for query in queries, result in response.queryResults do
     for ⟨mod, decl, isExactMatch⟩ in result do
-      let isDeclInEnv := query.env.contains decl
-      if ! isDeclInEnv && mod == query.env.mainModule then
+      let isDeclInEnv := query.ctx.env.contains decl
+      if ! isDeclInEnv && mod == query.ctx.env.mainModule then
         -- Don't offer any code actions for identifiers defined further down in the same file
         continue
       let insertion := query.determineInsertion decl
@@ -290,7 +306,7 @@ def handleUnknownIdentifierCodeAction
             edits := #[
               {
                 range := importInsertionRange
-                newText := s!"import {mod}\n"
+                newText := mkImportText query.ctx mod
               },
               insertion.edit
             ]
@@ -344,15 +360,15 @@ def handleResolveImportAllUnknownIdentifiersCodeAction?
   let mut imports : Std.HashSet Name := ∅
   for q in queries, result in response.queryResults do
     let some ⟨mod, decl, _⟩ := result.find? fun id =>
-        id.isExactMatch && ! q.env.contains id.decl
+        id.isExactMatch && ! q.ctx.env.contains id.decl
       | continue
-    if mod == q.env.mainModule then
+    if mod == q.ctx.env.mainModule then
       continue
     let insertion := q.determineInsertion decl
     if ! imports.contains mod then
       edits := edits.push {
         range := importInsertionRange
-        newText := s!"import {mod}\n"
+        newText := mkImportText q.ctx mod
       }
     edits := edits.push insertion.edit
     imports := imports.insert mod

src/Lean/Server/References.lean

@@ -86,19 +86,17 @@ def toLspRefInfo (i : RefInfo) : StateT Decls BaseIO Lsp.RefInfo := do
   let refToRefInfoLocation (ref : Reference) : StateT Decls BaseIO RefInfo.Location := do
     let parentDeclName? := ref.ci.parentDecl?
     let parentDeclNameString? := parentDeclName?.map (·.toString)
-    let .ok parentDeclInfo? ← 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.
-        let some parentDeclRanges := declRangeExt.find? (asyncMode := .local) (← getEnv) parentDeclName
-          | return none
-        return some <| .ofDeclarationRanges parentDeclRanges
-      -- we only use `CoreM` to get access to a `MonadEnv`, but these are currently all `IO`
-      | unreachable!
+    let parentDeclInfo? : Option DeclInfo := 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.
+      let parentDeclRanges ← declRangeExt.find? (asyncMode := .local) cmdEnv parentDeclName
+      return .ofDeclarationRanges parentDeclRanges
     if let some parentDeclNameString := parentDeclNameString? then
       if let some parentDeclInfo := parentDeclInfo? then
         modify (·.insert parentDeclNameString parentDeclInfo)

src/Lean/Server/Test/Runner.lean

@@ -692,9 +692,9 @@ partial def main (args : List String) : IO Unit := do
   let isProject := args[0]?.any (· == "-p")
   let (ipcCmd, ipcArgs) :=
     if isProject then
-      ("lake", #["serve", "--", "-DstderrAsMessages=false"])
+      ("lake", #["serve", "--", "-DstderrAsMessages=false", "-Dexperimental.module=true"])
     else
-      ("lean", #["--server", "-DstderrAsMessages=false"])
+      ("lean", #["--server", "-DstderrAsMessages=false", "-Dexperimental.module=true"])
   let path := if args.size == 1 then args[0]! else args[1]!
   let uri := s!"file:///{path}"
   -- We want `dbg_trace` tactics to write directly to stderr instead of being caught in reuse

tests/lean/interactive/unknownIdentifierCodeActions.lean

@@ -1,4 +1,8 @@
+module
 --^ waitForILeans
+
+public section
+
 #check Lean.Server.Test.Refs.test1
                                 --^ codeAction
 
@@ -48,3 +52,13 @@ example (f : Foobar) : Nat := f.veryLongAndHopefullyVeryUniqueBar
 
 example : Foobar := .veryLongAndHopefullyVeryUniqueFoobar
                                                        --^ codeAction
+
+def pubNonExposed : Lean.Server.Test.Refs.Test1
+                                             --^ codeAction
+  := Lean.Server.Test.Refs.Test1
+                              --^ codeAction
+
+public meta def pubMeta : Lean.Server.Test.Refs.Test1
+                                                   --^ codeAction
+  := Lean.Server.Test.Refs.Test1
+                              --^ codeAction