Merge pull request #4036 from retronym/topic/opt-tail-calls

SI-8893 Restore linear perf in TailCalls with nested matches

2 files changed, 39 insertions, 20 deletions

diff --git a/src/compiler/scala/tools/nsc/backend/icode/BasicBlocks.scala b/src/compiler/scala/tools/nsc/backend/icode/BasicBlocks.scala
index f9551697d2..ad1975ef23 100644
--- a/src/compiler/scala/tools/nsc/backend/icode/BasicBlocks.scala
+++ b/src/compiler/scala/tools/nsc/backend/icode/BasicBlocks.scala
@@ -300,14 +300,16 @@ trait BasicBlocks {
       if (!closed)
         instructionList = instructionList map (x => map.getOrElse(x, x))
       else
-        instrs.zipWithIndex collect {
-          case (oldInstr, i) if map contains oldInstr =>
-            // SI-6288 clone important here because `replaceInstruction` assigns
-            // a position to `newInstr`. Without this, a single instruction can
-            // be added twice, and the position last position assigned clobbers
-            // all previous positions in other usages.
-            val newInstr = map(oldInstr).clone()
-            code.touched |= replaceInstruction(i, newInstr)
+        instrs.iterator.zipWithIndex foreach {
+          case (oldInstr, i) =>
+            if (map contains oldInstr) {
+              // SI-6288 clone important here because `replaceInstruction` assigns
+              // a position to `newInstr`. Without this, a single instruction can
+              // be added twice, and the position last position assigned clobbers
+              // all previous positions in other usages.
+              val newInstr = map(oldInstr).clone()
+              code.touched |= replaceInstruction(i, newInstr)
+            }
         }
 
     ////////////////////// Emit //////////////////////
diff --git a/src/compiler/scala/tools/nsc/transform/TailCalls.scala b/src/compiler/scala/tools/nsc/transform/TailCalls.scala
index ef534f70fd..16ea3ea90f 100644
--- a/src/compiler/scala/tools/nsc/transform/TailCalls.scala
+++ b/src/compiler/scala/tools/nsc/transform/TailCalls.scala
@@ -129,6 +129,13 @@ abstract class TailCalls extends Transform {
       }
       override def toString = s"${method.name} tparams=$tparams tailPos=$tailPos label=$label label info=${label.info}"
 
+      final def noTailContext() = clonedTailContext(false)
+      final def yesTailContext() = clonedTailContext(true)
+      protected def clonedTailContext(tailPos: Boolean): TailContext = this match {
+        case _ if this.tailPos == tailPos => this
+        case clone: ClonedTailContext => clone.that.clonedTailContext(tailPos)
+        case _ => new ClonedTailContext(this, tailPos)
+      }
     }
 
     object EmptyTailContext extends TailContext {
@@ -174,7 +181,7 @@ abstract class TailCalls extends Transform {
       }
       def containsRecursiveCall(t: Tree) = t exists isRecursiveCall
     }
-    class ClonedTailContext(that: TailContext, override val tailPos: Boolean) extends TailContext {
+    class ClonedTailContext(val that: TailContext, override val tailPos: Boolean) extends TailContext {
       def method     = that.method
       def tparams    = that.tparams
       def methodPos  = that.methodPos
@@ -183,9 +190,6 @@ abstract class TailCalls extends Transform {
     }
 
     private var ctx: TailContext = EmptyTailContext
-    private def noTailContext()  = new ClonedTailContext(ctx, tailPos = false)
-    private def yesTailContext() = new ClonedTailContext(ctx, tailPos = true)
-
 
     override def transformUnit(unit: CompilationUnit): Unit = {
       try {
@@ -206,16 +210,16 @@ abstract class TailCalls extends Transform {
       finally this.ctx = saved
     }
 
-    def yesTailTransform(tree: Tree): Tree = transform(tree, yesTailContext())
-    def noTailTransform(tree: Tree): Tree = transform(tree, noTailContext())
+    def yesTailTransform(tree: Tree): Tree = transform(tree, ctx.yesTailContext())
+    def noTailTransform(tree: Tree): Tree = transform(tree, ctx.noTailContext())
     def noTailTransforms(trees: List[Tree]) = {
-      val nctx = noTailContext()
-      trees map (t => transform(t, nctx))
+      val nctx = ctx.noTailContext()
+      trees mapConserve (t => transform(t, nctx))
     }
 
     override def transform(tree: Tree): Tree = {
       /* A possibly polymorphic apply to be considered for tail call transformation. */
-      def rewriteApply(target: Tree, fun: Tree, targs: List[Tree], args: List[Tree]) = {
+      def rewriteApply(target: Tree, fun: Tree, targs: List[Tree], args: List[Tree], mustTransformArgs: Boolean = true) = {
         val receiver: Tree = fun match {
           case Select(qual, _)  => qual
           case _                => EmptyTree
@@ -223,7 +227,7 @@ abstract class TailCalls extends Transform {
         def receiverIsSame    = ctx.enclosingType.widen =:= receiver.tpe.widen
         def receiverIsSuper   = ctx.enclosingType.widen <:< receiver.tpe.widen
         def isRecursiveCall   = (ctx.method eq fun.symbol) && ctx.tailPos
-        def transformArgs     = noTailTransforms(args)
+        def transformArgs     = if (mustTransformArgs) noTailTransforms(args) else args
         def matchesTypeArgs   = ctx.tparams sameElements (targs map (_.tpe.typeSymbol))
 
         /* Records failure reason in Context for reporting.
@@ -265,6 +269,10 @@ abstract class TailCalls extends Transform {
         !(sym.hasAccessorFlag || sym.isConstructor)
       }
 
+      // intentionally shadowing imports from definitions for performance
+      val runDefinitions = currentRun.runDefinitions
+      import runDefinitions.{Boolean_or, Boolean_and}
+
       tree match {
         case ValDef(_, _, _, _) =>
           if (tree.symbol.isLazy && tree.symbol.hasAnnotation(TailrecClass))
@@ -312,8 +320,13 @@ abstract class TailCalls extends Transform {
           // the assumption is once we encounter a case, the remainder of the block will consist of cases
           // the prologue may be empty, usually it is the valdef that stores the scrut
           val (prologue, cases) = stats span (s => !s.isInstanceOf[LabelDef])
+          val transformedPrologue = noTailTransforms(prologue)
+          val transformedCases = transformTrees(cases)
+          val transformedStats =
+            if ((prologue eq transformedPrologue) && (cases eq transformedCases)) stats // allow reuse of `tree` if the subtransform was an identity
+            else transformedPrologue ++ transformedCases
           treeCopy.Block(tree,
-            noTailTransforms(prologue) ++ transformTrees(cases),
+            transformedStats,
             transform(expr)
           )
 
@@ -380,7 +393,7 @@ abstract class TailCalls extends Transform {
           if (res ne arg)
             treeCopy.Apply(tree, fun, res :: Nil)
           else
-            rewriteApply(fun, fun, Nil, args)
+            rewriteApply(fun, fun, Nil, args, mustTransformArgs = false)
 
         case Apply(fun, args) =>
           rewriteApply(fun, fun, Nil, args)
@@ -421,6 +434,10 @@ abstract class TailCalls extends Transform {
     def traverseNoTail(tree: Tree) = traverse(tree, maybeTailNew = false)
     def traverseTreesNoTail(trees: List[Tree]) = trees foreach traverseNoTail
 
+    // intentionally shadowing imports from definitions for performance
+    private val runDefinitions = currentRun.runDefinitions
+    import runDefinitions.{Boolean_or, Boolean_and}
+
     override def traverse(tree: Tree) = tree match {
       // we're looking for label(x){x} in tail position, since that means `a` is in tail position in a call `label(a)`
       case LabelDef(_, List(arg), body@Ident(_)) if arg.symbol == body.symbol =>