fix: load missing dynlibs individually with precompileModules and initialize lazily

src/lake/Lake/Build/Infos.lean

@@ -111,8 +111,9 @@ builtin_facet precompileImports : Module => Array Module
 /-- Shared library for `--load-dynlib`. -/
 builtin_facet dynlib : Module => Dynlib
 
-/-- A Lean library's Lean modules. -/
-builtin_facet modules : LeanLib => Array Module
+/-- A Lean library's Lean modules, topologically sorted
+(`A` comes before `B` when `B` imports `A`). -/
+builtin_facet modules : LeanLib => OrdModuleSet
 
 /-- The package's array of dependencies. -/
 builtin_facet deps : Package => Array Package

src/lake/Lake/Build/Library.lean

@@ -34,7 +34,7 @@ Collect the local modules of a library.
 That is, the modules from `getModuleArray` plus their local transitive imports.
 -/
 partial def LeanLib.recCollectLocalModules
-  (self : LeanLib) : FetchM (Job (Array Module))
+  (self : LeanLib) : FetchM (Job OrdModuleSet)
 := ensureJob do
   let mut col : ModuleCollection := {}
   for mod in (← self.getModuleArray) do
@@ -43,7 +43,7 @@ partial def LeanLib.recCollectLocalModules
     -- This is not considered a fatal error because we want the modules
     -- built to provide better error categorization in the monitor.
     logError s!"{self.name}: some modules have bad imports"
-  return Job.pure col.mods
+  return Job.pure ⟨col.modSet, col.mods⟩
 where
   go root col := do
     let mut col := col
@@ -83,7 +83,7 @@ public def LeanLib.leanArtsFacetConfig : LibraryFacetConfig leanArtsFacet :=
       ""
   withRegisterJob s!"{self.name}:static{suffix}" <| withCurrPackage self.pkg do
   let mods ← (← self.modules.fetch).await
-  let oJobs ← mods.flatMapM fun mod =>
+  let oJobs ← mods.toArray.flatMapM fun mod =>
     mod.nativeFacets shouldExport |>.mapM (·.fetch mod)
   let moreOJobs ← self.moreLinkObjs.mapM (·.fetchIn self.pkg)
   let libFile := if shouldExport then self.staticExportLibFile else self.staticLibFile
@@ -126,7 +126,7 @@ public def LeanLib.staticExportFacetConfig : LibraryFacetConfig staticExportFace
 def LeanLib.recBuildShared (self : LeanLib) : FetchM (Job Dynlib) := do
   withRegisterJob s!"{self.name}:shared" <| withCurrPackage self.pkg do
   let mods ← (← self.modules.fetch).await
-  let objJobs ← mods.flatMapM fun mod =>
+  let objJobs ← mods.toArray.flatMapM fun mod =>
     mod.nativeFacets true |>.mapM (·.fetch mod)
   let objJobs ← self.moreLinkObjs.foldlM (init := objJobs)
     (·.push <$> ·.fetchIn self.pkg)

src/lake/Lake/Build/Module.lean

