From e96b338cd96c591557f00595875b9ef79d77db79 Mon Sep 17 00:00:00 2001
From: tydeu <tydeu@hatpress.net>
Date: Thu, 7 Sep 2023 10:53:25 -0400
Subject: [PATCH] feat: `lake run <s>` for `s` in any pkg

---
 src/lake/Lake/CLI/Help.lean                  |  6 ++--
 src/lake/Lake/CLI/Main.lean                  | 31 ++++++++++----------
 src/lake/Lake/Config/Script.lean             |  2 ++
 src/lake/Lake/Config/Workspace.lean          |  4 +++
 src/lake/Lake/Load/Package.lean              | 11 ++++---
 src/lake/examples/scripts/.gitignore         |  2 +-
 src/lake/examples/scripts/dep/lakefile.lean  |  9 ++++++
 src/lake/examples/scripts/expected.out       |  6 +++-
 src/lake/examples/scripts/lake-manifest.json |  4 +++
 src/lake/examples/scripts/lakefile.lean      |  2 +-
 src/lake/examples/scripts/test.sh            | 12 ++++++--
 11 files changed, 62 insertions(+), 27 deletions(-)
 create mode 100644 src/lake/examples/scripts/dep/lakefile.lean
 create mode 100644 src/lake/examples/scripts/lake-manifest.json

diff --git a/src/lake/Lake/CLI/Help.lean b/src/lake/Lake/CLI/Help.lean
index 3b54130339..023295b0db 100644
--- a/src/lake/Lake/CLI/Help.lean
+++ b/src/lake/Lake/CLI/Help.lean
@@ -168,8 +168,8 @@ def helpScriptRun :=
 USAGE:
   lake script run [[<package>/]<script>] [<args>...]
 
-This command runs the given `script` from `package`, passing `args` to it.
-Defaults to the root package.
+This command runs the `script` of the workspace (or the specific `package`),
+passing `args` to it.
 
 A bare `lake run` command will run the default script(s) of the root package
 (with no arguments)."
@@ -180,7 +180,7 @@ def helpScriptDoc :=
 USAGE:
   lake script doc [<package>/]<script>
 
-Print the docstring of `script` in `package`. Defaults to the root package."
+Print the docstring of `script` in the workspace or the specific `package`."
 
 def helpServe :=
 "Start the Lean language server
diff --git a/src/lake/Lake/CLI/Main.lean b/src/lake/Lake/CLI/Main.lean
index 548f60774e..55ea5b90e6 100644
--- a/src/lake/Lake/CLI/Main.lean
+++ b/src/lake/Lake/CLI/Main.lean
@@ -181,10 +181,17 @@ def verifyInstall (opts : LakeOptions) : ExceptT CliError MainM PUnit := do
   let (leanInstall, _) ← opts.getInstall
   verifyLeanVersion leanInstall
 
-def parseScriptSpec (ws : Workspace) (spec : String) : Except CliError (Package × String) :=
+def parseScriptSpec (ws : Workspace) (spec : String) : Except CliError Script :=
   match spec.splitOn "/" with
-  | [script] => return (ws.root, script)
-  | [pkg, script] => return (← parsePackageSpec ws pkg, script)
+  | [scriptName] =>
+    match ws.findScript? scriptName with
+    | some script => return script
+    | none => throw <| CliError.unknownScript spec
+  | [pkg, scriptName] => do
+    let pkg ← parsePackageSpec ws pkg
+    match pkg.scripts.find? scriptName with
+    | some script => return script
+    | none => throw <| CliError.unknownScript spec
   | _ => throw <| CliError.invalidScriptSpec spec
 
 def parseTemplateSpec (spec : String) : Except CliError InitTemplate :=
@@ -220,11 +227,8 @@ protected nonrec def run : CliM PUnit := do
   let ws ← loadWorkspace config
   if let some spec ← takeArg? then
     let args ← takeArgs
-    let (pkg, scriptName) ← parseScriptSpec ws spec
-    if let some script := pkg.scripts.find? scriptName then
-        exit <| ← script.run args |>.run {opaqueWs := ws}
-    else do
-      throw <| CliError.unknownScript scriptName
+    let script ← parseScriptSpec ws spec
+    exit <| ← script.run args |>.run {opaqueWs := ws}
   else
     for script in ws.root.defaultScripts do
       exitIfErrorCode <| ← script.run [] |>.run {opaqueWs := ws}
