refactor: move inductive information in compiler from lazily filled local caches to eager persistent environment extensions

src/Lean/Compiler/LCNF/Irrelevant.lean

@@ -6,6 +6,7 @@ Authors: Henrik Böving
 module
 
 prelude
+public import Lean.EnvExtension
 public import Lean.Compiler.LCNF.CompilerM
 import Lean.Compiler.LCNF.BaseTypes
 import Lean.Compiler.LCNF.Util
@@ -38,23 +39,26 @@ public structure TrivialStructureInfo where
   deriving Inhabited, Repr
 
 /--
-Return `some fieldIdx` if `declName` is the name of an inductive datatype s.t.
+Returns `some fieldIdx` if `declName` is the name of an inductive datatype s.t.
 - It does not have builtin support in the runtime.
 - It has only one constructor.
 - This constructor has only one computationally relevant field.
+
+For declarations of the current module, the result is persisted in `infoExt` on first call. The
+information is computed eagerly for inductives by `compileInductives` in their defining module
+(needed because the computation walks constructor field types, which may be private).
 -/
 public def Irrelevant.hasTrivialStructure?
-    (cacheExt : CacheExtension Name (Option TrivialStructureInfo))
+    (infoExt : MapDeclarationExtension (Option TrivialStructureInfo))
     (trivialType : Expr → MetaM Bool) (declName : Name) : CoreM (Option TrivialStructureInfo) := do
-  match (← cacheExt.find? declName) with
-  | some info? => return info?
-  | none =>
-    let info? ← fillCache
-    cacheExt.insert declName info?
-    return info?
-where fillCache : CoreM (Option TrivialStructureInfo) := do
   if isRuntimeBuiltinType declName then return none
   let .inductInfo info ← getConstInfo declName | return none
+  if let some result := infoExt.find? (← getEnv) declName then return result
+  let result ← compute info
+  if (← getEnv).getModuleIdxFor? declName |>.isNone then
+    modifyEnv (infoExt.insert · declName result)
+  return result
+where compute (info : InductiveVal) : CoreM (Option TrivialStructureInfo) := do
   if info.isUnsafe || info.isRec then return none
   let [ctorName] := info.ctors | return none
   let ctorType ← getOtherDeclBaseType ctorName []
@@ -67,5 +71,4 @@ where fillCache : CoreM (Option TrivialStructureInfo) := do
       result := some { ctorName, fieldIdx := i, numParams := info.numParams }
   return result
 
-
 end Lean.Compiler.LCNF

src/Lean/Compiler/LCNF/Main.lean

@@ -9,6 +9,7 @@ import Lean.Compiler.Options
 import Lean.Compiler.IR
 import Lean.Compiler.LCNF.Passes
 import Lean.Compiler.LCNF.ToDecl
+import Lean.Compiler.LCNF.ToImpureType
 import Lean.Compiler.LCNF.Check
 import Lean.Meta.Match.MatcherInfo
 import Lean.Compiler.LCNF.SplitSCC
@@ -125,6 +126,17 @@ partial def run (declNames : Array Name) (baseOpts : Options) : CompilerM Unit :
   and it often creates a very deep recursion.
   Moreover, some declarations get very big during simplification.
   -/
+
+  if (← declNames.anyM isInductive) then
+    -- Eagerly compute and persist the cross-module inductive infos for these inductive types in
+    -- their defining module. The computation walks each constructor's field types, which can
+    -- reference constants from non-transitively (privately) imported modules; the defining module is
+    -- the only place where that walk is guaranteed to succeed, so consumers rely on the persisted
+    -- result. This could be postponed as well but the added complexity is likely not worth the
+    -- slight gain in rebuild avoidance/parallelism.
+    compileInductives declNames
+    return
+
   for declName in declNames do
     if let some fnName := Compiler.getImplementedBy? (← getEnv) declName then
       if !isDeclPublic (← getEnv) fnName then

src/Lean/Compiler/LCNF/MonoTypes.lean

