Handle branching control flow by creating a let-binding of the continuation.

Also, grok continue statement.

1 files changed, 28 insertions, 1 deletions

diff --git a/monkeypatch.hs b/monkeypatch.hs
index 14359fe..b3bbe2d 100644
--- a/monkeypatch.hs
+++ b/monkeypatch.hs
@@ -211,12 +211,36 @@ modifyOperand1 f (InfixApp l x op y) = InfixApp l (f x) op y
 informalize :: FormalLambda -> HS.Exp ()
 informalize (FormalLambda k x) = Lambda () [hspvar k] x
 
+
+
+-- Like applyComputation, but creates a let-binding rather than inlining the continuation.
+multiwayContinuation :: Computation FormalLambda -> Computation (HS.Exp ()) -> Computation (HS.Exp ())
+multiwayContinuation a@Computation{ comp = FormalLambda govar exp } b =
+    let k = uniqIdentifier "go" (foldr Map.union Map.empty [compFree a,compIntro a,compFree b,compIntro b])
+        vs = Map.keys $ compIntro a `Map.intersection` compFree b
+        matchgo (Var () (UnQual () (HS.Ident () v))) = v==govar
+        matchgo _                                    = False
+        subst x | matchgo x   = callsite
+                | otherwise   = x
+        callsite = foldl (App ()) (hsvar k) $ map hsvar vs
+        pats = map hspvar vs
+        letexpr = Let () (BDecls () [FunBind () [HS.Match () (HS.Ident () k) pats
+                                                             (UnGuardedRhs () (comp b)) Nothing]])
+                          (everywhere (mkT subst) exp)
+    in Computation
+            { compFree     = Map.union (compFree a) (compFree b) Map.\\ compIntro a
+            , compIntro    = compIntro a `Map.union` compIntro b
+            , compContinue = Nothing
+            , comp         = letexpr
+            }
+
+
 applyComputation :: Computation FormalLambda -> Computation (HS.Exp ()) -> Computation (HS.Exp ())
 applyComputation a@Computation{ comp = FormalLambda govar exp } b =
     let matchgo (Var () (UnQual () (HS.Ident () v))) = v==govar
         matchgo _                                    = False
     in case listify matchgo exp of
-        (_:_:_) -> error "TODO: Multiple go-refs; make let binding."
+        (_:_:_) -> multiwayContinuation a b
         _       -> Computation
             { compFree     = Map.union (compFree a) (compFree b) Map.\\ compIntro a
             , compIntro    = compIntro a `Map.union` compIntro b
@@ -564,6 +588,9 @@ grokStatement fe (CBlockStmt (CReturn (Just exp) _)) = do
     return $ fmap (\y -> FormalLambda k y) $ foldr applyComputation x' xs
 grokStatement fe (CBlockStmt (CReturn Nothing _)) =
     Just $ mkcomp $ FormalLambda "go" retUnit
+grokStatement fe (CBlockStmt (CCont _)) =
+    Just (mkcomp $ FormalLambda "go" $ hsvar " continue")
+        { compContinue = Just " continue" }
 grokStatement fe (CBlockStmt (CIf exp thn els _)) = do
     (xs,x) <- grokExpression fe exp
     let mkif0 = If () (comp x)