Simplify the adaptation of types of return expressions

Add `adaptTypeOfReturn` hook to `AnnotationCheckers`.
Move adaptation of types of return expressions from `addAnnotations`
to `typedReturn` via `adaptTypeOfReturn` hook.

This resolves an inconsistency where previously types could have
a plus marker without additional CPS annotations. This also adds
additional test cases.

8 files changed, 86 insertions, 11 deletions

diff --git a/src/compiler/scala/tools/nsc/symtab/AnnotationCheckers.scala b/src/compiler/scala/tools/nsc/symtab/AnnotationCheckers.scala
index d26c5d65a0..be2e520554 100644
--- a/src/compiler/scala/tools/nsc/symtab/AnnotationCheckers.scala
+++ b/src/compiler/scala/tools/nsc/symtab/AnnotationCheckers.scala
@@ -47,6 +47,13 @@ trait AnnotationCheckers {
      *  before. If the implementing class cannot do the adaptiong, it
      *  should return the tree unchanged.*/
     def adaptAnnotations(tree: Tree, mode: Int, pt: Type): Tree = tree
+
+    /** Adapt the type of a return expression. The decision of an annotation checker
+     *  whether the type should be adapted is based on the type of the expression
+     *  which is returned, as well as the result type of the method (pt).
+     *  By default, this method simply returns the passed `default` type.
+     */
+    def adaptTypeOfReturn(tree: Tree, pt: Type, default: => Type): Type = default
   }
 
   /** The list of annotation checkers that have been registered */
@@ -117,4 +124,23 @@ trait AnnotationCheckers {
     annotationCheckers.foldLeft(tree)((tree, checker) =>
       checker.adaptAnnotations(tree, mode, pt))
   }
+
+  /** Let a registered annotation checker adapt the type of a return expression.
+   *  Annotation checkers that cannot do the adaptation should simply return
+   *  the `default` argument.
+   *  
+   *  Note that the result is undefined if more than one annotation checker
+   *  returns an adapted type which is not a subtype of `default`.
+   */
+  def adaptTypeOfReturn(tree: Tree, pt: Type, default: => Type): Type = {
+    val adaptedTypes = annotationCheckers flatMap { checker =>
+      val adapted = checker.adaptTypeOfReturn(tree, pt, default)
+      if (!(adapted <:< default)) List(adapted)
+      else List()
+    }
+    adaptedTypes match {
+      case fst :: _ => fst
+      case List() => default
+    }
+  }
 }
diff --git a/src/compiler/scala/tools/nsc/typechecker/Typers.scala b/src/compiler/scala/tools/nsc/typechecker/Typers.scala
index 47f57a09ef..05b36d128e 100644
--- a/src/compiler/scala/tools/nsc/typechecker/Typers.scala
+++ b/src/compiler/scala/tools/nsc/typechecker/Typers.scala
@@ -3190,7 +3190,8 @@ trait Typers extends Modes {
               if (typed(expr).tpe.typeSymbol != UnitClass)
                 unit.warning(tree.pos, "enclosing method " + name + " has result type Unit: return value discarded")
             }
-            treeCopy.Return(tree, checkDead(expr1)) setSymbol enclMethod.owner setType NothingClass.tpe
+            treeCopy.Return(tree, checkDead(expr1)).setSymbol(enclMethod.owner)
+                                                   .setType(adaptTypeOfReturn(expr1, restpt.tpe, NothingClass.tpe))
           }
         }
       }
diff --git a/src/continuations/plugin/scala/tools/selectivecps/CPSAnnotationChecker.scala b/src/continuations/plugin/scala/tools/selectivecps/CPSAnnotationChecker.scala
index cde602bf83..249ae565cc 100644
--- a/src/continuations/plugin/scala/tools/selectivecps/CPSAnnotationChecker.scala
+++ b/src/continuations/plugin/scala/tools/selectivecps/CPSAnnotationChecker.scala
@@ -224,16 +224,32 @@ abstract class CPSAnnotationChecker extends CPSUtils with Modes {
         res
       } else if (retMode && !hasPlusMarker(tree.tpe) && annotsTree.isEmpty && annotsExpected.nonEmpty) {
         // add a marker annotation that will make tree.tpe behave as pt, subtyping wise
-        // tree will look like having no annotation
+        // tree will look like having any possible annotation
         
-        // note that we are only adding a plus marker if the method's result type is a CPS type
-        // (annotsExpected.nonEmpty == cpsParamAnnotation(pt).nonEmpty)
+        // note 1: we are only adding a plus marker if the method's result type is a cps type
+        //         (annotsExpected.nonEmpty == cpsParamAnnotation(pt).nonEmpty)
+        // note 2: we are not adding the expected cps annotations, since they will be added
+        //         by adaptTypeOfReturn (see below).
         val res = tree modifyType (_ withAnnotations List(newPlusMarker()))
         vprintln("adapted annotations (return) of " + tree + " to " + res.tpe)
         res
       } else tree
     }
 
