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,21 +39,13 @@ public structure TrivialStructureInfo where
   deriving Inhabited, Repr
 
 /--
-Return `some fieldIdx` if `declName` is the name of an inductive datatype s.t.
+Computes `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.
 -/
-public def Irrelevant.hasTrivialStructure?
-    (cacheExt : CacheExtension Name (Option TrivialStructureInfo))
+def Irrelevant.computeHasTrivialStructure?
     (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 info.isUnsafe || info.isRec then return none
@@ -67,5 +60,22 @@ where fillCache : CoreM (Option TrivialStructureInfo) := do
       result := some { ctorName, fieldIdx := i, numParams := info.numParams }
   return result
 
+/-- Eagerly computes and persists the trivial-structure info for `declName`; see `compileDecls`. -/
+public def Irrelevant.setHasTrivialStructure?
+    (infoExt : MapDeclarationExtension (Option TrivialStructureInfo))
+    (trivialType : Expr → MetaM Bool) (declName : Name) : CoreM Unit := do
+  unless (infoExt.find? (← getEnv) declName).isSome do
+    modifyEnv (infoExt.insert · declName (← computeHasTrivialStructure? trivialType declName))
+
+/-- Trivial-structure info for `declName` (`none` for non-inductives); requires `compileDecls` to
+have been run for inductive `declName`. -/
+public def Irrelevant.hasTrivialStructure?
+    (infoExt : MapDeclarationExtension (Option TrivialStructureInfo))
+    (declName : Name) : CoreM (Option TrivialStructureInfo) := do
+  if isRuntimeBuiltinType declName then return none
+  let .inductInfo _ ← getConstInfo declName | return none
+  let some info? := infoExt.find? (← getEnv) declName
+    | throwError "`{declName}` was not compiled; `compileDecls` must run on inductive types first"
+  return info?
 
 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,24 @@ 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)
+
+/-- Eagerly computes and persists the trivial-structure info of `declName`; see `compileDecls`. -/
+def setHasTrivialStructure? (declName : Name) : CoreM Unit :=
+  Irrelevant.setHasTrivialStructure? trivialStructureInfoExt
+    (fun type => Meta.isProp type <||> Meta.isTypeFormerType type) declName
 
 /--
 Return `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.
+
+Requires `compileDecls` to have been run for inductive `declName`.
 -/
-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 declName
 
 def getParamTypes (type : Expr) : Array Expr :=
   go type #[]
@@ -88,14 +94,21 @@ 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)
+
+/-- Eagerly computes and persists the mono type of `declName`; see `compileDecls`. -/
+def setOtherDeclMonoType (declName : Name) : CoreM Unit := do
+  unless (monoTypeExt.find? (← getEnv) declName).isSome do
+    modifyEnv (monoTypeExt.insert · declName (← toMonoType (← getOtherDeclBaseType declName [])))
 
+/--
+Returns the LCNF mono-phase type of `declName`, a declaration without associated code (constructor,
+inductive type, or foreign function). Requires `compileDecls` to have been run for it.
+-/
 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
+  let some type := monoTypeExt.find? (← getEnv) declName
+    | throwError "`{declName}` was not compiled; `compileDecls` must run on inductive types first"
+  return type
 
 end Lean.Compiler.LCNF

src/Lean/Compiler/LCNF/ToImpureType.lean

@@ -24,71 +24,105 @@ 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
+
+/-- Eagerly computes and persists the impure trivial-structure info of `declName`; see `compileDecls`. -/
+public def setHasTrivialImpureStructure? (declName : Name) : CoreM Unit :=
+  Irrelevant.setHasTrivialStructure? impureTrivialStructureInfoExt impureIrrelevantType declName
+
+public def hasTrivialImpureStructure? (declName : Name) : CoreM (Option TrivialStructureInfo) :=
+  Irrelevant.hasTrivialStructure? impureTrivialStructureInfoExt 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
 
+/-- Eagerly computes and persists the IR type of inductive `name`; see `compileDecls`. -/
+public def setImpureType (name : Name) : CoreM Unit := do
+  if (builtinImpureType? name).isSome then return
+  unless (impureTypeExt.find? (← getEnv) name).isSome do
+    modifyEnv (impureTypeExt.insert · name (← computeImpureType name))
+
+/--
+Returns the IR (impure) type representation of `name`. Requires `compileDecls` to have been run for
+inductive type `name`.
+-/
 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
+  let some type := impureTypeExt.find? (← getEnv) name
+    | throwError "`{name}` was not compiled; `compileDecls` must run on inductive types first"
+  return type
 
 def isAnyProducingType (type : Expr) : Bool :=
   match type with
@@ -149,16 +183,13 @@ 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)
 
-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
+/-- Eagerly computes and persists the layout of constructor `ctorName`; see `compileDecls`. -/
+public def setCtorLayout (ctorName : Name) : CoreM Unit := do
+  unless (ctorLayoutExt.find? (← getEnv) ctorName).isSome do
+    modifyEnv (ctorLayoutExt.insert · ctorName (← fillCache))
 where fillCache := do
   let .some (.ctorInfo ctorInfo) := (← getEnv).find? ctorName | unreachable!
   Meta.MetaM.run' <| Meta.forallTelescopeReducing ctorInfo.type fun params _ => do
@@ -240,4 +271,33 @@ where fillCache := do
       fieldInfo := fields
     }
 
+/--
+Returns the runtime layout of constructor `ctorName`. Requires `compileDecls` to have been run for
+its inductive type.
+-/
+public def getCtorLayout (ctorName : Name) : CoreM CtorLayout := do
+  let some info := ctorLayoutExt.find? (← getEnv) ctorName
+    | throwError "`{ctorName}` was not compiled; `compileDecls` must run on inductive types first"
+  return info
+
+/--
+Eagerly computes and persists all cross-module compiler caches for the inductive types `typeNames`
+(and their constructors) in their defining module; run from `compileDecls`.
+-/
+public def compileInductives (typeNames : Array Name) : CoreM Unit := do
+  let inductiveNames ← typeNames.filterM fun n => return (← getEnv).find? n matches some (.inductInfo _)
+  -- The readers are strict, so we fill in dependency phases across the whole (possibly mutual)
+  -- block: each phase only reads caches filled by an earlier phase or an imported module.
+  for typeName in inductiveNames do
+    setHasTrivialStructure? typeName
+    setHasTrivialImpureStructure? typeName
+  for typeName in inductiveNames do  -- reads the trivial-structure info above
+    setOtherDeclMonoType typeName
+    setImpureType typeName
+  for typeName in inductiveNames do  -- reads the type-level info above
+    let .inductInfo iv ← getConstInfo typeName | unreachable!
+    for ctorName in iv.ctors do
+      setOtherDeclMonoType ctorName
+      setCtorLayout 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