fix: realizeConst fixes (#7272)

Emerged and fixed while adding more `realizeConst` callers

src/Lean/CoreM.lean

@@ -535,7 +535,9 @@ opaque compileDeclsOld (env : Environment) (opt : @& Options) (decls : @& List N
 -- `ref?` is used for error reporting if available
 partial def compileDecls (decls : List Name) (ref? : Option Declaration := none)
     (logErrors := true) : CoreM Unit := do
-  if !Elab.async.get (← getOptions) then
+  -- When inside `realizeConst`, do compilation synchronously so that `_cstage*` constants are found
+  -- by the replay code
+  if !Elab.async.get (← getOptions) || (← getEnv).isRealizing then
     doCompile
     return
   let env ← getEnv

src/Lean/Environment.lean

@@ -439,11 +439,14 @@ private structure AsyncConsts where
 deriving Inhabited
 
 private def AsyncConsts.add (aconsts : AsyncConsts) (aconst : AsyncConst) : AsyncConsts :=
-  { aconsts with
+  let normalizedName := privateToUserName aconst.constInfo.name
+  if let some aconst' := aconsts.normalizedTrie.find? normalizedName then
+    panic! s!"AsyncConsts.add: duplicate normalized declaration name {aconst.constInfo.name} vs. {aconst'.constInfo.name}"
+  else { aconsts with
     size := aconsts.size + 1
     revList := aconst :: aconsts.revList
     map := aconsts.map.insert aconst.constInfo.name aconst
-    normalizedTrie := aconsts.normalizedTrie.insert (privateToUserName aconst.constInfo.name) aconst
+    normalizedTrie := aconsts.normalizedTrie.insert normalizedName aconst
   }
 
 private def AsyncConsts.find? (aconsts : AsyncConsts) (declName : Name) : Option AsyncConst :=
@@ -451,8 +454,9 @@ private def AsyncConsts.find? (aconsts : AsyncConsts) (declName : Name) : Option
 
 /-- Finds the constant in the collection that is a prefix of `declName`, if any. -/
 private def AsyncConsts.findPrefix? (aconsts : AsyncConsts) (declName : Name) : Option AsyncConst :=
-  -- as macro scopes are a strict suffix,
-  aconsts.normalizedTrie.findLongestPrefix? (privateToUserName declName.eraseMacroScopes)
+  -- as macro scopes are a strict suffix, we do not have to remove them before calling
+  -- `findLongestPrefix?`
+  aconsts.normalizedTrie.findLongestPrefix? (privateToUserName declName)
 
 /-- Context for `realizeConst` established by `enableRealizationsForConst`. -/
 private structure RealizationContext where
@@ -504,8 +508,8 @@ structure Environment where
   /-- Information about this asynchronous branch of the environment, if any. -/
   private asyncCtx?   : Option AsyncContext := none
   /--
-  Realized constants belonging to imported declarations. `none` only from `Environment.ofKernelEnv`,
-  which should never leak into general elaboration.
+  Realized constants belonging to imported declarations. Must be initialized by calling
+  `enableRealizationsForImports`.
   -/
   private realizedImportedConsts? : Option RealizationContext
   /--
@@ -644,6 +648,21 @@ def findConstVal? (env : Environment) (n : Name) : Option ConstantVal := do
     return asyncConst.constInfo.toConstantVal
   else env.findNoAsync n |>.map (·.toConstantVal)
 
+/--
+Allows `realizeConst` calls for imported declarations in all derived environment branches.
+Realizations will run using the given environment and options to ensure deterministic results.
+This function should be called directly after `setMainModule` to ensure that all realized constants
+use consistent private prefixes.
+-/
+def enableRealizationsForImports (env : Environment) (opts : Options) : BaseIO Environment :=
+  return { env with realizedImportedConsts? := some {
+      -- safety: `RealizationContext` is private
+      env := unsafe unsafeCast env
+      opts
+      constsRef := (← IO.mkRef {})
+    }
+  }
+
 /--
 Allows `realizeConst` calls for the given declaration in all derived environment branches.
 Realizations will run using the given environment and options to ensure deterministic results. Note
@@ -893,7 +912,10 @@ def imports (env : Environment) : Array Import :=
 def allImportedModuleNames (env : Environment) : Array Name :=
   env.header.moduleNames
 
-def setMainModule (env : Environment) (m : Name) : Environment :=
+def setMainModule (env : Environment) (m : Name) : Environment := Id.run do
+  if env.realizedImportedConsts?.isSome then
+    panic! "Environment.setMainModule: cannot set after `enableRealizationsForImports`"
+    return env
   env.modifyCheckedAsync ({ · with header.mainModule := m })
 
 def mainModule (env : Environment) : Name :=
@@ -1078,9 +1100,6 @@ def modifyState {σ : Type} (ext : EnvExtension σ) (env : Environment) (f : σ
         {if asyncCtx.realizing then "realization" else "async"} context '{asyncCtx.declPrefix}'"
     return { env with checkedWithoutAsync.extensions := unsafe ext.modifyStateImpl env.checkedWithoutAsync.extensions f }
   | .local =>
-    if let some asyncCtx := env.asyncCtx?.filter (·.realizing) then
-      panic! s!"Environment.modifyState: environment extension is marked as `local` but used in \
-        realization context '{asyncCtx.declPrefix}'"
     return { env with checkedWithoutAsync.extensions := unsafe ext.modifyStateImpl env.checkedWithoutAsync.extensions f }
   | _ =>
     if ext.replay?.isNone then
@@ -1692,14 +1711,6 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
        Safety: There are no concurrent accesses to `env` at this point. -/
     env ← unsafe Runtime.markPersistent env
   env ← finalizePersistentExtensions env s.moduleData opts
-  env := { env with
-    realizedImportedConsts? := some {
-      -- safety: `RealizationContext` is private
-      env := unsafe unsafeCast env
-      opts
-      constsRef := (← IO.mkRef {})
-    }
-  }
   if leakEnv then
     /- Ensure the final environment including environment extension states is
        marked persistent as documented.
@@ -1870,6 +1881,7 @@ def realizeConst (env : Environment) (forConst : Name) (constName : Name)
         -- allow realizations to recursively realize other constants for `forConst`. Do note that
         -- this allows for recursive realization of `constName` itself, which will deadlock.
         realizedLocalConsts := realizeEnv.realizedLocalConsts.insert forConst ctx
+        realizedImportedConsts? := env.realizedImportedConsts?
       }
       -- ensure realized constants are nested below `forConst` and that environment extension
       -- modifications know they are in an async context
@@ -1882,7 +1894,8 @@ def realizeConst (env : Environment) (forConst : Name) (constName : Name)
 
       -- find new constants incl. nested realizations, add current extension state, and compute
       -- closure
-      let consts := realizeEnv'.asyncConsts.revList.take (realizeEnv'.asyncConsts.size - realizeEnv.asyncConsts.size)
+      let numNewConsts := realizeEnv'.asyncConsts.size - realizeEnv.asyncConsts.size
+      let consts := realizeEnv'.asyncConsts.revList.take numNewConsts |>.reverse
       let consts := consts.map fun c =>
         if c.exts?.isNone then
           { c with exts? := some <| .pure realizeEnv'.checkedWithoutAsync.extensions }
@@ -1892,7 +1905,11 @@ def realizeConst (env : Environment) (forConst : Name) (constName : Name)
       prom.resolve (consts, replay, dyn)
       pure (consts, replay, dyn)
     return ({ env with
-      asyncConsts := consts.foldl (·.add) env.asyncConsts
+      asyncConsts := consts.foldl (init := env.asyncConsts) fun consts c =>
+        if consts.find? c.constInfo.name |>.isSome then
+          consts
+        else
+          consts.add c
       checked := env.checked.map replay
     }, dyn)
 where

src/Lean/Language/Lean.lean

@@ -425,6 +425,7 @@ where
         return { diagnostics, result? := none }
 
       let headerEnv := headerEnv.setMainModule setup.mainModuleName
+      let headerEnv ← headerEnv.enableRealizationsForImports setup.opts
       let mut traceState := default
       if trace.profiler.output.get? setup.opts |>.isSome then
         traceState := {

src/Lean/Meta/Basic.lean

@@ -2215,16 +2215,24 @@ private partial def setAllDiagRanges (snap : Language.SnapshotTree) (pos endPos
       task := (← BaseIO.mapTask (t := task.task) (setAllDiagRanges · pos endPos)) })
   }
 
+open Language
+
+private structure RealizeConstantResult where
+  snap       : SnapshotTree
+  error? : Option Exception
+deriving TypeName
+
 /--
 Makes the helper constant `constName` that is derived from `forConst` available in the environment.
 `enableRealizationsForConst forConst` must have been called first on this environment branch. If
 this is the first environment branch requesting `constName` to be realized (atomically), `realize`
 is called with the environment and options at the time of calling `enableRealizationsForConst` if
-`forConst` is from the current module and the state just after importing otherwise, thus helping
-achieve deterministic results despite the non-deterministic choice of which thread is tasked with
-realization. In other words, the state after calling `realizeConst` is *as if* `realize` had been
-called immediately after `enableRealizationsForConst forConst`, though the effects of this call are
-visible only after calling `realizeConst`. See below for more details on the replayed effects.
+`forConst` is from the current module and the state just after importing (when
+`enableRealizationsForImports` should be called) otherwise, thus helping achieve deterministic
+results despite the non-deterministic choice of which thread is tasked with realization. In other
+words, the state after calling `realizeConst` is *as if* `realize` had been called immediately after
+`enableRealizationsForConst forConst`, though the effects of this call are visible only after
+calling `realizeConst`. See below for more details on the replayed effects.
 
 `realizeConst` cannot check what other data is captured in the `realize` closure,
 so it is best practice to extract it into a separate function and pay close attention to the passed
@@ -2241,20 +2249,25 @@ to add `constName` to the environment, an appropriate diagnostic is reported to
 constants are added to the environment.
 -/
 def realizeConst (forConst : Name) (constName : Name) (realize : MetaM Unit) :
-    MetaM Unit := withTraceNode `Meta.realizeConst (fun _ => return constName) do
+    MetaM Unit := do
   let env ← getEnv
-  let coreCtx ← readThe Core.Context
-  -- these fields should be invariant throughout the file
-  let coreCtx := { fileName := coreCtx.fileName, fileMap := coreCtx.fileMap }
-  let (env, dyn) ← env.realizeConst forConst constName (realizeAndReport coreCtx)
-  if let some snap := dyn.get? Language.SnapshotTree then
-    let mut snap := snap
-    -- localize diagnostics
-    if let some range := (← getRef).getRange? then
-      let fileMap ← getFileMap
-      snap ← setAllDiagRanges snap (fileMap.toPosition range.start) (fileMap.toPosition range.stop)
-    Core.logSnapshotTask <| .finished (stx? := none) snap
-  setEnv env
+  if env.contains constName then
+    return
+  withTraceNode `Meta.realizeConst (fun _ => return constName) do
+    let coreCtx ← readThe Core.Context
+    -- these fields should be invariant throughout the file
+    let coreCtx := { fileName := coreCtx.fileName, fileMap := coreCtx.fileMap }
+    let (env, dyn) ← env.realizeConst forConst constName (realizeAndReport coreCtx)
+    if let some res := dyn.get? RealizeConstantResult then
+      let mut snap := res.snap
+      -- localize diagnostics
+      if let some range := (← getRef).getRange? then
+        let fileMap ← getFileMap
+        snap ← setAllDiagRanges snap (fileMap.toPosition range.start) (fileMap.toPosition range.stop)
+      Core.logSnapshotTask <| .finished (stx? := none) snap
+      if let some e := res.error? then
+        throw e
+    setEnv env
 where
   -- similar to `wrapAsyncAsSnapshot` but not sufficiently so to share code
   realizeAndReport (coreCtx : Core.Context) env opts := do
@@ -2267,14 +2280,20 @@ where
           realize
           if !(← getEnv).contains constName then
             throwError "Lean.Meta.realizeConst: {constName} was not added to the environment"
-        catch e : Exception =>
-          logError e.toMessageData
         finally
           addTraceAsMessages
     let res? ← act |>.run' |>.run coreCtx { env } |>.toBaseIO
     match res? with
-    | .ok ((output, ()), st) => pure (st.env, .mk (← Core.mkSnapshot output coreCtx st))
-    | .error _e => unreachable!; pure (env, .mk ({ diagnostics := .empty : Language.SnapshotLeaf}))
+    | .ok ((output, ()), st) => pure (st.env, .mk {
+      snap := (← Core.mkSnapshot output coreCtx st)
+      error? := none
+      : RealizeConstantResult
+    })
+    | .error e => pure (env, .mk {
+      snap := toSnapshotTree { diagnostics := .empty : Language.SnapshotLeaf}
+      error? := some e
+      : RealizeConstantResult
+    })
 
 end Meta
 

src/lake/Lake/Load/Lean/Elab.lean

@@ -58,6 +58,7 @@ def elabConfigFile (pkgDir : FilePath) (lakeOpts : NameMap String)
   let (header, parserState, messages) ← Parser.parseHeader inputCtx
   let (env, messages) ← processHeader header leanOpts inputCtx messages
   let env := env.setMainModule configModuleName
+  let env ← env.enableRealizationsForImports leanOpts
 
   -- Configure extensions
   let env := dirExt.setState env pkgDir

tests/lean/prvCtor.lean

@@ -1,33 +0,0 @@
-import Lean
-
-private def a := 10
-
-#check a
-
-structure Name (x : String) where
-  private mk ::
-  val : String := x
-  deriving Repr
-
-def n1 : Name "hello" := {}
-
-def n2 : Name "hello" := ⟨"hello"⟩
-
-def n3 : Name "hello" := Name.mk "hello"
-
-open Lean in
-#eval id (α := CoreM Unit) do
-  modifyEnv fun env => env.setMainModule `foo -- change module name to test `private`
-
-open Lean in
-#eval id (α := CoreM Unit) do
-  -- this implementation is no longer allowed because of a private constructor
-  modifyEnv fun env => { env with checked.header.mainModule := `foo }
-
-#check a -- Error
-
-def m1 : Name "hello" := {} -- Error
-
-def m2 : Name "hello" := ⟨"hello"⟩ -- Error
-
-def m3 : Name "hello" := Name.mk "hello"

tests/lean/prvCtor.lean.expected.out

@@ -1,12 +0,0 @@
-a : Nat
-prvCtor.lean:25:23-25:69: error: invalid {...} notation, constructor for `Lean.Environment` is marked as private
-prvCtor.lean:27:7-27:8: error: unknown identifier 'a'
-prvCtor.lean:29:25-29:27: error: overloaded, errors 
-  failed to synthesize
-    EmptyCollection (Name "hello")
-  
-  Additional diagnostic information may be available using the `set_option diagnostics true` command.
-  
-  invalid {...} notation, constructor for `Name` is marked as private
-prvCtor.lean:31:25-31:34: error: invalid ⟨...⟩ notation, constructor for `Name` is marked as private
-prvCtor.lean:33:25-33:40: error: unknown constant 'Name.mk'

tests/lean/run/matchEqsBug.lean

@@ -26,12 +26,7 @@ set_option trace.Meta.Match.matchEqs true
 test% f.match_1
 #check f.match_1.splitter
 
-def g.match_1.splitter := 4
-
 test% g.match_1
 #check g.match_1.eq_1
 #check g.match_1.eq_2
 #check g.match_1.splitter
-
-
-def bla.splitter := 5 -- ok

tests/pkg/prv/Prv.lean

@@ -2,4 +2,30 @@ import Prv.Foo
 
 #check { name := "leo", val := 15 : Foo }
 #check { name := "leo", val := 15 : Foo }.name
-#check { name := "leo", val := 15 : Foo }.val -- Error
+/-- error: field 'val' from structure 'Foo' is private -/
+#guard_msgs in
+#check { name := "leo", val := 15 : Foo }.val
+
+/-- error: unknown identifier 'a' -/
+#guard_msgs in
+#check a
+
+/--
+error: overloaded, errors ⏎
+  failed to synthesize
+    EmptyCollection (Name "hello")
+  ⏎
+  Additional diagnostic information may be available using the `set_option diagnostics true` command.
+  ⏎
+  invalid {...} notation, constructor for `Name` is marked as private
+-/
+#guard_msgs in
+def m1 : Name "hello" := {}
+
+/-- error: invalid ⟨...⟩ notation, constructor for `Name` is marked as private -/
+#guard_msgs in
+def m2 : Name "hello" := ⟨"hello"⟩
+
+/-- error: unknown constant 'Name.mk' -/
+#guard_msgs in
+def m3 : Name "hello" := Name.mk "hello"

tests/pkg/prv/Prv/Foo.lean

@@ -1,6 +1,13 @@
+private def a := 10
+
 structure Foo where
   private val : Nat
   name : String
 
 #check { name := "leo", val := 15 : Foo }
 #check { name := "leo", val := 15 : Foo }.val
+
+structure Name (x : String) where
+  private mk ::
+  val : String := x
+  deriving Repr

tests/pkg/prv/lean-toolchain

@@ -0,0 +1 @@
+lean4

tests/pkg/prv/test.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 
 rm -rf .lake/build
-lake build 2>&1 | grep 'error: .*: field.*private'
+lake build