+    // only adapt type if this return will
+    // (a) be removed (in tail position and method's result type (pt) is cps type), or
+    // (b) cause an error
+    override def adaptTypeOfReturn(tree: Tree, pt: Type, default: => Type): Type = {
+      // only adapt if method's result type (pt) is cps type
+      val annots = cpsParamAnnotation(pt)
+      if (annots.nonEmpty) {
+        // return type of `tree` without plus marker, but only if it doesn't have other cps annots
+        if (hasPlusMarker(tree.tpe) && !hasCpsParamTypes(tree.tpe))
+          tree.setType(removeAttribs(tree.tpe, MarkerCPSAdaptPlus))
+        tree.tpe
+      } else default
+    }
+
     def updateAttributesFromChildren(tpe: Type, childAnnots: List[AnnotationInfo], byName: List[Tree]): Type = {
       tpe match {
         // Would need to push annots into each alternative of overloaded type
@@ -486,13 +502,6 @@ abstract class CPSAnnotationChecker extends CPSUtils with Modes {
           }
           tpe
 
-        case ret @ Return(expr) =>
-          // only change type if this return will (a) be removed (in tail position) or (b) cause
-          // an error (not in tail position)
-          if (expr.tpe != null && hasPlusMarker(expr.tpe))
-            ret setType expr.tpe
-          ret.tpe
-
         case _ =>
           tpe
       }
diff --git a/src/continuations/plugin/scala/tools/selectivecps/SelectiveANFTransform.scala b/src/continuations/plugin/scala/tools/selectivecps/SelectiveANFTransform.scala
index 8ab9ba881f..501396b31c 100644
--- a/src/continuations/plugin/scala/tools/selectivecps/SelectiveANFTransform.scala
+++ b/src/continuations/plugin/scala/tools/selectivecps/SelectiveANFTransform.scala
@@ -43,6 +43,12 @@ abstract class SelectiveANFTransform extends PluginComponent with Transform with
         case If(cond, r1 @ Return(thenExpr), r2 @ Return(elseExpr)) =>
           treeCopy.If(tree, cond, transform(thenExpr), transform(elseExpr))
 
+        case If(cond, r1 @ Return(thenExpr), elseExpr) =>
+          treeCopy.If(tree, cond, transform(thenExpr), transform(elseExpr))
+
+        case If(cond, thenExpr, r2 @ Return(elseExpr)) =>
+          treeCopy.If(tree, cond, transform(thenExpr), transform(elseExpr))
+
         case If(cond, thenExpr, elseExpr) =>
           treeCopy.If(tree, cond, transform(thenExpr), transform(elseExpr))
 
diff --git a/test/files/continuations-run/t5314-with-if.check b/test/files/continuations-run/t5314-with-if.check
new file mode 100644
index 0000000000..7f8f011eb7
--- /dev/null
+++ b/test/files/continuations-run/t5314-with-if.check
@@ -0,0 +1 @@
+7
diff --git a/test/files/continuations-run/t5314-with-if.scala b/test/files/continuations-run/t5314-with-if.scala
new file mode 100644
index 0000000000..5840199a3c
--- /dev/null
+++ b/test/files/continuations-run/t5314-with-if.scala
@@ -0,0 +1,17 @@
+import scala.util.continuations._
+
+object Test extends App {
+
+  def foo(x:Int): Int @cps[Int] = 7
+
+  def bar(x:Int): Int @cps[Int] = {
+    val v = foo(x)
+    if (v > 0)
+      return v
+    else
+      return 10
+  }
+
+  println(reset { bar(10) })
+
+}
diff --git a/test/files/continuations-run/t5314.check b/test/files/continuations-run/t5314.check
index 4951e7caae..4b35d8e6d0 100644
--- a/test/files/continuations-run/t5314.check
+++ b/test/files/continuations-run/t5314.check
@@ -1,3 +1,7 @@
+7
+7
+7
+8
 8
 hi
 8
diff --git a/test/files/continuations-run/t5314.scala b/test/files/continuations-run/t5314.scala
index 0bdea19824..d611016ce4 100644
--- a/test/files/continuations-run/t5314.scala
+++ b/test/files/continuations-run/t5314.scala
@@ -29,6 +29,17 @@ object Test extends App {
 
   def nocps(x: Int): Int = { return x; x }
 
+  def foo2(x:Int): Int @cps[Int] = 7
+  def bar2(x:Int): Int @cps[Int] = { foo2(x); return 7 }
+  def bar3(x:Int): Int @cps[Int] = { foo2(x); if (x == 7) return 7 else return foo2(x) }
+  def bar4(x:Int): Int @cps[Int] = { foo2(x); if (x == 7) return 7 else foo2(x) }
+  def bar5(x:Int): Int @cps[Int] = { foo2(x); if (x == 7) return 7 else 8 }
+  println(reset { bar2(10) })
+  println(reset { bar3(10) })
+  println(reset { bar4(10) })
+  println(reset { bar5(10) })
+  
+  /* original test case */
   val repro = new ReturnRepro
   repro.caller
   repro.caller2