Added BYREF and BYVAL

3 files changed, 85 insertions, 25 deletions

diff --git a/lib/CodeGen/ArrayHandler.cc b/lib/CodeGen/ArrayHandler.cc
index 18b9253..5842f77 100644
--- a/lib/CodeGen/ArrayHandler.cc
+++ b/lib/CodeGen/ArrayHandler.cc
@@ -8,7 +8,7 @@
 using namespace llvm;
 using namespace cps;
 
-ArrayHandler::ArrayHandler(LLVMContext &C, IRBuilder<> &B, Module &M, std::map<std::string, AllocaInst*> &NV, RuntimeCheck &RC)
+ArrayHandler::ArrayHandler(LLVMContext &C, IRBuilder<> &B, Module &M, std::map<std::string, Value*> &NV, RuntimeCheck &RC)
     : TheContext(&C), Builder(&B), TheModule(&M), NamedValues(&NV), RuntimeChecker(RC) {
     setupExternalFunctions();
 }
@@ -125,7 +125,7 @@ Value* ArrayHandler::emitArrayAccess(ArrayAccessExprAST *Expr, CodeGen &CG) {
     
     Value *Offset = computeFlatIndex(Name, Indices);
     
-    AllocaInst *Alloca = (*NamedValues)[Name];
+    Value *Alloca = (*NamedValues)[Name];
     Value *RawPtr = Builder->CreateLoad(PointerType::getUnqual(*TheContext), Alloca);
     Value *IntPtr = Builder->CreateBitCast(RawPtr, PointerType::getUnqual(*TheContext));
     Value *ElemPtr = Builder->CreateGEP(Type::getInt64Ty(*TheContext), IntPtr, Offset);
@@ -154,7 +154,7 @@ void ArrayHandler::emitArrayAssign(ArrayAssignStmtAST *Stmt, CodeGen &CG) {
     Value *Val = CG.emitExpr(Stmt->getExpr());
     Value *Offset = computeFlatIndex(Name, Indices);
     
-    AllocaInst *Alloca = (*NamedValues)[Name];
+    Value *Alloca = (*NamedValues)[Name];
     Value *RawPtr = Builder->CreateLoad(PointerType::getUnqual(*TheContext), Alloca);
     Value *IntPtr = Builder->CreateBitCast(RawPtr, PointerType::getUnqual(*TheContext));
     Value *ElemPtr = Builder->CreateGEP(Type::getInt64Ty(*TheContext), IntPtr, Offset);
@@ -197,7 +197,7 @@ void ArrayHandler::emitPrintLoop(const std::string &Name, int CurrentDim, std::v
     if (CurrentDim == Meta.Rank) {
         Value *Offset = computeFlatIndex(Name, CurrentIndices);
         
-        AllocaInst *Alloca = (*NamedValues)[Name];
+        Value *Alloca = (*NamedValues)[Name];
         Value *RawPtr = Builder->CreateLoad(PointerType::getUnqual(*TheContext), Alloca);
         Value *IntPtr = Builder->CreateBitCast(RawPtr, PointerType::getUnqual(*TheContext));
         Value *ElemPtr = Builder->CreateGEP(Type::getInt64Ty(*TheContext), IntPtr, Offset);
diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index 6358699..bfbd8e0 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -70,7 +70,7 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
     }
     
     if (auto *Var = dynamic_cast<VariableExprAST*>(Expr)) {
-        AllocaInst *A = NamedValues[Var->getName()];
+        Value *A = NamedValues[Var->getName()];
         if (!A) {
             fprintf(stderr, "Error: Unknown variable name %s\n", Var->getName().c_str());
             return nullptr;
@@ -109,9 +109,34 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
     }
 
     if (auto *Call = dynamic_cast<CallExprAST*>(Expr)) {
+        Function *CalleeF = TheModule->getFunction(Call->getCallee());
         std::vector<Value*> Args;
-        for (const auto &Arg : Call->getArgs()) {
-            Args.push_back(emitExpr(Arg.get()));
+        
+        for (unsigned i = 0; i < Call->getArgs().size(); ++i) {
+            ExprAST *ArgExpr = Call->getArgs()[i].get();
+            bool isByRef = false;
+
+            if (CalleeF && i < CalleeF->arg_size()) {
+                if (CalleeF->getArg(i)->getType()->isPointerTy()) {
+                    isByRef = true;
+                }
+            }
+
+            if (isByRef) {
+                if (auto *Var = dynamic_cast<VariableExprAST*>(ArgExpr)) {
+                    Value *Ptr = NamedValues[Var->getName()];
+                    if (!Ptr) {
+                        fprintf(stderr, "Error: Unknown variable %s in BYREF call\n", Var->getName().c_str());
+                        return nullptr;
+                    }
+                    Args.push_back(Ptr);
+                } else {
+                     fprintf(stderr, "Error: BYREF argument must be a variable.\n");
+                     return nullptr;
+                }
+            } else {
+                Args.push_back(emitExpr(ArgExpr));
+            }
         }
         return FuncGen->emitCallExpr(Call, Args);
     }
@@ -203,7 +228,7 @@ void CodeGen::emitForStmt(ForStmtAST *Stmt) {
     Value *StartVal = emitExpr(Stmt->getStart());
     if (!StartVal) return;
 
-    AllocaInst *Alloca = NamedValues[VarName];
+    Value *Alloca = NamedValues[VarName];
     if (!Alloca) {
         fprintf(stderr, "Error: Unknown variable in FOR loop %s\n", VarName.c_str());
         return;
@@ -285,9 +310,34 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
     }
 
     else if (auto *Call = dynamic_cast<CallStmtAST*>(Stmt)) {
+        Function *CalleeF = TheModule->getFunction(Call->getCallee());
         std::vector<Value*> Args;
-        for (const auto &ArgExpr : Call->getArgs()) {
-            Args.push_back(emitExpr(ArgExpr.get()));
+        
+        for (unsigned i = 0; i < Call->getArgs().size(); ++i) {
+            ExprAST *ArgExpr = Call->getArgs()[i].get();
+            bool isByRef = false;
+            
+            if (CalleeF && i < CalleeF->arg_size()) {
+                if (CalleeF->getArg(i)->getType()->isPointerTy()) {
+                    isByRef = true;
+                }
+            }
+
+            if (isByRef) {
+                if (auto *Var = dynamic_cast<VariableExprAST*>(ArgExpr)) {
+                    Value *Ptr = NamedValues[Var->getName()];
+                    if (!Ptr) {
+                        fprintf(stderr, "Error: Unknown variable %s in BYREF call\n", Var->getName().c_str());
+                        return;
+                    }
+                    Args.push_back(Ptr);
+                } else {
+                     fprintf(stderr, "Error: BYREF argument must be a variable.\n");
+                     return;
+                }
+            } else {
+                Args.push_back(emitExpr(ArgExpr));
+            }
         }
         FuncGen->emitCallStmt(Call, Args);
         return;
@@ -311,7 +361,7 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
     else if (auto *Assign = dynamic_cast<AssignStmtAST*>(Stmt)) {
         Value *Val = emitExpr(Assign->getExpr());
         if (Val) {
-            AllocaInst *Alloca = NamedValues[Assign->getName()];
+            Value *Alloca = NamedValues[Assign->getName()];
             if (!Alloca) {
                 fprintf(stderr, "Error: Unknown variable name %s\n", Assign->getName().c_str());
                 return;
@@ -320,7 +370,7 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
         }
     }
     else if (auto *In = dynamic_cast<InputStmtAST*>(Stmt)) {
-        AllocaInst *Alloca = NamedValues[In->getName()];
+        Value *Alloca = NamedValues[In->getName()];
         if (Alloca) {
             std::vector<Value*> Args;
             Args.push_back(ScanfFormatStr);
diff --git a/lib/CodeGen/FunctionGen.cc b/lib/CodeGen/FunctionGen.cc
index b2adc08..a929943 100644
--- a/lib/CodeGen/FunctionGen.cc
+++ b/lib/CodeGen/FunctionGen.cc
@@ -13,27 +13,37 @@ Type *FunctionGen::getLLVMType(const std::string &TypeName) {
     return Type::getInt64Ty(Context);
 }
 
-void FunctionGen::createArgumentAllocas(Function *F, const std::vector<std::pair<std::string, std::string>> &Args) {
+void FunctionGen::createArgumentAllocas(Function *F, const std::vector<std::tuple<std::string, std::string, bool>> &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);
+        std::string ArgName = std::get<0>(Args[Idx]);
+        std::string ArgTypeStr = std::get<1>(Args[Idx]);
+        bool IsRef = std::get<2>(Args[Idx]);
         
         Value *ArgVal = &(*AI);
-        Builder.CreateStore(ArgVal, Alloca);
-        
-        NamedValues[ArgName] = Alloca; 
+        ArgVal->setName(ArgName);
+
+        if (IsRef) {
+            NamedValues[ArgName] = ArgVal;
+        } else {
+            Type *ArgType = getLLVMType(ArgTypeStr);
+            IRBuilder<> TmpB(&F->getEntryBlock(), F->getEntryBlock().begin());
+            AllocaInst *Alloca = TmpB.CreateAlloca(ArgType, nullptr, ArgName);
+            
+            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 *T = getLLVMType(std::get<1>(Arg));
+        if (std::get<2>(Arg)) {
+            T = T->getPointerTo();
+        }
+        ArgTypes.push_back(T);
     }
 
     Type *RetType = getLLVMType(Proto->getReturnType());
@@ -43,7 +53,7 @@ Function *FunctionGen::emitPrototype(PrototypeAST *Proto) {
     unsigned Idx = 0;
     for (auto &Arg : F->args()) {
         if (Idx < Proto->getArgs().size())
-            Arg.setName(Proto->getArgs()[Idx++].first);
+            Arg.setName(std::get<0>(Proto->getArgs()[Idx++]));
     }
     return F;
 }
@@ -65,7 +75,7 @@ Function *FunctionGen::emitFunctionDef(FunctionDefAST *FuncAST,
     BasicBlock *BB = BasicBlock::Create(Context, "entry", TheFunction);
     Builder.SetInsertPoint(BB);
 
-    std::map<std::string, llvm::AllocaInst*> OldNamedValues = NamedValues;
+    std::map<std::string, llvm::Value*> OldNamedValues = NamedValues;
     NamedValues.clear();
     
     createArgumentAllocas(TheFunction, Proto->getArgs());