Applicative instance for Computation.

1 files changed, 28 insertions, 20 deletions

diff --git a/monkeypatch.hs b/monkeypatch.hs
index d7c4282..a6c4db9 100644
--- a/monkeypatch.hs
+++ b/monkeypatch.hs
@@ -159,6 +159,19 @@ data Computation st = Computation
     }
  deriving (Eq,Ord,Functor)
 
+instance Applicative Computation where
+    pure = mkcomp
+    mf <*> ma = Computation
+                { compFree  = Map.union (compFree mf) (compFree ma)
+                , compIntro = Map.union (compIntro mf) (compIntro ma)
+                , compContinue = (if isJust (compContinue mf) && isJust (compContinue ma)
+                                    then trace "Warning: incompatible continue symbols."
+                                    else id)
+                                  $ mplus (compContinue mf) (compContinue ma)
+                , comp = comp mf $ comp ma
+                }
+
+
 mkcomp :: x -> Computation x
 mkcomp x = Computation Map.empty Map.empty Nothing x
 
@@ -323,9 +336,7 @@ grokExpression fe (CBinary op a b _) = do
              | otherwise          = infixOp
     -- trace ("intros("++hop++"): "++show (foldr Map.union Map.empty $ map compIntro as)) $ return ()
     -- TODO: Short-circuit boolean evaluation side-effects.
-    return $ (,) ss $ (mkcomp $ infx (comp ca) hop (comp cb))
-        { compFree  = compFree ca `Map.union` compFree cb
-        }
+    return $ (,) ss $ infx <$> ca <*> pure hop <*> cb
 grokExpression fe (CUnary CAdrOp (CVar cv0 _) _) = do
     let cv = identToString cv0
         hv = "p" ++ cv
@@ -348,9 +359,7 @@ grokExpression fe (CCond cond (Just thn) els _) = do
     (es,e) <- grokExpression fe els
     let tt = foldr applyComputation t ts
         ee = foldr applyComputation e es
-    return $ (,) cs $ fmap (\cnd -> If () cnd (comp tt) (comp ee)) c
-        { compFree = compFree ee `Map.union` compFree tt `Map.union` compFree c
-        }
+    return $ (,) cs $ If () <$> c <*> tt <*> ee
 grokExpression fe (CSizeofExpr expr _) = do
     (xs,x) <- grokExpression fe expr
     return $ (,) xs $ fmap (App () (hsvar "sizeOf")) x
@@ -379,16 +388,13 @@ grokExpression fe (C.CCall (CVar fn _) exps _) = do
     gs <- mapM (grokExpression fe) exps
     let ss = concatMap fst gs -- TODO: resolve variable name conflicts
         hv = "r" ++ identToString fn
-        cll = foldl (App ()) (hsvar (identToString fn)) $ map (comp . snd) gs
-        frees = foldr Map.union (Map.singleton (identToString fn) ()) (map (compFree . snd) gs)
-        k = uniqIdentifier "go" frees
-        s = Computation
-                { compFree = frees
-                , compIntro = Map.singleton hv ()
-                , compContinue = Nothing
-                , comp = FormalLambda k
-                            $ infixOp cll ">>="
-                                $ Lambda () [hspvar hv] (hsvar k)
+        -- cll = foldl (App ()) (hsvar (identToString fn)) $ map (comp . snd) gs
+        -- frees = foldr Map.union (Map.singleton (identToString fn) ()) (map (compFree . snd) gs)
+        fn' = identToString fn
+        cll = foldl (\f x -> App () <$> f <*> x) (mkcomp $ hsvar fn'){compFree = Map.singleton fn' ()} (map snd gs)
+        k = uniqIdentifier "go" (compFree s)
+        s = (fmap (\c -> FormalLambda k $ infixOp c ">>=" $ Lambda () [hspvar hv] (hsvar k)) cll)
+                { compIntro = Map.singleton hv ()
                 }
     return $ (,) (ss++[s]) (mkcomp $ hsvar hv)
         { compFree  = Map.singleton hv ()
@@ -556,13 +562,15 @@ grokStatement fe (CBlockStmt (CExpr (Just
     return $ fmap (FormalLambda k1) $ foldr applyComputation x' xs
 grokStatement fe (CBlockStmt (CExpr (Just
                     (C.CCall cvarfun exps _)) _)) = do
+    -- This case is technically not needed, but it makes slightly cleaner output
+    -- by avoiding a bind operation.
     fn <- cvarName cvarfun
     gs <- mapM (grokExpression fe) exps
-    let k = uniqIdentifier "go" frees
-        cll = foldl (App ()) (hsvar fn) $ map (comp . snd) gs
-        frees = foldr Map.union (Map.singleton fn ()) (map (compFree . snd) gs)
+    let k = uniqIdentifier "go" (compFree s1)
+        cll = foldl (\f x -> App () <$> f <*> x) (mkcomp $ hsvar fn){compFree = Map.singleton fn ()} $ map snd gs
+        s1 = fmap (`infixOp` ">>") cll
+        s = s1 <*> mkcomp (hsvar k)
         x = foldr applyComputation s $ concatMap fst gs
-        s = (mkcomp $ infixOp cll ">>" (hsvar k)) { compFree = frees }
     return $ fmap (FormalLambda k) x
 grokStatement fe (CBlockStmt (CExpr (Just
                 (CAssign CAssignOp cvarnew