feat: lake: precompileLibrary & precompileImports (#15015)

This PR adds two new configuration options that are subsets of
`precompileModules`: `precompileLibrary` for `lean_lib` targets and
`precompileImports` for all Lean configurations (e.g., settable on
`package`, `lean_lib`, or `lean_exe`). `precompileImports` compiles a
module's imports but not the module itself. `precompileLibrary` compiles
the whole library for importers, but the library's modules do not
compile their own imports during elaboration.

Closes #2757 with the finer-grain `precompileLibrary` instead of the
more general `precompilePackage`.

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: tydeu

src/lake/Lake/Build/Module.lean

@@ -111,10 +111,8 @@ def computePrecompileImportsAux
   (fileName : String) (imports : Array Module)
 : FetchM (Job (Array Module)) := do
   collectImportsAux fileName imports fun imp =>
-    if imp.shouldPrecompile then
-      (true, ·) <$> imp.transImports.fetch
-    else
-      (false, ·) <$> imp.precompileImports.fetch
+    -- `imp.shouldPrecompile` implies `imp.lib.shouldPrecompile`
+    (imp.lib.shouldPrecompile, ·) <$> imp.precompileImports.fetch
 
 /-- Recursively compute a module's precompiled imports. -/
 def Module.recComputePrecompileImports (mod : Module) : FetchM (Job (Array Module)) := ensureJob do
@@ -130,15 +128,20 @@ Modules from the same library are loaded individually, while modules
 from other libraries are loaded as part of the whole library.
 -/
 def Module.fetchImportLibs
-  (self : Module) (imps : Array Module) (compileSelf : Bool)
+  (self : Module) (imps : Array Module)
+  (precompileModules : Bool) (precompileImports : 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
-      let job ← imp.dynlib.fetch
-      return (libs, jobs.push job)
-    else if compileSelf || imp.shouldPrecompile then
+    else if self.lib.name = imp.lib.name then
+      -- The library as a whole cannot be loaded here, as it includes the module itself.
+      if precompileModules then
+        let job ← imp.dynlib.fetch
+        return (libs, jobs.push job)
+      else
+        return (libs, jobs)
+    else if precompileImports || imp.lib.shouldPrecompile then
       let jobs ← jobs.push <$> imp.lib.shared.fetch
       return (libs.insert imp.lib.name, jobs)
     else
@@ -155,7 +158,7 @@ def fetchImportLibs
   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
+    else if imp.lib.shouldPrecompile then
       let jobs ← jobs.push <$> imp.lib.shared.fetch
       return (libs.insert imp.lib.name, jobs)
     else
@@ -577,14 +580,14 @@ def Module.recFetchPreSetup (mod : Module) : FetchM (Job ModulePreSetup) := ensu
   Remark: It should be possible to avoid transitive imports here when the module
   itself is precompiled, but they are currently kept to preserve the "bad import" errors.
   -/
-  let precompileImports ← if mod.shouldPrecompile then
+  let precompileImports ← if mod.shouldPrecompileImports then
     mod.transImports.fetch else mod.precompileImports.fetch
   let precompileImports ← precompileImports.await
   let impLibsJob ← Job.collectArray (traceCaption := "import dynlibs") <$>
-    mod.fetchImportLibs precompileImports mod.shouldPrecompile
+    mod.fetchImportLibs precompileImports mod.shouldPrecompile mod.shouldPrecompileImports
 
   let externLibsJob ← Job.collectArray (traceCaption := "package external libraries") <$>
-    if mod.shouldPrecompile then mod.pkg.externLibs.mapM (·.dynlib.fetch) else pure #[]
+    if mod.shouldPrecompileImports then mod.pkg.externLibs.mapM (·.dynlib.fetch) else pure #[]
   let dynlibsJob ← mod.dynlibs.fetchIn mod.pkg "module dynlibs"
   let pluginsJob ← mod.plugins.fetchIn mod.pkg "module plugins"
 
@@ -1360,7 +1363,7 @@ def Module.recBuildDynlib (mod : Module) : FetchM (Job Dynlib) :=
   -- Fetch dependencies' dynlibs
   let libJobs ← id do
     let imps ← (← mod.imports.fetch).await
-    let libJobs ← mod.fetchImportLibs imps true
+    let libJobs ← mod.fetchImportLibs imps true true
     let libJobs ← mod.lib.moreLinkLibs.foldlM
       (·.push <$> ·.fetchIn mod.pkg) libJobs
     let libJobs ← mod.pkg.externLibs.foldlM
@@ -1439,14 +1442,14 @@ def setupEditedModule
   let impInfoJob ← fetchImportInfo fileName mod.pkg.keyName mod.name header
     (allowNonModules := mod.allowNonModules)
   let precompileImports ←
-    if mod.shouldPrecompile then
+    if mod.shouldPrecompileImports then
       (← computeTransImportsAux fileName localImports).await
     else
       (← computePrecompileImportsAux fileName localImports).await
   let impLibsJob ← Job.collectArray (traceCaption := "import dynlibs") <$>
-    mod.fetchImportLibs precompileImports mod.shouldPrecompile
+    mod.fetchImportLibs precompileImports mod.shouldPrecompile mod.shouldPrecompileImports
   let externLibsJob ← Job.collectArray (traceCaption := "package external libraries") <$>
-    if mod.shouldPrecompile then mod.pkg.externLibs.mapM (·.dynlib.fetch) else pure #[]
+    if mod.shouldPrecompileImports then mod.pkg.externLibs.mapM (·.dynlib.fetch) else pure #[]
   let dynlibsJob ← mod.dynlibs.fetchIn mod.pkg "module dynlibs"
   let pluginsJob ← mod.plugins.fetchIn mod.pkg "module plugins"
   extraDepJob.bindM (sync := true) fun _ => do

src/lake/Lake/Config/LeanConfig.lean

@@ -254,6 +254,14 @@ public configuration LeanConfig where
   and Lake will not catch it. Defaults to `none`.
   -/
   platformIndependent : Option Bool := none
+  /--
+  Whether to compile the imports of modules into native shared libraries that
+  are loaded when the module is built. This speeds up evaluation of metaprograms
+  and enables the interpreter to run functions marked `@[extern]`.
+
+  Defaults to `false`.
+  -/
+  precompileImports : Bool := false
   /-
   An array of dynamic library targets to load during the elaboration
   of a module (via `lean --load-dynlib`).

src/lake/Lake/Config/LeanLib.lean

@@ -104,6 +104,21 @@ Is true if either the package or the library have `precompileModules` set.
 @[inline] public def precompileModules (self : LeanLib) : Bool :=
   self.pkg.precompileModules || self.config.precompileModules
 
+/--
+Whether to precompile the imports of the library's modules.
+Is true if either the package or the library have `precompileImports` set,
+or if the library's modules are precompiled.
+-/
+@[inline] public def precompileImports (self : LeanLib) : Bool :=
+  self.precompileModules || self.pkg.precompileImports || self.config.precompileImports
+
+/--
+Whether to precompile the library for importers.
+Is true if the library has `precompileLibrary` set or its modules are precompiled.
+-/
+@[inline] public def shouldPrecompile (self : LeanLib) : Bool :=
+  self.precompileModules || self.config.precompileLibrary
+
 /--
 Whether to the library's Lean code is platform-independent.
 Returns the library's `platformIndependent` configuration if non-`none`.

src/lake/Lake/Config/LeanLibConfig.lean

@@ -72,11 +72,23 @@ public configuration LeanLibConfig (name : Name) extends LeanConfig where
   -/
   extraDepTargets : Array Name := #[]
 
+  /--
+  Whether to compile the library into a native shared library that is loaded
+  whenever one of its modules is imported by a module outside the library.
+  This speeds up evaluation of metaprograms and enables the interpreter to run
+  functions marked `@[extern]`.
+
+  Defaults to `false`.
+  -/
+  precompileLibrary : Bool := false
+
   /--
   Whether to compile each of the library's modules into a native shared library
   that is loaded whenever the module is imported. This speeds up evaluation of
   metaprograms and enables the interpreter to run functions marked `@[extern]`.
 
+  This implies `precompileImports` and `precompileLibrary`.
+
   Defaults to `false`.
   -/
   precompileModules : Bool := false

src/lake/Lake/Config/Module.lean

@@ -217,6 +217,9 @@ public def leanIncludeDir? (self : Module) : Option FilePath :=
 @[inline] public def platformIndependent (self : Module) : Option Bool :=
   self.lib.platformIndependent
 
+@[inline] public def shouldPrecompileImports (self : Module) : Bool :=
+  self.lib.precompileImports
+
 @[inline] public def shouldPrecompile (self : Module) : Bool :=
   self.lib.precompileModules
 

src/lake/Lake/Config/Package.lean

@@ -276,6 +276,10 @@ public def id? (self : Package) : Option PkgId :=
 @[inline] public def precompileModules (self : Package) : Bool :=
   self.config.precompileModules
 
+/-- The package's `precompileImports` configuration. -/
+@[inline] public def precompileImports (self : Package) : Bool :=
+  self.config.precompileImports
+
 /-- The package's `moreGlobalServerArgs` configuration. -/
 @[inline] public def moreGlobalServerArgs (self : Package) : Array String :=
   self.config.moreGlobalServerArgs

src/lake/schemas/lakefile-toml-schema.json

@@ -132,6 +132,11 @@
                     "type": "boolean",
                     "description": "Asserts whether Lake should assume Lean modules are platform-independent.\n\n* If `false`, Lake will add `System.Platform.target` to the module traces within the code unit (e.g., package or library). This will force Lean code to be re-elaborated on different platforms.\n\n* If `true`, Lake will exclude platform-dependent elements (e.g., precompiled modules, external libraries) from a module's trace, preventing re-elaboration on different platforms. Note that this will not effect modules outside the code unit in question. For example, a platform-independent package which depends on a platform-dependent library will still be platform-dependent.\n\n* If not set, Lake will construct traces as natural. That is, it will include platform-dependent artifacts in the trace if they module depends on them, but otherwise not force modules to be platform-dependent.\n\nThere is no check  for correctness here, so a configuration can lie and Lake will not catch it."
                 },
+                "precompileImports": {
+                    "type": "boolean",
+                    "default": false,
+                    "description": "Whether to compile the imports of modules into native shared libraries that are loaded when the module is built. This speeds up evaluation of metaprograms and enables the interpreter to run functions marked `@[extern]`."
+                },
                 "dynlibs": {
                     "type": "array",
                     "items": {
@@ -217,10 +222,15 @@
                             "default": [],
                             "description": "Build target identifiers of targets to build before the library's modules.\n\nFormat: `<target>[:<facet>]*`, where:\n- `<target>` is of the form `@<package>`, `@<package>/<targetName>`, +<module>` or `@<package>/+<module>`,\n- `<package>` may be empty to reference the default package,\n- `<facet>` is the name of a target facet."
                         },
+                        "precompileLibrary": {
+                            "type": "boolean",
+                            "default": false,
+                            "description": "Whether to compile the library into a native shared library that is loaded whenever one of its modules is imported by a module outside the library. This speeds up evaluation of metaprograms and enables the interpreter to run functions marked `@[extern]`."
+                        },
                         "precompileModules": {
                             "type": "boolean",
                             "default": false,
-                            "description": "Whether to compile each of the library's modules into a native shared library that is loaded whenever the module is imported. This speeds up evaluation of metaprograms and enables the interpreter to run functions marked `@[extern]`."
+                            "description": "Whether to compile each of the library's modules into a native shared library that is loaded whenever the module is imported. This speeds up evaluation of metaprograms and enables the interpreter to run functions marked `@[extern]`.\n\nThis implies `precompileImports` and `precompileLibrary`."
                         },
                         "defaultFacets": {
                             "type": "array",

tests/lake/tests/precompileLink/Lib.lean

@@ -0,0 +1 @@
+import Lib.Base

tests/lake/tests/precompileLink/Lib/Base.lean

@@ -0,0 +1,3 @@
+import LibDep
+
+builtin_initialize libGreetingRef : IO.Ref String ← IO.mkRef libGreeting

tests/lake/tests/precompileLink/LibDep.lean

@@ -0,0 +1 @@
+def libGreeting := "Hello"