Simplify `Evaluator.Mapping` and consolidate logic within the `Segments` data-type

10 files changed, 74 insertions, 74 deletions

diff --git a/core/src/main/scala/mill/define/Cross.scala b/core/src/main/scala/mill/define/Cross.scala
index 72424bde..691c5691 100644
--- a/core/src/main/scala/mill/define/Cross.scala
+++ b/core/src/main/scala/mill/define/Cross.scala
@@ -4,7 +4,6 @@ import scala.reflect.macros.blackbox
 
 
 object Cross{
-  def autoCast[A](x: Any): A = x.asInstanceOf[A]
   abstract class Implicit[T]{
     def make(v: Any, ctx: Module.Ctx): T
     def crossValues(v: Any): List[Any]
@@ -54,7 +53,7 @@ class Cross[T](cases: Any*)
     val sub = ci.make(
       c,
       ctx.copy(
-        segments0 = Segments(ctx.segments0.value :+ ctx.segment),
+        segments0 = ctx.segments0 ++ Seq(ctx.segment),
         segment = Segment.Cross(crossValues.reverse)
       )
     )
diff --git a/core/src/main/scala/mill/define/Module.scala b/core/src/main/scala/mill/define/Module.scala
index 9c830c1a..db1d511a 100644
--- a/core/src/main/scala/mill/define/Module.scala
+++ b/core/src/main/scala/mill/define/Module.scala
@@ -12,7 +12,19 @@ object Segment{
 }
 
 case class BasePath(value: Path)
-case class Segments(value: Seq[Segment])
+case class Segments(value: Segment*){
+  def ++(other: Seq[Segment]): Segments = Segments(value ++ other:_*)
+  def ++(other: Segments): Segments = Segments(value ++ other.value:_*)
+  def render = value match {
+    case Nil => ""
+    case Segment.Label(head) :: rest =>
+      val stringSegments = rest.map{
+        case Segment.Label(s) => "." + s
+        case Segment.Cross(vs) => "[" + vs.mkString(",") + "]"
+      }
+      head + stringSegments.mkString
+  }
+}
 /**
   * `Module` is a class meant to be extended by `trait`s *only*, in order to
   * propagate the implicit parameters forward to the final concrete
@@ -32,7 +44,7 @@ class Module(implicit ctx0: Module.Ctx) extends mill.moduledefs.Cacher{
   })
   implicit def millModuleBasePath: BasePath = BasePath(basePath)
   implicit def millModuleSegments: Segments = {
-    Segments(ctx.segments0.value :+ ctx.segment)
+    ctx.segments0 ++ Seq(ctx.segment)
   }
 }
 object Module{
@@ -43,7 +55,7 @@ object Module{
                  basePath: Path,
                  segments0: Segments,
                  overrides: Int){
-    def segments = Segments(segments0.value.drop(1))
+    def segments = Segments(segments0.value.drop(1):_*)
   }
   object Ctx{
     implicit def make(implicit millModuleEnclosing0: sourcecode.Enclosing,
@@ -67,5 +79,5 @@ class BaseModule(basePath: Path)
                  millName0: sourcecode.Name,
                  overrides0: Overrides)
   extends Module()(
-    Module.Ctx.make(implicitly, implicitly, implicitly, BasePath(basePath), Segments(Nil), implicitly)
+    Module.Ctx.make(implicitly, implicitly, implicitly, BasePath(basePath), Segments(), implicitly)
   )
\ No newline at end of file
diff --git a/core/src/main/scala/mill/define/Task.scala b/core/src/main/scala/mill/define/Task.scala
index 96fdbb8e..aa4d838a 100644
--- a/core/src/main/scala/mill/define/Task.scala
+++ b/core/src/main/scala/mill/define/Task.scala
@@ -221,7 +221,7 @@ object Caller {
 class TargetImpl[+T](t: Task[T],
                      ctx0: Module.Ctx,
                      val readWrite: RW[_]) extends Target[T] {
-  val ctx = ctx0.copy(segments0 = Segments(ctx0.segments0.value :+ ctx0.segment))
+  val ctx = ctx0.copy(segments0 = ctx0.segments0 ++ Seq(ctx0.segment))
   val inputs = Seq(t)
   def evaluate(args: Ctx) = args[T](0)
   override def toString = ctx.enclosing + "@" + Integer.toHexString(System.identityHashCode(this))
@@ -229,7 +229,7 @@ class TargetImpl[+T](t: Task[T],
 class Command[+T](t: Task[T],
                   ctx0: Module.Ctx,
                   val writer: W[_]) extends NamedTask[T] {
-  val ctx = ctx0.copy(segments0 = Segments(ctx0.segments0.value :+ ctx0.segment))
+  val ctx = ctx0.copy(segments0 = ctx0.segments0 ++ Seq(ctx0.segment))
   val inputs = Seq(t)
   def evaluate(args: Ctx) = args[T](0)
   override def asCommand = Some(this)
@@ -245,7 +245,7 @@ class Persistent[+T](t: Task[T],
 class Input[T](t: Task[T],
                ctx0: Module.Ctx,
                val readWrite: RW[_]) extends Target[T]{
-  val ctx = ctx0.copy(segments0 = Segments(ctx0.segments0.value :+ ctx0.segment))
+  val ctx = ctx0.copy(segments0 = ctx0.segments0 ++ Seq(ctx0.segment))
   val inputs = Seq(t)
   def evaluate(args: Ctx) = args[T](0)
   override def sideHash = util.Random.nextInt()
diff --git a/core/src/main/scala/mill/discover/Discovered.scala b/core/src/main/scala/mill/discover/Discovered.scala
index 49d2db4a..0728a027 100644
--- a/core/src/main/scala/mill/discover/Discovered.scala
+++ b/core/src/main/scala/mill/discover/Discovered.scala
@@ -5,6 +5,7 @@ import ammonite.main.Router
 import ammonite.main.Router.EntryPoint
 import mill.discover.Mirror.TargetPoint
 import mill.util.Ctx.Loader
+import mill.util.OSet
 
 import scala.language.experimental.macros
 import scala.reflect.macros.blackbox.Context
@@ -23,48 +24,32 @@ object Discovered {
   }
   case class Mapping[T](mirror: Mirror[T, T],
                         base: T){
-    val modulesToPaths = Mirror.traverse(base, mirror){ (mirror, segmentsRev) =>
-      val resolvedNode = mirror.node(
-        base,
-        segmentsRev.reverse.map{case Segment.Cross(vs) => vs.toList case _ => Nil}.toList
-      )
-      Seq(resolvedNode -> segmentsRev.reverse)
+    val modulesToMirrors = Mirror.traverseNode[T, T, (Any, Mirror[_, _])](base, mirror){
+      (mirror, segments, resolvedNode) => Seq(resolvedNode -> mirror)
     }.toMap
 
-    val modulesToMirrors = Mirror.traverse[T, T, (Any, Mirror[_, _])](base, mirror){ (mirror, segmentsRev) =>
-      val resolvedNode = mirror.node(
-        base,
-        segmentsRev.reverse.map{case Segment.Cross(vs) => vs.toList case _ => Nil}.toList
-      )
-      Seq(resolvedNode -> mirror)
-    }.toMap
-
-    val targetsToSegments = Mirror.traverse(base, mirror){ (mirror, segmentsRev) =>
-      val resolvedNode = mirror.node(
-        base,
-        segmentsRev.reverse.map{case Segment.Cross(vs) => vs.toList case _ => Nil}.toList
-      )
-      for(target <- mirror.targets) yield {
-        target.asInstanceOf[TargetPoint[Any, Any]].run(resolvedNode) -> (segmentsRev.reverse :+ Segment.Label(target.label))
+    val targets = OSet.from(
+      Mirror.traverseNode(base, mirror){ (mirror, segmentsRev, resolvedNode) =>
+        for(target <- mirror.targets)
+        yield target.asInstanceOf[TargetPoint[Any, Any]].run(resolvedNode)
       }
+    )
 
-    }.toMap
-
-    val segmentsToCommands = Mirror.traverse[T, T, (Seq[Segment], EntryPoint[_])](base, mirror){ (mirror, segmentsRev) =>
+    val segmentsToCommands = Mirror.traverse[T, T, (Segments, EntryPoint[_])](base, mirror){ (mirror, segments) =>
       for(command <- mirror.commands)
-      yield (segmentsRev.reverse :+ Segment.Label(command.name)) -> command
+      yield segments ++ Seq(Segment.Label(command.name)) -> command
     }.toMap
 
-    val segmentsToTargets = targetsToSegments.map(_.swap)
+    val segmentsToTargets = targets.map(t => (t.ctx.segments, t)).toMap
+    val targetsToSegments = segmentsToTargets.map(_.swap)
   }
 
-  def consistencyCheck[T](mapping: Discovered.Mapping[T]) = {
+  def consistencyCheck[T](mapping: Discovered.Mapping[T]): OSet[Segments] = {
     val mapping2 = Discovered.Mapping(mapping.mirror, mapping.base)
-
     for{
-      (t1, t2) <- mapping2.targetsToSegments.zip(mapping.targetsToSegments)
-      if t1._1 ne t2._1
-    } yield t1._2
+      (t1, t2) <- mapping2.targets.zip(mapping.targets)
+      if t1 ne t2
+    } yield t1.ctx.segments
   }
 
 
diff --git a/core/src/main/scala/mill/discover/Mirror.scala b/core/src/main/scala/mill/discover/Mirror.scala
index 3443535e..bb6e44ad 100644
--- a/core/src/main/scala/mill/discover/Mirror.scala
+++ b/core/src/main/scala/mill/discover/Mirror.scala
@@ -1,6 +1,6 @@
 package mill.discover
 
-import mill.define.{Segment, Target, Task}
+import mill.define.{Segment, Segments, Target, Task}
 import ammonite.main.Router.EntryPoint
 
 import scala.language.experimental.macros
@@ -20,22 +20,23 @@ case class Mirror[-T, V](node: (T, List[List[Any]]) => V,
                          crossChildren: Option[(V => List[List[Any]], Mirror[T, _])])
 
 object Mirror{
-  def renderSelector(selector: Seq[Segment]) = selector.toList match {
-    case Nil => ""
-    case Segment.Label(head) :: rest =>
-      val stringSegments = rest.map{
-        case Segment.Label(s) => "." + s
-        case Segment.Cross(vs) => "[" + vs.mkString(",") + "]"
-      }
-      head + stringSegments.mkString
-  }
 
 
+  def traverseNode[T, V, R](t: T, hierarchy: Mirror[T, V])
+                           (f: (Mirror[T, _], => Segments, => Any) => Seq[R]): Seq[R] = {
+    traverse(t, hierarchy){ (mirror, segmentsRev) =>
+      lazy val resolvedNode = mirror.node(
+        t,
+        segmentsRev.value.map{case Segment.Cross(vs) => vs.toList case _ => Nil}.toList
+      )
+      f(mirror, segmentsRev, resolvedNode)
+    }
+  }
   def traverse[T, V, R](t: T, hierarchy: Mirror[T, V])
-                       (f: (Mirror[T, _], => Seq[Segment]) => Seq[R]): Seq[R] = {
+                       (f: (Mirror[T, _], => Segments) => Seq[R]): Seq[R] = {
     def rec[C](segmentsRev: List[Segment], h: Mirror[T, C]): Seq[R]= {
       val crossValues = segmentsRev.map{case Segment.Cross(vs) => vs case _ => Nil}
-      val self = f(h, segmentsRev)
+      val self = f(h, Segments(segmentsRev.reverse:_*))
       self ++
       h.children.flatMap{case (label, c) => rec(Segment.Label(label) :: segmentsRev, c)} ++
       h.crossChildren.toSeq.flatMap{
diff --git a/core/src/main/scala/mill/eval/Evaluator.scala b/core/src/main/scala/mill/eval/Evaluator.scala
index ff2d64f2..efb10d36 100644
--- a/core/src/main/scala/mill/eval/Evaluator.scala
+++ b/core/src/main/scala/mill/eval/Evaluator.scala
@@ -13,7 +13,7 @@ import mill.util._
 import scala.collection.mutable
 import scala.util.control.NonFatal
 case class Labelled[T](target: NamedTask[T],
-                       segments: Seq[Segment]){
+                       segments: Segments){
   def format = target match{
     case t: Target[T] => Some(t.readWrite.asInstanceOf[upickle.default.ReadWriter[T]])
     case _ => None
@@ -41,7 +41,7 @@ class Evaluator[T](val workspacePath: Path,
     val topoSorted = Graph.topoSorted(transitive)
     val sortedGroups = Graph.groupAroundImportantTargets(topoSorted){
       case t: NamedTask[Any]   =>
-        val segments = t.ctx.segments.value
+        val segments = t.ctx.segments
         val (finalTaskOverrides, enclosing) = t match{
           case t: Target[_] => mapping.segmentsToTargets(segments).ctx.overrides -> t.ctx.enclosing
           case c: mill.define.Command[_] => mapping.segmentsToCommands(segments).overrides -> c.ctx.enclosing
@@ -125,7 +125,7 @@ class Evaluator[T](val workspacePath: Path,
 
           case _ =>
 
-            val Seq(first, rest @_*) = labelledTarget.segments
+            val Seq(first, rest @_*) = labelledTarget.segments.value
             val msgParts = Seq(first.asInstanceOf[Segment.Label].value) ++ rest.map{
               case Segment.Label(s) => "." + s
               case Segment.Cross(s) => "[" + s.mkString(",") + "]"
@@ -256,11 +256,11 @@ object Evaluator{
                    dest: Path,
                    meta: Path,
                    log: Path)
-  def makeSegmentStrings(segments: Seq[Segment]) = segments.flatMap{
+  def makeSegmentStrings(segments: Segments) = segments.value.flatMap{
     case Segment.Label(s) => Seq(s)
     case Segment.Cross(values) => values.map(_.toString)
   }
-  def resolveDestPaths(workspacePath: Path, segments: Seq[Segment]): Paths = {
+  def resolveDestPaths(workspacePath: Path, segments: Segments): Paths = {
     val segmentStrings = makeSegmentStrings(segments)
     val targetPath = workspacePath / segmentStrings
     Paths(targetPath, targetPath / 'dest, targetPath / "meta.json", targetPath / 'log)
diff --git a/core/src/main/scala/mill/main/ReplApplyHandler.scala b/core/src/main/scala/mill/main/ReplApplyHandler.scala
index 0bcf8e6e..7ce00f67 100644
--- a/core/src/main/scala/mill/main/ReplApplyHandler.scala
+++ b/core/src/main/scala/mill/main/ReplApplyHandler.scala
@@ -65,13 +65,13 @@ class ReplApplyHandler(pprinter0: pprint.PPrinter, evaluator: Evaluator[_]) exte
         ))
 
       )
-    case t: mill.define.Target[_] if evaluator.mapping.targetsToSegments.contains(t) =>
+    case t: mill.define.Target[_] if evaluator.mapping.targets.contains(t) =>
       val seen = mutable.Set.empty[Task[_]]
-      def rec(t: Task[_]): Seq[Seq[Segment]] = {
+      def rec(t: Task[_]): Seq[Segments] = {
         if (seen(t)) Nil // do nothing
         else t match {
-          case t: Target[_] if evaluator.mapping.targetsToSegments.contains(t) =>
-            Seq(evaluator.mapping.targetsToSegments(t))
+          case t: Target[_] if evaluator.mapping.targets.contains(t) =>
+            Seq(t.ctx.segments)
           case _ =>
             seen.add(t)
             t.inputs.flatMap(rec)
@@ -79,11 +79,10 @@ class ReplApplyHandler(pprinter0: pprint.PPrinter, evaluator: Evaluator[_]) exte
       }
       pprint.Tree.Lazy(ctx =>
         Iterator(t.ctx.enclosing, ":", t.ctx.lineNum.toString, "\n", ctx.applyPrefixColor("Inputs:").toString) ++
-        t.inputs.iterator.flatMap(rec).map("\n    " + Mirror.renderSelector(_))
+        t.inputs.iterator.flatMap(rec).map("\n    " + _.render)
       )
 
   }
-  ammonite.main.Cli
   val pprinter = pprinter0.copy(
     additionalHandlers = millHandlers orElse pprinter0.additionalHandlers
   )
diff --git a/core/src/main/scala/mill/main/Resolve.scala b/core/src/main/scala/mill/main/Resolve.scala
index b8968091..35b07a26 100644
--- a/core/src/main/scala/mill/main/Resolve.scala
+++ b/core/src/main/scala/mill/main/Resolve.scala
@@ -1,6 +1,6 @@
 package mill.main
 
-import mill.define.{Segment, Task}
+import mill.define.{Segment, Segments, Task}
 import mill.define.Task.TaskModule
 import mill.discover.Mirror
 import ammonite.main.Router
@@ -43,8 +43,8 @@ object Resolve {
         def command = invokeCommand(hierarchy, last)
 
         command orElse target orElse runDefault.headOption.flatten match{
-          case None =>  Left("Cannot resolve task " + Mirror.renderSelector(
-            (Segment.Label(last) :: revSelectorsSoFar).reverse)
+          case None =>  Left("Cannot resolve task " +
+            Segments((Segment.Label(last) :: revSelectorsSoFar).reverse:_*).render
           )
           // Contents of `either` *must* be a `Task`, because we only select
           // methods returning `Task` in the discovery process
@@ -60,7 +60,7 @@ object Resolve {
               case (label, child) if label == singleLabel => child
             } match{
               case Some(child) => resolve(tail, child, obj, rest, remainingCrossSelectors, newRevSelectorsSoFar)
-              case None => Left("Cannot resolve module " + Mirror.renderSelector(newRevSelectorsSoFar.reverse))
+              case None => Left("Cannot resolve module " + Segments(newRevSelectorsSoFar.reverse:_*).render)
             }
 
           case Segment.Cross(cross) =>
@@ -69,7 +69,7 @@ object Resolve {
             if (crossOptions.contains(cross)){
               resolve(tail, childMirror, obj, rest, remainingCrossSelectors, newRevSelectorsSoFar)
             }else{
-              Left("Cannot resolve cross " + Mirror.renderSelector(newRevSelectorsSoFar.reverse))
+              Left("Cannot resolve cross " + Segments(newRevSelectorsSoFar.reverse:_*).render)
             }
 
 
diff --git a/core/src/main/scala/mill/main/RunScript.scala b/core/src/main/scala/mill/main/RunScript.scala
index e27d4e1c..e79349ad 100644
--- a/core/src/main/scala/mill/main/RunScript.scala
+++ b/core/src/main/scala/mill/main/RunScript.scala
@@ -146,7 +146,7 @@ object RunScript{
       (for((k, fs) <- evaluated.failing.items()) yield {
         val ks = k match{
           case Left(t) => t.toString
-          case Right(t) => Mirror.renderSelector(t.segments.toList)
+          case Right(t) => t.segments.render
         }
         val fss = fs.map{
           case Result.Exception(t, outerStack) =>
@@ -161,10 +161,9 @@ object RunScript{
         val json = for(t <- Seq(target)) yield {
           t match {
             case t: mill.define.NamedTask[_] =>
-              val segments = t.ctx.segments.value
-              val jsonFile = Evaluator.resolveDestPaths(
-                evaluator.workspacePath, segments
-              ).meta
+              val jsonFile = Evaluator
+                .resolveDestPaths(evaluator.workspacePath, t.ctx.segments)
+                .meta
               val metadata = upickle.json.read(jsonFile.toIO)
               Some(metadata(1))
 
@@ -206,7 +205,7 @@ object RunScript{
   def consistencyCheck[T](mapping: Discovered.Mapping[T]): Either[String, Unit] = {
     val consistencyErrors = Discovered.consistencyCheck(mapping)
     if (consistencyErrors.nonEmpty) {
-      Left(s"Failed Discovered.consistencyCheck: ${consistencyErrors.map(Mirror.renderSelector)}")
+      Left(s"Failed Discovered.consistencyCheck: ${consistencyErrors.map(_.render)}")
     } else {
       Right(())
     }
diff --git a/core/src/main/scala/mill/util/OSet.scala b/core/src/main/scala/mill/util/OSet.scala
index ce5d573b..52b47cde 100644
--- a/core/src/main/scala/mill/util/OSet.scala
+++ b/core/src/main/scala/mill/util/OSet.scala
@@ -16,9 +16,11 @@ trait OSet[V] extends TraversableOnce[V]{
   def flatMap[T](f: V => TraversableOnce[T]): OSet[T]
   def map[T](f: V => T): OSet[T]
   def filter(f: V => Boolean): OSet[V]
+  def withFilter(f: V => Boolean): OSet[V]
   def collect[T](f: PartialFunction[V, T]): OSet[T]
   def zipWithIndex: OSet[(V, Int)]
   def reverse: OSet[V]
+  def zip[T](other: OSet[T]): OSet[(V, T)]
 }
 
 object OSet{
@@ -66,6 +68,7 @@ object OSet{
       for(i <- items) if (f(i)) output.append(i)
       output
     }
+    def withFilter(f: V => Boolean): OSet[V] = filter(f)
 
     def collect[T](f: PartialFunction[V, T]) = this.filter(f.isDefinedAt).map(x => f(x))
 
@@ -79,6 +82,8 @@ object OSet{
 
     def reverse = OSet.from(indexed.reverseIterator)
 
+    def zip[T](other: OSet[T]) = OSet.from(items.zip(other.items))
+
     // Members declared in scala.collection.GenTraversableOnce
     def isTraversableAgain: Boolean = items.isTraversableAgain
     def toIterator: Iterator[V] = items.toIterator