@@ -14,18 +14,18 @@ public section
 
 namespace Lean.Compiler.LCNF
 
-private builtin_initialize trivialStructureInfoExt : CacheExtension Name (Option TrivialStructureInfo) ←
-  CacheExtension.register
+private builtin_initialize trivialStructureInfoExt : MapDeclarationExtension (Option TrivialStructureInfo) ←
+  mkMapDeclarationExtension (asyncMode := .sync)
 
 /--
-Return `some fieldIdx` if `declName` is the name of an inductive datatype s.t.
+Returns `some fieldIdx` if `declName` is the name of an inductive datatype s.t.
 - It does not have builtin support in the runtime.
 - It has only one constructor.
 - This constructor has only one computationally relevant field.
 -/
-def hasTrivialStructure? (declName : Name) : CoreM (Option TrivialStructureInfo) := do
-  let irrelevantType type := Meta.isProp type <||> Meta.isTypeFormerType type
-  Irrelevant.hasTrivialStructure? trivialStructureInfoExt irrelevantType declName
+def hasTrivialStructure? (declName : Name) : CoreM (Option TrivialStructureInfo) :=
+  Irrelevant.hasTrivialStructure? trivialStructureInfoExt
+    (fun type => Meta.isProp type <||> Meta.isTypeFormerType type) declName
 
 def getParamTypes (type : Expr) : Array Expr :=
   go type #[]
@@ -88,14 +88,19 @@ State for the environment extension used to save the LCNF mono phase type for de
 that do not have code associated with them.
 Example: constructors, inductive types, foreign functions.
 -/
-builtin_initialize monoTypeExt : CacheExtension Name Expr ← CacheExtension.register
+builtin_initialize monoTypeExt : MapDeclarationExtension Expr ←
+  mkMapDeclarationExtension (asyncMode := .sync)
 
+/--
+Returns the LCNF mono-phase type of `declName`, a declaration without associated code (constructor,
+inductive type, foreign function, or noncomputable definition). For declarations of the current
+module, the result is persisted in `monoTypeExt` on first call.
+-/
 def getOtherDeclMonoType (declName : Name) : CoreM Expr := do
-  match (← monoTypeExt.find? declName) with
-  | some type => return type
-  | none =>
-    let type ← toMonoType (← getOtherDeclBaseType declName [])
-    monoTypeExt.insert declName type
-    return type
+  if let some type := monoTypeExt.find? (← getEnv) declName then return type
+  let type ← toMonoType (← getOtherDeclBaseType declName [])
+  if (← getEnv).getModuleIdxFor? declName |>.isNone then
+    modifyEnv (monoTypeExt.insert · declName type)
+  return type
 
 end Lean.Compiler.LCNF

src/Lean/Compiler/LCNF/ToImpureType.lean

@@ -24,71 +24,97 @@ def impureTypeForEnum (numCtors : Nat) : Expr :=
   else
     ImpureType.tagged
 
-builtin_initialize impureTypeExt : CacheExtension Name Expr ←
-  CacheExtension.register
+/--
+Maps each inductive type to its IR (impure) type representation. Populated eagerly in the type's
+defining module by `setImpureType` (driven from `compileDecls`); consumers read it via
+`nameToImpureType`. Persisted across modules because computing an inductive's impure type walks each
+constructor's field types, which can reference constants from non-transitively (privately) imported
+modules.
+-/
+builtin_initialize impureTypeExt : MapDeclarationExtension Expr ←
+  mkMapDeclarationExtension (asyncMode := .sync)
 
 builtin_initialize impureTrivialStructureInfoExt :
-    CacheExtension Name (Option TrivialStructureInfo) ←
-  CacheExtension.register
+    MapDeclarationExtension (Option TrivialStructureInfo) ←
+  mkMapDeclarationExtension (asyncMode := .sync)
 
 /--
 The idea of this function is the same as in `ToMono`, however the notion of "irrelevancy" has
 changed because we now have the `void` type which can only be erased in impure context and thus at
 earliest at the conversion from mono to impure.
 -/
