Merge pull request #1614 from alexarchambault/better-revolver-input

Better revolver output

Author: alexarchambault

modules/cli/src/main/scala/scala/cli/commands/WatchUtil.scala

@@ -1,5 +1,7 @@
 package scala.cli.commands
 
+import scala.annotation.tailrec
+
 object WatchUtil {
 
   lazy val isDevMode: Boolean =
@@ -9,18 +11,47 @@ object WatchUtil {
   def waitMessage(message: String): String = {
     // Both short cuts actually always work, but Ctrl+C also exits mill in mill watch mode.
     val shortCut = if (isDevMode) "Ctrl+D" else "Ctrl+C"
-    val gray     = "\u001b[90m"
-    val reset    = Console.RESET
-    s"$gray$message, press $shortCut to exit, or press Enter to re-run.$reset"
+    gray(s"$message, press $shortCut to exit, or press Enter to re-run.")
+  }
+
+  private def gray(message: String): String = {
+    val gray  = "\u001b[90m"
+    val reset = Console.RESET
+    s"$gray$message$reset"
   }
 
   def printWatchMessage(): Unit =
     System.err.println(waitMessage("Watching sources"))
 
-  def waitForCtrlC(onPressEnter: () => Unit = () => ()): Unit = {
+  def printWatchWhileRunningMessage(): Unit =
+    System.err.println(gray("Watching sources while your program is running."))
+
+  def waitForCtrlC(
+    onPressEnter: () => Unit = () => (),
+    shouldReadInput: () => Boolean = () => true
+  ): Unit = synchronized {
+
+    @tailrec
+    def readNextChar(): Int =
+      if (shouldReadInput())
+        try System.in.read()
+        catch {
+          case _: InterruptedException =>
+            // Actually never called, as System.in.read isn't interruptible…
+            // That means we sometimes read input when we shouldn't.
+            readNextChar()
+        }
+      else {
+        try wait()
+        catch {
+          case _: InterruptedException =>
+        }
+        readNextChar()
+      }
+
     var readKey = -1
     while ({
-      readKey = System.in.read()
+      readKey = readNextChar()
       readKey != -1
     })
       if (readKey == '\n')

modules/cli/src/main/scala/scala/cli/commands/run/Run.scala

@@ -149,17 +149,22 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
             val onExitProcess = process.onExit().thenApply { p1 =>
               val retCode = p1.exitValue()
               onExitOpt.foreach(_())
-              if (retCode != 0)
-                if (allowTerminate)
+              (retCode, allowTerminate) match {
+                case (0, true) =>
+                case (0, false) =>
+                  val gray  = "\u001b[90m"
+                  val reset = Console.RESET
+                  System.err.println(s"${gray}Program exited with return code $retCode.$reset")
+                case (_, true) =>
                   sys.exit(retCode)
-                else {
+                case (_, false) =>
                   val red      = Console.RED
                   val lightRed = "\u001b[91m"
                   val reset    = Console.RESET
                   System.err.println(
                     s"${red}Program exited with return code $lightRed$retCode$red.$reset"
                   )
-                }
+              }
             }
 
             Some((process, onExitProcess))
@@ -191,7 +196,9 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
       )
 
     if (options.sharedRun.watch.watchMode) {
-      var processOpt = Option.empty[(Process, CompletableFuture[_])]
+      var processOpt      = Option.empty[(Process, CompletableFuture[_])]
+      var shouldReadInput = false
+      var mainThreadOpt   = Option.empty[Thread]
       val watcher = Build.watch(
         inputs,
         initialBuildOptions,
@@ -202,16 +209,23 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
         buildTests = false,
         partial = None,
         actionableDiagnostics = actionableDiagnostics,
-        postAction = () => WatchUtil.printWatchMessage()
+        postAction = () =>
+          if (processOpt.exists(_._1.isAlive()))
+            WatchUtil.printWatchWhileRunningMessage()
+          else
+            WatchUtil.printWatchMessage()
       ) { res =>
         for ((process, onExitProcess) <- processOpt) {
           onExitProcess.cancel(true)
           ProcUtil.interruptProcess(process, logger)
         }
         res.orReport(logger).map(_.main).foreach {
           case s: Build.Successful =>
-            for ((proc, _) <- processOpt) // If the process doesn't exit, send SIGKILL
-              if (proc.isAlive) ProcUtil.forceKillProcess(proc, logger)
+            for ((proc, _) <- processOpt if proc.isAlive)
+              // If the process doesn't exit, send SIGKILL
+              ProcUtil.forceKillProcess(proc, logger)
+            shouldReadInput = false
+            mainThreadOpt.foreach(_.interrupt())
             val maybeProcess = maybeRun(
               s,
               allowTerminate = false,
@@ -221,18 +235,34 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
             )
               .orReport(logger)
               .flatten
+              .map {
+                case (proc, onExit) =>
+                  if (options.sharedRun.watch.restart)
+                    onExit.thenApply { _ =>
+                      shouldReadInput = true
+                      mainThreadOpt.foreach(_.interrupt())
+                    }
+                  (proc, onExit)
+              }
             s.copyOutput(options.shared)
             if (options.sharedRun.watch.restart)
               processOpt = maybeProcess
-            else
+            else {
               for ((proc, onExit) <- maybeProcess)
                 ProcUtil.waitForProcess(proc, onExit)
+              shouldReadInput = true
+              mainThreadOpt.foreach(_.interrupt())
+            }
           case _: Build.Failed =>
             System.err.println("Compilation failed")
         }
       }
-      try WatchUtil.waitForCtrlC(() => watcher.schedule())
-      finally watcher.dispose()
+      mainThreadOpt = Some(Thread.currentThread())
+      try WatchUtil.waitForCtrlC(() => watcher.schedule(), () => shouldReadInput)
+      finally {
+        mainThreadOpt = None
+        watcher.dispose()
+      }
     }
     else {
       val builds =