@@ -236,13 +240,10 @@ protected def doc : CliM PUnit := do
   let config ← mkLoadConfig (← getThe LakeOptions)
   noArgsRem do
     let ws ← loadWorkspace config
-    let (pkg, scriptName) ← parseScriptSpec ws spec
-    if let some script := pkg.scripts.find? scriptName then
-      match script.doc? with
-      | some doc => IO.println doc
-      | none => throw <| CliError.missingScriptDoc scriptName
-    else
-      throw <| CliError.unknownScript scriptName
+    let script ← parseScriptSpec ws spec
+    match script.doc? with
+    | some doc => IO.println doc
+    | none => throw <| CliError.missingScriptDoc script.name
 
 protected def help : CliM PUnit := do
   IO.println <| helpScript <| (← takeArg?).getD ""
diff --git a/src/lake/Lake/Config/Script.lean b/src/lake/Lake/Config/Script.lean
index f0cbd7cb91..79b2fb9080 100644
--- a/src/lake/Lake/Config/Script.lean
+++ b/src/lake/Lake/Config/Script.lean
@@ -26,6 +26,8 @@ A package `Script` is a `ScriptFn` definition that is
 indexed by a `String` key and can be be run by `lake run <key> [-- <args>]`.
 -/
 structure Script where
+  /-- The full name of the `Script` (e.g., `pkg/script`). -/
+  name : String
   fn : ScriptFn
   doc? : Option String
   deriving Inhabited
diff --git a/src/lake/Lake/Config/Workspace.lean b/src/lake/Lake/Config/Workspace.lean
index 8bbba8fc14..14cdc9b318 100644
--- a/src/lake/Lake/Config/Workspace.lean
+++ b/src/lake/Lake/Config/Workspace.lean
@@ -72,6 +72,10 @@ def addPackage (pkg : Package) (self : Workspace) : Workspace :=
 @[inline] def findPackage? (name : Name) (self : Workspace) : Option (NPackage name) :=
   self.packageMap.find? name
 
+/-- Try to find a script in the workspace with the given name. -/
+def findScript? (script : Name) (self : Workspace) : Option Script :=
+  self.packageArray.findSome? (·.scripts.find? script)
+
 /-- Check if the module is local to any package in the workspace. -/
 def isLocalModule (mod : Name) (self : Workspace) : Bool :=
   self.packageMap.any fun _ pkg => pkg.isLocalModule mod
diff --git a/src/lake/Lake/Load/Package.lean b/src/lake/Lake/Load/Package.lean
index 06a42e2721..74082c108d 100644
--- a/src/lake/Lake/Load/Package.lean
+++ b/src/lake/Lake/Load/Package.lean
@@ -63,12 +63,15 @@ Load the remainder of a `Package`
 from its configuration environment after resolving its dependencies.
 -/
 def Package.finalize (self : Package) (deps : Array Package) : LogIO Package := do
-  let env := self.configEnv; let opts := self.leanOpts
+  let env := self.configEnv
+  let opts := self.leanOpts
+  let strName := self.name.toString (escape := false)
 
   -- Load Script, Facet, & Target Configurations
-  let scripts : NameMap Script ← mkTagMap env scriptAttr fun name => do
-    let fn ← IO.ofExcept <| evalConstCheck env opts ScriptFn ``ScriptFn name
-    return {fn, doc? := (← findDocString? env name)}
+  let scripts : NameMap Script ← mkTagMap env scriptAttr fun scriptName => do
+    let name := strName ++ "/" ++ scriptName.toString (escape := false)
+    let fn ← IO.ofExcept <| evalConstCheck env opts ScriptFn ``ScriptFn scriptName
+    return {name, fn, doc? := ← findDocString? env scriptName}
   let defaultScripts ← defaultScriptAttr.getAllEntries env |>.mapM fun name =>
     if let some script := scripts.find? name then pure script else
       error s!"package is missing script `{name}` marked as a default"
