LLVMCodeBuilder: Properly destroy finished coroutines

Author: adazem009

src/engine/internal/llvm/llvmcodebuilder.cpp

@@ -1363,7 +1363,7 @@ std::shared_ptr<ExecutableCode> LLVMCodeBuilder::finalize()
     if (m_warp)
         m_builder.CreateRet(m_builder.getInt1(true));
     else
-        m_builder.CreateRet(coro->createResume(resumeFunc->getArg(0)));
+        m_builder.CreateRet(coro->createResume(resumeFunc, resumeFunc->getArg(0)));
 
     verifyFunction(resumeFunc);
 

src/engine/internal/llvm/llvmcoroutine.cpp

@@ -105,10 +105,32 @@ void LLVMCoroutine::createSuspend()
     m_builder->SetInsertPoint(resumeBranch);
 }
 
-llvm::Value *LLVMCoroutine::createResume(llvm::Value *coroHandle)
+llvm::Value *LLVMCoroutine::createResume(llvm::Function *function, llvm::Value *coroHandle)
 {
+    llvm::LLVMContext &ctx = m_builder->getContext();
+    llvm::Function *coroDone = llvm::Intrinsic::getDeclaration(m_module, llvm::Intrinsic::coro_done);
+
+    llvm::Value *ret = m_builder->CreateAlloca(m_builder->getInt1Ty());
+    llvm::Value *done = m_builder->CreateCall(coroDone, { coroHandle });
+    done = m_builder->CreateCall(coroDone, { coroHandle });
+
+    llvm::BasicBlock *destroyBranch = llvm::BasicBlock::Create(ctx, "", function);
+    llvm::BasicBlock *resumeBranch = llvm::BasicBlock::Create(ctx, "", function);
+    llvm::BasicBlock *nextBranch = llvm::BasicBlock::Create(ctx, "", function);
+    m_builder->CreateCondBr(done, destroyBranch, resumeBranch);
+
+    m_builder->SetInsertPoint(destroyBranch);
+    m_builder->CreateCall(llvm::Intrinsic::getDeclaration(m_module, llvm::Intrinsic::coro_destroy), { coroHandle });
+    m_builder->CreateBr(nextBranch);
+
+    m_builder->SetInsertPoint(resumeBranch);
     m_builder->CreateCall(llvm::Intrinsic::getDeclaration(m_module, llvm::Intrinsic::coro_resume), { coroHandle });
-    return m_builder->CreateCall(llvm::Intrinsic::getDeclaration(m_module, llvm::Intrinsic::coro_done), { coroHandle });
+    done = m_builder->CreateCall(coroDone, { coroHandle });
+    m_builder->CreateStore(done, ret);
+    m_builder->CreateCondBr(done, destroyBranch, nextBranch);
+
+    m_builder->SetInsertPoint(nextBranch);
+    return m_builder->CreateLoad(m_builder->getInt1Ty(), ret);
 }
 
 void LLVMCoroutine::end()

src/engine/internal/llvm/llvmcoroutine.h

@@ -22,7 +22,7 @@ class LLVMCoroutine
         llvm::Value *didSuspendVar() const;
 
         void createSuspend();
-        llvm::Value *createResume(llvm::Value *coroHandle);
+        llvm::Value *createResume(llvm::Function *function, llvm::Value *coroHandle);
         void end();
 
     private: