First pass at cleaning up evaluator caching system

6 files changed, 45 insertions, 44 deletions

diff --git a/main/src/mill/Main.scala b/main/src/mill/Main.scala
index e05423cd..c8e445c9 100644
--- a/main/src/mill/Main.scala
+++ b/main/src/mill/Main.scala
@@ -6,6 +6,7 @@ import ammonite.main.Cli
 import ammonite.main.Cli.{formatBlock, genericSignature, replSignature}
 import ammonite.ops._
 import ammonite.util.Util
+import mill.eval.Evaluator
 import mill.main.MainRunner
 
 object Main {
@@ -21,12 +22,12 @@ object Main {
     System.exit(if(result) 0 else 1)
   }
   def main0(args: Array[String],
-            mainRunner: Option[(Cli.Config, MainRunner)],
+            stateCache: Option[Evaluator.State],
             mainInteractive: Boolean,
             watchInterrupted: () => Boolean,
             stdin: InputStream,
             stdout: PrintStream,
-            stderr: PrintStream): (Boolean, Option[(Cli.Config, MainRunner)]) = {
+            stderr: PrintStream): (Boolean, Option[Evaluator.State]) = {
     import ammonite.main.Cli
 
     val removed = Set("predef-code", "home", "no-home-predef")
@@ -73,18 +74,15 @@ object Main {
           config.copy(home = pwd / "out" / ".ammonite", colored = Some(mainInteractive)),
           stdout, stderr, stdin,
           watchInterrupted,
-          mainRunner match{
-            case Some((c, mr)) if c.copy(storageBackend = null) == cliConfig.copy(storageBackend = null) =>
-              mr.lastEvaluator
-            case _ => None
-          }
+          stateCache
         )
 
+
         if (repl){
           runner.printInfo("Loading...")
-          (runner.watchLoop(isRepl = true, printing = false, _.run()), Some(cliConfig -> runner))
+          (runner.watchLoop(isRepl = true, printing = false, _.run()), stateCache)
         } else {
-          (runner.runScript(pwd / "build.sc", leftoverArgs), Some(cliConfig -> runner))
+          (runner.runScript(pwd / "build.sc", leftoverArgs), stateCache)
         }
     }
   }
diff --git a/main/src/mill/ServerClient.scala b/main/src/mill/ServerClient.scala
index 9a2139b1..383a8865 100644
--- a/main/src/mill/ServerClient.scala
+++ b/main/src/mill/ServerClient.scala
@@ -5,6 +5,7 @@ import java.nio.channels.FileChannel
 import java.util
 
 import ammonite.main.Cli
