From b53bc9001e10605b27e525efdd0515e5bd4e5f6c Mon Sep 17 00:00:00 2001
From: Lorenzo Gabriele <lorenzolespaul@gmail.com>
Date: Tue, 22 Feb 2022 09:51:18 +0100
Subject: [PATCH] Follow up refactor of importTree in RunScript

- Remove extra comma in Evaluator toString
---
 main/core/src/mill/eval/Evaluator.scala |  2 +-
 main/src/mill/main/RunScript.scala      | 52 +++++++++++++------------
 2 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/main/core/src/mill/eval/Evaluator.scala b/main/core/src/mill/eval/Evaluator.scala
index de31ef13465..153955a563b 100644
--- a/main/core/src/mill/eval/Evaluator.scala
+++ b/main/core/src/mill/eval/Evaluator.scala
@@ -726,7 +726,7 @@ class Evaluator private (
        |  env = $env,
        |  failFast = $failFast,
        |  threadCount = $threadCount,
-       |  importTree = $importTree,
+       |  importTree = $importTree
        |)""".stripMargin
   }
 
diff --git a/main/src/mill/main/RunScript.scala b/main/src/mill/main/RunScript.scala
index f2b2cdcc623..49881e4c9dd 100644
--- a/main/src/mill/main/RunScript.scala
+++ b/main/src/mill/main/RunScript.scala
@@ -3,7 +3,7 @@ package mill.main
 import java.nio.file.NoSuchFileException
 import ammonite.runtime.SpecialClassLoader
 import ammonite.util.Util.CodeSource
-import ammonite.util.{Name, Res, Util}
+import ammonite.util.{Name, Res, ScriptOutput, Util}
 import mill.define
 import mill.define._
 import mill.eval.{Evaluator, EvaluatorPaths}
@@ -66,29 +66,6 @@ object RunScript {
             val eval =
               for (rootModule <- evaluateRootModule(wd, path, interp, log))
                 yield {
-                  val importTreeMap = mutable.Map.empty[String, Seq[String]]
-                  interp.alreadyLoadedFiles.foreach { case (a, b) =>
-                    val filePath = AmmoniteUtils.normalizeAmmoniteImportPath(a.filePathPrefix)
-                    val importPaths = b.blockInfo.flatMap { b =>
-                      val relativePath = b.hookInfo.trees.map { t =>
-                        val prefix = t.prefix
-                        val mappings = t.mappings.toSeq.flatMap(_.map(_._1))
-                        prefix ++ mappings
-                      }
-                      relativePath.collect {
-                        case "$file" :: tail =>
-                          val concatenated = filePath.init ++ tail
-                          AmmoniteUtils.normalizeAmmoniteImportPath(concatenated)
-                      }
-                    }
-                    def toCls(segments: Seq[String]): String = segments.mkString(".")
-                    val key = toCls(filePath)
-                    val toAppend = importPaths.map(toCls)
-                    importTreeMap(key) = importTreeMap.getOrElse(key, Seq.empty) ++ toAppend
-                  }
-
-                  val importTree = GraphUtils.linksToScriptNodeGraph(importTreeMap)
-
                   EvaluatorState(
                     rootModule,
                     rootModule.getClass.getClassLoader.asInstanceOf[
@@ -97,7 +74,7 @@ object RunScript {
                     mutable.Map.empty[Segments, (Int, Any)],
                     interp.watchedValues.toSeq,
                     systemProperties.keySet,
-                    importTree
+                    importTree(interp.alreadyLoadedFiles)
                   )
                 }
             (eval, interp.watchedValues)
@@ -132,6 +109,31 @@ object RunScript {
     sig.forall { case (p, l) => p.poll() == l }
   }
 
+  private def importTree(alreadyLoadedFiles: collection.Map[CodeSource, ScriptOutput.Metadata]): Seq[ScriptNode] = {
+    val importTreeMap = mutable.Map.empty[String, Seq[String]]
+    alreadyLoadedFiles.foreach { case (a, b) =>
+      val filePath = AmmoniteUtils.normalizeAmmoniteImportPath(a.filePathPrefix)
+      val importPaths = b.blockInfo.flatMap { b =>
+        val relativePath = b.hookInfo.trees.map { t =>
+          val prefix = t.prefix
+          val mappings = t.mappings.toSeq.flatMap(_.map(_._1))
+          prefix ++ mappings
+        }
+        relativePath.collect {
+          case "$file" :: tail =>
+            val concatenated = filePath.init ++ tail
+            AmmoniteUtils.normalizeAmmoniteImportPath(concatenated)
+        }
+      }
+      def toCls(segments: Seq[String]): String = segments.mkString(".")
+      val key = toCls(filePath)
+      val toAppend = importPaths.map(toCls)
+      importTreeMap(key) = importTreeMap.getOrElse(key, Seq.empty) ++ toAppend
+    }
+
+    GraphUtils.linksToScriptNodeGraph(importTreeMap)
+  }
+
   def evaluateRootModule(
       wd: os.Path,
       path: os.Path,