-public def hasTrivialImpureStructure? (declName : Name) : CoreM (Option TrivialStructureInfo) := do
+def impureIrrelevantType (type : Expr) : MetaM Bool := do
   let isVoidType type := do
     let type ← Meta.whnfD type
     return type matches .proj ``Subtype 0 (.app (.const ``Void.nonemptyType []) _)
-  let irrelevantType type :=
-    Meta.isProp type <||> Meta.isTypeFormerType type <||> isVoidType type
-  Irrelevant.hasTrivialStructure? impureTrivialStructureInfoExt irrelevantType declName
+  Meta.isProp type <||> Meta.isTypeFormerType type <||> isVoidType type
+
+public def hasTrivialImpureStructure? (declName : Name) : CoreM (Option TrivialStructureInfo) :=
+  Irrelevant.hasTrivialStructure? impureTrivialStructureInfoExt impureIrrelevantType declName
 
+/--
+IR representations that are fixed independently of the environment: builtin scalar types and the
+compiler's pseudo-constants (`lcErased`/`lcVoid`, which are not inductives). These never need to be
+persisted in `impureTypeExt`.
+-/
+def builtinImpureType? : Name → Option Expr
+  | ``UInt8 => some ImpureType.uint8
+  | ``UInt16 => some ImpureType.uint16
+  | ``UInt32 => some ImpureType.uint32
+  | ``UInt64 => some ImpureType.uint64
+  | ``USize => some ImpureType.usize
+  | ``Float => some ImpureType.float
+  | ``Float32 => some ImpureType.float32
+  | ``lcErased => some ImpureType.erased
+  -- `Int` is specified as an inductive type with two constructors that have relevant arguments,
+  -- but it has the same runtime representation as `Nat` and thus needs to be special-cased here.
+  | ``Int => some ImpureType.tobject
+  | ``lcVoid => some ImpureType.void
+  | _ => none
+
+/--
+Computes the IR (impure) type of `name` from scratch by inspecting its constructors. For inductives
+this walks the constructor field types, so it must run in the defining module, where those types
+(including private ones) are accessible.
+-/
+def computeImpureType (name : Name) : CoreM Expr := do
+  if let some type := builtinImpureType? name then return type
+  let env ← Lean.getEnv
+  let some (.inductInfo inductiveVal) := env.find? name | return ImpureType.tobject
+  let ctorNames := inductiveVal.ctors
+  let numCtors := ctorNames.length
+  let mut numScalarCtors := 0
+  for ctorName in ctorNames do
+    let some (.ctorInfo ctorInfo) := env.find? ctorName | unreachable!
+    let hasRelevantField ← Meta.MetaM.run' <|
+                           Meta.forallTelescope ctorInfo.type fun params _ => do
+      for field in params[ctorInfo.numParams...*] do
+        let fieldType ← field.fvarId!.getType
+        let lcnfFieldType ← toLCNFType fieldType
+        let monoFieldType ← toMonoType lcnfFieldType
+        if !monoFieldType.isErased then return true
+      return false
+    if !hasRelevantField then
+      numScalarCtors := numScalarCtors + 1
+  if numScalarCtors == numCtors then
+    return impureTypeForEnum numCtors
+  else if numScalarCtors == 0 then
+    return ImpureType.object
+  else
+    return ImpureType.tobject
+
+/--
+Returns the IR (impure) type representation of `name`. For inductive types of the current module,
+the result is persisted in `impureTypeExt` on first call.
+-/
 public def nameToImpureType (name : Name) : CoreM Expr := do
