Grok more if-then cases.

1 files changed, 16 insertions, 2 deletions

diff --git a/monkeypatch.hs b/monkeypatch.hs
index 98ebe38..897323d 100644
--- a/monkeypatch.hs
+++ b/monkeypatch.hs
@@ -380,15 +380,29 @@ grokStatement (CBlockStmt (CExpr (Just
             $ infixOp (hsvar fn) ">>="
                 $ Lambda () [hspvar v] (hsvar k)
         }
-grokStatement (CBlockStmt (CIf exp (CCompound [] stmts _) Nothing _)) = do
+grokStatement (CBlockStmt (CIf exp thn els _)) = do
     (xs,x) <- grokExpression exp
+    let mkif0 = If () (comp x)
+    (mkif,stmts) <- case (thn,els) of
+
+        (CCompound [] stmts _, Nothing                                               ) -> Just (mkif0, stmts)
+        (CCompound [] stmts _, Just (CExpr Nothing _)                                ) -> Just (mkif0, stmts)
+        (CCompound [] stmts _, Just (CCompound [] [ CBlockStmt (CExpr Nothing _) ] _)) -> Just (mkif0, stmts)
+
+        (CExpr Nothing _                              ,Just (CCompound [] stmts _)) -> Just (flip mkif0, stmts)
+        (CCompound [] [CBlockStmt (CExpr Nothing _)] _,Just (CCompound [] stmts _)) -> Just (flip mkif0, stmts)
+        (CExpr Nothing _                              ,Just e@(CExpr (Just _) _))   -> Just (flip mkif0, [CBlockStmt e])
+        (CCompound [] [CBlockStmt (CExpr Nothing _)] _,Just e@(CExpr (Just _) _))   -> Just (flip mkif0, [CBlockStmt e])
+
+        _ -> Nothing -- TODO
+
     ss <- sequence $ map grokStatement stmts
     let s = foldr applyComputation (Computation Map.empty Map.empty (hsvar k)) ss
         k = uniqIdentifier "go" (Map.union (compFree x) (compFree s))
     return $ flip (foldr applyComputation) xs Computation
         { compFree = compFree x `Map.union` compFree s
         , compIntro = compIntro s
-        , comp = Lambda () [hspvar k] $ If () (comp x) (comp s) (hsvar k)
+        , comp = Lambda () [hspvar k] $ mkif (comp s) (hsvar k)
         }
 grokStatement (CBlockStmt (CExpr (Just (CUnary CPostIncOp (CMember (CVar cv0 _) fld True _) _)) _)) = do
     let k1 = uniqIdentifier "go" (varmap [fieldlbl,v])