@@ -124,40 +124,58 @@ def Module.recComputePrecompileImports (mod : Module) : FetchM (Job (Array Modul
 public def Module.precompileImportsFacetConfig : ModuleFacetConfig precompileImportsFacet :=
   mkFacetJobConfig recComputePrecompileImports (buildable := false)
 
+/-- Whether this module is included in `lib.modules` for its parent library `lib`.
+This is false of "orphaned" modules not imported by their parent library's root module. -/
+def Module.containedInLibModules (mod : Module) : FetchM Bool := do
+  return (← (← mod.lib.modules.fetch).await).contains mod
+
 /--
-Computes the transitive dynamic libraries of a module's imports.
-Modules from the same library are loaded individually, while modules
-from other libraries are loaded as part of the whole library.
--/
+Computes the dynamic libraries of `self`'s transitive imports `imps`.
+Modules from the same library as `self` are loaded individually,
+while modules from other libraries are
+- loaded through their library's shared target if they are included in this target,
+- otherwise also individually. -/
 def Module.fetchImportLibs
   (self : Module) (imps : Array Module) (compileSelf : Bool)
 : FetchM (Array (Job Dynlib)) := do
   let (_, jobs) ← imps.foldlM (init := (({} : NameSet), #[])) fun (libs, jobs) imp => do
-    if libs.contains imp.lib.name then
-      return (libs, jobs)
-    else if compileSelf && self.lib.name = imp.lib.name then
+    if compileSelf && self.lib.name = imp.lib.name then
       let job ← imp.dynlib.fetch
       return (libs, jobs.push job)
     else if compileSelf || imp.shouldPrecompile then
-      let jobs ← jobs.push <$> imp.lib.shared.fetch
-      return (libs.insert imp.lib.name, jobs)
+      if ← imp.containedInLibModules then
+        if libs.contains imp.lib.name then
+          return (libs, jobs)
+        else
+          let job ← imp.lib.shared.fetch
+          return (libs.insert imp.lib.name, jobs.push job)
+      else
+        -- TODO: consider linting for orphaned modules
+        let job ← imp.dynlib.fetch
+        return (libs, jobs.push job)
     else
       return (libs, jobs)
   return jobs
 
 /--
-Fetches the library dynlibs of a list of non-local imports.
-Modules are loaded as part of their whole library.
+Fetches the dynamic libraries of a list `mods` of non-local imports.
+Modules are loaded through their library's shared target if they are included in this target,
+otherwise individually.
 -/
 def fetchImportLibs
   (mods : Array Module) : FetchM (Job (Array Dynlib))
 := do
   let (_, jobs) ← mods.foldlM (init := (({} : NameSet), #[])) fun (libs, jobs) imp => do
-    if libs.contains imp.lib.name then
-      return (libs, jobs)
-    else if imp.shouldPrecompile then
-      let jobs ← jobs.push <$> imp.lib.shared.fetch
-      return (libs.insert imp.lib.name, jobs)
+    if imp.shouldPrecompile then
+      if ← imp.containedInLibModules then
+        if libs.contains imp.lib.name then
+          return (libs, jobs)
+        else
+          let job ← imp.lib.shared.fetch
+          return (libs.insert imp.lib.name, jobs.push job)
+      else
+        let job ← imp.dynlib.fetch
+        return (libs, jobs.push job)
     else
       return (libs, jobs)
   return Job.collectArray jobs "import dynlibs"
@@ -206,17 +224,19 @@ def computeModuleDeps
   -/
   let impLibs ← mkLoadOrder impLibs
   let mut dynlibs := externLibs ++ dynlibs
-  let mut plugins := plugins
   for impLib in impLibs do
-    if impLib.plugin then
-      plugins := plugins.push impLib
-    else
-      dynlibs := dynlibs.push impLib
+    /- Load as dynlib (not plugin) because:
+    - imported modules will be initialized in `importModules` anyway; and
+    - since imports from a different `pkg` are loaded together through `pkg:shared`,
+      passing `pkg:shared` as a plugin would initialize too much,
+      namely all modules in `pkg` rather than just the ones we have imported. -/
+    dynlibs := dynlibs.push impLib
   /-
   On MacOS, Lake must be loaded as a plugin for
   `import Lake` to work with precompiled modules.
   https://github.com/leanprover/lean4/issues/7388
   -/
+  let mut plugins := plugins
   if Platform.isOSX && !(plugins.isEmpty && dynlibs.isEmpty) then
     plugins := plugins.push (← getLakeInstall).sharedDynlib
   return {dynlibs, plugins}

src/lake/Lake/Util/OrdHashSet.lean

@@ -30,6 +30,9 @@ public instance : EmptyCollection (OrdHashSet α) := ⟨empty⟩
 public def mkEmpty (size : Nat) : OrdHashSet α :=
   ⟨∅, .mkEmpty size⟩
 
+public def contains (self : OrdHashSet α) (a : α) : Bool :=
+  self.toHashSet.contains a
+
 public def insert (self : OrdHashSet α) (a : α) : OrdHashSet α :=
   if self.toHashSet.contains a then
     self

tests/lake/tests/precompileLink/test.sh

@@ -3,6 +3,7 @@ source ../common.sh
 
 ./clean.sh
 
+PKG=precompileArgs
 
 # Test that precompilation works with a Lake import
 # https://github.com/leanprover/lean4/issues/7388
@@ -13,8 +14,8 @@ test_run -v build LakeTest
 test_run -v exe orderTest
 
 # Test that transitively importing a precompiled module
-# from a non-precompiled module works
-test_not_out '"plugins":[]' -v setup-file ImportDownstream.lean
+# from a non-precompiled module loads the correct shared object.
+test_out "${PKG}_Foo.$SHARED_LIB_EXT" -v setup-file ImportDownstream.lean
 test_run -v build Downstream
 
 # Test that `moreLinkArgs` are included when linking precompiled modules
@@ -23,7 +24,6 @@ test_maybe_err "-lBogus" build -KlinkArgs=-lBogus
 ./clean.sh
 
 # Test that dynlibs are part of the module trace unless `platformIndependent` is set
-PKG=precompileArgs
 test_run build -R
 echo foo > .lake/build/lib/lean/${PKG}_Foo_Bar.$SHARED_LIB_EXT
 test_err "Building Foo" build --rehash

tests/lake/tests/precompileModules-detached/Downstream/ImportDetached.lean

@@ -0,0 +1 @@
+import Upstream.Detached

tests/lake/tests/precompileModules-detached/Downstream/ImportImportDetached.lean

@@ -0,0 +1,3 @@
+import Downstream.ImportDetached
+
+#eval detachedValue

tests/lake/tests/precompileModules-detached/ImportDetached.lean

@@ -0,0 +1,3 @@
+import Upstream.Detached
+
+#eval detachedValue

tests/lake/tests/precompileModules-detached/Upstream/Detached.lean

@@ -0,0 +1,4 @@
+/-! A module in the library `Upstream` not imported by `Upstream.lean`. -/
+
+initialize detachedValue : Nat ←
+  return 1234

tests/lake/tests/precompileModules-detached/clean.sh

@@ -0,0 +1 @@
+rm -rf .lake lake-manifest.json produced.out

tests/lake/tests/precompileModules-detached/lakefile.toml

@@ -0,0 +1,8 @@
+name = "precompileModules-detached"
+precompileModules = true
+
+[[lean_lib]]
+name = "Upstream"
+
+[[lean_lib]]
+name = "Downstream"

tests/lake/tests/precompileModules-detached/test.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+source ../common.sh
+
+./clean.sh
+
+# Test that a precompiled module absent from its parent library's shared target
+# is loaded as an individual dylib when elaborating a downstream module that imports it.
+PKG="precompileModules_x2ddetached"
+test_out "${PKG}_Upstream_Detached.$SHARED_LIB_EXT" -v setup-file ImportDetached.lean
+test_out "${PKG}_Upstream_Detached.$SHARED_LIB_EXT" -v setup-file Downstream/ImportDetached.lean
+
+# This segfaults if the individual module dynlib isn't loaded.
+test_run build -R Downstream.ImportImportDetached
+
+# cleanup
+rm -f produced.out

tests/lake/tests/precompileModules-pkgInit/Downstream.lean

@@ -0,0 +1,2 @@
+import Downstream.A
+import Downstream.B

tests/lake/tests/precompileModules-pkgInit/Downstream/A.lean

@@ -0,0 +1 @@
+import Upstream.Detached

tests/lake/tests/precompileModules-pkgInit/ExternalMod.lean

@@ -0,0 +1 @@
+import Downstream.B

tests/lake/tests/precompileModules-pkgInit/Upstream/Detached.lean

@@ -0,0 +1,4 @@
+/-! A module in the library `Upstream` not imported by `Upstream.lean`. -/
+
+initialize detachedValue : Nat ←
+  return 1234

tests/lake/tests/precompileModules-pkgInit/clean.sh

@@ -0,0 +1 @@
+rm -rf .lake lake-manifest.json produced.out

tests/lake/tests/precompileModules-pkgInit/lakefile.toml

@@ -0,0 +1,8 @@
+name = "precompileModules-pkgInit"
+precompileModules = true
+
+[[lean_lib]]
+name = "Upstream"
+
+[[lean_lib]]
+name = "Downstream"

tests/lake/tests/precompileModules-pkgInit/test.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+source ../common.sh
+
+./clean.sh
+
+# Test that elaborating `ExternalMod` succeeds.
+# Prior to https://github.com/leanprover/lean4/pull/14326,
+# we used to load `Downstream` as a plugin rather than a dynlib.
+# This would try to initialize `Downstream`,
+# and then (transitively through `Downstream.A`) `Upstream.Detached`.
+# But the native symbol to initialize `Upstream.Detached` is not in `Upstream:shared`,
+# so we would get a crash.
+test_run -v lean ExternalMod.lean
+
+# cleanup
+rm -f produced.out