Added FUNCTION and PROCEDURE

5 files changed, 132 insertions, 1 deletions

diff --git a/include/cps/CodeGen.h b/include/cps/CodeGen.h
index 57f0899..832ba7a 100644
--- a/include/cps/CodeGen.h
+++ b/include/cps/CodeGen.h
@@ -4,6 +4,7 @@
 #include "llvm/IR/Module.h"
 #include "cps/AST.h"
 #include "cps/RuntimeCheck.h"
+#include "cps/FunctionGen.h"
 #include <map>
 #include <memory>
 
@@ -19,6 +20,7 @@ class CodeGen {
     
     std::unique_ptr<ArrayHandler> Arrays;
     std::unique_ptr<RuntimeCheck> RuntimeChecker;
+    std::unique_ptr<FunctionGen> FuncGen;
 
     llvm::FunctionCallee PrintfFunc;
     llvm::FunctionCallee ScanfFunc;
diff --git a/include/cps/FunctionAST.h b/include/cps/FunctionAST.h
new file mode 100644
index 0000000..155bd8c
--- /dev/null
+++ b/include/cps/FunctionAST.h
@@ -0,0 +1,76 @@
+#pragma once
+#include "cps/AST.h"
+#include <vector>
+#include <string>
+#include <memory>
+
+namespace cps {
+
+class PrototypeAST {
+    std::string Name;
+    std::vector<std::pair<std::string, std::string>> Args;
+    std::string ReturnType;
+    bool IsExternal;
+
+public:
+    PrototypeAST(const std::string &Name, 
+                 std::vector<std::pair<std::string, std::string>> Args, 
+                 const std::string &ReturnType,
+                 bool IsExternal = false)
+        : Name(Name), Args(std::move(Args)), ReturnType(ReturnType), IsExternal(IsExternal) {}
+
+    const std::string &getName() const { return Name; }
+    const std::vector<std::pair<std::string, std::string>> &getArgs() const { return Args; }
+    const std::string &getReturnType() const { return ReturnType; }
+    bool isExternal() const { return IsExternal; }
+};
+
+class FunctionDefAST : public StmtAST {
+    std::unique_ptr<PrototypeAST> Proto;
+    std::vector<std::unique_ptr<StmtAST>> Body;
+
+public:
+    FunctionDefAST(std::unique_ptr<PrototypeAST> Proto, 
+                   std::vector<std::unique_ptr<StmtAST>> Body)
+        : Proto(std::move(Proto)), Body(std::move(Body)) {}
+
+    PrototypeAST *getProto() const { return Proto.get(); }
+    const std::vector<std::unique_ptr<StmtAST>> &getBody() const { return Body; }
+};
+
+class CallExprAST : public ExprAST {
+    std::string Callee;
+    std::vector<std::unique_ptr<ExprAST>> Args;
+
+public:
+    CallExprAST(const std::string &Callee, std::vector<std::unique_ptr<ExprAST>> Args)
+        : Callee(Callee), Args(std::move(Args)) {}
+
+    const std::string &getCallee() const { return Callee; }
+    const std::vector<std::unique_ptr<ExprAST>> &getArgs() const { return Args; }
+};
+
+class CallStmtAST : public StmtAST {
+    std::string Callee;
+    std::vector<std::unique_ptr<ExprAST>> Args;
+
+public:
+    CallStmtAST(const std::string &Callee, std::vector<std::unique_ptr<ExprAST>> Args)
+        : Callee(Callee), Args(std::move(Args)) {}
+
+    const std::string &getCallee() const { return Callee; }
+    const std::vector<std::unique_ptr<ExprAST>> &getArgs() const { return Args; }
+};
+
+
+class ReturnStmtAST : public StmtAST {
+    std::unique_ptr<ExprAST> RetVal;
+
+public:
+    ReturnStmtAST(std::unique_ptr<ExprAST> RetVal = nullptr) 
+        : RetVal(std::move(RetVal)) {}
+        
+    ExprAST *getRetVal() const { return RetVal.get(); }
+};
+
+} // namespace cps
diff --git a/include/cps/FunctionGen.h b/include/cps/FunctionGen.h
new file mode 100644
index 0000000..c1a6ebd
--- /dev/null
+++ b/include/cps/FunctionGen.h
@@ -0,0 +1,39 @@
+#pragma once
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Module.h"
+#include "cps/FunctionAST.h"
+#include <map>
+#include <vector>
+#include <functional>
+#include <string>
+
+namespace cps {
+
+class FunctionGen {
+    llvm::LLVMContext &Context;
+    llvm::Module &Module;
+    llvm::IRBuilder<> &Builder;
+    
+    std::map<std::string, llvm::AllocaInst*> &NamedValues;
+
+    void createArgumentAllocas(llvm::Function *F, const std::vector<std::pair<std::string, std::string>> &Args);
+
+public:
+    FunctionGen(llvm::LLVMContext &C, llvm::Module &M, llvm::IRBuilder<> &B, 
+                std::map<std::string, llvm::AllocaInst*> &NV)
+        : Context(C), Module(M), Builder(B), NamedValues(NV) {}
+
+    llvm::Type *getLLVMType(const std::string &TypeName);
+
+    llvm::Function *emitPrototype(PrototypeAST *Proto);
+
+    llvm::Function *emitFunctionDef(FunctionDefAST *FuncAST, const std::function<void(StmtAST*)> &StmtEmitter);
+
+    llvm::Value *emitCallExpr(CallExprAST *Call, const std::vector<llvm::Value*> &Args);
+    
+    void emitCallStmt(CallStmtAST *Call, const std::vector<llvm::Value*> &Args);
+
+    void emitReturn(ReturnStmtAST *Ret, llvm::Value *RetVal);
+};
+
+}
diff --git a/include/cps/Lexer.h b/include/cps/Lexer.h
index ccbda98..cd34435 100644
--- a/include/cps/Lexer.h
+++ b/include/cps/Lexer.h
@@ -41,7 +41,15 @@ enum Token {
     tok_le = -203, // <=
     tok_ge = -204, // >=
     
-    tok_colon = -300 // :
+    tok_colon = -300, // :
+
+    tok_function = -400, 
+    tok_endfunction = -401, 
+    tok_return = -402, 
+    tok_returns = -403,
+    tok_procedure = -404, 
+    tok_endprocedure = -405, 
+    tok_call = -406  
 };
 
 class Lexer {
diff --git a/include/cps/Parser.h b/include/cps/Parser.h
index f5be655..f328141 100644
--- a/include/cps/Parser.h
+++ b/include/cps/Parser.h
@@ -31,6 +31,12 @@ class Parser {
     std::unique_ptr<StmtAST> ParseForStmt();
     
     std::unique_ptr<StmtAST> ParseDeclare();
+
+    std::unique_ptr<StmtAST> ParseFunction();
+    std::unique_ptr<StmtAST> ParseProcedure();
+    std::unique_ptr<StmtAST> ParseCallStmt();
+    std::unique_ptr<StmtAST> ParseReturnStmt();
+    std::vector<std::pair<std::string, std::string>> ParsePrototypeArgs();
     
 public:
     Parser(Lexer &L);