3 files changed, 13 insertions, 14 deletions

diff --git a/node/ExprNode.cc b/node/ExprNode.cc
index 93f0068..f4a26a8 100644
--- a/node/ExprNode.cc
+++ b/node/ExprNode.cc
@@ -95,14 +95,13 @@ FunctionCallNode::FunctionCallNode(Context* context, const std::string& name,
 
 Value FunctionCallNode::CodeGen() {
     TypeManager* manager = context_->GetTypeManager();
-    std::shared_ptr<FunctionType> type = context_->GetFunctionType(name_);
+    Function function = context_->FindFunction(name_);
+    std::shared_ptr<FunctionType> type = std::get<0>(function);
     if (type == nullptr) {
         Log(LogLevel::PCC_ERROR, "call to undefined function %s",
             name_.c_str());
         return nullptr;
     }
-    llvm::FunctionCallee fp = context_->GetModule()->getOrInsertFunction(
-        name_, llvm::cast<llvm::FunctionType>(type->GetLLVMType()));
     const std::vector<std::shared_ptr<Type>>& ArgTypes = type->GetArgTypes();
     std::vector<Value> ArgValues(args_.size());
     auto ArgTypeIt = ArgTypes.begin();
@@ -121,6 +120,6 @@ Value FunctionCallNode::CodeGen() {
         ++ArgTypeIt, ++ArgValueIt;
     }
     type_ = type->GetReturnType();
-    return context_->GetBuilder()->CreateCall(fp, ArgValues);
+    return context_->GetBuilder()->CreateCall(std::get<1>(function), ArgValues);
 }
 } // namespace pcc
diff --git a/node/FunctionNode.cc b/node/FunctionNode.cc
index f6810fd..736007b 100644
--- a/node/FunctionNode.cc
+++ b/node/FunctionNode.cc
@@ -12,12 +12,12 @@ namespace pcc {
 FunctionNode::FunctionNode(
     Context* context, std::string&& name, std::list<Declaration>&& arguments,
     std::list<std::shared_ptr<DeclNode>>&& LocalDeclarations,
-    std::shared_ptr<Type> type, std::shared_ptr<StatementListNode> body)
+    std::shared_ptr<Type> ReturnType, std::shared_ptr<StatementListNode> body)
     : BaseNode(context)
     , name_(std::move(name))
     , arguments_(std::move(arguments))
     , LocalDeclarations_(std::move(LocalDeclarations))
-    , type_(type)
+    , ReturnType_(ReturnType)
     , locals_(context->GetGlobals())
     , body_(body) {
     for (auto&& decl : LocalDeclarations_) {
@@ -33,13 +33,13 @@ Value FunctionNode::CodeGen() {
     std::transform(arguments_.cbegin(), arguments_.cend(), params.begin(),
                    [](const Declaration& decl) { return std::get<1>(decl); });
     std::shared_ptr<FunctionType> type =
-        manager->CreateFunctionType(type_, params);
-    context_->AddFunctionType(name_, type);
+        manager->CreateFunctionType(ReturnType_, params);
     llvm::FunctionType* LLVMType =
         llvm::cast<llvm::FunctionType>(type->GetLLVMType());
     llvm::Function* function = llvm::Function::Create(
         LLVMType, llvm::GlobalValue::ExternalLinkage, name_, *module);
     function->setCallingConv(llvm::CallingConv::C);
+    context_->AddFunction(name_, {type, function});
     llvm::BasicBlock* block =
         llvm::BasicBlock::Create(GlobalLLVMContext, name_, function);
     Log(LogLevel::PCC_INFO, "generating function %s of type %s", name_.c_str(),
@@ -57,17 +57,17 @@ Value FunctionNode::CodeGen() {
         decl->CodeGen();
     }
     Value rv = nullptr;
-    bool IsVoid = type_ == manager->GetBuiltinType(BuiltinType::VOID);
+    bool IsVoid = ReturnType_ == manager->GetBuiltinType(BuiltinType::VOID);
     if (IsVoid == false) {
         // create a variable with the same name as the function for return value
-        rv = builder->CreateAlloca(type_->GetLLVMType());
-        locals_.Add(name_, {type_, rv});
+        rv = builder->CreateAlloca(ReturnType_->GetLLVMType());
+        locals_.Add(name_, {ReturnType_, rv});
     }
     // push the function scope variables
     context_->PushScope(&locals_);
     body_->CodeGen();
     if (IsVoid == false) {
-        builder->CreateRet(builder->CreateLoad(type_->GetLLVMType(), rv));
+        builder->CreateRet(builder->CreateLoad(ReturnType_->GetLLVMType(), rv));
     } else {
         builder->CreateRetVoid();
     }
diff --git a/node/FunctionNode.hh b/node/FunctionNode.hh
index 415c2cc..d317402 100644
--- a/node/FunctionNode.hh
+++ b/node/FunctionNode.hh
@@ -13,7 +13,7 @@ class FunctionNode : public BaseNode {
     FunctionNode(Context* context, std::string&& name,
                  std::list<Declaration>&& arguments,
                  std::list<std::shared_ptr<DeclNode>>&& LocalDeclarations,
-                 std::shared_ptr<Type> type,
+                 std::shared_ptr<Type> ReturnType,
                  std::shared_ptr<StatementListNode> body);
     virtual Value CodeGen() override;
 
@@ -21,7 +21,7 @@ class FunctionNode : public BaseNode {
     std::string name_;
     std::list<Declaration> arguments_;
     std::list<std::shared_ptr<DeclNode>> LocalDeclarations_;
-    std::shared_ptr<Type> type_;
+    std::shared_ptr<Type> ReturnType_;
     VariableList locals_;
     std::shared_ptr<StatementListNode> body_;
 };