Fix LLVM coroutine memory leaks

Author: adazem009

src/engine/internal/llvm/llvmexecutablecode.cpp

@@ -86,15 +86,23 @@ bool LLVMExecutableCode::runPredicate(ExecutionContext *context)
 void LLVMExecutableCode::kill(ExecutionContext *context)
 {
     LLVMExecutionContext *ctx = getContext(context);
-    ctx->setCoroutineHandle(nullptr);
+    // NOTE: Do not destroy coroutine, it will be destroyed later
+    // (really bad things happen when a coroutine destroys itself...)
+    // The code is finished and the coroutine won't be used again until reset() is called
     ctx->setFinished(true);
     ctx->setPromise(nullptr);
 }
 
 void LLVMExecutableCode::reset(ExecutionContext *context)
 {
     LLVMExecutionContext *ctx = getContext(context);
-    ctx->setCoroutineHandle(nullptr);
+
+    if (ctx->coroutineHandle()) {
+        // TODO: Can a script reset itself? (really bad things happen when a coroutine destroys itself...)
+        m_ctx->destroyCoroutine(ctx->coroutineHandle());
+        ctx->setCoroutineHandle(nullptr);
+    }
+
     ctx->setFinished(false);
     ctx->setPromise(nullptr);
 }
@@ -124,7 +132,7 @@ std::shared_ptr<ExecutionContext> LLVMExecutableCode::createExecutionContext(Thr
     }
 
     m_resumeFunction = m_ctx->lookupFunction<ResumeFunctionType>(m_resumeFunctionName);
-    return std::make_shared<LLVMExecutionContext>(thread);
+    return std::make_shared<LLVMExecutionContext>(m_ctx, thread);
 }
 
 LLVMExecutionContext *LLVMExecutableCode::getContext(ExecutionContext *context)

src/engine/internal/llvm/llvmexecutioncontext.cpp

@@ -1,14 +1,27 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include "llvmexecutioncontext.h"
+#include "llvmcompilercontext.h"
 
 using namespace libscratchcpp;
 
-LLVMExecutionContext::LLVMExecutionContext(Thread *thread) :
-    ExecutionContext(thread)
+LLVMExecutionContext::LLVMExecutionContext(LLVMCompilerContext *compilerCtx, Thread *thread) :
+    ExecutionContext(thread),
+    m_compilerCtx(compilerCtx)
 {
 }
 
+LLVMExecutionContext::~LLVMExecutionContext()
+{
+    if (m_coroutineHandle) {
+        if (m_compilerCtx) {
+            m_compilerCtx->destroyCoroutine(m_coroutineHandle);
+            m_coroutineHandle = nullptr;
+        } else
+            assert(false);
+    }
+}
+
 void *LLVMExecutionContext::coroutineHandle() const
 {
     return m_coroutineHandle;

src/engine/internal/llvm/llvmexecutioncontext.h

@@ -7,10 +7,13 @@
 namespace libscratchcpp
 {
 
+class LLVMCompilerContext;
+
 class LLVMExecutionContext : public ExecutionContext
 {
     public:
-        LLVMExecutionContext(Thread *thread);
+        LLVMExecutionContext(LLVMCompilerContext *compilerCtx, Thread *thread);
+        ~LLVMExecutionContext();
 
         void *coroutineHandle() const;
         void setCoroutineHandle(void *newCoroutineHandle);
@@ -19,6 +22,7 @@ class LLVMExecutionContext : public ExecutionContext
         void setFinished(bool newFinished);
 
     private:
+        LLVMCompilerContext *m_compilerCtx = nullptr;
         void *m_coroutineHandle = nullptr;
         bool m_finished = false;
 };

test/llvm/llvmcodebuilder_test.cpp

@@ -70,8 +70,9 @@ class LLVMCodeBuilderTest : public testing::Test
         void createBuilder(Target *target, BlockPrototype *procedurePrototype, Compiler::CodeType codeType = Compiler::CodeType::Script)
         {
             if (m_contexts.find(target) == m_contexts.cend() || !target)
-                m_contexts[target] = std::make_unique<LLVMCompilerContext>(&m_engine, target);
+                m_contexts[target] = std::make_shared<LLVMCompilerContext>(&m_engine, target);
 
+            m_contextList.push_back(m_contexts[target]);
             m_builder = std::make_unique<LLVMCodeBuilder>(m_contexts[target].get(), procedurePrototype, codeType);
         }
 
@@ -433,7 +434,8 @@ class LLVMCodeBuilderTest : public testing::Test
             ASSERT_THAT(testing::internal::GetCapturedStdout(), Eq(expected)) << quotes << v.toString() << quotes;
         };
 
-        std::unordered_map<Target *, std::unique_ptr<LLVMCompilerContext>> m_contexts;
+        std::unordered_map<Target *, std::shared_ptr<LLVMCompilerContext>> m_contexts;
+        std::vector<std::shared_ptr<LLVMCompilerContext>> m_contextList;
         std::unique_ptr<LLVMCodeBuilder> m_builder;
         std::shared_ptr<BlockPrototype> m_procedurePrototype;
         EngineMock m_engine;
@@ -2854,6 +2856,34 @@ TEST_F(LLVMCodeBuilderTest, Yield)
     ASSERT_EQ(testing::internal::GetCapturedStdout(), expected2);
     ASSERT_TRUE(code->isFinished(ctx.get()));
 
+    // Terminate unfinished coroutine
+    EXPECT_CALL(m_target, isStage()).Times(3);
+    testing::internal::CaptureStdout();
+    code->reset(ctx.get());
+    code->run(ctx.get());
+    ASSERT_EQ(testing::internal::GetCapturedStdout(), expected1);
+    ASSERT_FALSE(code->isFinished(ctx.get()));
+
+    code->kill(ctx.get());
+    ASSERT_TRUE(code->isFinished(ctx.get()));
+
+    // Reset unfinished coroutine
+    EXPECT_CALL(m_target, isStage()).Times(3);
+    testing::internal::CaptureStdout();
+    code->reset(ctx.get());
+    code->run(ctx.get());
+    ASSERT_EQ(testing::internal::GetCapturedStdout(), expected1);
+    ASSERT_FALSE(code->isFinished(ctx.get()));
+
+    code->reset(ctx.get());
+
+    EXPECT_CALL(m_target, isStage()).Times(3);
+    testing::internal::CaptureStdout();
+    code->reset(ctx.get());
+    code->run(ctx.get());
+    ASSERT_EQ(testing::internal::GetCapturedStdout(), expected1);
+    ASSERT_FALSE(code->isFinished(ctx.get())); // leave unfinished coroutine
+
     // With warp
     createBuilder(true);
     build();

test/llvm/llvmexecutioncontext_test.cpp

@@ -1,12 +1,14 @@
 #include <scratchcpp/thread.h>
 #include <engine/internal/llvm/llvmexecutioncontext.h>
+#include <engine/internal/llvm/llvmcompilercontext.h>
 #include <gtest/gtest.h>
 
 using namespace libscratchcpp;
 
 TEST(LLVMExecutionContextTest, Constructor)
 {
     Thread thread(nullptr, nullptr, nullptr);
-    LLVMExecutionContext ctx(&thread);
+    LLVMCompilerContext compilerCtx(nullptr, nullptr);
+    LLVMExecutionContext ctx(&compilerCtx, &thread);
     ASSERT_EQ(ctx.thread(), &thread);
 }