Implement non-local returns

Non-local returns are now implemented.

5 files changed, 126 insertions, 0 deletions

diff --git a/src/dotty/tools/dotc/Compiler.scala b/src/dotty/tools/dotc/Compiler.scala
index 76cf10428..1742fb1b8 100644
--- a/src/dotty/tools/dotc/Compiler.scala
+++ b/src/dotty/tools/dotc/Compiler.scala
@@ -68,6 +68,7 @@ class Compiler {
            new LazyVals,
            new Memoize,
            new LinkScala2ImplClasses,
+           new NonLocalReturns,
            new CapturedVars, // capturedVars has a transformUnit: no phases should introduce local mutable vars here
            new Constructors,
            new FunctionalInterfaces,
diff --git a/src/dotty/tools/dotc/ast/tpd.scala b/src/dotty/tools/dotc/ast/tpd.scala
index b05d23107..4e9940ac4 100644
--- a/src/dotty/tools/dotc/ast/tpd.scala
+++ b/src/dotty/tools/dotc/ast/tpd.scala
@@ -161,6 +161,10 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
   def Bind(sym: TermSymbol, body: Tree)(implicit ctx: Context): Bind =
     ta.assignType(untpd.Bind(sym.name, body), sym)
 
+  /** A pattern corrsponding to `sym: tpe` */ 
+  def BindTyped(sym: TermSymbol, tpe: Type)(implicit ctx: Context): Bind =
+    Bind(sym, Typed(Underscore(tpe), TypeTree(tpe)))
+
   def Alternative(trees: List[Tree])(implicit ctx: Context): Alternative =
     ta.assignType(untpd.Alternative(trees), trees)
 
diff --git a/src/dotty/tools/dotc/core/Definitions.scala b/src/dotty/tools/dotc/core/Definitions.scala
index 7ed0a26e0..502b42987 100644
--- a/src/dotty/tools/dotc/core/Definitions.scala
+++ b/src/dotty/tools/dotc/core/Definitions.scala
@@ -326,6 +326,7 @@ class Definitions {
     lazy val Product_productArity = ProductClass.requiredMethod(nme.productArity)
     lazy val Product_productPrefix = ProductClass.requiredMethod(nme.productPrefix)
   lazy val LanguageModuleClass          = ctx.requiredModule("dotty.language").moduleClass.asClass
+  lazy val NonLocalReturnControlClass   = ctx.requiredClass("scala.runtime.NonLocalReturnControl")
 
   // Annotation base classes
   lazy val AnnotationClass              = ctx.requiredClass("scala.annotation.Annotation")
diff --git a/src/dotty/tools/dotc/transform/NonLocalReturns.scala b/src/dotty/tools/dotc/transform/NonLocalReturns.scala
new file mode 100644
index 000000000..d0a2f2ca7
--- /dev/null
+++ b/src/dotty/tools/dotc/transform/NonLocalReturns.scala
@@ -0,0 +1,88 @@
+package dotty.tools.dotc
+package transform
+
+import core._
+import Contexts._, Symbols._, Types._, Flags._, Decorators._, StdNames._, Constants._, Phases._
+import TreeTransforms._
+import ast.Trees._
+import collection.mutable
+
+/** Implement non-local returns using NonLocalReturnControl exceptions.
+ */
+class NonLocalReturns extends MiniPhaseTransform { thisTransformer =>
+  override def phaseName = "nonLocalReturns"
+
+  import ast.tpd._
+
+  override def runsAfter: Set[Class[_ <: Phase]] = Set(classOf[ElimByName])
+
+  private def ensureConforms(tree: Tree, pt: Type)(implicit ctx: Context) =
+    if (tree.tpe <:< pt) tree
+    else Erasure.Boxing.adaptToType(tree, pt)
+
+  /** The type of a non-local return expression with given argument type */
+  private def nonLocalReturnExceptionType(argtype: Type)(implicit ctx: Context) =
+    defn.NonLocalReturnControlClass.typeRef.appliedTo(argtype)
+
+  /** A hashmap from method symbols to non-local return keys */
+  private val nonLocalReturnKeys = mutable.Map[Symbol, TermSymbol]()
+
+  /** Return non-local return key for given method */
+  private def nonLocalReturnKey(meth: Symbol)(implicit ctx: Context) =
+    nonLocalReturnKeys.getOrElseUpdate(meth,
+      ctx.newSymbol(
+        meth, ctx.freshName("nonLocalReturnKey").toTermName, Synthetic, defn.ObjectType, coord = meth.pos))
+
+  /** Generate a non-local return throw with given return expression from given method.
+   *  I.e. for the method's non-local return key, generate:
+   *
+   *    throw new NonLocalReturnControl(key, expr)
+   *  todo: maybe clone a pre-existing exception instead?
+   *  (but what to do about exceptions that miss their targets?)
+   */
+  private def nonLocalReturnThrow(expr: Tree, meth: Symbol)(implicit ctx: Context) =
+    Throw(
+      New(
+        defn.NonLocalReturnControlClass.typeRef,
+        ref(nonLocalReturnKey(meth)) :: ensureConforms(expr, defn.ObjectType) :: Nil))
+
+  /** Transform (body, key) to:
+   *
+   *  {
+   *    val key = new Object()
+   *    try {
+   *      body
+   *    } catch {
+   *      case ex: NonLocalReturnControl =>
+   *        if (ex.key().eq(key)) ex.value().asInstanceOf[T]
+   *        else throw ex
+   *    }
+   *  }
+   */
+  private def nonLocalReturnTry(body: Tree, key: TermSymbol, meth: Symbol)(implicit ctx: Context) = {
+    val keyDef = ValDef(key, New(defn.ObjectType, Nil))
+    val nonLocalReturnControl = defn.NonLocalReturnControlClass.typeRef
+    val ex = ctx.newSymbol(meth, nme.ex, EmptyFlags, nonLocalReturnControl, coord = body.pos)
+    val pat = BindTyped(ex, nonLocalReturnControl)
+    val rhs = If(
+        ref(ex).select(nme.key).appliedToNone.select(nme.eq).appliedTo(ref(key)),
+        ensureConforms(ref(ex).select(nme.value), meth.info.finalResultType),
+        Throw(ref(ex)))
+    val catches = CaseDef(pat, EmptyTree, rhs) :: Nil
+    val tryCatch = Try(body, catches, EmptyTree)
+    Block(keyDef :: Nil, tryCatch)
+  }
+
+  def isNonLocalReturn(ret: Return)(implicit ctx: Context) =
+    ret.from.symbol != ctx.owner.enclosingMethod || ctx.owner.is(Lazy) // Lazy needed?
+
+  override def transformDefDef(tree: DefDef)(implicit ctx: Context, info: TransformerInfo): Tree =
+    nonLocalReturnKeys.remove(tree.symbol) match {
+      case Some(key) => cpy.DefDef(tree)(rhs = nonLocalReturnTry(tree.rhs, key, tree.symbol))
+      case _ => tree
+    }
+
+  override def transformReturn(tree: Return)(implicit ctx: Context, info: TransformerInfo): Tree =
+    if (isNonLocalReturn(tree)) nonLocalReturnThrow(tree.expr, tree.from.symbol).withPos(tree.pos)
+    else tree
+}
diff --git a/tests/run/nonLocalReturns.scala b/tests/run/nonLocalReturns.scala
new file mode 100644
index 000000000..a425ac723
--- /dev/null
+++ b/tests/run/nonLocalReturns.scala
@@ -0,0 +1,32 @@
+object Test {
+
+  def foo(xs: List[Int]): Int = {
+    xs.foreach(x => return x)
+    0
+  }
+
+  def bar(xs: List[Int]): Int = {
+    lazy val y = if (xs.isEmpty) return -1 else xs.head
+    y
+  }
+
+  def baz(x: Int): Int =
+    byName { return -2; 3 }
+
+  def byName(x: => Int): Int = x
+
+  def bam(): Int = { // no non-local return needed here
+    val foo = {
+      return -3
+      3
+    }
+    foo
+  }
+
+  def main(args: Array[String]) = {
+    assert(foo(List(1, 2, 3)) == 1)
+    assert(bar(Nil) == -1)
+    assert(baz(3) == -2)
+    assert(bam() == -3)
+  }
+}