From 346da2c29cfe687c47f7dec929b3279795cc141e Mon Sep 17 00:00:00 2001
From: tydeu <tydeu@hatpress.net>
Date: Sat, 15 Apr 2023 20:07:47 -0400
Subject: [PATCH] feat: bare `lake run` default scripts

---
 .gitignore                     |  1 +
 Lake/CLI/Help.lean             |  5 ++++-
 Lake/CLI/Main.lean             | 19 +++++++++++--------
 Lake/Config/Package.lean       |  5 +++++
 Lake/Config/Script.lean        |  5 +++--
 Lake/DSL/Attributes.lean       |  6 ++++++
 Lake/DSL/Script.lean           | 19 ++++++++++++++-----
 Lake/Load/Package.lean         |  5 ++++-
 Lake/Util/Exit.lean            |  4 ++++
 examples/scripts/expected.out  | 17 +++++++++++++++++
 examples/scripts/lakefile.lean |  7 +++++++
 examples/scripts/test.sh       | 21 +++++++++++++--------
 12 files changed, 89 insertions(+), 25 deletions(-)
 create mode 100644 examples/scripts/expected.out

diff --git a/.gitignore b/.gitignore
index 1d60dbb262..59ebc302c5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
 /build
+produced.out
 result*
 # Do not commit the flake lockfile to avoid having to maintain it
 flake.lock
diff --git a/Lake/CLI/Help.lean b/Lake/CLI/Help.lean
index e3cc18e448..b738900017 100644
--- a/Lake/CLI/Help.lean
+++ b/Lake/CLI/Help.lean
@@ -163,7 +163,10 @@ USAGE:
   lake script run [<package>/]<script> [<args>...]
 
 This command runs the given `script` from `package`, passing `args` to it.
-Defaults to the root package."
+Defaults to the root package.
+
+A bare `run` command will run the default script(s) of the root package
+(with no arguments)."
 
 def helpScriptDoc :=
 "Print a script's docstring
diff --git a/Lake/CLI/Main.lean b/Lake/CLI/Main.lean
index 1cc079412e..fdbf91c32f 100644
--- a/Lake/CLI/Main.lean
+++ b/Lake/CLI/Main.lean
@@ -204,16 +204,19 @@ protected def list : CliM PUnit := do
 
 protected nonrec def run : CliM PUnit := do
   processOptions lakeOption
-  let spec ← takeArg "script name"; let args ← takeArgs
   let config ← mkLoadConfig (← getThe LakeOptions)
   let ws ← loadWorkspace config
-  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
+  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
+  else
+    for script in ws.root.defaultScripts do
+      exitIfErrorCode <| ← script.run [] |>.run {opaqueWs := ws}
+    exit 0
 
 protected def doc : CliM PUnit := do
   processOptions lakeOption
diff --git a/Lake/Config/Package.lean b/Lake/Config/Package.lean
index 1999b4ea65..5001bb26a7 100644
--- a/Lake/Config/Package.lean
+++ b/Lake/Config/Package.lean
@@ -205,6 +205,11 @@ structure Package where
   defaultTargets : Array Name := #[]
   /-- Scripts for the package. -/
   scripts : NameMap Script := {}
+  /--
+  The names of the package's scripts run by default
+  (i.e., on a bare `lake run` of the package).
+  -/
+  defaultScripts : Array Script := #[]
 
 instance : Nonempty Package :=
   have : Inhabited Environment := Classical.inhabited_of_nonempty inferInstance
diff --git a/Lake/Config/Script.lean b/Lake/Config/Script.lean
index f9f1bb08aa..f0cbd7cb91 100644
--- a/Lake/Config/Script.lean
+++ b/Lake/Config/Script.lean
@@ -3,6 +3,7 @@ Copyright (c) 2021 Mac Malone. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Mac Malone
 -/
+import Lake.Util.Exit
 import Lake.Config.Context
 
 namespace Lake
@@ -18,7 +19,7 @@ The type of a `Script`'s function.
 Similar to the `main` function's signature, except that its monad is
 also equipped with information about the Lake configuration.
 -/
-abbrev ScriptFn := (args : List String) → ScriptM UInt32
+abbrev ScriptFn := (args : List String) → ScriptM ExitCode
 
 /--
 A package `Script` is a `ScriptFn` definition that is
@@ -29,5 +30,5 @@ structure Script where
   doc? : Option String
   deriving Inhabited
 
-def Script.run (args : List String) (self : Script) : ScriptM UInt32 :=
+def Script.run (args : List String) (self : Script) : ScriptM ExitCode :=
   self.fn args
