revamp loop feature

now CBT and builds pass their file names to the current build
via the context. The build then simply blocks until any file changes.
Then it returns with a special exit code, which the bash script picks
up and restarts CBT. Thats works well for looping over project files.
It works less well for looping over builds and CBT itself. For this
a build has to success once, so that the .cbt-loop.tmp file exists.
Then looping works for cbt and builds, but the file list is not
updated in case of compile errors, etc.

Fixes
- https://github.com/cvogt/cbt/issues/406
- https://github.com/cvogt/cbt/issues/405
- https://github.com/cvogt/cbt/issues/202
- https://github.com/cvogt/cbt/issues/50
- https://github.com/cvogt/cbt/issues/22

We should improve for 1.0 in https://github.com/cvogt/cbt/issues/419
to handle looping over build files and cbt itself smarter.

15 files changed, 202 insertions, 109 deletions

diff --git a/.gitignore b/.gitignore
index 5a580e1..a924fed 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,3 +15,4 @@ node_modules
 examples/dotty-example/_site
 examples/scalajs-plain-example/server/public/generated
 examples/scalajs-react-example/server/public/generated
+.cbt-loop.tmp
diff --git a/cbt b/cbt
index 3a856ec..1f3d581 100755
--- a/cbt
+++ b/cbt
@@ -199,6 +199,7 @@ stage1 () {
 			log "Running JVM directly" "$@"
 			# JVM options to improve startup time. See https://github.com/cvogt/cbt/pull/262
 			java $JAVA_OPTS $DEBUG -Xmx6072m -Xss10M -XX:MaxJavaStackTraceDepth=-1 -XX:+TieredCompilation -XX:TieredStopAtLevel=1 -Xverify:none -cp $NAILGUN$TARGET cbt.NailgunLauncher $(time_taken) "$CWD" "$@"
+			exitCode=$?
 		else
 			log "Running via background process (nailgun)" "$@"
 			for i in 0 1 2 3 4 5 6 7 8 9; do
@@ -222,17 +223,52 @@ stage1 () {
 			done
 			log "Running CBT via Nailgun." "$@"
 			$NG cbt.NailgunLauncher $(time_taken) "$CWD" "$@"
+			exitCode=$?
 		fi
-		exitCode=$?
 		log "Done running CBT." "$@"
 	fi
 }
 
+
+loop=1
+case "$1" in
+	"loop") loop=0
+esac
+case "$2" in
+	"loop") loop=0
+esac
+
+CBT_SIGNALS_LOOPING=253
+USER_PRESSED_CTRL_C=130
+
+CBT_LOOP_FILE="$CWD/target/.cbt-loop.tmp"
 while true; do
 	stage1 "$@"
-	if [ ! "$1" = "loop" ]; then
+	if [ ! $loop -eq 0 ] || [ $exitCode -eq $USER_PRESSED_CTRL_C ]; then
+		log "not looping, exiting" "$@"
 		break
+	else
+		if [ ! $exitCode -eq $CBT_SIGNALS_LOOPING ]; then
+			log "exitCode $exitCode" "$@"
+			which fswatch >/dev/null 2>/dev/null
+			export fswatch_installed=$?
+			if [ -f "$CBT_LOOP_FILE" ]; then
+				if [ $fswatch_installed -eq 0 ]; then
+					# fswatch allows looping over CBT's sources itself
+					log "fswatch found. looping." "$@"
+					files=($(cat "$CBT_LOOP_FILE"))
+					fswatch --one-event "${files[@]}"
+				else
+					log "fswatch not installed, stopping cbt" "$@"
+					break
+				fi
+			else
+				log "no $CBT_LOOP_FILE file, stopping cbt" "$@"
+				break
+			fi
+		fi
 	fi
+
 	echo "======= Restarting CBT =======" 1>&2
 done
 