-  match (← impureTypeExt.find? name) with
-  | some type => return type
-  | none =>
-    let type ← fillCache
-    impureTypeExt.insert name type
-    return type
-where fillCache : CoreM Expr := do
-    match name with
-    | ``UInt8 => return ImpureType.uint8
-    | ``UInt16 => return ImpureType.uint16
-    | ``UInt32 => return ImpureType.uint32
-    | ``UInt64 => return ImpureType.uint64
-    | ``USize => return ImpureType.usize
-    | ``Float => return ImpureType.float
-    | ``Float32 => return ImpureType.float32
-    | ``lcErased => return ImpureType.erased
-    -- `Int` is specified as an inductive type with two constructors that have relevant arguments,
-    -- but it has the same runtime representation as `Nat` and thus needs to be special-cased here.
-    | ``Int => return ImpureType.tobject
-    | ``lcVoid => return ImpureType.void
-    | _ =>
-      let env ← Lean.getEnv
-      let some (.inductInfo inductiveVal) := env.find? name | return ImpureType.tobject
-      let ctorNames := inductiveVal.ctors
-      let numCtors := ctorNames.length
-      let mut numScalarCtors := 0
-      for ctorName in ctorNames do
-        let some (.ctorInfo ctorInfo) := env.find? ctorName | unreachable!
-        let hasRelevantField ← Meta.MetaM.run' <|
-                               Meta.forallTelescope ctorInfo.type fun params _ => do
-          for field in params[ctorInfo.numParams...*] do
-            let fieldType ← field.fvarId!.getType
-            let lcnfFieldType ← toLCNFType fieldType
-            let monoFieldType ← toMonoType lcnfFieldType
-            if !monoFieldType.isErased then return true
-          return false
-        if !hasRelevantField then
-          numScalarCtors := numScalarCtors + 1
-      if numScalarCtors == numCtors then
-        return impureTypeForEnum numCtors
-      else if numScalarCtors == 0 then
-        return ImpureType.object
-      else
-        return ImpureType.tobject
+  if let some type := builtinImpureType? name then return type
+  let some (.inductInfo _) := (← getEnv).find? name | return ImpureType.tobject
+  if let some type := impureTypeExt.find? (← getEnv) name then return type
+  let type ← computeImpureType name
+  if (← getEnv).getModuleIdxFor? name |>.isNone then
+    modifyEnv (impureTypeExt.insert · name type)
+  return type
 
 def isAnyProducingType (type : Expr) : Bool :=
   match type with
@@ -149,16 +175,19 @@ public structure CtorLayout where
   fieldInfo : Array CtorFieldInfo
   deriving Inhabited
 
-builtin_initialize ctorLayoutExt : CacheExtension Name CtorLayout ←
-  CacheExtension.register
+builtin_initialize ctorLayoutExt : MapDeclarationExtension CtorLayout ←
+  mkMapDeclarationExtension (asyncMode := .sync)
 
+/--
+Returns the runtime layout of constructor `ctorName`. For constructors of the current module, the
+result is persisted in `ctorLayoutExt` on first call.
+-/
 public def getCtorLayout (ctorName : Name) : CoreM CtorLayout := do
-  match (← ctorLayoutExt.find? ctorName) with
-  | some info => return info
-  | none =>
-    let info ← fillCache
-    ctorLayoutExt.insert ctorName info
-    return info
+  if let some info := ctorLayoutExt.find? (← getEnv) ctorName then return info
+  let info ← fillCache
+  if (← getEnv).getModuleIdxFor? ctorName |>.isNone then
+    modifyEnv (ctorLayoutExt.insert · ctorName info)
+  return info
 where fillCache := do
   let .some (.ctorInfo ctorInfo) := (← getEnv).find? ctorName | unreachable!
   Meta.MetaM.run' <| Meta.forallTelescopeReducing ctorInfo.type fun params _ => do
@@ -240,4 +269,21 @@ where fillCache := do
       fieldInfo := fields
     }
 
