Switched grokStatement to use StateT.

2 files changed, 40 insertions, 36 deletions

diff --git a/Unique.hs b/Unique.hs
index 13ae7cf..d0d3eb1 100644
--- a/Unique.hs
+++ b/Unique.hs
@@ -1,5 +1,6 @@
 module Unique
     ( UniqueFactory
+    , freshUniques
     , Unique
     , uniqueSymbol
     , substituteUnique
diff --git a/monkeypatch.hs b/monkeypatch.hs
index 3b93e5b..ceeaf25 100644
--- a/monkeypatch.hs
+++ b/monkeypatch.hs
@@ -127,6 +127,9 @@ capitalize xs = concatMap (cap . drop 1) gs
         gs = groupBy (\a b -> b/='_') $ '_':xs
         cap (c:cs) = toUpper c : cs
 
+mb :: Functor m => m a -> StateT t m a
+mb m = StateT $ \s ->  fmap (, s) m
+
 transField :: CDeclaration t -> [(TH.Name, TH.Bang, TH.Type)]
 transField (CDecl [CTypeSpec (CTypeDef ctyp _)] vars _)
     = do
@@ -345,10 +348,10 @@ isGlobalRef fe sym = fromMaybe False $ do
 -- expression.
 grokExpression :: FunctionEnvironment
                   -> CExpression a
-                  -> Maybe ([Computation FormalLambda], Computation (HS.Exp ()))
+                  -> StateT UniqueFactory Maybe ([Computation FormalLambda], Computation (HS.Exp ()))
 grokExpression fe (CVar cv _) =
     let v = identToString cv
-    in Just $
+    in return $
         if isGlobalRef fe v
             then let k = uniqIdentifier "go" (varmap [v,hv])
                      s = Computation
@@ -367,9 +370,9 @@ grokExpression fe (CVar cv _) =
                 { compFree  = Map.singleton (identToString cv) ()
                 }
 grokExpression fe (CConst (CIntConst n _)) =
-    Just $ (,) [] $ mkcomp $ Lit () (Int () (getCInteger n) (show n))
+    return $ (,) [] $ mkcomp $ Lit () (Int () (getCInteger n) (show n))
 grokExpression fe (CConst (CStrConst s _)) =
-    Just $ (,) [] $ mkcomp $ Lit () (HS.String () (getCString s) (getCString s))
+    return $ (,) [] $ mkcomp $ Lit () (HS.String () (getCString s) (getCString s))
 grokExpression fe (CBinary op a b _) = do
     (as,ca) <- grokExpression fe a
     (bs0,cb0) <- grokExpression fe b
@@ -432,8 +435,8 @@ grokExpression fe (C.CCall fn exps u) = grokCall fe True (C.CCall fn exps u)
 grokExpression fe (CStatExpr (CCompound idents xs _) _) = do
     let (y,ys) = splitAt 1 (reverse xs)
     y' <- case y of
-                [CBlockStmt (CExpr mexp ni)] -> Just $ CBlockStmt (CReturn mexp ni)
-                _                            -> Just (head y) -- Nothing FIXME
+                [CBlockStmt (CExpr mexp ni)] -> return $ CBlockStmt (CReturn mexp ni)
+                _                            -> return (head y) -- Nothing FIXME
     gs <- mapM (grokStatement fe) (reverse $ y' : ys)
     let s0 = foldr applyComputation (mkcomp retUnit) gs
         s1 = fmap (\xp -> Paren () xp) s0
@@ -451,7 +454,7 @@ grokExpression fe (CStatExpr (CCompound idents xs _) _) = do
         { compFree  = Map.singleton hv ()
         }
 grokExpression fe (CAssign CAssignOp cvar expr _) = do
-    v <- cvarName cvar
+    v <- mb $ cvarName cvar
     (ss,x) <- grokExpression fe expr
     let k = uniqIdentifier "go" (Map.insert v () $ foldr (\s m -> compFree s `Map.union` compIntro s `Map.union` m) Map.empty ss)
         s = x
@@ -462,7 +465,7 @@ grokExpression fe (CAssign CAssignOp cvar expr _) = do
             }
     return $ (,) (ss ++ [s]) $ mkcomp (hsvar v)
 grokExpression fe (CMember cvar fld isptr _) = do
-    v <- cvarName cvar
+    v <- mb $ cvarName cvar
     let fieldlbl = identToString fld
         hv = v ++ fieldlbl
         e = App () (App () (hsvar "get")
@@ -473,12 +476,12 @@ grokExpression fe (CMember cvar fld isptr _) = do
         s = (FormalLambda k <$> fmap (($ Lambda () [hspvar hv] (hsvar k)) . (`infixOp` ">>=")) e')
                 { compIntro = Map.singleton hv () }
     return $ (,) [s] (mkcomp $ hsvar hv){ compFree = Map.singleton hv () }
-grokExpression fe _ = Nothing
+grokExpression fe _ = mzero
 
 grokCall :: FunctionEnvironment
             -> Bool
             -> CExpression a
-            -> Maybe ([Computation FormalLambda], Computation (HS.Exp ()))
+            -> StateT UniqueFactory Maybe ([Computation FormalLambda], Computation (HS.Exp ()))
 grokCall fe wantsRet (C.CCall (CVar fn _) exps _) = do
     gs <- mapM (grokExpression fe) exps
     let ss = concatMap fst gs -- TODO: resolve variable name conflicts
@@ -519,14 +522,14 @@ grokCall fe wantsRet (C.CCall fnx@(CMember cvar fld isptr _) exps _) = do
     return $ (,) (ss++[s]) (mkcomp $ hsvar hv)
         { compFree  = Map.singleton hv ()
         }
-grokCall _ _ _ = Nothing
+grokCall _ _ _ = mzero
 
 
 grokInitialization :: Foldable t1 =>
                       FunctionEnvironment
                       -> t1 (CDeclarationSpecifier t2)
                       -> (Maybe (CDeclarator a1), CInitializer a2)
-                      -> Maybe (Computation FormalLambda)
+                      -> StateT UniqueFactory Maybe (Computation FormalLambda)
 grokInitialization fe _ (Just (CDeclr (Just cv0) _ _ _ _),CInitExpr exp _) = do
     let v = identToString cv0
     (xs,x) <- grokExpression fe exp
@@ -547,7 +550,7 @@ grokInitialization fe ts (Just (CDeclr (Just cv0) _ _ _ _),CInitList exps _) = d
                 forM exps $ \(ms,initexpr) -> do
                     case initexpr of
                         CInitExpr ie _ -> (grokExpression fe) ie >>= \g -> return (ms,g)
-                        _              -> Nothing
+                        _              -> mzero
             let assigns = do
                     (ms,(ss,x)) <- gs
                     let k2 = uniqIdentifier "gopoo" (compFree ret)
@@ -577,8 +580,8 @@ grokInitialization fe ts (Just (CDeclr (Just cv0) _ _ _ _),CInitList exps _) = d
                 k = uniqIdentifier "go" Map.empty -- (compFree ret) TODO
                 ret = foldr applyComputation (mkcomp $ hsvar k) $ newstruct : assigns
             return $ fmap (FormalLambda k) ret
-        _ -> Nothing
-grokInitialization _ _ _ = Nothing
+        _ -> mzero
+grokInitialization _ _ _ = mzero
 
 hasBool :: HS.Type () -> Bool
 hasBool = (1 <=) . gcount (mkQ False (\t -> case t of { HS.Ident () "Bool" -> True; _ -> False }))
@@ -591,33 +594,33 @@ promote fe y@(Lit () (Int () n _)) | (n==0 || n==1) &&  hasBool (fe Map.! "") =
 promote _ y = y
 
 
-grokStatement :: FunctionEnvironment -> CCompoundBlockItem a -> Maybe (Computation FormalLambda)
+grokStatement :: FunctionEnvironment -> CCompoundBlockItem a -> StateT UniqueFactory Maybe (Computation FormalLambda)
 grokStatement fe (CBlockStmt (CReturn (Just exp) _)) = do
     (xs,x) <- grokExpression fe exp
     let k = uniqIdentifier "go" (compFree x `Map.union` compIntro x)
         x' = fmap (\y -> App () (hsvar "return") $ promote (fnArgs fe) y) x
     return $ fmap (\y -> FormalLambda k y) $ foldr applyComputation x' xs
 grokStatement fe (CBlockStmt (CReturn Nothing _)) =
-    Just $ mkcomp $ FormalLambda "go" retUnit
+    return $ mkcomp $ FormalLambda "go" retUnit
 grokStatement fe (CBlockStmt (CCont _)) =
-    Just (mkcomp $ FormalLambda "go" $ hsvar " continue")
+    return (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)
     (mkif,stmts) <- case (thn,els) of
 
-        (CCompound [] stmts _, Nothing                                               ) -> Just (mkif0, stmts)
-        (stmt                , Nothing                                               ) -> Just (mkif0, [CBlockStmt stmt])
-        (CCompound [] stmts _, Just (CExpr Nothing _)                                ) -> Just (mkif0, stmts)
-        (CCompound [] stmts _, Just (CCompound [] [ CBlockStmt (CExpr Nothing _) ] _)) -> Just (mkif0, stmts)
+        (CCompound [] stmts _, Nothing                                               ) -> return (mkif0, stmts)
+        (stmt                , Nothing                                               ) -> return (mkif0, [CBlockStmt stmt])
+        (CCompound [] stmts _, Just (CExpr Nothing _)                                ) -> return (mkif0, stmts)
+        (CCompound [] stmts _, Just (CCompound [] [ CBlockStmt (CExpr Nothing _) ] _)) -> return (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])
+        (CExpr Nothing _                              ,Just (CCompound [] stmts _)) -> return (flip mkif0, stmts)
+        (CCompound [] [CBlockStmt (CExpr Nothing _)] _,Just (CCompound [] stmts _)) -> return (flip mkif0, stmts)
+        (CExpr Nothing _                              ,Just e@(CExpr (Just _) _))   -> return (flip mkif0, [CBlockStmt e])
+        (CCompound [] [CBlockStmt (CExpr Nothing _)] _,Just e@(CExpr (Just _) _))   -> return (flip mkif0, [CBlockStmt e])
 
-        _ -> trace ("Unhandled if: "++show (fmap (const LT) thn)) $ Nothing -- TODO
+        _ -> trace ("Unhandled if: "++show (fmap (const LT) thn)) $ mzero -- TODO
 
     ss <- sequence $ map (grokStatement fe) stmts
     let s = foldr applyComputation (mkcomp $ hsvar k) ss
@@ -631,8 +634,8 @@ grokStatement fe (CBlockStmt (CIf exp thn els _)) = do
 grokStatement fe (CBlockStmt (CExpr (Just (C.CCall (CVar (C.Ident "__assert_fail" _ _) _) xs _)) _)) = do
     x <- case xs of
             (CConst (CStrConst msg _):_) -> let s = getCString msg
-                                            in Just $ mkcomp $ Lit () (HS.String () s s)
-            _                            -> Nothing
+                                            in return $ mkcomp $ Lit () (HS.String () s s)
+            _                            -> mzero
     let k = uniqIdentifier "go" (compFree x `Map.union` compIntro x)
         x' = fmap (\y -> App () (hsvar "error") y) x
     return $ fmap (FormalLambda k) x'
@@ -640,7 +643,7 @@ grokStatement fe (CBlockStmt (CExpr (Just
                 (CAssign CAssignOp
                    (CMember cvar fld isptr _) expr _)) _)) = do
     (xs,x) <- grokExpression fe expr
-    v <- cvarName cvar
+    v <- mb $ cvarName cvar
     let fieldlbl = identToString fld
         k1 = uniqIdentifier "go" (compFree x)
         fieldinit = comp x
@@ -662,8 +665,8 @@ grokStatement fe (CBlockStmt (CExpr (Just (C.CCall cvarfun exps a)) _)) = do
 grokStatement fe (CBlockStmt (CExpr (Just
                 (CAssign CAssignOp cvarnew
                     (C.CCall cvarfun [] _) _)) _)) = do
-    v <- cvarName cvarnew
-    fn <- cvarName cvarfun
+    v <- mb $ cvarName cvarnew
+    fn <- mb $ cvarName cvarfun
     let k = uniqIdentifier "go" (varmap [v,fn])
     return Computation
         { compFree = Map.singleton fn ()
@@ -690,7 +693,7 @@ grokStatement fe (CBlockStmt (CExpr (Just (CUnary CPostIncOp (CMember (CVar cv0
             }
 grokStatement fe (CBlockStmt (CExpr mexpr _)) = do
     -- trace ("CExpr statement: " ++ take 50 (show $ fmap (fmap $ const ()) mexpr)) $ return ()
-    (ss,pre) <- maybe (Just $ (,) [] $ mkcomp id)
+    (ss,pre) <- maybe (return $ (,) [] $ mkcomp id)
                       (let -- Discard pure value since we are interested only in side-effects.
                            discard = const $ mkcomp id
                            -- Alternate: keep pure-value using `seq` operator.
@@ -721,7 +724,7 @@ grokStatement fe (CBlockStmt (CWhile cond (CCompound [] bdy _) isDoWhile _)) = d
         x = foldr applyComputation c' ss -- continue function
         vs = [] -- Map.keys $ compIntro g
     return $ fmap (FormalLambda "fin") $ fmap (factorOutFunction "continue" vs (comp x) " continue") g
-grokStatement fe _ = Nothing
+grokStatement fe _ = mzero
 
 isFunctionDecl :: CExternalDeclaration a -> Bool
 isFunctionDecl (CDeclExt (CDecl _ [(Just (CDeclr _ [CFunDeclr _ _ _] _ _ _),_,_)] _)) = True
@@ -851,7 +854,7 @@ transpile o fname incs (CTranslUnit edecls _) = do
                     then do printHeader
                             forM_ bdy $ \d -> putStrLn $ ppShow $ cleanTree d
                     else do
-                        let mhask = do
+                        let mhask = (`evalStateT` freshUniques) $ do
                                 xs <- mapM (grokStatement fe) bdy
                                 return $ foldr applyComputation (mkcomp retUnit) xs
                         case mhask of
@@ -861,7 +864,7 @@ transpile o fname incs (CTranslUnit edecls _) = do
                                 printHeader
                                 forM_ bdy $ \d -> do
                                     putStrLn $ " C: " ++ show (pretty d)
-                                    case grokStatement fe d of
+                                    case grokStatement fe d `evalStateT` freshUniques of
 
                                         Just hd -> do putStrLn $ "fr: " ++ intercalate " " (Map.keys (compFree hd))
                                                       putStrLn $ "HS: " ++ HS.prettyPrint (informalize $ comp hd)