Rearrange code

src/Compiler.hs

@@ -3,12 +3,12 @@
 
 module Compiler (compile) where
 
+import           Auxiliary           (snoc)
 import           Control.Monad.State (StateT, execStateT, gets, modify)
 import           Data.Map            (Map)
 import qualified Data.Map            as Map
-import           Data.Tuple.Extra    (second)
+import           Data.Tuple.Extra    (dupe, first, second)
 import           Grammar.ErrM        (Err)
-import           Grammar.Print       (printTree)
 import           LlvmIr              (LLVMIr (..), LLVMType (..),
                                       LLVMValue (..), Visibility (..),
                                       llvmIrToString)
@@ -32,11 +32,11 @@ data FunctionInfo = FunctionInfo
 
 -- | Adds a instruction to the CodeGenerator state
 emit :: LLVMIr -> CompilerState ()
-emit l = modify (\t -> t{instructions = instructions t ++ [l]})
+emit l = modify $ \t -> t { instructions = snoc l $ instructions t }
 
 -- | Increases the variable counter in the CodeGenerator state
 increaseVarCount :: CompilerState ()
-increaseVarCount = modify (\t -> t{variableCount = variableCount t + 1})
+increaseVarCount = modify $ \t -> t { variableCount = variableCount t + 1 }
 
 -- | Returns the variable count from the CodeGenerator state
 getVarCount :: CompilerState Integer
@@ -46,38 +46,47 @@ getVarCount = gets variableCount
 getNewVar :: CompilerState Integer
 getNewVar = increaseVarCount >> getVarCount
 
-{- | Produces a map of functions infos from a list of binds,
+-- | Produces a map of functions infos from a list of binds,
-  which contains useful data for code generation.
+-- which contains useful data for code generation.
--}
 getFunctions :: [Bind] -> Map Id FunctionInfo
-getFunctions xs =
+getFunctions bs = Map.fromList $ map go bs
-    Map.fromList $
+  where
-        map
+    go (Bind id args _) =
-            ( \(Bind id args _) ->
+        (id, FunctionInfo { numArgs=length args, arguments=args })
-                ( id
-                , FunctionInfo
-                    { numArgs = length args
-                    , arguments = args
-                    }
-                )
-            )
-            xs
 
