Added FUNCTION and PROCEDURE

5 files changed, 354 insertions, 11 deletions

diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index 9dc88af..6358699 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -17,6 +17,8 @@ CodeGen::CodeGen() {
     
     Arrays = std::make_unique<ArrayHandler>(*TheContext, *Builder, *TheModule, NamedValues, *RuntimeChecker);
     
+    FuncGen = std::make_unique<FunctionGen>(*TheContext, *TheModule, *Builder, NamedValues);
+
     SetupExternalFunctions();
 }
 
@@ -105,6 +107,15 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
             return nullptr;
         }
     }
+
+    if (auto *Call = dynamic_cast<CallExprAST*>(Expr)) {
+        std::vector<Value*> Args;
+        for (const auto &Arg : Call->getArgs()) {
+            Args.push_back(emitExpr(Arg.get()));
+        }
+        return FuncGen->emitCallExpr(Call, Args);
+    }
+
     return nullptr;
 }
 
@@ -261,6 +272,36 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
         return;
     }
 
+    if (auto *FuncDef = dynamic_cast<FunctionDefAST*>(Stmt)) {
+        BasicBlock *SavedBlock = Builder->GetInsertBlock();
+
+        FuncGen->emitFunctionDef(FuncDef, [this](StmtAST *S) { 
+            this->emitStmt(S); 
+        });
+
+        if (SavedBlock) Builder->SetInsertPoint(SavedBlock);
+        
+        return;
+    }
+
+    else if (auto *Call = dynamic_cast<CallStmtAST*>(Stmt)) {
+        std::vector<Value*> Args;
+        for (const auto &ArgExpr : Call->getArgs()) {
+            Args.push_back(emitExpr(ArgExpr.get()));
+        }
+        FuncGen->emitCallStmt(Call, Args);
+        return;
+    }
+
+    else if (auto *Ret = dynamic_cast<ReturnStmtAST*>(Stmt)) {
+        Value *RetVal = nullptr;
+        if (Ret->getRetVal()) {
+            RetVal = emitExpr(Ret->getRetVal());
+        }
+        FuncGen->emitReturn(Ret, RetVal);
+        return;
+    }
+
     if (auto *Decl = dynamic_cast<DeclareStmtAST*>(Stmt)) {
         Function *TheFunction = Builder->GetInsertBlock()->getParent();
         AllocaInst *Alloca = CreateEntryBlockAlloca(TheFunction, Decl->getName());
@@ -313,4 +354,3 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
         emitForStmt(ForStmt);
     }
 }
