improvements to cps exception handling.

5 files changed, 163 insertions, 38 deletions

diff --git a/src/continuations/library/scala/util/continuations/ControlContext.scala b/src/continuations/library/scala/util/continuations/ControlContext.scala
index 9ed83a7e71..37e84e7d59 100644
--- a/src/continuations/library/scala/util/continuations/ControlContext.scala
+++ b/src/continuations/library/scala/util/continuations/ControlContext.scala
@@ -14,7 +14,7 @@ private class cpsMinus extends Annotation // implementation detail
 
 
 
-@serializable final class ControlContext[+A,-B,+C](val fun: (A => B, Throwable => B) => C, val x: A) {
+@serializable final class ControlContext[+A,-B,+C](val fun: (A => B, Exception => B) => C, val x: A) {
 
   /*
     final def map[A1](f: A => A1): ControlContext[A1,B,C] = {
@@ -29,9 +29,14 @@ private class cpsMinus extends Annotation // implementation detail
 
   final def map[A1](f: A => A1): ControlContext[A1,B,C] = {
     if (fun eq null)
-      new ControlContext(null, f(x))
+      try {
+        new ControlContext(null, f(x)) // TODO: only alloc if f(x) != x
+      } catch {
+        case ex: Exception =>
+          new ControlContext((k: A1 => B, thr: Exception => B) => thr(ex).asInstanceOf[C], null.asInstanceOf[A1])
+      }
     else
-      new ControlContext({ (k: A1 => B, thr: Throwable => B) =>
+      new ControlContext({ (k: A1 => B, thr: Exception => B) =>
         fun( { (x:A) =>
           var done = false
           try {
@@ -39,7 +44,7 @@ private class cpsMinus extends Annotation // implementation detail
             done = true
             k(res)
           } catch {
-            case ex if !done =>
+            case ex: Exception if !done =>
               thr(ex)
           }
         }, thr)
@@ -52,9 +57,14 @@ private class cpsMinus extends Annotation // implementation detail
 
   /*@inline*/ final def flatMap[A1,B1,C1<:B](f: (A => ControlContext[A1,B1,C1])): ControlContext[A1,B1,C] = {
     if (fun eq null)
-      f(x).asInstanceOf[ControlContext[A1,B1,C]]
+      try {
+        f(x).asInstanceOf[ControlContext[A1,B1,C]]
+      } catch {
+        case ex: Exception =>
+          new ControlContext((k: A1 => B1, thr: Exception => B1) => thr(ex).asInstanceOf[C], null.asInstanceOf[A1])
+      }
     else
-      new ControlContext({ (k: A1 => B1, thr: Throwable => B1) =>
+      new ControlContext({ (k: A1 => B1, thr: Exception => B1) =>
         fun( { (x:A) =>
           var done = false
           try {
@@ -63,16 +73,16 @@ private class cpsMinus extends Annotation // implementation detail
             val res: C1 = ctxR.foreachFull(k, thr) // => B1
             res
           } catch {
-            case ex if !done =>
+            case ex: Exception if !done =>
               thr(ex).asInstanceOf[B] // => B NOTE: in general this is unsafe!
           }                           // However, the plugin will not generate offending code
-        }, thr.asInstanceOf[Throwable=>B]) // => B
+        }, thr.asInstanceOf[Exception=>B]) // => B
       }, null.asInstanceOf[A1])
   }
 
   final def foreach(f: A => B) = foreachFull(f, throw _)
 
-  def foreachFull(f: A => B, g: Throwable => B): C = {
+  def foreachFull(f: A => B, g: Exception => B): C = {
     if (fun eq null)
       f(x).asInstanceOf[C]
     else
@@ -85,12 +95,12 @@ private class cpsMinus extends Annotation // implementation detail
 
   // need filter or other functions?
 
-  final def flatCat[A1>:A,B1<:B,C1>:C<:B1](pf: PartialFunction[Throwable, ControlContext[A1,B1,C1]]): ControlContext[A1,B1,C1] = {
+  final def flatMapCatch[A1>:A,B1<:B,C1>:C<:B1](pf: PartialFunction[Exception, ControlContext[A1,B1,C1]]): ControlContext[A1,B1,C1] = {
     if (fun eq null)
       this
     else {
-      val fun1 = (ret1: A1 => B1, thr1: Throwable => B1) => {
-        val thr: Throwable => B1 = { t: Throwable =>
+      val fun1 = (ret1: A1 => B1, thr1: Exception => B1) => {
+        val thr: Exception => B1 = { t: Exception =>
           var captureExceptions = true
           try {
             if (pf.isDefinedAt(t)) {
@@ -102,14 +112,50 @@ private class cpsMinus extends Annotation // implementation detail
               thr1(t) // Throw => B1
             }
           } catch {
-            case t1 if captureExceptions => thr1(t1) // => E2
+            case t1: Exception if captureExceptions => thr1(t1) // => E2
           }
         }
-        foreachFull(ret1, thr)// fun(ret1, thr)  // => B
+        fun(ret1, thr)// fun(ret1, thr)  // => B
       }
       new ControlContext(fun1, null.asInstanceOf[A1])
     }
   }
 
-  // TODO: finally
+  final def mapFinally(f: () => Unit): ControlContext[A,B,C] = {
+    if (fun eq null) {
+      try {
+        f()
+        this
+      } catch {
+        case ex: Exception =>
+          new ControlContext((k: A => B, thr: Exception => B) => thr(ex).asInstanceOf[C], null.asInstanceOf[A])
+      }
+    } else {
+      val fun1 = (ret1: A => B, thr1: Exception => B) => {
+        val ret: A => B = { x: A =>
+          var captureExceptions = true
+          try {
+            f()
+            captureExceptions = false
+            ret1(x)
+          } catch {
+            case t1: Exception if captureExceptions => thr1(t1)
+          }
+        }
+        val thr: Exception => B = { t: Exception =>
+          var captureExceptions = true
+          try {
+            f()
+            captureExceptions = false
+            thr1(t)
+          } catch {
+            case t1: Exception if captureExceptions => thr1(t1)
+          }
+        }
+        fun(ret, thr1)
+      }
+      new ControlContext(fun1, null.asInstanceOf[A])
+    }
+  }
+
 }
diff --git a/src/continuations/library/scala/util/continuations/package.scala b/src/continuations/library/scala/util/continuations/package.scala
index 1fa89da2b4..aa4681a0cc 100644
--- a/src/continuations/library/scala/util/continuations/package.scala
+++ b/src/continuations/library/scala/util/continuations/package.scala
@@ -55,7 +55,7 @@ package object continuations {
   }
 
   def shiftR[A,B,C](fun: (A => B) => C): ControlContext[A,B,C] = {
-    new ControlContext((f:A=>B,g:Throwable=>B) => fun(f), null.asInstanceOf[A])
+    new ControlContext((f:A=>B,g:Exception=>B) => fun(f), null.asInstanceOf[A])
   }
 
   def reifyR[A,B,C](ctx: => ControlContext[A,B,C]): ControlContext[A,B,C] = {
diff --git a/src/continuations/plugin/scala/tools/selectivecps/CPSAnnotationChecker.scala b/src/continuations/plugin/scala/tools/selectivecps/CPSAnnotationChecker.scala
index a1112fa3c8..500f102790 100644
--- a/src/continuations/plugin/scala/tools/selectivecps/CPSAnnotationChecker.scala
+++ b/src/continuations/plugin/scala/tools/selectivecps/CPSAnnotationChecker.scala
@@ -438,6 +438,8 @@ abstract class CPSAnnotationChecker extends CPSUtils {
             val atp0::atp1::Nil = ann.atp.normalize.typeArgs
             if (!(atp0 =:= atp1))
               throw new TypeError("only simple cps types allowed in try/catch blocks (found: " + tpe1 + ")")
+            if (!finalizer.isEmpty) // no finalizers allowed. see explanation in SelectiveCPSTransform
+              reporter.error(tree.pos, "try/catch blocks that use continuations cannot have finalizers")
           }
           tpe1
 
diff --git a/src/continuations/plugin/scala/tools/selectivecps/CPSUtils.scala b/src/continuations/plugin/scala/tools/selectivecps/CPSUtils.scala
index 3e1f05e731..566f175183 100644
--- a/src/continuations/plugin/scala/tools/selectivecps/CPSUtils.scala
+++ b/src/continuations/plugin/scala/tools/selectivecps/CPSUtils.scala
@@ -11,7 +11,7 @@ trait CPSUtils {
 
   var cpsEnabled = false
   val verbose: Boolean = System.getProperty("cpsVerbose", "false") == "true"
-  def vprintln(x: =>Any): Unit = if (verbose) println(x)
+  @inline def vprintln(x: =>Any): Unit = if (verbose) println(x)
 
 
   lazy val MarkerCPSSym = definitions.getClass("scala.util.continuations.cpsSym")
diff --git a/src/continuations/plugin/scala/tools/selectivecps/SelectiveCPSTransform.scala b/src/continuations/plugin/scala/tools/selectivecps/SelectiveCPSTransform.scala
index 7f1610e8cc..3c35149636 100644
--- a/src/continuations/plugin/scala/tools/selectivecps/SelectiveCPSTransform.scala
+++ b/src/continuations/plugin/scala/tools/selectivecps/SelectiveCPSTransform.scala
@@ -117,36 +117,107 @@ abstract class SelectiveCPSTransform extends PluginComponent with
           }
 
       case Try(block, catches, finalizer) =>
-        // currently duplicates the catch block into a partial function
+        // currently duplicates the catch block into a partial function.
         // this is kinda risky, but we don't expect there will be lots
-        // of try/catches inside catch blocks (exp. blowup)
+        // of try/catches inside catch blocks (exp. blowup unlikely).
+
+        // CAVEAT: finalizers are surprisingly tricky!
+        // the problem is that they cannot easily be removed
+        // from the regular control path and hence will
+        // also be invoked after creating the Context object.
+
+        /*
+        object Test {
+          def foo1 = {
+            throw new Exception("in sub")
+            shift((k:Int=>Int) => k(1))
+            10
+          }
+          def foo2 = {
+            shift((k:Int=>Int) => k(2))
+            20
+          }
+          def foo3 = {
+            shift((k:Int=>Int) => k(3))
+            throw new Exception("in sub")
+            30
+          }
+          def foo4 = {
+            shift((k:Int=>Int) => 4)
+            throw new Exception("in sub")
+            40
+          }
+          def bar(x: Int) = try {
+            if (x == 1)
+              foo1
+            else if (x == 2)
+              foo2
+            else if (x == 3)
+              foo3
+            else //if (x == 4)
+              foo4
+          } catch {
+            case _ =>
+              println("exception")
+              0
+          } finally {
+            println("done")
+          }
+        }
+
+        reset(Test.bar(1)) // should print: exception,done,0
+        reset(Test.bar(2)) // should print: done,20 <-- but prints: done,done,20
+        reset(Test.bar(3)) // should print: exception,done,0 <-- but prints: done,exception,done,0
+        reset(Test.bar(4)) // should print: 4 <-- but prints: done,4
+        */
 
         val block1 = transform(block)
         val catches1 = transformCaseDefs(catches)
         val finalizer1 = transform(finalizer)
 
-        if (block1.tpe.typeSymbol.tpe <:< Context.tpe) {
-          //println("CPS Transform: " + tree)
-
-          val pos = catches.head.pos
+        if (hasAnswerTypeAnn(tree.tpe)) {
+          //vprintln("CPS Transform: " + tree + "/" + tree.tpe + "/" + block1.tpe)
 
-          val (stms, expr) = block1 match {
+          val (stms, expr1) = block1 match {
             case Block(stms, expr) => (stms, expr)
             case expr => (Nil, expr)
           }
 
-          val arg = currentOwner.newValueParameter(pos, "$ex").setInfo(ThrowableClass.tpe)
-          val catches2 = catches1 map (duplicateTree(_).asInstanceOf[CaseDef])
-          val rhs = Match(Ident(arg), catches2)
-          val fun = Function(List(ValDef(arg)), rhs)
-          val expr2 = localTyper.typed(atPos(pos) { Apply(Select(expr, expr.tpe.member("flatCat")), List(fun)) })
-          val block2 = treeCopy.Block(block1, stms, expr2)
+          val expr2 = if (catches.nonEmpty) {
+            val pos = catches.head.pos
+            val arg = currentOwner.newValueParameter(pos, "$ex").setInfo(ThrowableClass.tpe)
+            val catches2 = catches1 map (duplicateTree(_).asInstanceOf[CaseDef])
+            val rhs = Match(Ident(arg), catches2)
+            val fun = Function(List(ValDef(arg)), rhs)
+            val expr2 = localTyper.typed(atPos(pos) { Apply(Select(expr1, expr1.tpe.member("flatMapCatch")), List(fun)) })
 
-          arg.owner = fun.symbol
-          val chown = new ChangeOwnerTraverser(currentOwner, fun.symbol)
-          chown.traverse(rhs)
+            arg.owner = fun.symbol
+            val chown = new ChangeOwnerTraverser(currentOwner, fun.symbol)
+            chown.traverse(rhs)
 
-          treeCopy.Try(tree, block2, catches1, finalizer1)
+            expr2
+          } else
+            expr1
+
+/*
+          disabled for now - see notes above
+
+          val expr3 = if (!finalizer.isEmpty) {
+            val pos = finalizer.pos
+            val finalizer2 = duplicateTree(finalizer1)
+            val fun = Function(List(), finalizer2)
+            val expr3 = localTyper.typed(atPos(pos) { Apply(Select(expr2, expr2.tpe.member("mapFinally")), List(fun)) })
+
+            val chown = new ChangeOwnerTraverser(currentOwner, fun.symbol)
+            chown.traverse(finalizer2)
+
+            expr3
+          } else
+            expr2
+*/
+
+
+          treeCopy.Try(tree, treeCopy.Block(block1, stms, expr2), catches1, finalizer1)
         } else {
           treeCopy.Try(tree, block1, catches1, finalizer1)
         }
@@ -204,8 +275,10 @@ abstract class SelectiveCPSTransform extends PluginComponent with
 
                 val body2 = localTyper.typed(atPos(vd.symbol.pos) { body })
 
-                if ((body2.tpe == null) || !(body2.tpe.typeSymbol.tpe <:< Context.tpe)) {
-                  println(body2 + "/" + body2.tpe)
+                // in theory it would be nicer to look for an @cps annotation instead
+                // of testing for Context
+                if ((body2.tpe == null) || !(body2.tpe.typeSymbol == Context)) {
+                  //println(body2 + "/" + body2.tpe)
                   unit.error(rhs.pos, "cannot compute type for CPS-transformed function result")
                 }
                 body2
@@ -228,7 +301,7 @@ abstract class SelectiveCPSTransform extends PluginComponent with
                 var methodName = "map"
 
                 if (body.tpe != null) {
-                  if (body.tpe.typeSymbol.tpe <:< Context.tpe)
+                  if (body.tpe.typeSymbol == Context)
                     methodName = "flatMap"
                 }
                 else
@@ -244,9 +317,13 @@ abstract class SelectiveCPSTransform extends PluginComponent with
               try {
                 if (specialCaseTrivial) {
                   log("will optimize possible tail call: " + bodyExpr)
+
+                  // FIXME: flatMap impl has become more complicated due to
+                  // exceptions. do we need to put a try/catch in the then part??
+
                   // val ctx = <rhs>
                   // if (ctx.isTrivial)
-                  //   val <lhs> = ctx.getTrivialValue; ...
+                  //   val <lhs> = ctx.getTrivialValue; ...    <--- TODO: try/catch ??? don't bother for the moment...
                   // else
                   //   ctx.flatMap { <lhs> => ... }
                   val ctxSym = currentOwner.newValue(vd.symbol.name + "$shift").setInfo(rhs1.tpe)