Restrict transpilation to functions.

1 files changed, 32 insertions, 14 deletions

diff --git a/monkeypatch.hs b/monkeypatch.hs
index c09fcc7..121809a 100644
--- a/monkeypatch.hs
+++ b/monkeypatch.hs
@@ -149,6 +149,11 @@ grokExpression (CVar cv _) = Just Computation
     , compIntro = Map.empty
     , comp      = hsvar (identToString cv)
     }
+grokExpression (CConst (CIntConst n _)) = Just Computation
+    { compFree  = Map.empty
+    , compIntro = Map.empty
+    , comp      = Lit () (Int () (getCInteger n) (show n))
+    }
 grokExpression (CBinary CNeqOp a b _) = do
     ca <- grokExpression a
     cb <- grokExpression b
@@ -204,7 +209,6 @@ grokStatement (CBlockStmt (CReturn (Just exp) _)) = do
     x <- grokExpression exp
     let k = uniqIdentifier "go" (compFree x `Map.union` compIntro x)
     return $ fmap (\y -> Lambda () [hspvar k] $ App () (hsvar "return") y) x
-
 grokStatement (CBlockStmt (CExpr (Just
                 (CAssign CAssignOp cvarnew
                     (C.CCall cvarfun [] _) _)) _)) = do
@@ -236,6 +240,11 @@ grokStatement (CBlockDecl (CDecl (t:_) (v:vs) _)) = do
         }
 grokStatement _ = Nothing
 
+isFunctionDecl (CDeclExt (CDecl _ [(Just (CDeclr _ [CFunDeclr _ _ _] _ _ _),_,_)] _)) = True
+isFunctionDecl (CFDefExt (CFunDef _ _ _ (CCompound [] _ _) _)) = True
+isFunctionDecl _ = False
+
+
 
 transpile :: C2HaskellOptions -> FilePath -> IncludeStack -> CTranslationUnit NodeInfo -> IO ()
 transpile o fname incs (CTranslUnit edecls _) = do
@@ -256,19 +265,28 @@ transpile o fname incs (CTranslUnit edecls _) = do
                         i <- m
                         maybe [] (return . identToString) i
             -- mapM_ (putStrLn . show . pretty) (symbolSource sym)
-            forM_ (take 1 h) $ \hh -> do
-                putStrLn . HS.prettyPrint $ changeType makeFunctionUseIO hh
-            putStrLn $ unwords (hname:as) ++ " ="
-            let bdy = concat $ take 1 $ dropWhile null $ map body $ symbolSource sym
-            if oPrettyTree o
-                then forM_ bdy $ \d -> putStrLn $ ppShow $ everywhere (mkT eraseNodeInfo) $ d
-                else do
-                    let mhask = do
-                            xs <- sequence $ map grokStatement bdy
-                            return $ foldr applyComputation (Computation Map.empty Map.empty hsopUnit) xs
-                    case mhask of
-                        Just hask -> mapM_ (putStrLn . ("  "++)) $ lines $ HS.prettyPrint $ comp hask
-                        Nothing -> forM_ bdy $ \d -> do
+            let mprintHeader = do
+                    hh <- listToMaybe h
+                    guard (isJust (oSelectFunction o) || isFunctionDecl c)
+                    Just $ do
+                        -- putStrLn $ show (fmap (const LT) c)
+                        putStrLn . HS.prettyPrint $ changeType makeFunctionUseIO hh
+                        putStrLn $ unwords (hname:as) ++ " ="
+            forM_ mprintHeader $ \printHeader -> do
+                let bdy = concat $ take 1 $ dropWhile null $ map body $ symbolSource sym
+                if oPrettyTree o
+                    then do printHeader
+                            forM_ bdy $ \d -> putStrLn $ ppShow $ everywhere (mkT eraseNodeInfo) $ d
+                    else do
+                        let mhask = do
+                                xs <- sequence $ map grokStatement bdy
+                                return $ foldr applyComputation (Computation Map.empty Map.empty hsopUnit) xs
+                        case mhask of
+                            Just hask -> do printHeader
+                                            mapM_ (putStrLn . ("  "++)) $ lines $ HS.prettyPrint $ comp hask
+                            Nothing -> forM_ (oSelectFunction o) $ \_ -> do
+                                printHeader
+                                forM_ bdy $ \d -> do
                                     putStrLn . show . pretty $ d
                                     mapM_ (putStrLn . HS.prettyPrint . comp) (grokStatement d)
     return ()