improve cps handling of if-then-else. no review.

1 files changed, 33 insertions, 29 deletions

diff --git a/src/continuations/plugin/scala/tools/selectivecps/SelectiveANFTransform.scala b/src/continuations/plugin/scala/tools/selectivecps/SelectiveANFTransform.scala
index 9f1d3dfcd6..b383227243 100644
--- a/src/continuations/plugin/scala/tools/selectivecps/SelectiveANFTransform.scala
+++ b/src/continuations/plugin/scala/tools/selectivecps/SelectiveANFTransform.scala
@@ -173,40 +173,44 @@ abstract class SelectiveANFTransform extends PluginComponent with Transform with
 
           (Nil, tree1, cpsA)
 
-        case If(cond, thenp, elsep) =>
-
-          val (condStats, condVal, spc) = transInlineValue(cond, cpsA)
-
-          val (cpsA2, cpsR2) = (spc, linearize(spc, getAnswerTypeAnn(tree.tpe)))
-//          val (cpsA2, cpsR2) = (None, getAnswerTypeAnn(tree.tpe))
-          val thenVal = transExpr(thenp, cpsA2, cpsR2)
-          val elseVal = transExpr(elsep, cpsA2, cpsR2)
-
-          // check that then and else parts agree (not necessary any more, but left as sanity check)
-          if (cpsR.isDefined) {
-            if (elsep == EmptyTree)
-              unit.error(tree.pos, "always need else part in cps code")
-          }
-          if (hasAnswerTypeAnn(thenVal.tpe) != hasAnswerTypeAnn(elseVal.tpe)) {
-            unit.error(tree.pos, "then and else parts must both be cps code or neither of them")
-          }
+          case If(cond, thenp, elsep) =>
+            /* possible situations:
+            cps before (cpsA)
+            cps in condition (spc)  <-- synth flag set if *only* here!
+            cps in (one or both) branches */
+            val (condStats, condVal, spc) = transInlineValue(cond, cpsA)
+            val (cpsA2, cpsR2) = if (hasSynthMarker(tree.tpe))
+              (spc, linearize(spc, getAnswerTypeAnn(tree.tpe))) else
+              (None, getAnswerTypeAnn(tree.tpe)) // if no cps in condition, branches must conform to tree.tpe directly
+            val thenVal = transExpr(thenp, cpsA2, cpsR2)
+            val elseVal = transExpr(elsep, cpsA2, cpsR2)
+
+            // check that then and else parts agree (not necessary any more, but left as sanity check)
+            if (cpsR.isDefined) {
+              if (elsep == EmptyTree)
+                unit.error(tree.pos, "always need else part in cps code")
+            }
+            if (hasAnswerTypeAnn(thenVal.tpe) != hasAnswerTypeAnn(elseVal.tpe)) {
+              unit.error(tree.pos, "then and else parts must both be cps code or neither of them")
+            }
 
-          (condStats, updateSynthFlag(treeCopy.If(tree, condVal, thenVal, elseVal)), spc)
+            (condStats, updateSynthFlag(treeCopy.If(tree, condVal, thenVal, elseVal)), spc)
 
-        case Match(selector, cases) =>
+          case Match(selector, cases) =>
 
-          val (selStats, selVal, spc) = transInlineValue(selector, cpsA)
-          val (cpsA2, cpsR2) = (spc, linearize(spc, getAnswerTypeAnn(tree.tpe)))
-//          val (cpsA2, cpsR2) = (None, getAnswerTypeAnn(tree.tpe))
+            val (selStats, selVal, spc) = transInlineValue(selector, cpsA)
+            val (cpsA2, cpsR2) = if (hasSynthMarker(tree.tpe))
+              (spc, linearize(spc, getAnswerTypeAnn(tree.tpe))) else
+              (None, getAnswerTypeAnn(tree.tpe))
 
-          val caseVals = for {
-            cd @ CaseDef(pat, guard, body) <- cases
-            bodyVal = transExpr(body, cpsA2, cpsR2)
-          } yield {
-            treeCopy.CaseDef(cd, transform(pat), transform(guard), bodyVal)
-          }
+            val caseVals = for {
+              cd @ CaseDef(pat, guard, body) <- cases
+              bodyVal = transExpr(body, cpsA2, cpsR2)
+            } yield {
+              treeCopy.CaseDef(cd, transform(pat), transform(guard), bodyVal)
+            }
 
-          (selStats, updateSynthFlag(treeCopy.Match(tree, selVal, caseVals)), spc)
+            (selStats, updateSynthFlag(treeCopy.Match(tree, selVal, caseVals)), spc)
 
 
         case ldef @ LabelDef(name, params, rhs) =>