feat: add lake shake command (#11921)

This PR adds `lake shake` as a built-in Lake command, moving the shake
functionality from `script/Shake.lean` into the Lake CLI.

## Motivation

Per discussion with @Kha and @tydeu, having shake as a top-level Lake
command is preferable to `lake exe shake` because:
- Avoids the awkwardness of accessing core tools via `lake exe`
- Compiles shake into the Lake binary, avoiding lakefile issues
- No benefit to lazy compilation on user machines for this tool

## Changes

- Move shake logic from `script/Shake.lean` to
`src/lake/Lake/CLI/Shake.lean`
- Add `lake shake` command dispatch in `Lake/CLI/Main.lean`
- Add help text in `Lake/CLI/Help.lean`
- Remove the standalone shake executable from `script/lakefile.toml`

## Usage

```
lake shake [OPTIONS] [<MODULE>...]
```

See `lake shake --help` for full documentation.

🤖 Prepared with Claude Code

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Mac Malone <mac@lean-fro.org>
Co-authored-by: Sebastian Ullrich <sebasti@nullri.ch>

script/lakefile.toml

@@ -3,9 +3,3 @@ name = "scripts"
 [[lean_exe]]
 name = "modulize"
 root = "Modulize"
-
-[[lean_exe]]
-name = "shake"
-root = "Shake"
-# needed by `Lake.loadWorkspace`
-supportInterpreter = true

src/lake/Lake/Build/Context.lean

@@ -36,10 +36,18 @@ Whether the build should show progress information.
 public def BuildConfig.showProgress (cfg : BuildConfig) : Bool :=
   (cfg.noBuild ∧ cfg.verbosity == .verbose) ∨ cfg.verbosity != .quiet
 
+/-- Mutable reference of registered build jobs. -/
+@[expose] -- for codegen
+public def JobQueue := IO.Ref (Array OpaqueJob)
+
+/-- Returns a new empty job queue. -/
+@[inline] public def mkJobQueue : BaseIO JobQueue :=
+  IO.mkRef #[]
+
 /-- A Lake context with a build configuration and additional build data. -/
 public structure BuildContext extends BuildConfig, Context where
   leanTrace : BuildTrace
-  registeredJobs : IO.Ref (Array OpaqueJob)
+  registeredJobs : JobQueue
 
 /-- A transformer to equip a monad with a `BuildContext`. -/
 public abbrev BuildT := ReaderT BuildContext

src/lake/Lake/Build/Run.lean

@@ -22,6 +22,7 @@ open System
 namespace Lake
 
 /-- Create a fresh build context from a workspace and a build configuration. -/
+@[deprecated "Deprecated without replacement." (since := "2025-01-08")]
 public def mkBuildContext (ws : Workspace) (config : BuildConfig) : BaseIO BuildContext := do
   return {
     opaqueWs := ws,
@@ -37,7 +38,7 @@ private def Monitor.spinnerFrames :=
 
 /-- Context of the Lake build monitor. -/
 private structure MonitorContext where
-  jobs : IO.Ref (Array OpaqueJob)
+  jobs : JobQueue
   out : IO.FS.Stream
   outLv : LogLevel
   failLv : LogLevel
@@ -49,6 +50,9 @@ private structure MonitorContext where
   /-- How often to poll jobs (in milliseconds). -/
   updateFrequency : Nat
 
+@[inline] def MonitorContext.logger (ctx : MonitorContext) : MonadLog IO :=
+  .stream ctx.out ctx.outLv ctx.useAnsi
+
 /-- State of the Lake build monitor. -/
 private structure MonitorState where
   jobNo : Nat := 0
@@ -189,12 +193,51 @@ private def main (init : Array OpaqueJob) : MonitorM PUnit := do
 
 end Monitor
 
+/-- **For internal use only.** -/
 public structure MonitorResult where
   didBuild : Bool
   failures : Array String
   numJobs : Nat
 
+@[inline] def MonitorResult.isOk (self : MonitorResult) : Bool :=
+  self.failures.isEmpty
+
+def mkMonitorContext (cfg : BuildConfig) (jobs : JobQueue) : BaseIO MonitorContext := do
+  let out ← cfg.out.get
+  let useAnsi ← cfg.ansiMode.isEnabled out
+  let outLv := cfg.outLv
+  let failLv := cfg.failLv
+  let isVerbose := cfg.verbosity = .verbose
+  let showProgress := cfg.showProgress
+  let minAction := if isVerbose then .unknown else .fetch
+  let showOptional := isVerbose
+  let showTime := isVerbose || !useAnsi
+  let updateFrequency := 100
+  return {
+    jobs, out, failLv, outLv, minAction, showOptional
+    useAnsi, showProgress, showTime, updateFrequency
+  }
+
+def monitorJobs'
+  (ctx : MonitorContext)
+  (initJobs : Array OpaqueJob)
+  (initFailures : Array String := #[])
+  (resetCtrl : String := "")
+: BaseIO MonitorResult := do
+  let s := {
+    resetCtrl
+    lastUpdate := ← IO.monoMsNow
+    failures := initFailures
+  }
+  let (_,s) ← Monitor.main initJobs |>.run ctx s
+  return {
+    failures := s.failures
+    numJobs := s.totalJobs
+    didBuild := s.didBuild
+  }
+
 /-- The job monitor function. An auxiliary definition for `runFetchM`. -/
+@[inline, deprecated "Deprecated without replacement" (since := "2025-01-08")]
 public def monitorJobs
   (initJobs : Array OpaqueJob)
   (jobs : IO.Ref (Array OpaqueJob))
@@ -210,99 +253,148 @@ public def monitorJobs
     jobs, out, failLv, outLv, minAction, showOptional
     useAnsi, showProgress, showTime, updateFrequency
   }
-  let s := {
-    resetCtrl
-    lastUpdate := ← IO.monoMsNow
-    failures := initFailures
-  }
-  let (_,s) ← Monitor.main initJobs |>.run ctx s
-  return {
-    failures := s.failures
-    numJobs := s.totalJobs
-    didBuild := s.didBuild
-  }
+  monitorJobs' ctx initJobs initFailures resetCtrl
 
 /-- Exit code to return if `--no-build` is set and a build is required. -/
 public def noBuildCode : ExitCode := 3
 
+def Workspace.saveOutputs
+  [logger : MonadLog IO] (ws : Workspace)
+  (out : IO.FS.Stream) (outputsFile : FilePath) (isVerbose : Bool)
+: IO Unit := do
+  unless ws.isRootArtifactCacheEnabled do
+    logWarning s!"{ws.root.prettyName}: \
+      the artifact cache is not enabled for this package, so the artifacts described \
+      by the mappings produced by `-o` will not necessarily be available in the cache."
+  if let some ref := ws.root.outputsRef? then
+    match (← (← ref.get).writeFile outputsFile {}) with
+    | .ok _ log =>
+      if !log.isEmpty && isVerbose then
+        print! out "There were issues saving input-to-output mappings from the build:\n"
+        log.replay
+    | .error _ log =>
+      print! out "Failed to save input-to-output mappings from the build.\n"
+      if isVerbose then
+        log.replay
+  else
+    print! out "Workspace missing input-to-output mappings from build. (This is likely a bug in Lake.)\n"
+
+def reportResult (cfg : BuildConfig) (out : IO.FS.Stream) (result : MonitorResult) : BaseIO Unit := do
+  if result.failures.isEmpty then
+    if cfg.showProgress && cfg.showSuccess then
+      let numJobs := result.numJobs
+      let jobs := if numJobs == 1 then "1 job" else s!"{numJobs} jobs"
+      if cfg.noBuild then
+        print! out s!"All targets up-to-date ({jobs}).\n"
+      else
+        print! out s!"Build completed successfully ({jobs}).\n"
+  else
+    print! out "Some required targets logged failures:\n"
+    result.failures.forM (print! out s!"- {·}\n")
+    flush out
+
+structure BuildResult (α : Type u) extends MonitorResult where
+  out : Except String α
+
+instance : CoeOut (BuildResult α) MonitorResult := ⟨BuildResult.toMonitorResult⟩
+
+@[inline] def BuildResult.isOk (self : BuildResult α) : Bool :=
+  self.out.isOk
+
+def monitorJob (ctx : MonitorContext) (job : Job α) : BaseIO (BuildResult α) := do
+  let result ← monitorJobs' ctx #[job]
+  if result.isOk then
+    if let some a ← job.wait? then
+      return {toMonitorResult := result, out := .ok a}
+    else
+      -- Computation job failed but was unreported in the monitor. This should be impossible.
+      return {toMonitorResult := result, out := .error <|
+        "uncaught top-level build failure (this is likely a bug in Lake)"}
+  else
+    return {toMonitorResult := result, out := .error "build failed"}
+
+def monitorFetchM
+  (mctx : MonitorContext) (bctx : BuildContext) (build : FetchM α)
+  (caption := "job computation")
+: BaseIO (BuildResult α) := do
+  let compute := Job.async build (caption := caption)
+  let job ← compute.run.run'.run bctx |>.run nilTrace
+  monitorJob mctx job
+
+def Workspace.finalizeBuild
+  (ws : Workspace) (cfg : BuildConfig) (ctx : MonitorContext) (result : BuildResult α)
+: IO α := do
+  reportResult cfg ctx.out result
+  if let some outputsFile := cfg.outputsFile? then
+    ws.saveOutputs (logger := ctx.logger) ctx.out outputsFile (cfg.verbosity matches .verbose)
+  if cfg.noBuild && result.didBuild then
+    IO.Process.exit noBuildCode.toUInt8
+  else
+    IO.ofExcept result.out
+
+def mkBuildContext' (ws : Workspace) (cfg : BuildConfig) (jobs : JobQueue) : BuildContext where
+  opaqueWs := ws
+  toBuildConfig := cfg
+  registeredJobs := jobs
+  leanTrace := .ofHash (pureHash ws.lakeEnv.leanGithash)
+    s!"Lean {Lean.versionStringCore}, commit {ws.lakeEnv.leanGithash}"
+
 /--
 Run a build function in the Workspace's context using the provided configuration.
 Reports incremental build progress and build logs. In quiet mode, only reports
 failing build jobs (e.g., when using `-q` or non-verbose `--no-build`).
 -/
 public def Workspace.runFetchM
-  (ws : Workspace) (build : FetchM α) (cfg : BuildConfig := {})
+  (ws : Workspace) (build : FetchM α) (cfg : BuildConfig := {}) (caption := "job computation")
 : IO α := do
-  -- Configure
-  let out ← cfg.out.get
-  let useAnsi ← cfg.ansiMode.isEnabled out
-  let outLv := cfg.outLv
-  let failLv := cfg.failLv
-  let isVerbose := cfg.verbosity = .verbose
-  let showProgress := cfg.showProgress
-  let showSuccess := cfg.showSuccess
-  let ctx ← mkBuildContext ws cfg
-  -- Job Computation
-  let caption := "job computation"
-  let compute := Job.async build (caption := caption)
-  let job ← compute.run.run'.run ctx |>.run nilTrace
-  -- Job Monitor
-  let minAction := if isVerbose then .unknown else .fetch
-  let showOptional := isVerbose
-  let showTime := isVerbose || !useAnsi
-  let {failures, numJobs, didBuild} ← monitorJobs #[job] ctx.registeredJobs
-    out failLv outLv minAction showOptional useAnsi showProgress showTime
-  -- Save input-to-output mappings
-  if let some outputsFile := cfg.outputsFile? then
-    let logger := .stream out outLv useAnsi
-    unless ws.isRootArtifactCacheEnabled do
-      logger.logEntry <| .warning s!"{ws.root.prettyName}: \
-        the artifact cache is not enabled for this package, so the artifacts described \
-        by the mappings produced by `-o` will not necessarily be available in the cache."
-    if let some ref := ws.root.outputsRef? then
-      match (← (← ref.get).writeFile outputsFile {}) with
-      | .ok _ log =>
-        if !log.isEmpty && isVerbose then
-          print! out "There were issues saving input-to-output mappings from the build:\n"
-          log.replay (logger := logger)
-      | .error _ log =>
-        print! out "Failed to save input-to-output mappings from the build.\n"
-        if isVerbose then
-          log.replay (logger := logger)
+  let jobs ← mkJobQueue
+  let mctx ← mkMonitorContext cfg jobs
+  let bctx := mkBuildContext' ws cfg jobs
+  let result ← monitorFetchM mctx bctx build caption
+  ws.finalizeBuild cfg mctx result
+
+def monitorBuild
+  (mctx : MonitorContext) (bctx : BuildContext) (build : FetchM (Job α))
+  (caption := "job computation")
+: BaseIO (BuildResult α) := do
+  let result ← monitorFetchM mctx bctx build caption
+   match result.out with
+  | .ok job =>
+    if let some a ← job.wait? then
+      return {result with out := .ok a}
     else
-      print! out "Workspace missing input-to-output mappings from build. (This is likely a bug in Lake.)\n"
-  -- Report
-  let isNoBuild := cfg.noBuild
-  if failures.isEmpty then
-    let some a ← job.wait?
-      | print! out "Uncaught top-level build failure (this is likely a bug in Lake).\n"
-        error "build failed"
-    if showProgress && showSuccess then
-      let jobs := if numJobs == 1 then "1 job" else s!"{numJobs} jobs"
-      if isNoBuild then
-        print! out s!"All targets up-to-date ({jobs}).\n"
-      else
-        print! out s!"Build completed successfully ({jobs}).\n"
-    return a
-  else
-    print! out "Some required targets logged failures:\n"
-    failures.forM (print! out s!"- {·}\n")
-    flush out
-    if isNoBuild && didBuild then
-      IO.Process.exit noBuildCode.toUInt8
-    else
-      error "build failed"
+      -- Job failed but was unreported in the monitor. It was likely not properly registered.
+      return {result with out := .error <|
+        "uncaught top-level build failure (this is likely a bug in the build script)"}
+  | .error e =>
+    return {result with out := .error e}
+
+/--
+Returns whether a build is needed to validate `build`. Does not report on the attempted build.
+
+This is equivalent to checking whether `lake build --no-build` exits with code 0.
+-/
+@[inline] public def Workspace.checkNoBuild
+  (ws : Workspace) (build : FetchM (Job α))
+: BaseIO Bool := do
+  let jobs ← mkJobQueue
+  let cfg := {noBuild := true}
+  let mctx ← mkMonitorContext cfg jobs
+  let bctx := mkBuildContext' ws cfg jobs
+  let result ← monitorBuild mctx bctx build
+  return result.isOk && !result.didBuild
 
 /-- Run a build function in the Workspace's context and await the result. -/
 @[inline] public def Workspace.runBuild
   (ws : Workspace) (build : FetchM (Job α)) (cfg : BuildConfig := {})
 : IO α := do
-  let job ← ws.runFetchM build cfg
-  let some a ← job.wait? | error "build failed"
-  return a
+  let jobs ← mkJobQueue
+  let mctx ← mkMonitorContext cfg jobs
+  let bctx := mkBuildContext' ws cfg jobs
+  let result ← monitorBuild mctx bctx build
+  ws.finalizeBuild cfg mctx result
 
 /-- Produce a build job in the Lake monad's workspace and await the result. -/
 @[inline] public def runBuild
   (build : FetchM (Job α)) (cfg : BuildConfig := {})
-: LakeT IO α := do
-  (← getWorkspace).runBuild build cfg
+: LakeT IO α := do (← getWorkspace).runBuild build cfg

src/lake/Lake/CLI.lean

@@ -13,6 +13,7 @@ public import Lake.CLI.Help
 public import Lake.CLI.Init
 public import Lake.CLI.Main
 public import Lake.CLI.Serve
+public import Lake.CLI.Shake
 public import Lake.CLI.Translate
 public import Lake.CLI.Translate.Lean
 public import Lake.CLI.Translate.Toml

src/lake/Lake/CLI/Help.lean

@@ -30,6 +30,7 @@ COMMANDS:
   lint                  lint the package using the configured lint driver
   check-lint            check if there is a properly configured lint driver
   clean                 remove build outputs
+  shake                 minimize imports in source files
   env <cmd> <args>...   execute a command in Lake's environment
   lean <file>           elaborate a Lean file in Lake's context
   update                update dependencies and save them to the manifest
@@ -310,6 +311,44 @@ USAGE:
 If no package is specified, deletes the build directories of every package in
 the workspace. Otherwise, just deletes those of the specified packages."
 
+def helpShake :=
+"Minimize imports in Lean source files
+
+USAGE:
+  lake shake [OPTIONS] [<MODULE>...]
+
+Checks the current project for unused imports by analyzing generated `.olean`
+files to deduce required imports and ensuring that every import contributes
+some constant or other elaboration dependency.
+
+ARGUMENTS:
+  <MODULE>              A module path like `Mathlib`. All files transitively
+                        reachable from the provided module(s) will be checked.
+                        If not specified, uses the package's default targets.
+
+OPTIONS:
+  --force               Skip the `lake build --no-build` sanity check
+  --keep-implied        Preserve imports implied by other imports
+  --keep-prefix         Prefer parent module imports over specific submodules
+  --keep-public         Preserve all `public` imports for API stability
+  --add-public          Add new imports as `public` if they were in the
+                        original public closure
+  --explain             Show which constants require each import
+  --fix                 Apply suggested fixes directly to source files
+  --gh-style            Output in GitHub problem matcher format
+
+ANNOTATIONS:
+  Source files can contain special comments to control shake behavior:
+
+  * `module -- shake: keep-downstream`
+    Preserves this module in all downstream modules
+
+  * `module -- shake: keep-all`
+    Preserves all existing imports in this module
+
+  * `import X -- shake: keep`
+    Preserves this specific import"
+
 def helpCacheCli :=
 "Manage the Lake cache
 
@@ -557,6 +596,7 @@ public def help : (cmd : String) → String
 | "lint"                => helpLint
 | "check-lint"          => helpCheckLint
 | "clean"               => helpClean
+| "shake"               => helpShake
 | "script"              => helpScriptCli
 | "scripts"             => helpScriptList
 | "run"                 => helpScriptRun

src/lake/Lake/CLI/Main.lean

@@ -10,6 +10,7 @@ public import Init.System.IO
 public import Lake.Util.Exit
 public import Lake.Load.Config
 public import Lake.CLI.Error
+public import Lake.CLI.Shake
 import Lake.Version
 import Lake.Build.Run
 import Lake.Build.Targets
@@ -74,6 +75,7 @@ public structure LakeOptions where
   toolchain? : Option String := none
   rev? : Option String := none
   maxRevs : Nat := 100
+  shake : Shake.Args := {}
 
 def LakeOptions.outLv (opts : LakeOptions) : LogLevel :=
   opts.outLv?.getD opts.verbosity.minLogLv
@@ -299,6 +301,16 @@ def lakeLongOption : (opt : String) → CliM PUnit
 | "--"            => do
   let subArgs ← takeArgs
   modifyThe LakeOptions ({· with subArgs})
+-- Shake options
+| "--keep-implied" => modifyThe LakeOptions ({· with shake.keepImplied := true})
+| "--keep-prefix" => modifyThe LakeOptions ({· with shake.keepPrefix := true})
+| "--keep-public" => modifyThe LakeOptions ({· with shake.keepPublic := true})
+| "--add-public" => modifyThe LakeOptions ({· with shake.addPublic := true})
+| "--force" => modifyThe LakeOptions ({· with shake.force := true})
+| "--gh-style" => modifyThe LakeOptions ({· with shake.githubStyle := true})
+| "--explain" => modifyThe LakeOptions ({· with shake.explain := true})
+| "--trace" => modifyThe LakeOptions ({· with shake.trace := true})
+| "--fix" => modifyThe LakeOptions ({· with shake.fix := true})
 | opt             =>  throw <| CliError.unknownLongOption opt
 
 def lakeOption :=
@@ -358,7 +370,6 @@ def parseTemplateLangSpec (spec : String) : Except CliError (InitTemplate × Con
   | [tmp] => return (← parseTemplateSpec tmp, default)
   | _ => return default
 
-
 /-! ## Commands -/
 
 namespace lake
@@ -756,6 +767,31 @@ protected def clean : CliM PUnit := do
       | some pkg => pure pkg
     pkgs.forM (·.clean)
 
+/-- The `lake shake` command: minimize imports in Lean source files. -/
+protected def shake : CliM PUnit := do
+  processOptions lakeOption
+  let opts ← getThe LakeOptions
+  let config ← mkLoadConfig opts
+  let ws ← loadWorkspace config
+  -- Get remaining arguments as module names
+  let mods := (← takeArgs).toArray.map (·.toName)
+  -- Get default target modules from workspace if no modules specified
+  let mods := if mods.isEmpty then ws.defaultTargetRoots else mods
+  if h : 0 < mods.size then
+    let args := {opts.shake with mods}
+    unless args.force do
+      let specs ← parseTargetSpecs ws []
+      let upToDate ← ws.checkNoBuild (buildSpecs specs)
+      unless upToDate do
+        error "there are out of date oleans; run `lake build` or fetch them from a cache first"
+    -- Run shake with workspace search paths
+    Lean.searchPathRef.set ws.augmentedLeanPath
+    let exitCode ← Shake.run args h ws.augmentedLeanSrcPath
+    if exitCode != 0 then
+      exit exitCode
+  else
+    error "no modules specified and there are no applicable default targets"
+
 protected def script : CliM PUnit := do
   if let some cmd ← takeArg? then
     processLeadingOptions lakeOption -- between `lake script <cmd>` and args
@@ -910,6 +946,7 @@ def lakeCli : (cmd : String) → CliM PUnit
 | "lint"                => lake.lint
 | "check-lint"          => lake.checkLint
 | "clean"               => lake.clean
+| "shake"               => lake.shake
 | "script"              => lake.script
 | "scripts"             => lake.script.list
 | "run"                 => lake.script.run

src/lake/Lake/CLI/Shake.lean

@@ -5,12 +5,13 @@ Authors: Mario Carneiro, Sebastian Ullrich
 -/
 module
 
+prelude
+public import Init.Prelude
+public import Init.System.IO
+public import Lean.Util.Path
 import Lean.Environment
 import Lean.ExtraModUses
-
-import Lake.CLI.Main
 import Lean.Parser.Module
-import Lake.Load.Workspace
 
 /-! # Shake: A Lean import minimizer
 
@@ -20,84 +21,12 @@ ensuring that every import is used to contribute some constant or other elaborat
 recorded by `recordExtraModUse` and friends.
 -/
 
-/-- help string for the command line interface -/
-def help : String := "Lean project tree shaking tool
-Usage: lake exe shake [OPTIONS] <MODULE>..
-
-Arguments:
-  <MODULE>
-    A module path like `Mathlib`. All files transitively reachable from the
-    provided module(s) will be checked.
-
-Options:
-  --force
-    Skips the `lake build --no-build` sanity check
-
-  --keep-implied
-    Preserves existing imports that are implied by other imports and thus not technically needed
-    anymore
-
-  --keep-prefix
-    If an import `X` would be replaced in favor of a more specific import `X.Y...` it implies,
-    preserves the original import instead. More generally, prefers inserting `import X` even if it
-    was not part of the original imports as long as it was in the original transitive import closure
-    of the current module.
-
-  --keep-public
-    Preserves all `public` imports to avoid breaking changes for external downstream modules
-
-  --add-public
-    Adds new imports as `public` if they have been in the original public closure of that module.
-    In other words, public imports will not be removed from a module unless they are unused even
-    in the private scope, and those that are removed will be re-added as `public` in downstream
-    modules even if only needed in the private scope there. Unlike `--keep-public`, this may
-    introduce breaking changes but will still limit the number of inserted imports.
-
-  --explain
-    Gives constants explaining why each module is needed
-
-  --fix
-    Apply the suggested fixes directly. Make sure you have a clean checkout
-    before running this, so you can review the changes.
-
-  --gh-style
-    Outputs messages that can be parsed by `gh-problem-matcher-wrap`
-
-Annotations:
-  The following annotations can be added to Lean files in order to configure the behavior of
-  `shake`. Only the substring `shake: ` directly followed by a directive is checked for, so multiple
-  directives can be mixed in one line such as `-- shake: keep-downstream, shake: keep-all`, and they
-  can be surrounded by arbitrary comments such as `-- shake: keep (metaprogram output dependency)`.
-
-  * `module -- shake: keep-downstream`:
-    Preserves this module in all (current) downstream modules, adding new imports of it if needed.
-
-  * `module -- shake: keep-all`:
-    Preserves all existing imports in this module as is. New imports now needed because of upstream
-    changes may still be added.
-
-  * `import X -- shake: keep`:
-    Preserves this specific import in the current module. The most common use case is to preserve a
-    public import that will be needed in downstream modules to make sense of the output of a
-    metaprogram defined in this module. For example, if a tactic is defined that may synthesize a
-    reference to a theorem when run, there is no way for `shake` to detect this by itself and the
-    module of that theorem should be publicly imported and annotated with `keep` in the tactic's
-    module.
-    ```
-    public import X  -- shake: keep (metaprogram output dependency)
-
-    ...
-
-    elab \"my_tactic\" : tactic => do
-      ... mkConst ``f -- `f`, defined in `X`, may appear in the output of this tactic
-    ```
-"
-
 open Lean
 
-/-- The parsed CLI arguments. See `help` for more information -/
-structure Args where
-  help : Bool := false
+namespace Lake.Shake
+
+/-- The parsed CLI arguments for shake. -/
+public structure Args where
   keepImplied : Bool := false
   keepPrefix : Bool := false
   keepPublic : Bool := false
@@ -640,7 +569,7 @@ def visitModule (pkg : Name) (srcSearchPath : SearchPath)
         if toRemove.any fun imp => imp == decodeImport stx then
           let pos := inputCtx.fileMap.toPosition stx.raw.getPos?.get!
           println! "{path}:{pos.line}:{pos.column+1}: warning: unused import \
-            (use `lake exe shake --fix` to fix this, or `lake exe shake --update` to ignore)"
+            (use `lake shake --fix` to fix this, or `lake shake --update` to ignore)"
       if !toAdd.isEmpty then
         -- we put the insert message on the beginning of the last import line
         let pos := inputCtx.fileMap.toPosition endHeader.offset
@@ -685,76 +614,31 @@ def visitModule (pkg : Name) (srcSearchPath : SearchPath)
         run j
     for i in toAdd do run i
 
-/-- Convert a list of module names to a bitset of module indexes -/
-def toBitset (s : State) (ns : List Name) : Bitset :=
-  ns.foldl (init := ∅) fun c name =>
-    match s.env.getModuleIdxFor? name with
-    | some i => c ∪ {i}
-    | none => c
-
 local instance : Ord Import where
   compare :=
     let _ := @lexOrd
     compareOn fun imp => (!imp.isExported, imp.module.toString)
 
-/-- The main entry point. See `help` for more information on arguments. -/
-public def main (args : List String) : IO UInt32 := do
-  initSearchPath (← findSysroot)
-  -- Parse the arguments
-  let rec parseArgs (args : Args) : List String → Args
-    | [] => args
-    | "--help" :: rest => parseArgs { args with help := true } rest
-    | "--keep-implied" :: rest => parseArgs { args with keepImplied := true } rest
-    | "--keep-prefix" :: rest => parseArgs { args with keepPrefix := true } rest
-    | "--keep-public" :: rest => parseArgs { args with keepPublic := true } rest
-    | "--add-public" :: rest => parseArgs { args with addPublic := true } rest
-    | "--force" :: rest => parseArgs { args with force := true } rest
-    | "--fix" :: rest => parseArgs { args with fix := true } rest
-    | "--explain" :: rest => parseArgs { args with explain := true } rest
-    | "--trace" :: rest => parseArgs { args with trace := true } rest
-    | "--gh-style" :: rest => parseArgs { args with githubStyle := true } rest
-    | "--" :: rest => { args with mods := args.mods ++ rest.map (·.toName) }
-    | other :: rest => parseArgs { args with mods := args.mods.push other.toName } rest
-  let args := parseArgs {} args
+/--
+Run the shake analysis with the given arguments.
 
-  -- Bail if `--help` is passed
-  if args.help then
-    IO.println help
-    IO.Process.exit 0
+Assumes Lean's search path has already been properly configured.
+-/
+public def run (args : Args) (h : 0 < args.mods.size)
+    (srcSearchPath : SearchPath := {}) : IO UInt32 := do
 
-  if !args.force then
-    if (← IO.Process.output { cmd := "lake", args := #["build", "--no-build"] }).exitCode != 0 then
-      IO.println "There are out of date oleans. Run `lake build` or `lake exe cache get` first"
-      IO.Process.exit 1
-
-  -- Determine default module(s) to run shake on
-  let defaultTargetModules : Array Name ← try
-    let (elanInstall?, leanInstall?, lakeInstall?) ← Lake.findInstall?
-    let config ← Lake.MonadError.runEIO <| Lake.mkLoadConfig { elanInstall?, leanInstall?, lakeInstall? }
-    let some workspace ← Lake.loadWorkspace config |>.toBaseIO
-      | throw <| IO.userError "failed to load Lake workspace"
-    let defaultTargetModules := workspace.root.defaultTargets.flatMap fun target =>
-      if let some lib := workspace.root.findLeanLib? target then
-        lib.roots
-      else if let some exe := workspace.root.findLeanExe? target then
-        #[exe.config.root]
-      else
-        #[]
-    pure defaultTargetModules
-  catch _ =>
-    pure #[]
-
-  let srcSearchPath ← getSrcSearchPath
   -- the list of root modules
-  let mods := if args.mods.isEmpty then defaultTargetModules else args.mods
+  let mods := args.mods
   -- Only submodules of `pkg` will be edited or have info reported on them
-  let pkg := mods[0]!.components.head!
+  let pkg := mods[0].getRoot
 
   -- Load all the modules
   let imps := mods.map ({ module := · })
   let (_, s) ← importModulesCore imps (isExported := true) |>.run
   let s := s.markAllExported
   let mut env ← finalizeImport s (isModule := true) imps {} (leakEnv := false) (loadExts := false)
+  if env.header.moduleData.any (!·.isModule) then
+    throw <| .userError "`lake shake` only works with `module`s currently"
   -- the one env ext we want to initialize
   let is := indirectModUseExt.toEnvExtension.getState env
   let newState ← indirectModUseExt.addImportedFn is.importedEntries { env := env, opts := {} }

src/lake/Lake/Config/Workspace.lean

@@ -52,6 +52,16 @@ public instance : Nonempty Workspace :=
 
 public hydrate_opaque_type OpaqueWorkspace Workspace
 
+/-- Returns the names of the root modules of the package's default targets. -/
+public def Package.defaultTargetRoots (self : Package) : Array Lean.Name :=
+  self.defaultTargets.flatMap fun target =>
+    if let some lib := self.findLeanLib? target then
+      lib.roots
+    else if let some exe := self.findLeanExe? target then
+      #[exe.root.name]
+    else
+      #[]
+
 namespace Workspace
 
 /-- **For internal use.** Whether this workspace is Lean itself.  -/
@@ -102,6 +112,10 @@ public def isRootArtifactCacheEnabled (ws : Workspace) : Bool :=
 @[inline] public def serverOptions (self : Workspace) : LeanOptions :=
   self.root.moreServerOptions
 
+/-- Returns the names of the root modules of the workpace root's default targets. -/
+@[inline] public def defaultTargetRoots (self : Workspace) : Array Lean.Name :=
+  self.root.defaultTargetRoots
+
 /-- The workspace's Lake manifest. -/
 @[inline] public def manifestFile (self : Workspace) : FilePath :=
   self.root.manifestFile

tests/lake/tests/shake/.gitignore

@@ -0,0 +1,8 @@
+.lake/
+lake-manifest.json
+lake-packages/
+produced.out
+# Working files copied from input/
+lakefile.toml
+Main.lean
+Lib/

tests/lake/tests/shake/clean.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+rm -rf .lake lake-manifest.json lake-packages
+rm -f lakefile.toml Main.lean
+rm -rf Lib

tests/lake/tests/shake/expected/Main.lean

@@ -0,0 +1,6 @@
+module
+
+import Lib.A
+
+-- Does not use Lib.B, uses Lib.A privately only
+def myValue : Nat := valueA

tests/lake/tests/shake/input/Lib/A.lean

@@ -0,0 +1,3 @@
+module
+
+public def valueA : Nat := 42

tests/lake/tests/shake/input/Lib/B.lean

@@ -0,0 +1,3 @@
+module
+
+public def valueB : Nat := 99

tests/lake/tests/shake/input/Main.lean

@@ -0,0 +1,7 @@
+module
+
+public import Lib.A
+public import Lib.B
+
+-- Does not use Lib.B, uses Lib.A privately only
+def myValue : Nat := valueA

tests/lake/tests/shake/input/lakefile.toml

@@ -0,0 +1,10 @@
+name = "test"
+version = "0.1.0"
+defaultTargets = ["Main"]
+
+[[lean_lib]]
+name = "Lib"
+globs = ["Lib.**"]
+
+[[lean_lib]]
+name = "Main"

tests/lake/tests/shake/test.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+source ../common.sh
+
+./clean.sh
+
+# Test the `lake shake` command
+
+# Copy input project to working directory
+cp -r input/* .
+
+# Build the project first (shake needs .olean files)
+test_run build
+
+# Run shake to check for unused imports
+# Shake exits with code 1 when issues are found, which is expected here
+lake_out shake Main || true
+match_pat 'remove.*Lib.B' produced.out
+
+# Test --fix mode: apply the fixes and verify the result
+./clean.sh
+cp -r input/* .
+test_run build
+test_run shake --fix Main
+
+# Verify Main.lean matches expected (Lib.B import removed)
+check_diff expected/Main.lean Main.lean