diff --git a/Lake/DSL/Attributes.lean b/Lake/DSL/Attributes.lean
index e9c2424054..3f0ec39945 100644
--- a/Lake/DSL/Attributes.lean
+++ b/Lake/DSL/Attributes.lean
@@ -17,6 +17,12 @@ initialize packageDepAttr : OrderedTagAttribute ←
 initialize scriptAttr : OrderedTagAttribute ←
   registerOrderedTagAttribute `script "mark a definition as a Lake script"
 
+initialize defaultScriptAttr : OrderedTagAttribute ←
+  registerOrderedTagAttribute `default_script "mark a Lake script as the package's default"
+    fun name => do
+      unless (← getEnv <&> (scriptAttr.hasTag · name)) do
+        throwError "attribute `default_script` can only be used on a `script`"
+
 initialize leanLibAttr : OrderedTagAttribute ←
   registerOrderedTagAttribute `lean_lib "mark a definition as a Lake Lean library target configuration"
 
diff --git a/Lake/DSL/Script.lean b/Lake/DSL/Script.lean
index 889f1f9c8b..0bfa583cdd 100644
--- a/Lake/DSL/Script.lean
+++ b/Lake/DSL/Script.lean
@@ -13,14 +13,23 @@ open Lean Parser Command
 syntax scriptDeclSpec :=
   ident (ppSpace simpleBinder)? (declValSimple <|> declValDo)
 
+/--
+Define a new Lake script for the package. Has two forms:
+
+```lean
+script «script-name» (args) do /- ... -/
+script «script-name» (args) := ...
+```
+-/
 scoped syntax (name := scriptDecl)
-(docComment)? "script " scriptDeclSpec : command
+(docComment)?  optional(Term.attributes) "script " scriptDeclSpec : command
 
 @[macro scriptDecl]
 def expandScriptDecl : Macro
-| `($[$doc?]? script $id:ident $[$args?]? do $seq $[$wds?]?) => do
-  `($[$doc?]? script $id:ident $[$args?]? := do $seq $[$wds?]?)
-| `($[$doc?]? script $id:ident $[$args?]? := $defn $[$wds?]?) => do
+| `($[$doc?]? $[$attrs?]? script $id:ident $[$args?]? do $seq $[$wds?]?) => do
+  `($[$doc?]? $[$attrs?]? script $id:ident $[$args?]? := do $seq $[$wds?]?)
+| `($[$doc?]? $[$attrs?]? script $id:ident $[$args?]? := $defn $[$wds?]?) => do
   let args ← expandOptSimpleBinder args?
-  `($[$doc?]? @[«script»] def $id : ScriptFn := fun $args => $defn $[$wds?]?)
+  let attrs := #[← `(Term.attrInstance| «script»)] ++ expandAttrs attrs?
+  `($[$doc?]? @[$attrs,*] def $id : ScriptFn := fun $args => $defn $[$wds?]?)
 | stx => Macro.throwErrorAt stx "ill-formed script declaration"
diff --git a/Lake/Load/Package.lean b/Lake/Load/Package.lean
index dab1ce1a6f..a206c1011f 100644
--- a/Lake/Load/Package.lean
+++ b/Lake/Load/Package.lean
@@ -69,6 +69,9 @@ def Package.finalize (self : Package) (deps : Array Package) : LogIO Package :=
   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 defaultScripts ← defaultScriptAttr.ext.getState 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"
   let leanLibConfigs ← IO.ofExcept <| mkTagMap env leanLibAttr fun name =>
     evalConstCheck env opts LeanLibConfig ``LeanLibConfig name
   let leanExeConfigs ← IO.ofExcept <| mkTagMap env leanExeAttr fun name =>
@@ -95,7 +98,7 @@ def Package.finalize (self : Package) (deps : Array Package) : LogIO Package :=
   return {self with
     opaqueDeps := deps.map (.mk ·)
     leanLibConfigs, leanExeConfigs, externLibConfigs
-    opaqueTargetConfigs, defaultTargets, scripts
+    opaqueTargetConfigs, defaultTargets, scripts, defaultScripts
   }
 
 /--
diff --git a/Lake/Util/Exit.lean b/Lake/Util/Exit.lean
index d4b0405088..2ce22f9a23 100644
--- a/Lake/Util/Exit.lean
+++ b/Lake/Util/Exit.lean
@@ -15,3 +15,7 @@ export MonadExit (exit)
 
 instance [MonadLift m n] [MonadExit m] : MonadExit n where
   exit rc := liftM (m := m) <| exit rc
+
+/-- Exit with `ExitCode` if it is not 0. Otherwise, continue. -/
+@[inline] def exitIfErrorCode [Pure m] [MonadExit m]  (rc : ExitCode) : m Unit :=
+  if rc != 0 then exit rc else pure ()
diff --git a/examples/scripts/expected.out b/examples/scripts/expected.out
new file mode 100644
index 0000000000..f48636d82f
--- /dev/null
+++ b/examples/scripts/expected.out
@@ -0,0 +1,17 @@
+scripts/dismiss
+scripts/greet
+Hello, world!
+Hello, me!
+Display a greeting
+
+USAGE:
+  lake run greet [name]
+
+Greet the entity with the given name. Otherwise, greet the whole world.
+
+error: unknown script nonexistant
+error: unknown script nonexistant
+scripts/dismiss
+scripts/greet
+Hello, world!
+Goodbye!
diff --git a/examples/scripts/lakefile.lean b/examples/scripts/lakefile.lean
index 513df8ec90..5d292e08fa 100644
--- a/examples/scripts/lakefile.lean
+++ b/examples/scripts/lakefile.lean
@@ -11,9 +11,16 @@ USAGE:
 
 Greet the entity with the given name. Otherwise, greet the whole world.
 -/
+@[default_script]
 script greet (args) do
   if h : 0 < args.length then
     IO.println s!"Hello, {args.get ⟨0, h⟩}!"
   else
     IO.println "Hello, world!"
   return 0
+
+
+@[default_script]
+script dismiss do
+  IO.println "Goodbye!"
+  return 0
diff --git a/examples/scripts/test.sh b/examples/scripts/test.sh
index 6c8fd9ffe0..c327b6226f 100755
--- a/examples/scripts/test.sh
+++ b/examples/scripts/test.sh
@@ -1,9 +1,14 @@
-set -ex
+#!/usr/bin/env bash
+set -exo pipefail
+
 LAKE=${LAKE:-../../build/bin/lake}
-$LAKE script list
-$LAKE run scripts/greet
-$LAKE script run greet me
-$LAKE script doc greet
-$LAKE script run nonexistant && false || true
-$LAKE script doc nonexistant && false || true
-$LAKE scripts
+$LAKE script list | tee produced.out
+$LAKE run scripts/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 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
+
+diff --strip-trailing-cr expected.out produced.out