+import mill.eval.Evaluator
 import mill.main.MainRunner
 
 class ServerClient(lockBase: String){
@@ -199,7 +200,7 @@ class ProbeThread(lockChannel: FileChannel,
   var running = true
   def run() = {
     while({
-      Thread.sleep(1)
+      Thread.sleep(3)
       lockChannel.tryLock() match{
         case null => true && running
         case locked =>
@@ -219,7 +220,7 @@ class Server(lockBase: String) extends ServerClient(lockBase){
   var currentMeta: OutputStream = System.err
   val lockFile = new RandomAccessFile(lockBase + "/lock", "rw")
   val channel = lockFile.getChannel
-  var mainRunner = Option.empty[(Cli.Config, MainRunner)]
+  var stateCache = Option.empty[Evaluator.State]
 
   def run() = {
     val originalStdout = System.out
@@ -278,9 +279,9 @@ class Server(lockBase: String) extends ServerClient(lockBase){
     originalStdout.println("Parsed Args " + args.toList)
     try {
       originalStdout.println("Running Main")
-      val (_, mr) = mill.Main.main0(
+      val (_, newStateCache) = mill.Main.main0(
         args,
-        mainRunner,
+        stateCache,
         interactive,
         () => {
           channel.tryLock()  match{
@@ -296,9 +297,9 @@ class Server(lockBase: String) extends ServerClient(lockBase){
         new PrintStream(new ProxyOutputStream(currentOutErr, currentMeta, 1), true)
       )
       originalStdout.println("Finished Main")
-      mainRunner = mr
-    } catch{case MainRunner.WatchInterrupted(mr) =>
-      mainRunner = Some((mr.config, mr))
+      stateCache = newStateCache
+    } catch{case MainRunner.WatchInterrupted(sc) =>
+      stateCache = sc
     } finally{
 //      lockChannel.close()
 //      lockFile.close()
diff --git a/main/src/mill/main/MainRunner.scala b/main/src/mill/main/MainRunner.scala
index f6aec9c8..8ec018c6 100644
--- a/main/src/mill/main/MainRunner.scala
+++ b/main/src/mill/main/MainRunner.scala
@@ -12,7 +12,7 @@ import mill.util.PrintLogger
 import mill.main.RunScript
 import upickle.Js
 object MainRunner{
-  case class WatchInterrupted(mr: MainRunner) extends Exception
+  case class WatchInterrupted(stateCache: Option[Evaluator.State]) extends Exception
 }
 /**
   * Customized version of [[ammonite.MainRunner]], allowing us to run Mill
@@ -24,12 +24,13 @@ class MainRunner(val config: ammonite.main.Cli.Config,
                  errPrintStream: PrintStream,
                  stdIn: InputStream,
                  interruptWatch: () => Boolean,
-                 var lastEvaluator: Option[(Seq[(Path, Long)], Evaluator[Any])] = None)
+                 stateCache0: Option[Evaluator.State] = None)
   extends ammonite.MainRunner(
     config, outprintStream, errPrintStream,
     stdIn, outprintStream, errPrintStream
   ){
 
+  var stateCache  = stateCache0
   override def watchAndWait(watched: Seq[(Path, Long)]) = {
     printInfo(s"Watching for changes to ${watched.length} files... (Ctrl-C to exit)")
     def statAll() = watched.forall{ case (file, lastMTime) =>
@@ -37,7 +38,7 @@ class MainRunner(val config: ammonite.main.Cli.Config,
     }
 
     while(statAll()) {
-      if (interruptWatch()) throw WatchInterrupted(this)
+      if (interruptWatch()) throw WatchInterrupted(stateCache)
       Thread.sleep(100)
     }
   }
@@ -52,7 +53,7 @@ class MainRunner(val config: ammonite.main.Cli.Config,
           scriptPath,
           mainCfg.instantiateInterpreter(),
           scriptArgs,
-          lastEvaluator,
+          stateCache,
           new PrintLogger(
             colors != ammonite.util.Colors.BlackWhite,
             colors,
@@ -66,9 +67,10 @@ class MainRunner(val config: ammonite.main.Cli.Config,
           case Res.Success(data) =>
             val (eval, evaluationWatches, res) = data
 
-            lastEvaluator = Some((interpWatched, eval))
+            val watched = interpWatched ++ evaluationWatches
+            stateCache = Some(Evaluator.State(eval.rootModule, eval.classLoaderSig, eval.workerCache, watched))
 
-            (Res(res), interpWatched ++ evaluationWatches)
+            (Res(res), watched)
           case _ => (result, interpWatched)
         }
       }
diff --git a/main/src/mill/main/ReplApplyHandler.scala b/main/src/mill/main/ReplApplyHandler.scala
index eac0c7f0..5681d75b 100644
--- a/main/src/mill/main/ReplApplyHandler.scala
+++ b/main/src/mill/main/ReplApplyHandler.scala
@@ -19,7 +19,6 @@ object ReplApplyHandler{
         ammonite.ops.pwd / 'out,
         ammonite.ops.pwd / 'out,
         rootModule,
-        discover,
         new mill.util.PrintLogger(
           colors != ammonite.util.Colors.BlackWhite,
           colors,
@@ -45,7 +44,7 @@ object ReplApplyHandler{
         else
           ctx.applyPrefixColor("\nChildren:").toString +:
             m.millInternal.reflect[mill.Module].map("\n    ." + _.millOuterCtx.segment.pathSegments.mkString("."))) ++
-        (evaluator.discover.value.get(m.getClass) match{
+        (evaluator.rootModule.millDiscover.value.get(m.getClass) match{
           case None => Nil
           case Some(commands) =>
             ctx.applyPrefixColor("\nCommands:").toString +: commands.map{c =>
diff --git a/main/src/mill/main/RunScript.scala b/main/src/mill/main/RunScript.scala
index 1e7f977a..9a6700fd 100644
--- a/main/src/mill/main/RunScript.scala
+++ b/main/src/mill/main/RunScript.scala
@@ -14,6 +14,7 @@ import mill.util.{EitherOps, Logger, ParseArgs, Watched}
 import mill.util.Strict.Agg
 import upickle.Js
 
+import scala.collection.mutable
 import scala.reflect.ClassTag
 
 /**
@@ -26,32 +27,33 @@ object RunScript{
                 path: Path,
                 instantiateInterpreter: => Either[(Res.Failing, Seq[(Path, Long)]), ammonite.interp.Interpreter],
                 scriptArgs: Seq[String],
-                lastEvaluator: Option[(Seq[(Path, Long)], Evaluator[Any])],
+                stateCache: Option[Evaluator.State],
                 log: Logger)
   : (Res[(Evaluator[Any], Seq[(Path, Long)], Either[String, Seq[Js.Value]])], Seq[(Path, Long)]) = {
 
-    val (evalRes, interpWatched) = lastEvaluator match{
-      case Some((prevInterpWatchedSig, prevEvaluator))
-        if watchedSigUnchanged(prevInterpWatchedSig) =>
-
-        (Res.Success(prevEvaluator), prevInterpWatchedSig)
-
+    val (evalState, interpWatched) = stateCache match{
+      case Some(s) if watchedSigUnchanged(s.classLoaderSig) => Res.Success(s) -> s.watched
       case _ =>
         instantiateInterpreter match{
           case Left((res, watched)) => (res, watched)
           case Right(interp) =>
             interp.watch(path)
             val eval =
-              for((mapping, discover) <- evaluateMapping(wd, path, interp))
-              yield new Evaluator[Any](
-                wd / 'out, wd / 'out, mapping, discover, log,
-                mapping.getClass.getClassLoader.asInstanceOf[SpecialClassLoader].classpathSignature
+              for((rootModule, discover) <- evaluateMapping(wd, path, interp))
+              yield Evaluator.State(
+                rootModule,
+                rootModule.getClass.getClassLoader.asInstanceOf[SpecialClassLoader].classpathSignature,
+                mutable.Map.empty[Segments, (Int, Any)],
+                interp.watchedFiles
               )
-
             (eval, interp.watchedFiles)
         }
     }
 
+    val evalRes =
+      for(s <- evalState)
+      yield new Evaluator[Any](wd / 'out, wd / 'out, s.rootModule, log, s.classLoaderSig, s.workerCache)
+
     val evaluated = for{
       evaluator <- evalRes
       (evalWatches, res) <- Res(evaluateTasks(evaluator, scriptArgs, multiSelect = false))
@@ -142,7 +144,7 @@ object RunScript{
         val selected = selectors.map { case (scopedSel, sel) =>
           for(res <- prepareResolve(evaluator, scopedSel, sel))
           yield {
-            val (rootModule, discover, crossSelectors) = res
+            val (rootModule, crossSelectors) = res
 
 
             try {
@@ -153,7 +155,7 @@ object RunScript{
               // is not currently supported
               mill.eval.Evaluator.currentEvaluator.set(evaluator)
               resolver.resolve(
-                sel.value.toList, rootModule, discover,
+                sel.value.toList, rootModule, rootModule.millDiscover,
                 args, crossSelectors.toList, Nil
               )
             } finally {
@@ -169,7 +171,7 @@ object RunScript{
 
   def resolveRootModule[T](evaluator: Evaluator[T], scopedSel: Option[Segments]) = {
     scopedSel match {
-      case None => Right((evaluator.rootModule, evaluator.discover))
+      case None => Right(evaluator.rootModule)
       case Some(scoping) =>
         for {
           moduleCls <-
@@ -179,19 +181,18 @@ object RunScript{
             case rootModule: ExternalModule => Right(rootModule)
             case _ => Left("Class " + scoping.render + " is not an external module")
           }
-        } yield (rootModule, rootModule.millDiscover)
+        } yield rootModule
     }
   }
 
   def prepareResolve[T](evaluator: Evaluator[T], scopedSel: Option[Segments], sel: Segments) = {
-    for (res <- resolveRootModule(evaluator, scopedSel))
+    for (rootModule<- resolveRootModule(evaluator, scopedSel))
     yield {
-      val (rootModule, discover) = res
       val crossSelectors = sel.value.map {
         case Segment.Cross(x) => x.toList.map(_.toString)
         case _ => Nil
       }
-      (rootModule, discover, crossSelectors)
+      (rootModule, crossSelectors)
     }
   }
 
diff --git a/main/test/src/mill/util/TestEvaluator.scala b/main/test/src/mill/util/TestEvaluator.scala
index a5be0488..b35ebeb8 100644
--- a/main/test/src/mill/util/TestEvaluator.scala
+++ b/main/test/src/mill/util/TestEvaluator.scala
@@ -30,7 +30,7 @@ class TestEvaluator[T <: TestUtil.BaseModule](module: T)
 
   val logger = DummyLogger
 //  val logger = new PrintLogger(true, ammonite.util.Colors.Default, System.out, System.out, System.err)
-  val evaluator = new Evaluator(outPath, TestEvaluator.externalOutPath, module, discover, logger)
+  val evaluator = new Evaluator(outPath, TestEvaluator.externalOutPath, module, logger)
 
   def apply[T](t: Task[T]): Either[Result.Failing[T], (T, Int)] = {
     val evaluated = evaluator.evaluate(Agg(t))