-
diff --git a/lib/CodeGen/FunctionGen.cc b/lib/CodeGen/FunctionGen.cc
new file mode 100644
index 0000000..b2adc08
--- /dev/null
+++ b/lib/CodeGen/FunctionGen.cc
@@ -0,0 +1,125 @@
+#include "cps/FunctionGen.h"
+#include "llvm/IR/Verifier.h"
+#include <iostream>
+
+using namespace llvm;
+using namespace cps;
+
+Type *FunctionGen::getLLVMType(const std::string &TypeName) {
+    if (TypeName == "INTEGER") return Type::getInt64Ty(Context);
+    if (TypeName == "BOOLEAN") return Type::getInt1Ty(Context);
+    if (TypeName == "REAL")    return Type::getDoubleTy(Context);
+    if (TypeName == "VOID")    return Type::getVoidTy(Context);
+    return Type::getInt64Ty(Context);
+}
+
+void FunctionGen::createArgumentAllocas(Function *F, const std::vector<std::pair<std::string, std::string>> &Args) {
+    Function::arg_iterator AI = F->arg_begin();
+    for (unsigned Idx = 0, E = Args.size(); Idx != E; ++Idx, ++AI) {
+        std::string ArgName = Args[Idx].first;
+        std::string ArgTypeStr = Args[Idx].second;
+        
+        Type *ArgType = getLLVMType(ArgTypeStr);
+        IRBuilder<> TmpB(&F->getEntryBlock(), F->getEntryBlock().begin());
+        AllocaInst *Alloca = TmpB.CreateAlloca(ArgType, nullptr, ArgName);
+        
+        Value *ArgVal = &(*AI);
+        Builder.CreateStore(ArgVal, Alloca);
+        
+        NamedValues[ArgName] = Alloca; 
+    }
+}
+
+Function *FunctionGen::emitPrototype(PrototypeAST *Proto) {
+    std::vector<Type*> ArgTypes;
+    for (const auto &Arg : Proto->getArgs()) {
+        ArgTypes.push_back(getLLVMType(Arg.second));
+    }
+
+    Type *RetType = getLLVMType(Proto->getReturnType());
+    FunctionType *FT = FunctionType::get(RetType, ArgTypes, false);
+    Function *F = Function::Create(FT, Function::ExternalLinkage, Proto->getName(), &Module);
+
+    unsigned Idx = 0;
+    for (auto &Arg : F->args()) {
+        if (Idx < Proto->getArgs().size())
+            Arg.setName(Proto->getArgs()[Idx++].first);
+    }
+    return F;
+}
+
+Function *FunctionGen::emitFunctionDef(FunctionDefAST *FuncAST, 
+                                       const std::function<void(StmtAST*)> &StmtEmitter) {
+    PrototypeAST *Proto = FuncAST->getProto();
+    Function *TheFunction = Module.getFunction(Proto->getName());
+
+    if (!TheFunction)
+        TheFunction = emitPrototype(Proto);
+
+    if (!TheFunction) return nullptr;
+    if (!TheFunction->empty()) {
+        fprintf(stderr, "Error: Function %s cannot be redefined.\n", Proto->getName().c_str());
+        return nullptr;
+    }
+
+    BasicBlock *BB = BasicBlock::Create(Context, "entry", TheFunction);
+    Builder.SetInsertPoint(BB);
+
+    std::map<std::string, llvm::AllocaInst*> OldNamedValues = NamedValues;
+    NamedValues.clear();
+    
+    createArgumentAllocas(TheFunction, Proto->getArgs());
+
+    for (const auto &Stmt : FuncAST->getBody()) {
+        StmtEmitter(Stmt.get());
+    }
+
+    if (!Builder.GetInsertBlock()->getTerminator()) {
+        if (Proto->getReturnType() == "VOID")
+            Builder.CreateRetVoid();
+        else
+            Builder.CreateRet(Constant::getNullValue(TheFunction->getReturnType()));
+    }
+
+    verifyFunction(*TheFunction);
+    NamedValues = OldNamedValues;
+    return TheFunction;
+}
+
+static Value *GenerateCall(llvm::Module &Module, llvm::IRBuilder<> &Builder, llvm::LLVMContext &Context,
+                    const std::string &CalleeName, const std::vector<llvm::Value*> &Args) {
+    Function *CalleeF = Module.getFunction(CalleeName);
+    if (!CalleeF) {
+        std::vector<Type*> ArgTypes;
+        for (auto *Val : Args) ArgTypes.push_back(Val->getType());
+        FunctionType *FT = FunctionType::get(Type::getInt64Ty(Context), ArgTypes, false);
+        CalleeF = Function::Create(FT, Function::ExternalLinkage, CalleeName, &Module);
+    }
+
+    if (CalleeF->arg_size() != Args.size()) {
+        fprintf(stderr, "Error: Incorrect # arguments passed to %s\n", CalleeName.c_str());
+        return nullptr;
+    }
+
+    if (CalleeF->getReturnType()->isVoidTy()) {
+        return Builder.CreateCall(CalleeF, Args);
+    }
+
+    return Builder.CreateCall(CalleeF, Args, "calltmp");
+}
+
+llvm::Value *FunctionGen::emitCallExpr(CallExprAST *Call, const std::vector<llvm::Value*> &Args) {
+    return GenerateCall(Module, Builder, Context, Call->getCallee(), Args);
+}
+
+void FunctionGen::emitCallStmt(CallStmtAST *Call, const std::vector<llvm::Value*> &Args) {
+    GenerateCall(Module, Builder, Context, Call->getCallee(), Args);
+}
+
+void FunctionGen::emitReturn(ReturnStmtAST *Ret, llvm::Value *RetVal) {
+    if (RetVal) {
+        Builder.CreateRet(RetVal);
+    } else {
+        Builder.CreateRetVoid();
+    }
+}
diff --git a/lib/Lexer/Lexer.cc b/lib/Lexer/Lexer.cc
index 51e27dd..6854cc9 100644
--- a/lib/Lexer/Lexer.cc
+++ b/lib/Lexer/Lexer.cc
@@ -21,6 +21,9 @@ int Lexer::gettok() {
 
         if (IdentifierStr == "DECLARE") return tok_declare;
         if (IdentifierStr == "INTEGER") return tok_integer_kw;
+        if (IdentifierStr == "BOOLEAN") return tok_integer_kw; // TODO: implement real BOOLEAN
+        if (IdentifierStr == "REAL")    return tok_integer_kw; // TODO: implement real REAL
+        
         if (IdentifierStr == "INPUT") return tok_input;
         if (IdentifierStr == "OUTPUT") return tok_output;
         
@@ -44,6 +47,14 @@ int Lexer::gettok() {
         if (IdentifierStr == "ARRAY") return tok_array;
         if (IdentifierStr == "OF") return tok_of;
 
+        if (IdentifierStr == "FUNCTION") return tok_function;
+        if (IdentifierStr == "ENDFUNCTION") return tok_endfunction;
+        if (IdentifierStr == "PROCEDURE") return tok_procedure;
+        if (IdentifierStr == "ENDPROCEDURE") return tok_endprocedure;
+        if (IdentifierStr == "RETURN") return tok_return;
+        if (IdentifierStr == "RETURNS") return tok_returns;
+        if (IdentifierStr == "CALL") return tok_call;
+
         return tok_identifier;
     }
 
@@ -93,6 +104,15 @@ int Lexer::gettok() {
         return tok_colon;
     }
 
+    if (LastChar == '(') { LastChar = getchar(); return '('; }
+    if (LastChar == ')') { LastChar = getchar(); return ')'; }
+    if (LastChar == '[') { LastChar = getchar(); return '['; }
+    if (LastChar == ']') { LastChar = getchar(); return ']'; }
+    if (LastChar == ',') { LastChar = getchar(); return ','; }
+    if (LastChar == '+') { LastChar = getchar(); return '+'; }
+    if (LastChar == '-') { LastChar = getchar(); return '-'; }
+    if (LastChar == '*') { LastChar = getchar(); return '*'; }
+
     if (LastChar == '/') {
         int NextChar = getchar();
         if (NextChar == '/') {
diff --git a/lib/Parser/Parser.cc b/lib/Parser/Parser.cc
index b9396ce..563dd9b 100644
--- a/lib/Parser/Parser.cc
+++ b/lib/Parser/Parser.cc
@@ -1,4 +1,5 @@
 #include "cps/Parser.h"
+#include "cps/FunctionAST.h"
 
 using namespace cps;
 
@@ -40,24 +41,40 @@ std::unique_ptr<ExprAST> Parser::ParseIdentifierExpr() {
     std::string IdName = Lex.IdentifierStr;
     int Line = Lex.getLine();
     getNextToken();
+
+    if (CurTok == '(') {
+        getNextToken();
+        std::vector<std::unique_ptr<ExprAST>> Args;
+        if (CurTok != ')') {
+            while (true) {
+                if (auto Arg = ParseExpression())
+                    Args.push_back(std::move(Arg));
+                else
+                    return nullptr;
+
+                if (CurTok == ')') break;
+
+                if (CurTok != ',') {
+                    fprintf(stderr, "Error: Expected ')' or ',' in function call.\n");
+                    return nullptr;
+                }
+                getNextToken();
+            }
+        }
+        getNextToken();
+        return std::make_unique<CallExprAST>(IdName, std::move(Args));
+    }
     
     if (CurTok == '[') {
         getNextToken();
         std::vector<std::unique_ptr<ExprAST>> Indices;
-        
         while (true) {
             auto Exp = ParseExpression();
             if (!Exp) return nullptr;
             Indices.push_back(std::move(Exp));
-            
-            if (CurTok == ']') {
-                break;
-            }
-            if (CurTok == ',') {
-                getNextToken();
-                continue;
-            }
-            fprintf(stderr, "Error: Expected ',' or ']' in array index\n");
+            if (CurTok == ']') break;
+            if (CurTok == ',') { getNextToken(); continue; }
+            fprintf(stderr, "Error: Expected ',' or ']'\n");
             return nullptr;
         }
         getNextToken(); 
@@ -407,6 +424,18 @@ std::unique_ptr<StmtAST> Parser::ParseStatement() {
     else if (CurTok == tok_for) {
         return ParseForStmt();
     }
+    else if (CurTok == tok_function) {
+        return ParseFunction();
+    }
+    else if (CurTok == tok_procedure) {
+        return ParseProcedure();
+    }
+    else if (CurTok == tok_call) {
+        return ParseCallStmt();
+    }
+    else if (CurTok == tok_return) {
+        return ParseReturnStmt();
+    }
 
     return nullptr;
 }
diff --git a/lib/Parser/Parser_Function.cc b/lib/Parser/Parser_Function.cc
new file mode 100644
index 0000000..bdb4ee7
--- /dev/null
+++ b/lib/Parser/Parser_Function.cc
@@ -0,0 +1,129 @@
+#include "cps/Parser.h"
+#include "cps/FunctionAST.h"
+
+using namespace cps;
+
+std::vector<std::pair<std::string, std::string>> Parser::ParsePrototypeArgs() {
+    std::vector<std::pair<std::string, std::string>> Args;
+    if (CurTok != '(') return Args;
+    getNextToken();
+
+    while (CurTok != ')') {
+        if (CurTok == -1000 || CurTok == -1001) getNextToken();
+
+        if (CurTok != tok_identifier) {
+            fprintf(stderr, "Error: Expected argument name\n");
+            return Args;
+        }
+        std::string Name = Lex.IdentifierStr;
+        getNextToken();
+
+        if (CurTok != tok_colon) {
+            fprintf(stderr, "Error: Expected ':' after argument name\n");
+            return Args;
+        }
+        getNextToken();
+
+        std::string Type = "INTEGER";
+        if (CurTok == tok_integer_kw) {
+            getNextToken();
+        } else if (CurTok == tok_identifier) {
+            Type = Lex.IdentifierStr; 
+            getNextToken();
+        } else {
+             fprintf(stderr, "Error: Expected type (e.g. INTEGER) for argument '%s'\n", Name.c_str());
+             return Args;
+        }
+
+        Args.push_back({Name, Type});
+
+        if (CurTok == ')') break;
+        if (CurTok != ',') {
+            fprintf(stderr, "Error: Expected ',' or ')'\n");
+            return Args;
+        }
+        getNextToken();
+    }
+    getNextToken();
+    return Args;
+}
+
+
+std::unique_ptr<StmtAST> Parser::ParseFunction() {
+    getNextToken();
+    std::string Name = Lex.IdentifierStr;
+    getNextToken();
+
+    auto Args = ParsePrototypeArgs();
+
+    std::string RetType = "INTEGER";
+    if (CurTok == tok_returns) { 
+        getNextToken();
+        getNextToken();
+    }
+
+    auto Proto = std::make_unique<PrototypeAST>(Name, std::move(Args), RetType);
+
+    std::vector<std::unique_ptr<StmtAST>> Body;
+    while (CurTok != tok_endfunction && CurTok != tok_eof) {
+        auto Stmt = ParseStatement();
+        if (Stmt) Body.push_back(std::move(Stmt));
+        else if (CurTok != tok_eof && CurTok != tok_endfunction) getNextToken();
+    }
+    getNextToken();
+
+    return std::make_unique<FunctionDefAST>(std::move(Proto), std::move(Body));
+}
+
+std::unique_ptr<StmtAST> Parser::ParseProcedure() {
+    getNextToken();
+    std::string Name = Lex.IdentifierStr;
+    getNextToken();
+
+    auto Args = ParsePrototypeArgs();
+    auto Proto = std::make_unique<PrototypeAST>(Name, std::move(Args), "VOID");
+
+    std::vector<std::unique_ptr<StmtAST>> Body;
+    while (CurTok != tok_endprocedure && CurTok != tok_eof) {
+        auto Stmt = ParseStatement();
+        if (Stmt) Body.push_back(std::move(Stmt));
+        else if (CurTok != tok_eof && CurTok != tok_endprocedure) getNextToken();
+    }
+    getNextToken();
+
+    return std::make_unique<FunctionDefAST>(std::move(Proto), std::move(Body));
+}
+
+std::unique_ptr<StmtAST> Parser::ParseCallStmt() {
+    getNextToken();
+    std::string Callee = Lex.IdentifierStr;
+    getNextToken();
+
+    std::vector<std::unique_ptr<ExprAST>> Args;
+    if (CurTok == '(') {
+        getNextToken();
+        while (CurTok != ')') {
+            auto Arg = ParseExpression();
+            if (Arg) Args.push_back(std::move(Arg));
+            
+            if (CurTok == ')') break;
+            if (CurTok != ',') {
+                fprintf(stderr, "Error: Expected ',' or ')' in call\n");
+                return nullptr;
+            }
+            getNextToken();
+        }
+        getNextToken();
+    }
+    
+    return std::make_unique<CallStmtAST>(Callee, std::move(Args));
+}
+
+std::unique_ptr<StmtAST> Parser::ParseReturnStmt() {
+    getNextToken();
+    std::unique_ptr<ExprAST> Expr = nullptr;
+    if (CurTok != tok_endif && CurTok != tok_else && CurTok != tok_endfunction && CurTok != tok_endprocedure) {
+        Expr = ParseExpression();
+    }
+    return std::make_unique<ReturnStmtAST>(std::move(Expr));
+}