From 19f91909a8051da28b7d071604bdf037f208875c Mon Sep 17 00:00:00 2001
From: "A. R. Shajii" <ars@ars.me>
Date: Sun, 12 Feb 2023 11:32:20 -0500
Subject: [PATCH] Update module init codegen

---
 codon/cir/llvm/llvisitor.cpp | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/codon/cir/llvm/llvisitor.cpp b/codon/cir/llvm/llvisitor.cpp
index 17c2adc6..aa074633 100644
--- a/codon/cir/llvm/llvisitor.cpp
+++ b/codon/cir/llvm/llvisitor.cpp
@@ -1013,6 +1013,14 @@ void LLVMVisitor::writeToPythonExtension(const PyModule &pymod,
       M->getOrInsertFunction("PyModule_AddObject", i32, ptr, ptr, ptr).getCallee());
   pyModuleAddObject->setDoesNotThrow();
 
+  auto *pyIncRef = llvm::cast<llvm::Function>(
+      M->getOrInsertFunction("Py_IncRef", B->getVoidTy(), ptr).getCallee());
+  pyIncRef->setDoesNotThrow();
+
+  auto *pyDecRef = llvm::cast<llvm::Function>(
+      M->getOrInsertFunction("Py_DecRef", B->getVoidTy(), ptr).getCallee());
+  pyDecRef->setDoesNotThrow();
+
   auto *pyModuleInit = llvm::cast<llvm::Function>(
       M->getOrInsertFunction("PyInit_" + pymod.name, ptr).getCallee());
   auto *block = llvm::BasicBlock::Create(*context, "entry", pyModuleInit);
@@ -1050,7 +1058,7 @@ void LLVMVisitor::writeToPythonExtension(const PyModule &pymod,
     seqassertn(it != typeVars.end(), "type not found");
     auto *typeVar = it->second;
 
-    // TODO incref typeVar
+    B->CreateCall(pyIncRef, typeVar);
     auto *status =
         B->CreateCall(pyModuleAddObject, {mod, pyString(pytype.name), typeVar});
     fail = llvm::BasicBlock::Create(*context, "failure", pyModuleInit);
@@ -1058,8 +1066,8 @@ void LLVMVisitor::writeToPythonExtension(const PyModule &pymod,
     B->CreateCondBr(B->CreateICmpSLT(status, zero32), fail, block);
 
     B->SetInsertPoint(fail);
-    // TODO decref typeVar
-    // TODO decref mod
+    B->CreateCall(pyDecRef, typeVar);
+    B->CreateCall(pyDecRef, mod);
     B->CreateRet(null);
 
     B->SetInsertPoint(block);