diff --git a/src/lake/examples/scripts/.gitignore b/src/lake/examples/scripts/.gitignore
index c6270e433c..b96c2ef548 100644
--- a/src/lake/examples/scripts/.gitignore
+++ b/src/lake/examples/scripts/.gitignore
@@ -1 +1 @@
-/lakefile.olean
+lakefile.olean
diff --git a/src/lake/examples/scripts/dep/lakefile.lean b/src/lake/examples/scripts/dep/lakefile.lean
new file mode 100644
index 0000000000..6d859dd9ef
--- /dev/null
+++ b/src/lake/examples/scripts/dep/lakefile.lean
@@ -0,0 +1,9 @@
+import Lake
+open Lake DSL
+
+package dep
+
+@[default_script]
+script hello do
+  IO.println "Hello from Dep!"
+  return 0
diff --git a/src/lake/examples/scripts/expected.out b/src/lake/examples/scripts/expected.out
index f48636d82f..47508d273a 100644
--- a/src/lake/examples/scripts/expected.out
+++ b/src/lake/examples/scripts/expected.out
@@ -1,3 +1,4 @@
+dep/hello
 scripts/dismiss
 scripts/greet
 Hello, world!
@@ -5,12 +6,15 @@ Hello, me!
 Display a greeting
 
 USAGE:
-  lake run greet [name]
+lake run greet [name]
 
 Greet the entity with the given name. Otherwise, greet the whole world.
 
+Hello from Dep!
+Hello from Dep!
 error: unknown script nonexistant
 error: unknown script nonexistant
+dep/hello
 scripts/dismiss
 scripts/greet
 Hello, world!
diff --git a/src/lake/examples/scripts/lake-manifest.json b/src/lake/examples/scripts/lake-manifest.json
new file mode 100644
index 0000000000..cd998b70db
--- /dev/null
+++ b/src/lake/examples/scripts/lake-manifest.json
@@ -0,0 +1,4 @@
+{"version": 5,
+ "packagesDir": "lake-packages",
+ "packages":
+ [{"path": {"opts": {}, "name": "dep", "inherited": false, "dir": ".\\dep"}}]}
diff --git a/src/lake/examples/scripts/lakefile.lean b/src/lake/examples/scripts/lakefile.lean
index 9552d46561..7208145c2d 100644
--- a/src/lake/examples/scripts/lakefile.lean
+++ b/src/lake/examples/scripts/lakefile.lean
@@ -2,6 +2,7 @@ import Lake
 open Lake DSL
 
 package scripts
+require dep from "dep"
 
 /--
 Display a greeting
@@ -19,7 +20,6 @@ script greet (args) do
     IO.println "Hello, world!"
   return 0
 
-
 @[default_script]
 script dismiss do
   IO.println "Goodbye!"
diff --git a/src/lake/examples/scripts/test.sh b/src/lake/examples/scripts/test.sh
index ba14283628..1356f62753 100755
--- a/src/lake/examples/scripts/test.sh
+++ b/src/lake/examples/scripts/test.sh
@@ -1,15 +1,23 @@
 #!/usr/bin/env bash
 set -exo pipefail
 
+# make prefix `/` behave on MSYS2
+[ "$OSTYPE" == "msys" ] && export MSYS2_ARG_CONV_EXCL=*
+
+./clean.sh
 LAKE=${LAKE:-../../build/bin/lake}
+
+$LAKE update
 $LAKE script list | tee produced.out
-$LAKE run scripts/greet | tee -a produced.out
+$LAKE run /greet | tee -a produced.out
 $LAKE script run greet me | tee -a produced.out
 $LAKE script doc greet | tee -a produced.out
+$LAKE script run hello | tee -a produced.out
+$LAKE script run dep/hello | tee -a produced.out
 ($LAKE script run nonexistant 2>&1 | tee -a produced.out) && false || true
 ($LAKE script doc nonexistant 2>&1 | tee -a produced.out) && false || true
 $LAKE scripts | tee -a produced.out
 $LAKE run | tee -a produced.out
 
-# TODO: Figure out why indent in USAGE disappears in the CI
+# FIXME: Figure out why indent in USAGE disappears in the CI
 diff --ignore-all-space --strip-trailing-cr expected.out produced.out