-{- | Compiles an AST and produces a LLVM Ir string.
+
-  An easy way to actually "compile" this output is to
+
-  Simply pipe it to LLI
+initCodeGenerator :: [Bind] -> CodeGenerator
--}
+initCodeGenerator scs = CodeGenerator { instructions = defaultStart
-compile :: Program -> Err String
+                                      , functions = getFunctions scs
-compile (Program prg) = do
-    let s =
-            CodeGenerator
-                { instructions = defaultStart
-                , functions = getFunctions prg
                                       , variableCount = 0
                                       }
-    ins <- instructions <$> execStateT (goDef prg) s
+
-    pure $ llvmIrToString ins
+-- | Compiles an AST and produces a LLVM Ir string.
+-- An easy way to actually "compile" this output is to
+-- Simply pipe it to lli
+compile :: Program -> Err String
+compile (Program scs) = do
+    let codegen = initCodeGenerator scs
+    llvmIrToString . instructions <$> execStateT (compileScs scs) codegen
+
+compileScs :: [Bind] -> CompilerState ()
+compileScs []                             = pure ()
+compileScs (Bind (name, t) args exp : xs) = do
+    emit $ UnsafeRaw "\n"
+    emit . Comment $ show name <> ": " <> show exp
+    let args' = map (second type2LlvmType) args
+    emit $ Define (type2LlvmType t_return) name args'
+    functionBody <- exprToValue exp
+    if name == "main"
+        then mapM_ emit $ mainContent functionBody
+        else emit $ Ret I64 functionBody
+    emit DefineEnd
+    modify $ \s -> s { variableCount = 0 }
+    compileScs xs
   where
+    t_return = snd $ partitionType (length args) t
+
 mainContent :: LLVMValue -> [LLVMIr]
 mainContent var =
     [ UnsafeRaw $
@@ -97,51 +106,25 @@ compile (Program prg) = do
     ]
 
 defaultStart :: [LLVMIr]
-    defaultStart =
+defaultStart = [ UnsafeRaw "@.str = private unnamed_addr constant [3 x i8] c\"%i\n\", align 1\n"
-        [ Comment (show $ printTree (Program prg))
-        , UnsafeRaw "@.str = private unnamed_addr constant [3 x i8] c\"%i\n\", align 1\n"
                , UnsafeRaw "declare i32 @printf(ptr noalias nocapture, ...)\n"
                ]
 
-    goDef :: [Bind] -> CompilerState ()
+compileExp :: Exp -> CompilerState ()
-    goDef [] = return ()
+compileExp = \case
-    goDef (Bind (name, t) args exp : xs) = do
+  EInt i        -> emitInt i
-        emit $ UnsafeRaw "\n"
+  EAdd t e1 e2  -> emitAdd t e1 e2
-        emit $ Comment $ show name <> ": " <> show exp
+  EId (name, _) -> emitIdent name
-        emit $ Define (type2LlvmType t_return) name (map (second type2LlvmType) args)
+  EApp t e1 e2  -> emitApp t e1 e2
-        functionBody <- exprToValue exp
+  EAbs t ti e   -> emitAbs t ti e
-        if name == "main"
+  ELet bind e   -> emitLet bind e
-            then mapM_ emit (mainContent functionBody)
-            else emit $ Ret I64 functionBody
-        emit DefineEnd
-        modify (\s -> s{variableCount = 0})
-        goDef xs
-      where
-        t_return = snd $ partitionType (length args) t
-
-    go :: Exp -> CompilerState ()
-    go (EInt int)      = emitInt int
-    go (EAdd t e1 e2)  = emitAdd t e1 e2
-    go (EId (name, _)) = emitIdent name
-    go (EApp t e1 e2)  = emitApp t e1 e2
-    go (EAbs t ti e)   = emitAbs t ti e
-    go (ELet bind e)   = emitLet bind e
-    -- go (ESub e1 e2)  = emitSub e1 e2
-    -- go (EMul e1 e2)  = emitMul e1 e2
-    -- go (EDiv e1 e2)  = emitDiv e1 e2
-    -- go (EMod e1 e2)  = emitMod e1 e2
 
 --- aux functions ---
 emitAbs :: Type -> Id -> Exp -> CompilerState ()
-    emitAbs _t tid e = do
+emitAbs _t tid e = emit . Comment $ "Lambda escaped previous stages: \\" <> show tid <> " . " <> show e
-        emit . Comment $
+
-            "Lambda escaped previous stages: \\" <> show tid <> " . " <> show e
 emitLet :: Bind -> Exp -> CompilerState ()
-    emitLet b e = do
+emitLet b e = emit . Comment $ concat [ "ELet ("
-        emit $
-            Comment $
-                concat
-                    [ "ELet ("
                                       , show b
                                       , " = "
                                       , show e
@@ -160,10 +143,9 @@ compile (Program prg) = do
                 args <- traverse exprToValue newStack
                 vs <- getNewVar
                 funcs <- gets functions
-                    let vis = case Map.lookup id funcs of
+                let visibility = maybe Local (const Global) $ Map.lookup id funcs
-                            Nothing -> Local
+                    args'      = map (first valueGetType . dupe) args
-                            Just _  -> Global
+                    call       = Call (type2LlvmType t) visibility name args'
-                    let call = Call (type2LlvmType t) vis name ((\x -> (valueGetType x, x)) <$> args)
                 emit $ SetVariable (Ident $ show vs) call
             x -> do
                 emit . Comment $ "The unspeakable happened: "
@@ -236,22 +218,26 @@ compile (Program prg) = do
 --     emit $ Sub I64 v1 v2
 
 exprToValue :: Exp -> CompilerState LLVMValue
-    exprToValue (EInt i) = return $ VInteger i
+exprToValue = \case
-    exprToValue (EId id@(name, t)) = do
+    EInt i -> pure $ VInteger i
+
+    EId id@(name, t) -> do
         funcs <- gets functions
         case Map.lookup id funcs of
             Just fi -> do
                 if numArgs fi == 0
                     then do
                         vc <- getNewVar
-                        emit $ SetVariable (Ident $ show vc) (Call (type2LlvmType t) Global name [])
+                        emit $ SetVariable (Ident $ show vc)
-                        return $ VIdent (Ident $ show vc) (type2LlvmType t)
+                            (Call (type2LlvmType t) Global name [])
-                    else return $ VFunction name Global (type2LlvmType t)
+                        pure $ VIdent (Ident $ show vc) (type2LlvmType t)
-            Nothing -> return $ VIdent name (type2LlvmType t)
+                    else pure $ VFunction name Global (type2LlvmType t)
-    exprToValue e = do
+            Nothing -> pure $ VIdent name (type2LlvmType t)
-        go e
+
+    e -> do
+        compileExp e
         v <- getVarCount
-        return $ VIdent (Ident $ show v) (getType e)
+        pure $ VIdent (Ident $ show v) (getType e)
 
 type2LlvmType :: Type -> LLVMType
 type2LlvmType = \case

src/LlvmIr.hs

@@ -68,9 +68,8 @@ instance Show Visibility where
     show Local  = "%"
     show Global = "@"
 
-{- | Represents a LLVM "value", as in an integer, a register variable,
+-- | Represents a LLVM "value", as in an integer, a register variable,
-  or a string contstant
+-- or a string contstant
--}
 data LLVMValue
     = VInteger Integer
     | VIdent Ident LLVMType
@@ -127,10 +126,8 @@ llvmIrToString = go 0
                 _         -> (i, i)
         insToString n x <> go i' xs
 
-{- | Converts a LLVM inststruction to a String, allowing for printing etc.
+-- | Converts a LLVM inststruction to a String, allowing for printing etc.
-  The integer represents the indentation
+-- The integer represents the indentation
--}
-{- FOURMOLU_DISABLE -}
 insToString :: Int -> LLVMIr -> String
 insToString i l =
     replicate i '\t' <> case l of
@@ -205,4 +202,3 @@ llvmIrToString = go 0
                 ]
         (Comment s) -> "; " <> s <> "\n"
         (Variable (Ident id)) -> "%" <> id
-{- FOURMOLU_ENABLE -}