1 files changed, 39 insertions, 40 deletions

diff --git a/stage2/Stage2.scala b/stage2/Stage2.scala
index d43c4f6..33a67ac 100644
--- a/stage2/Stage2.scala
+++ b/stage2/Stage2.scala
@@ -1,6 +1,6 @@
 package cbt
 import java.io._
-import java.util._
+import java.util.{Set=>_,_}
 
 object Stage2 extends Stage2Base{
   def getBuild(context: Context) = {
@@ -13,19 +13,14 @@ object Stage2 extends Stage2Base{
     import paths._
     val lib = new Lib(args.logger)
     logger.stage2(s"Stage2 start")
-    val loop = args.args.lift(0) == Some("loop")
 
-    val taskIndex = if (loop) {
-      1
-    } else {
-      0
-    }
-    val task = args.args.lift( taskIndex )
+    val task = args.args.lift( 0 )
 
+    import scala.collection.JavaConverters._
     val context: Context = new ContextImplementation(
       args.cwd,
       args.cwd,
-      args.args.drop( taskIndex +1 ).toArray,
+      args.args.drop( 1 ).toArray,
       logger.enabledLoggers.toArray,
       logger.start,
       args.stage2LastModified,
@@ -36,38 +31,42 @@ object Stage2 extends Stage2Base{
       args.cbtHome,
       args.cbtHome,
       args.compatibilityTarget,
-      null
+      null,
+      NailgunLauncher.compatibilitySourceFiles.asScala.toArray[File]
+        ++ NailgunLauncher.nailgunLauncherSourceFiles.asScala.toArray[File]
+        ++ NailgunLauncher.stage1SourceFiles.asScala.toArray[File]
+        ++ args.stage2sourceFiles.toArray[File]
     )
-    val first = lib.loadRoot( context )
-    val build = first.finalBuild( context.cwd )
-
-    val res =
-      if (loop) {
-        // TODO: this should allow looping over task specific files, like test files as well
-        val triggerFiles = first.triggerLoopFiles.map(lib.realpath)
-        val triggerCbtFiles = Seq( nailgun, stage1, stage2 ).map(lib.realpath _)
-        val allTriggerFiles = triggerFiles ++ triggerCbtFiles
-
-        logger.loop("Looping change detection over:\n - "++allTriggerFiles.mkString("\n - "))
-
-        lib.watch(allTriggerFiles){
-          case file if triggerCbtFiles.exists(file.toString startsWith _.toString) =>
-            logger.loop("Change is in CBT's own source code.")
-            logger.loop("Restarting CBT.")
-            scala.util.control.Breaks.break
-
-          case file if triggerFiles.exists(file.toString startsWith _.toString) =>
-            val build = lib.loadRoot(context).finalBuild( context.cwd )
-            logger.loop(s"Re-running $task for " ++ build.show)
-            lib.callReflective(build, task, context)
-        }
-        ExitCode.Success
-      } else {
-        val code = lib.callReflective(build, task, context)
-        logger.stage2(s"Stage2 end")
-        code
+    def loop( code: ExitCode, files: () => Set[File] ): ExitCode = {
+      code match {
+        case c@ExitCode(253 | 130) => c // CBT signals loop | user pressed ctrl+C
+        case c if !task.contains("loop") => c
+        case c =>
+          // this allows looping over broken builds
+          lib.watch{ files }()
+          c
       }
-
-    res
+    }
+    val code = lib.trapExitCode{
+      val first = lib.loadRoot( context )
+      val build = first.finalBuild( context.cwd )
+      val code = lib.callReflective(build, task, context)
+      if( !context.loopFile.exists ){
+        loop( code, () => build.triggerLoopFiles )
+      }
+      code
+    }
+    if( context.loopFile.exists ){
+      loop(
+        code,
+        () => {
+          val files = context.loopFile.readAsString.split("\n").map(new File(_)).toSet
+          logger.loop("Looping change detection over:\n - "++files.mkString("\n - "))
+          files
+        }
+      )
+    }
+    logger.stage2(s"Stage2 end with exit code "+code.integer)
+    code
   }
 }