fix: inline/specialize may only refer to publicly imported decls for now (#10494)

This PR resolves a potential bad interaction between the compiler and
the module system where references to declarations not imported are
brought into scope by inlining or specializing. We now proactively check
that declarations to be inlined/specialized only reference public
imports. The intention is to later resolve this limitation by moving out
compilation into a separate build step with its own import/incremental
system.

Author: Kha

src/Init/Data/Array/Lex/Basic.lean

@@ -9,8 +9,8 @@ prelude
 public import Init.Core
 import Init.Data.Array.Basic
 import Init.Data.Nat.Lemmas
-import Init.Data.Range.Polymorphic.Iterators
-import Init.Data.Range.Polymorphic.Nat
+public import Init.Data.Range.Polymorphic.Iterators
+public import Init.Data.Range.Polymorphic.Nat
 import Init.Data.Iterators.Consumers
 
 public section

src/Init/Data/Array/Subarray.lean

@@ -7,6 +7,7 @@ module
 
 prelude
 public import Init.GetElem
+public import Init.Data.Array.Basic
 import Init.Data.Array.GetLit
 public import Init.Data.Slice.Basic
 

src/Init/Data/ToString/Name.lean

@@ -7,7 +7,7 @@ module
 
 prelude
 public import Init.Meta
-import Init.Data.String.Extra
+public import Init.Data.String.Extra
 
 /-!
 Here we give the. implementation of `Name.toString`. There is also a private implementation in

src/Lean/Compiler/LCNF/Main.lean

@@ -112,7 +112,7 @@ def run (declNames : Array Name) : CompilerM (Array IR.Decl) := withAtLeastMaxRe
   -- Check meta accesses now before optimizations may obscure references. This check should stay in
   -- `lean` if some compilation is moved out.
   for decl in decls do
-    checkMeta (isMeta (← getEnv) decl.name) decl
+    checkMeta decl
   let decls := markRecDecls decls
   let manager ← getPassManager
   let isCheckEnabled := compiler.check.get (← getOptions)

src/Lean/Compiler/LCNF/Passes.lean

@@ -85,6 +85,7 @@ def builtinPassManager : PassManager := {
     -/
     simp { etaPoly := true, inlinePartial := true, implementedBy := true } (occurrence := 1),
     eagerLambdaLifting,
+    checkTemplateVisibility,
     specialize,
     simp (occurrence := 2),
     cse (shouldElimFunDecls := false) (occurrence := 1),

src/Lean/Compiler/LCNF/Visibility.lean

@@ -57,20 +57,21 @@ partial def markDeclPublicRec (phase : Phase) (decl : Decl) : CompilerM Unit :=
             markDeclPublicRec phase refDecl
 
 /-- Checks whether references in the given declaration adhere to phase distinction. -/
-partial def checkMeta (isMeta : Bool) (origDecl : Decl) : CompilerM Unit := do
+partial def checkMeta (origDecl : Decl) : CompilerM Unit := do
   if !(← getEnv).header.isModule then
     return
+  let isMeta := isMeta (← getEnv) origDecl.name
   -- If the meta decl is public, we want to ensure it can only refer to public meta imports so that
   -- references to private imports cannot escape the current module. In particular, we check that
   -- decls with relevant global attrs are public (`Lean.ensureAttrDeclIsMeta`).
   let isPublic := !isPrivateName origDecl.name
-  go isPublic origDecl |>.run' {}
-where go (isPublic : Bool) (decl : Decl) : StateT NameSet CompilerM Unit := do
+  go isMeta isPublic origDecl |>.run' {}
+where go (isMeta isPublic : Bool) (decl : Decl) : StateT NameSet CompilerM Unit := do
   decl.value.forCodeM fun code =>
     for ref in collectUsedDecls code do
       if (← get).contains ref then
         continue
-      modify (·.insert decl.name)
+      modify (·.insert ref)
       if isMeta && isPublic then
         if let some modIdx := (← getEnv).getModuleIdxFor? ref then
           if (← getEnv).header.modules[modIdx]?.any (!·.isExported) then
@@ -85,7 +86,7 @@ where go (isPublic : Bool) (decl : Decl) : StateT NameSet CompilerM Unit := do
         -- *their* references in this case. We also need to do this for non-auxiliary defs in case a
         -- public meta def tries to use a private meta import via a local private meta def :/ .
         if let some refDecl ← getLocalDecl? ref then
-          go isPublic refDecl
+          go isMeta isPublic refDecl
 
 /--
 Checks meta availability just before `evalConst`. This is a "last line of defense" as accesses
@@ -104,13 +105,48 @@ where go (decl : Decl) : StateT NameSet (Except String) Unit :=
     for ref in collectUsedDecls code do
       if (← get).contains ref then
         continue
-      modify (·.insert decl.name)
+      modify (·.insert ref)
       if let some localDecl := baseExt.getState env |>.find? ref then
         go localDecl
       else
         if getIRPhases env ref == .runtime then
           throw s!"Cannot evaluate constant `{declName}` as it uses `{ref}` which is neither marked nor imported as `meta`"
 
+/--
+Checks that imports necessary for inlining/specialization are public as otherwise we may run into
+unknown declarations at the point of inlining/specializing. This is a limitation that we want to
+lift in the future by moving main compilation into a different process that can use a different
+import/incremental system.
+-/
+partial def checkTemplateVisibility : Pass where
+  occurrence := 0
+  phase := .base
+  name := `checkTemplateVisibility
+  run decls := do
+    if (← getEnv).header.isModule then
+      for decl in decls do
+        -- A private template-like decl cannot directly be used by a different module. If it could be used
+        -- indirectly via a public template-like, we do a recursive check when checking the latter.
+        if !isPrivateName decl.name && (← decl.isTemplateLike) then
+          let isMeta := isMeta (← getEnv) decl.name
+          go isMeta decl decl |>.run' {}
+    return decls
+where go (isMeta : Bool) (origDecl decl : Decl) : StateT NameSet CompilerM Unit := do
+  decl.value.forCodeM fun code =>
+    for ref in collectUsedDecls code do
+      if (← get).contains ref then
+        continue
+      modify (·.insert decl.name)
+      if let some localDecl := baseExt.getState (← getEnv) |>.find? ref then
+        -- check transitively through local decls
+        if isPrivateName localDecl.name && (← localDecl.isTemplateLike) then
+          go isMeta origDecl localDecl
+      else if let some modIdx := (← getEnv).getModuleIdxFor? ref then
+        if (← getEnv).header.modules[modIdx]?.any (!·.isExported) then
+          throwError "Cannot compile inline/specializing declaration `{.ofConstName origDecl.name}` as \
+            it uses `{.ofConstName ref}` of module `{(← getEnv).header.moduleNames[modIdx]!}` \
+            which must be imported publicly. This limitation may be lifted in the future."
+
 def inferVisibility (phase : Phase) : Pass where
   occurrence := 0
   phase

src/Lean/Data/Array.lean

@@ -8,8 +8,8 @@ module
 prelude
 public import Init.Prelude
 import Init.Data.Stream
-import Init.Data.Range.Polymorphic.Nat
-import Init.Data.Range.Polymorphic.Iterators
+public import Init.Data.Range.Polymorphic.Nat
+public import Init.Data.Range.Polymorphic.Iterators
 
 namespace Array
 

src/Lean/DocString/Add.lean

@@ -14,7 +14,7 @@ import Lean.Elab.DocString
 import Lean.DocString.Extension
 import Lean.DocString.Links
 import Lean.Parser.Types
-import Lean.DocString.Parser
+public import Lean.DocString.Parser
 import Lean.ResolveName
 public import Lean.Elab.Term.TermElabM
 import Std.Data.HashMap

src/Lean/Elab/DocString/Builtin/Parsing.lean

@@ -6,7 +6,7 @@ Author: David Thrane Christiansen
 module
 prelude
 import Lean.Elab.DocString
-import Lean.Parser.Extension
+public import Lean.Parser.Extension
 public import Lean.Parser.Types
 
 namespace Lean.Doc

src/Lean/Elab/Util.lean

@@ -6,7 +6,7 @@ Authors: Leonardo de Moura
 module
 
 prelude
-import Lean.Parser.Extension
+public import Lean.Parser.Extension
 meta import Lean.Parser.Command
 public import Lean.KeyedDeclsAttribute
 public import Lean.Elab.Exception