diff --git a/compatibility/Context.java b/compatibility/Context.java
index 389d401..1657ef0 100644
--- a/compatibility/Context.java
+++ b/compatibility/Context.java
@@ -32,6 +32,9 @@ public interface Context{
   public abstract File cbtRootHome(); // REMOVE
   public abstract File compatibilityTarget(); // maybe replace this with search in the classloader for it?
   public abstract BuildInterface parentBuildOrNull();
+  public default File[] triggerLoopFilesArray(){
+    return new File[0]; // REMOVE default value on next compatibility breaking release
+  }
 
   // deprecated methods
   @java.lang.Deprecated
diff --git a/doc/docs.md b/doc/docs.md
new file mode 100644
index 0000000..5a34e5b
--- /dev/null
+++ b/doc/docs.md
@@ -0,0 +1,8 @@
+
+### Clearing the screen during task looking
+```
+override def loop{
+  lib.clearScreen
+  super.compile
+}
+```
diff --git a/nailgun_launcher/NailgunLauncher.java b/nailgun_launcher/NailgunLauncher.java
index d89c764..958b052 100644
--- a/nailgun_launcher/NailgunLauncher.java
+++ b/nailgun_launcher/NailgunLauncher.java
@@ -51,6 +51,10 @@ public class NailgunLauncher{
     ClassLoader.getSystemClassLoader().getParent()
   );
 
+  public static List<File> compatibilitySourceFiles;
+  public static List<File> nailgunLauncherSourceFiles;
+  public static List<File> stage1SourceFiles;
+
   public static void main( String[] args ) throws Throwable {
     long _start = System.currentTimeMillis();
     if(args[0].equals("check-alive")){
@@ -80,7 +84,6 @@ public class NailgunLauncher{
     }
     // ---------------------
 
-    _assert(System.getenv("CBT_HOME") != null, "environment variable CBT_HOME not defined");
     String CBT_HOME = System.getenv("CBT_HOME");
     String cache = CBT_HOME + "/cache/";
     String compatibilityTarget = CBT_HOME + "/compatibility/" + TARGET;
@@ -132,7 +135,8 @@ public class NailgunLauncher{
     final String compatibilityTarget, final ClassLoaderCache classLoaderCache
   ) throws Throwable {
     _assert(TARGET != null, "environment variable TARGET not defined");
-    String nailgunTarget = cbtHome + "/" + NAILGUN + TARGET;
+    String nailgunSources = cbtHome + "/" + NAILGUN;
+    String nailgunTarget = nailgunSources + TARGET;
     String stage1Sources = cbtHome + "/" + STAGE1;
     String stage1Target = stage1Sources + TARGET;
     File compatibilitySources = new File(cbtHome + "/compatibility");
@@ -142,13 +146,20 @@ public class NailgunLauncher{
     ClassLoader rootClassLoader = new CbtURLClassLoader( new URL[]{}, ClassLoader.getSystemClassLoader().getParent() ); // wrap for caching
     EarlyDependencies earlyDeps = new EarlyDependencies(mavenCache, mavenUrl, classLoaderCache, rootClassLoader);
 
+    nailgunLauncherSourceFiles = new ArrayList<File>();
+    for( File f: new File(nailgunSources).listFiles() ){
+      if( f.isFile() && f.toString().endsWith(".java") ){
+        nailgunLauncherSourceFiles.add(f);
+      }
+    }
+
     long nailgunLauncherLastModified = new File( nailgunTarget + "../classes.last-success" ).lastModified();
 
     long compatibilityLastModified;
     if(!compatibilityTarget.startsWith(cbtHome)){
       compatibilityLastModified = new File( compatibilityTarget + "../classes.last-success" ).lastModified();
     } else {
-      List<File> compatibilitySourceFiles = new ArrayList<File>();
+      compatibilitySourceFiles = new ArrayList<File>();
       for( File f: compatibilitySources.listFiles() ){
         if( f.isFile() && f.toString().endsWith(".java") ){
           compatibilitySourceFiles.add(f);
@@ -174,7 +185,7 @@ public class NailgunLauncher{
       append( append( nailgunClasspathArray, compatibilityTarget ), stage1Target );
     String stage1Classpath = classpath( stage1ClasspathArray );
 
-    List<File> stage1SourceFiles = new ArrayList<File>();
+    stage1SourceFiles = new ArrayList<File>();
     for( File f: new File(stage1Sources).listFiles() ){
       if( f.isFile() && f.toString().endsWith(".scala") ){
         stage1SourceFiles.add(f);
diff --git a/stage1/ContextImplementation.scala b/stage1/ContextImplementation.scala
index b263ef4..90d9d5f 100644
--- a/stage1/ContextImplementation.scala
+++ b/stage1/ContextImplementation.scala
@@ -15,7 +15,8 @@ class ContextImplementation(
   override val cbtHome: File,
   override val cbtRootHome: File,
   override val compatibilityTarget: File,
-  override val parentBuildOrNull: BuildInterface
+  override val parentBuildOrNull: BuildInterface,
+  override val triggerLoopFilesArray: Array[File]
 ) extends Context{
   @deprecated("this method is replaced by workingDirectory","")
   def projectDirectory = workingDirectory
diff --git a/stage1/Stage1.scala b/stage1/Stage1.scala
index 714ed65..85ec0e9 100644
--- a/stage1/Stage1.scala
+++ b/stage1/Stage1.scala
@@ -42,7 +42,8 @@ class Stage2Args(
   val stage2LastModified: Long,
   val cache: File,
   val cbtHome: File,
-  val compatibilityTarget: File
+  val compatibilityTarget: File,
+  val stage2sourceFiles: Seq[File]
 )(
   implicit val transientCache: java.util.Map[AnyRef,AnyRef], val classLoaderCache: ClassLoaderCache, val logger: Logger
 ){
@@ -57,7 +58,7 @@ object Stage1{
   def getBuild( _context: java.lang.Object, buildStage1: BuildStage1Result ) = {
     val context = _context.asInstanceOf[Context]
     val logger = new Logger( context.enabledLoggers, buildStage1.start )
-    val (cbtLastModified, classLoader) = buildStage2(
+    val (_, cbtLastModified, classLoader) = buildStage2(
       buildStage1,
       context.cbtHome,
       context.cache
@@ -76,7 +77,9 @@ object Stage1{
 
   def buildStage2(
     buildStage1: BuildStage1Result, cbtHome: File, cache: File
-  )(implicit transientCache: java.util.Map[AnyRef,AnyRef], classLoaderCache: ClassLoaderCache, logger: Logger): (Long, ClassLoader) = {
+  )(
+    implicit transientCache: java.util.Map[AnyRef,AnyRef], classLoaderCache: ClassLoaderCache, logger: Logger
+  ): (Seq[File], Long, ClassLoader) = {
 
     import buildStage1._
 
@@ -149,7 +152,7 @@ object Stage1{
       )
     }
 
-    ( stage2LastModified, stage2ClassLoader )
+    ( stage2sourceFiles, stage2LastModified, stage2ClassLoader )
   }
 
   def run(
@@ -166,7 +169,7 @@ object Stage1{
     implicit val transientCache: java.util.Map[AnyRef,AnyRef] = new java.util.HashMap
     implicit val classLoaderCache = new ClassLoaderCache( persistentCache )
 
-    val (stage2LastModified, classLoader) = buildStage2( buildStage1, cbtHome, cache )
+    val (stage2sourceFiles, stage2LastModified, classLoader) = buildStage2( buildStage1, cbtHome, cache )
 
     val stage2Args = new Stage2Args(
       new File( args.args(0) ),
@@ -175,7 +178,8 @@ object Stage1{
       stage2LastModified = stage2LastModified,
       cache,
       cbtHome,
-      new File(buildStage1.compatibilityClasspath)
+      new File(buildStage1.compatibilityClasspath),
+      stage2sourceFiles
     )
 
     logger.stage1(s"Run Stage2")
@@ -193,7 +197,7 @@ object Stage1{
         case _ => ExitCode.Success
       }
     ).integer
-    logger.stage1(s"Stage1 end")
+    logger.stage1(s"Stage1 end with exit code " + exitCode)
     return exitCode;
   }
 }
diff --git a/stage1/Stage1Lib.scala b/stage1/Stage1Lib.scala
index 01115d4..f701c72 100644
--- a/stage1/Stage1Lib.scala
+++ b/stage1/Stage1Lib.scala
@@ -349,20 +349,22 @@ ${sourceFiles.sorted.mkString(" \\\n")}
     res
   }
 
-  def trapExitCode( code: => ExitCode ): ExitCode = {
+  def trapExitCodeOrValue[T]( result: => T ): Either[ExitCode,T] = {
     val trapExitCodeBefore = TrapSecurityManager.trapExitCode().get
     try{
       TrapSecurityManager.trapExitCode().set(true)
-      code
+      Right( result )
     } catch {
       case CatchTrappedExitCode(exitCode) =>
         logger.stage1(s"caught exit code $exitCode")
-        exitCode
+        Left( exitCode )
     } finally {
       TrapSecurityManager.trapExitCode().set(trapExitCodeBefore)
     }
   }
 
+  def trapExitCode( code: => ExitCode ): ExitCode = trapExitCodeOrValue(code).merge
+
   def ScalaDependency(
     groupId: String, artifactId: String, version: String, classifier: Classifier = Classifier.none,
     scalaMajorVersion: String, verifyHash: Boolean = true
diff --git a/stage1/cbt.scala b/stage1/cbt.scala
index 062e11d..05737d0 100644
--- a/stage1/cbt.scala
+++ b/stage1/cbt.scala
@@ -87,9 +87,9 @@ object `package`{
   implicit class BuildInterfaceExtensions(build: BuildInterface){
     import build._
     // TODO: if every build has a method triggers a callback if files change
-    // then we wouldn't need this and could provide this method from a 
+    // then we wouldn't need this and could provide this method from a
     // plugin rather than hard-coding trigger files stuff in cbt
-    def triggerLoopFiles: Seq[File] = triggerLoopFilesArray.to
+    def triggerLoopFiles: Set[File] = triggerLoopFilesArray.to
   }
   implicit class ArtifactInfoExtensions(subject: ArtifactInfo){
     import subject._
@@ -121,6 +121,9 @@ object `package`{
     def scalaVersion = Option(scalaVersionOrNull)
     def parentBuild = Option(parentBuildOrNull)
     def cbtLastModified: scala.Long = subject.cbtLastModified
+    def triggerLoopFiles: Set[File] = triggerLoopFilesArray.toSet[File]
+
+    private[cbt] def loopFile = cwd / "target/.cbt-loop.tmp"
 
     def copy(
       workingDirectory: File = workingDirectory,
@@ -129,7 +132,8 @@ object `package`{
       cbtLastModified: Long = cbtLastModified,
       scalaVersion: Option[String] = scalaVersion,
       cbtHome: File = cbtHome,
-      parentBuild: Option[BuildInterface] = None
+      parentBuild: Option[BuildInterface] = None,
+      triggerLoopFiles: Set[File] = Set()
     ): Context = new ContextImplementation(
       workingDirectory,
       cwd,
@@ -144,7 +148,8 @@ object `package`{
       cbtHome,
       cbtRootHome,
       compatibilityTarget,
-      parentBuild.getOrElse(null)
+      parentBuild.getOrElse(null),
+      (triggerLoopFiles ++ triggerLoopFilesArray.toSet[File]).toArray
     )
   }
 }
diff --git a/stage2/BasicBuild.scala b/stage2/BasicBuild.scala
index 68bda15..bc6f9d9 100644
--- a/stage2/BasicBuild.scala
+++ b/stage2/BasicBuild.scala
@@ -2,6 +2,7 @@ package cbt
 
 import java.io._
 import java.net._
+import java.nio.file._
 
 class BasicBuild(final val context: Context) extends BaseBuild
 trait BaseBuild extends BuildInterface with DependencyImplementation with TriggerLoop with SbtDependencyDsl{
@@ -124,7 +125,11 @@ trait BaseBuild extends BuildInterface with DependencyImplementation with Trigge
     pathToNestedBuild: _*
   )
 
-  def triggerLoopFiles: Seq[File] = sources ++ transitiveDependencies.collect{ case b: TriggerLoop => b.triggerLoopFiles }.flatten
+  def triggerLoopFiles: Set[File] = (
+    context.triggerLoopFiles
+    ++ sources
+    ++ transitiveDependencies.collect{ case b: TriggerLoop => b.triggerLoopFiles }.flatten
+  )
 
   def localJars: Seq[File] =
     Seq(projectDirectory ++ "/lib")
@@ -318,4 +323,17 @@ trait BaseBuild extends BuildInterface with DependencyImplementation with Trigge
   final def crossScalaVersionsArray = Array(scalaVersion)
 
   def publish: Seq[URL] = Seq()
+
+  def loop = {
+    lib.callReflective(this, context.args.headOption, context.copy(args=context.args.drop(1)))
+    val files = triggerLoopFiles
+    lib.watch{ () =>
+      logger.loop("Looping change detection over:\n - "++files.mkString("\n - "))
+      files
+    }()
+    context.loopFile.getParentFile.mkdirs
+    lib.write( context.loopFile, files.mkString("\n"), StandardOpenOption.CREATE )
+
+    ExitCode(253) // signal bash script to restart
+  }
 }
diff --git a/stage2/BuildBuild.scala b/stage2/BuildBuild.scala
index 9d3458d..9811db5 100644
--- a/stage2/BuildBuild.scala
+++ b/stage2/BuildBuild.scala
@@ -35,9 +35,10 @@ trait BuildBuildWithoutEssentials extends BaseBuild{
     s"You can't extend ${lib.buildBuildClassName} in: " + projectDirectory + "/" + lib.buildDirectoryName
   )
 
-  protected final val managedContext = context.copy(
+  protected def managedContext = context.copy(
     workingDirectory = managedBuildDirectory,
-    parentBuild=Some(this)
+    parentBuild=Some(this),
+    triggerLoopFiles = triggerLoopFiles
   )
 
   override def dependencies =
@@ -108,7 +109,7 @@ trait BuildBuildWithoutEssentials extends BaseBuild{
         throw new Exception(s"Your ${lib.buildClassName} class needs to extend BaseBuild in: "+projectDirectory, e)
     }
   }
-  override def triggerLoopFiles = super.triggerLoopFiles ++ managedBuild.triggerLoopFiles
+
   @deprecated("use finalbuild(File)","")
   override def finalBuild: BuildInterface = finalBuild( context.cwd )
   override def finalBuild( current: File ): BuildInterface = {
diff --git a/stage2/BuildDependency.scala b/stage2/BuildDependency.scala
index d7520d8..9e35c41 100644
--- a/stage2/BuildDependency.scala
+++ b/stage2/BuildDependency.scala
@@ -12,7 +12,7 @@ sealed abstract class ProjectProxy extends Ha{
 */
 trait TriggerLoop extends DependencyImplementation{
   final def triggerLoopFilesArray = triggerLoopFiles.toArray
-  def triggerLoopFiles: Seq[File]
+  def triggerLoopFiles: Set[File]
 }
 /** You likely want to use the factory method in the BasicBuild class instead of this. */
 object DirectoryDependency{
@@ -52,7 +52,7 @@ object DirectoryDependency{
 /*
 case class DependencyOr(first: DirectoryDependency, second: JavaDependency) extends ProjectProxy with DirectoryDependencyBase{
   val isFirst = new File(first.projectDirectory).exists
-  def triggerLoopFiles = if(isFirst) first.triggerLoopFiles else Seq()
+  def triggerLoopFiles = if(isFirst) first.triggerLoopFiles else Set()
   protected val delegate = if(isFirst) first else second
 }
 */
diff --git a/stage2/Lib.scala b/stage2/Lib.scala
index 46668b7..5e35ea7 100644
--- a/stage2/Lib.scala
+++ b/stage2/Lib.scala
@@ -160,7 +160,7 @@ final class Lib(val logger: Logger) extends Stage1Lib(logger){
       val name = NameTransformer.decode(taskName)
       logger.lib("Calling task " ++ taskName.toString)
       taskMethods(obj.getClass).get(name).map{ method =>
-        Option(method.invoke(obj) /* null in case of Unit */ ).getOrElse(().asInstanceOf[AnyRef]) match {
+        Option(trapExitCodeOrValue(method.invoke(obj)).merge /* null in case of Unit */ ).getOrElse(().asInstanceOf[AnyRef]) match {
           case code if code.getClass.getSimpleName == "ExitCode" =>
             // FIXME: ExitCode needs to be part of the compatibility interfaces
             Seq((None, Some(ExitCode(Stage0Lib.get(code,"integer").asInstanceOf[Int])), None))
@@ -480,49 +480,50 @@ final class Lib(val logger: Logger) extends Stage1Lib(logger){
 
 
   // code for continuous compile
-  def watch(files: Seq[File])(action: PartialFunction[File, Unit]): Unit = {
+  def watch[T]( files: () => Set[File] )(
+    action: Seq[File] => Option[T] = (f: Seq[File]) => Some(f)
+  ): T = {
     import com.barbarysoftware.watchservice._
     import scala.collection.JavaConversions._
-    val watcher = WatchService.newWatchService
-
-    val realFiles = files.map(realpath)
-
-    realFiles.map{
-      // WatchService can only watch folders
-      case file if file.isFile => dirname(file)
-      case file => file
-    }.distinct.map{ file =>
-      val watchableFile = new WatchableFile(file)
-      val key = watchableFile.register(
-        watcher,
-        StandardWatchEventKind.ENTRY_CREATE,
-        StandardWatchEventKind.ENTRY_DELETE,
-        StandardWatchEventKind.ENTRY_MODIFY
-      )
-    }
-
-    scala.util.control.Breaks.breakable{
-      while(true){
-        logger.loop("Waiting for file changes...")
-        logger.loop("Waiting for file changes...2")
-        Option(watcher.take).map{
-          key =>
-          val changedFiles = key
-            .pollEvents
-            .toVector
-            .filterNot(_.kind == StandardWatchEventKind.OVERFLOW)
-            .map(_.context.toString)
-            // make sure we don't react on other files changed
-            // in the same folder like the files we care about
-            .filter{ name => realFiles.exists(name startsWith _.toString) }
-            .map(new File(_))
-
-          changedFiles.foreach( f => logger.loop( "Changed: " ++ f.toString ) )
-          changedFiles.collect(action)
-          key.reset
-        }
+    Iterator.continually{
+      val watcher = WatchService.newWatchService
+      val realFiles = files().map(realpath)
+      realFiles.map{
+        // WatchService can only watch folders
+        case file if file.isFile => dirname(file)
+        case file => file
+      }.map{ file =>
+        val watchableFile = new WatchableFile(file)
+        val key = watchableFile.register(
+          watcher,
+          StandardWatchEventKind.ENTRY_CREATE,
+          StandardWatchEventKind.ENTRY_DELETE,
+          StandardWatchEventKind.ENTRY_MODIFY
+        )
       }
-    }
+      Option( watcher.take ) -> realFiles
+    }.collect{
+      case (Some(key),f) => key -> f
+    }.map{ case (key, realFiles) =>
+      logger.loop("Waiting for file changes...")
+      val changedFiles = key
+        .pollEvents
+        .toVector
+        .filterNot(_.kind == StandardWatchEventKind.OVERFLOW)
+        .map(_.context.toString)
+        // make sure we don't react on other files changed
+        // in the same folder like the files we care about
+        .filter{ name => realFiles.exists(name startsWith _.toString) }
+        .map(new File(_))
+
+        changedFiles.foreach( f => logger.loop( "Changed: " ++ f.toString ) )
+      val res = action(changedFiles)
+      key.reset
+      res
+    }.filterNot(_.isEmpty)
+     .take(1)
+     .toList
+     .head.get
   }
 
   def findInnerMostModuleDirectory(directory: File): File = {
@@ -559,4 +560,6 @@ final class Lib(val logger: Logger) extends Stage1Lib(logger){
 
     ( results.map(_.right.toOption).flatten.flatten, results.map(_.left.toOption).flatten )
   }
+
+  def clearScreen = System.err.println( (27.toChar +: "[2J").mkString )
 }
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
   }
 }
diff --git a/test/test.scala b/test/test.scala
index 45315ce..4d1d7fb 100644
--- a/test/test.scala
+++ b/test/test.scala
@@ -133,7 +133,8 @@ object Main{
         cbtHome,
         cbtHome,
         cbtHome ++ "/compatibilityTarget",
-        null
+        null,
+        Array()
       )
 
       val b = new BasicBuild(noContext){