feat: support multiple top level expressions

lib/Codegen/Codegen.hs

@@ -6,7 +6,9 @@
 module Codegen.Codegen where
 
 import qualified Ast.Types as AT
+import qualified Control.Monad as CM
 import qualified Control.Monad.Fix as Fix
+import qualified Data.ByteString.Short as BS
 import qualified LLVM.AST as AST
 import qualified LLVM.AST.Constant as C
 import qualified LLVM.AST.IntegerPredicate as IP
@@ -16,16 +18,32 @@ import qualified LLVM.IRBuilder.Module as M
 import qualified LLVM.IRBuilder.Monad as IRM
 
 -- | Type alias for the monad stack used for code generation.
-type MonadCodegen m = (IRM.MonadIRBuilder m, Fix.MonadFix m)
+type MonadCodegen m = (IRM.MonadIRBuilder m, M.MonadModuleBuilder m, Fix.MonadFix m)
+
+-- | Converts a parameter name to a pair of type and parameter name.
+-- The `toParamType` function takes a string and returns a pair of type and parameter name.
+-- The type is always `i32` (32-bit integer), and the parameter name is the input string.
+toParamType :: String -> (T.Type, M.ParameterName)
+toParamType param =
+  ( T.i32,
+    M.ParameterName $ BS.pack $ map (fromIntegral . fromEnum) param
+  )
 
 -- | Generates LLVM code for a given abstract syntax tree (AST).
 -- The `codegen` function takes an AST and returns the corresponding LLVM module.
 codegen :: AT.AST -> AST.Module
-codegen ast = M.buildModule "generated" $ do
-  M.function "$$generated" [] T.i32 $ \_ -> do
-    result <- case ast of
-      AT.AST exprs -> generateExpr $ head exprs
-    I.ret result
+codegen (AT.AST exprs) = M.buildModule "$$generated" $ do
+  IRM.runIRBuilderT IRM.emptyIRBuilder $ mapM_ generateTopLevel exprs
+
+-- | Generates LLVM code for a top-level expression.
+-- The `generateTopLevel` function takes an expression and generates the corresponding LLVM code.
+generateTopLevel :: (MonadCodegen m) => AT.Expr -> m ()
+generateTopLevel expr = case expr of
+  AT.Define name body -> CM.void $
+    M.function (AST.mkName name) [] T.i32 $ \_ -> do
+      result <- generateExpr body
+      I.ret result
+  _ -> error ("Unsupported top-level expression: " ++ show expr)
 
 -- | Maps binary operators to LLVM instructions.
 binaryOps :: [(AT.Operation, AST.Operand -> AST.Operand -> (IRM.MonadIRBuilder m) => m AST.Operand)]
@@ -81,4 +99,20 @@ generateExpr expr = case expr of
     pure $ AST.ConstantOperand $ C.Int 1 (if b then 1 else 0)
   AT.Op op e1 e2 -> generateOp op e1 e2
   AT.If cond then_ else_ -> generateIf cond then_ else_
+  AT.Call (AT.Lambda params body) args -> do
+    uniqueName <- IRM.freshName "lambda"
+    func <- M.function
+      uniqueName
+      [toParamType param | param <- params]
+      T.i32
+      $ \_ -> do
+        result <- generateExpr body
+        I.ret result
+    args' <- mapM generateExpr args
+    I.call func [(arg', []) | arg' <- args']
+  AT.Call func args -> do
+    func' <- generateExpr func
+    args' <- mapM generateExpr args
+    I.call func' [(arg, []) | arg <- args']
+  AT.Var name -> pure $ AST.LocalReference T.i32 $ AST.mkName (name ++ "_0")
   _ -> error ("Unimplemented expression type: " ++ show expr)