+/--
+Eagerly computes and persists all cross-module compiler caches for the inductive types `typeNames`
+(and their constructors) in their defining module; run from `compileDecls`. The readers compute on
+miss, so mutually-recursive inductives resolve their cross-references without explicit phasing.
+-/
+public def compileInductives (typeNames : Array Name) : CoreM Unit := do
+  let inductiveNames ← typeNames.filterM fun n => return (← getEnv).find? n matches some (.inductInfo _)
+  for typeName in inductiveNames do
+    discard <| hasTrivialStructure? typeName
+    discard <| hasTrivialImpureStructure? typeName
+    discard <| getOtherDeclMonoType typeName
+    discard <| nameToImpureType typeName
+    let .inductInfo iv ← getConstInfo typeName | unreachable!
+    for ctorName in iv.ctors do
+      discard <| getOtherDeclMonoType ctorName
+      discard <| getCtorLayout ctorName
+
 end Lean.Compiler.LCNF

src/Lean/Elab/ComputedFields.lean

@@ -104,16 +104,18 @@ def validateComputedFields : M Unit := do
     if indices.any (ty.containsFVar ·.fvarId!) then
       throwError "computed field {cf}'s type must not depend on indices{indentExpr ty}"
 
-def mkImplType : M Unit := do
+def mkImplType : M Name := do
   let {name, isUnsafe, type, ctors, levelParams, numParams, lparams, params, compFieldVars, ..} ← read
+  let name := name ++ `_impl
   addDecl <| .inductDecl levelParams numParams
     (isUnsafe := isUnsafe) -- Note: inlining is disabled with unsafe inductives
-    [{ name := name ++ `_impl, type,
+    [{ name, type,
        ctors := ← ctors.mapM fun ctor => do
          forallTelescope (← inferType (mkAppN (mkConst ctor lparams) params)) fun fields retTy => do
-           let retTy := mkAppN (mkConst (name ++ `_impl) lparams) retTy.getAppArgs
+           let retTy := mkAppN (mkConst name lparams) retTy.getAppArgs
            let type ← mkForallFVars (params ++ (if ← isScalarField ctor then #[] else compFieldVars) ++ fields) retTy
            return { name := ctor ++ `_impl, type } }]
+  return name
 
 def overrideCasesOn : M Unit := do
   let {name, numIndices, ctors, lparams, params, compFieldVars, ..} ← read
@@ -218,7 +220,8 @@ def mkComputedFieldOverrides (declName : Name) (compFields : Array Name) : MetaM
       let ctx := { ind with lparams, params, compFields, compFieldVars, indices, val }
       ReaderT.run (r := ctx) do
         validateComputedFields
-        mkImplType
+        let ty ← mkImplType
+        Lean.compileDecls #[ty]
         overrideCasesOn
         overrideConstructors
         overrideComputedFields

src/Lean/Elab/MutualInductive.lean

@@ -1566,6 +1566,8 @@ private def elabInductiveViews (vars : Array Expr) (elabs : Array InductiveElabS
   Term.withDeclName view0.declName do withRef ref do
   withExporting (isExporting := !isPrivateName view0.declName) do
     let res ← mkInductiveDecl vars elabs
+    -- Must happen *before* `sizeOf` etc are compiled below
+    Lean.compileDecls (elabs.map (·.view.declName))
     -- This might be too coarse, consider reconsidering on construction-by-construction basis
     withoutExporting (when := view0.ctors.any (isPrivateName ·.declName)) do
       mkAuxConstructions (elabs.map (·.view.declName))
@@ -1587,6 +1589,7 @@ private def elabFlatInductiveViews (vars : Array Expr) (elabs : Array InductiveE
   withExporting (isExporting := !isPrivateName view0.declName) do
     withElaboratedHeaders vars elabs <| mkInductiveDeclCore (fun context =>
      mkFlatInductive context.views context.indFVars context.levelParams context.numVars context.numParams context.indTypes)
+    Lean.compileDecls (elabs.map (·.view.declName))
     -- This might be too coarse, consider reconsidering on construction-by-construction basis
     withoutExporting (when := view0.ctors.any (isPrivateName ·.declName)) do
       mkAuxConstructions (elabs.map (·.view.declName))

stage0/src/stdlib_flags.h

@@ -1,5 +1,7 @@
 #include "util/options.h"
 
+// Dear CI, please update stage0
+
 namespace lean {
 options get_default_options() {
     options opts;