Add stop thread LLVM instruction

Author: adazem009

src/engine/internal/icodebuilder.h

@@ -99,6 +99,7 @@ class ICodeBuilder
         virtual void yield() = 0;
 
         virtual void createStop() = 0;
+        virtual void createThreadStop() = 0;
         virtual void createStopWithoutSync() = 0;
 
         virtual void createProcedureCall(BlockPrototype *prototype, const Compiler::Args &args) = 0;

src/engine/internal/llvm/instructions/control.cpp

@@ -60,6 +60,10 @@ ProcessResult Control::process(LLVMInstruction *ins)
             ret.next = buildStop(ins);
             break;
 
+        case LLVMInstruction::Type::ThreadStop:
+            ret.next = buildThreadStop(ins);
+            break;
+
         case LLVMInstruction::Type::StopWithoutSync:
             ret.next = buildStopWithoutSync(ins);
             break;
@@ -340,6 +344,17 @@ LLVMInstruction *Control::buildStop(LLVMInstruction *ins)
     return buildStopWithoutSync(ins);
 }
 
+LLVMInstruction *Control::buildThreadStop(LLVMInstruction *ins)
+{
+    m_utils.syncVariables();
+    m_builder.CreateBr(m_utils.endThreadBranch());
+
+    llvm::BasicBlock *nextBranch = llvm::BasicBlock::Create(m_utils.llvmCtx(), "", m_utils.function());
+    m_builder.SetInsertPoint(nextBranch);
+
+    return ins->next;
+}
+
 LLVMInstruction *Control::buildStopWithoutSync(LLVMInstruction *ins)
 {
     m_builder.CreateBr(m_utils.endBranch());

src/engine/internal/llvm/instructions/control.h

@@ -27,6 +27,7 @@ class Control : public InstructionGroup
         LLVMInstruction *buildBeginLoopCondition(LLVMInstruction *ins);
         LLVMInstruction *buildEndLoop(LLVMInstruction *ins);
         LLVMInstruction *buildStop(LLVMInstruction *ins);
+        LLVMInstruction *buildThreadStop(LLVMInstruction *ins);
         LLVMInstruction *buildStopWithoutSync(LLVMInstruction *ins);
 };
 

src/engine/internal/llvm/instructions/procedures.cpp

@@ -66,9 +66,17 @@ LLVMInstruction *Procedures::buildCallProcedure(LLVMInstruction *ins)
             args.push_back(m_utils.createValue(arg.second));
     }
 
+    // Call the procedure
     llvm::Value *handle = m_builder.CreateCall(m_utils.functions().resolveFunction(name, type), args);
 
+    // Check for end thread sentinel value
+    llvm::BasicBlock *nextBranch = llvm::BasicBlock::Create(llvmCtx, "", function);
+    llvm::Value *endThread = m_builder.CreateICmpEQ(handle, m_utils.threadEndSentinel());
+    m_builder.CreateCondBr(endThread, m_utils.endThreadBranch(), nextBranch);
+    m_builder.SetInsertPoint(nextBranch);
+
     if (!m_utils.warp() && !ins->procedurePrototype->warp()) {
+        // Handle suspend
         llvm::BasicBlock *suspendBranch = llvm::BasicBlock::Create(llvmCtx, "", function);
         llvm::BasicBlock *nextBranch = llvm::BasicBlock::Create(llvmCtx, "", function);
         m_builder.CreateCondBr(m_builder.CreateIsNull(handle), nextBranch, suspendBranch);

src/engine/internal/llvm/llvmbuildutils.cpp

@@ -161,8 +161,9 @@ void LLVMBuildUtils::init(llvm::Function *function, BlockPrototype *procedurePro
     reloadVariables();
     reloadLists();
 
-    // Create end branch
+    // Create end branches
     m_endBranch = llvm::BasicBlock::Create(m_llvmCtx, "end", m_function);
+    m_endThreadBranch = llvm::BasicBlock::Create(m_llvmCtx, "endThread", m_function);
 }
 
 void LLVMBuildUtils::end(LLVMInstruction *lastInstruction, LLVMRegister *lastConstant)
@@ -184,9 +185,9 @@ void LLVMBuildUtils::end(LLVMInstruction *lastInstruction, LLVMRegister *lastCon
     syncVariables();
     m_builder.CreateBr(m_endBranch);
 
+    // End branch
     m_builder.SetInsertPoint(m_endBranch);
 
-    // End the script function
     llvm::PointerType *pointerType = llvm::PointerType::get(llvm::Type::getInt8Ty(m_llvmCtx), 0);
 
     switch (m_codeType) {
@@ -216,6 +217,31 @@ void LLVMBuildUtils::end(LLVMInstruction *lastInstruction, LLVMRegister *lastCon
                 m_builder.CreateRet(castValue(lastConstant, Compiler::StaticType::Bool));
             break;
     }
+
+    // Thread end branch (stop the entire thread, including procedure callers)
+    m_builder.SetInsertPoint(m_endThreadBranch);
+
+    switch (m_codeType) {
+        case Compiler::CodeType::Script:
+            // Return a sentinel value (special pointer) to terminate any procedure callers
+            if (m_warp)
+                m_builder.CreateRet(threadEndSentinel());
+            else if (m_procedurePrototype)
+                m_coroutine->endWithSentinel(threadEndSentinel());
+            else {
+                // There's no need to return the sentinel value in standard scripts because they don't have any callers
+                m_coroutine->endWithSentinel(threadEndSentinel());
+                // m_coroutine->end();
+            }
+
+            break;
+
+        case Compiler::CodeType::Reporter:
+        case Compiler::CodeType::HatPredicate:
+            // Procedures cannot be called by these scripts, so we don't have to return the sentinel value
+            m_builder.CreateBr(m_endBranch);
+            break;
+    }
 }
 
 LLVMCompilerContext *LLVMBuildUtils::compilerCtx() const
@@ -323,6 +349,17 @@ llvm::BasicBlock *LLVMBuildUtils::endBranch() const
     return m_endBranch;
 }
 
+llvm::BasicBlock *LLVMBuildUtils::endThreadBranch() const
+{
+    return m_endThreadBranch;
+}
+
+llvm::Value *LLVMBuildUtils::threadEndSentinel() const
+{
+    llvm::PointerType *pointerType = llvm::PointerType::get(llvm::Type::getInt8Ty(m_llvmCtx), 0);
+    return m_builder.CreateIntToPtr(m_builder.getInt64(1), pointerType, "threadEndSentinel");
+}
+
 BlockPrototype *LLVMBuildUtils::procedurePrototype() const
 {
     return m_procedurePrototype;

src/engine/internal/llvm/llvmbuildutils.h

@@ -56,6 +56,9 @@ class LIBSCRATCHCPP_TEST_EXPORT LLVMBuildUtils
         size_t stringCount() const;
 
         llvm::BasicBlock *endBranch() const;
+        llvm::BasicBlock *endThreadBranch() const;
+
+        llvm::Value *threadEndSentinel() const;
 
         BlockPrototype *procedurePrototype() const;
         bool warp() const;
@@ -166,6 +169,7 @@ class LIBSCRATCHCPP_TEST_EXPORT LLVMBuildUtils
         llvm::Value *m_functionIdValue = nullptr;
 
         llvm::BasicBlock *m_endBranch = nullptr;
+        llvm::BasicBlock *m_endThreadBranch = nullptr;
 
         llvm::StructType *m_valueDataType = nullptr;
         llvm::StructType *m_stringPtrType = nullptr;

src/engine/internal/llvm/llvmcodebuilder.cpp

@@ -567,6 +567,12 @@ void LLVMCodeBuilder::createStop()
     m_instructions.addInstruction(ins);
 }
 
+void LLVMCodeBuilder::createThreadStop()
+{
+    auto ins = std::make_shared<LLVMInstruction>(LLVMInstruction::Type::ThreadStop, m_loopCondition);
+    m_instructions.addInstruction(ins);
+}
+
 void LLVMCodeBuilder::createStopWithoutSync()
 {
     auto ins = std::make_shared<LLVMInstruction>(LLVMInstruction::Type::StopWithoutSync, m_loopCondition);

src/engine/internal/llvm/llvmcodebuilder.h

@@ -113,6 +113,7 @@ class LIBSCRATCHCPP_TEST_EXPORT LLVMCodeBuilder : public ICodeBuilder
         void yield() override;
 
         void createStop() override;
+        void createThreadStop() override;
         void createStopWithoutSync() override;
 
         void createProcedureCall(BlockPrototype *prototype, const Compiler::Args &args) override;

src/engine/internal/llvm/llvmcoroutine.cpp

@@ -35,15 +35,23 @@ LLVMCoroutine::LLVMCoroutine(llvm::Module *module, llvm::IRBuilder<> *builder, l
 
     // Begin
     m_handle = builder->CreateCall(coroBegin, { coroIdRet, alloc });
+
     m_didSuspendVar = builder->CreateAlloca(builder->getInt1Ty(), nullptr, "didSuspend");
     builder->CreateStore(builder->getInt1(false), m_didSuspendVar);
+
+    m_sentinelVar = builder->CreateAlloca(pointerType, nullptr, "sentinel");
+    builder->CreateStore(nullPointer, m_sentinelVar);
+
     llvm::BasicBlock *entry = builder->GetInsertBlock();
 
     // Create suspend branch
     m_suspendBlock = llvm::BasicBlock::Create(ctx, "suspend", func);
     builder->SetInsertPoint(m_suspendBlock);
+
+    llvm::Value *sentinelValue = builder->CreateLoad(pointerType, m_sentinelVar);
+    llvm::Value *sentinelIsNull = builder->CreateIsNull(sentinelValue);
     builder->CreateCall(coroEnd, { m_handle, builder->getInt1(false), llvm::ConstantTokenNone::get(ctx) });
-    builder->CreateRet(m_handle);
+    builder->CreateRet(builder->CreateSelect(sentinelIsNull, m_handle, sentinelValue));
 
     // Create free branches
     m_freeMemRetBlock = llvm::BasicBlock::Create(ctx, "freeMemRet", func);
@@ -63,6 +71,15 @@ LLVMCoroutine::LLVMCoroutine(llvm::Module *module, llvm::IRBuilder<> *builder, l
     llvm::Value *needFree = builder->CreateIsNotNull(mem);
     builder->CreateCondBr(needFree, freeBranch, m_suspendBlock);
 
+    // Create final suspend point
+    m_finalSuspendBlock = llvm::BasicBlock::Create(ctx, "finalSuspend", m_function);
+
+    m_builder->SetInsertPoint(m_finalSuspendBlock);
+    llvm::Value *suspendResult = m_builder->CreateCall(llvm::Intrinsic::getDeclaration(m_module, llvm::Intrinsic::coro_suspend), { llvm::ConstantTokenNone::get(ctx), m_builder->getInt1(true) });
+    llvm::SwitchInst *sw = m_builder->CreateSwitch(suspendResult, m_suspendBlock, 2);
+    sw->addCase(m_builder->getInt8(0), m_freeMemRetBlock);
+    sw->addCase(m_builder->getInt8(1), m_cleanupBlock);
+
     builder->SetInsertPoint(entry);
 }
 
@@ -135,20 +152,11 @@ llvm::Value *LLVMCoroutine::createResume(llvm::Module *module, llvm::IRBuilder<>
 
 void LLVMCoroutine::end()
 {
-    llvm::LLVMContext &ctx = m_builder->getContext();
-
-    // Add final suspend point
-    llvm::BasicBlock *endBranch = llvm::BasicBlock::Create(ctx, "end", m_function);
-    llvm::BasicBlock *finalSuspendBranch = llvm::BasicBlock::Create(ctx, "finalSuspend", m_function);
-    m_builder->CreateCondBr(m_builder->CreateLoad(m_builder->getInt1Ty(), m_didSuspendVar), finalSuspendBranch, endBranch);
-
-    m_builder->SetInsertPoint(finalSuspendBranch);
-    llvm::Value *suspendResult = m_builder->CreateCall(llvm::Intrinsic::getDeclaration(m_module, llvm::Intrinsic::coro_suspend), { llvm::ConstantTokenNone::get(ctx), m_builder->getInt1(true) });
-    llvm::SwitchInst *sw = m_builder->CreateSwitch(suspendResult, m_suspendBlock, 2);
-    sw->addCase(m_builder->getInt8(0), endBranch); // unreachable
-    sw->addCase(m_builder->getInt8(1), m_cleanupBlock);
+    m_builder->CreateCondBr(m_builder->CreateLoad(m_builder->getInt1Ty(), m_didSuspendVar), m_finalSuspendBlock, m_freeMemRetBlock);
+}
 
-    // Jump to "free and return" branch
-    m_builder->SetInsertPoint(endBranch);
-    m_builder->CreateBr(m_freeMemRetBlock);
+void LLVMCoroutine::endWithSentinel(llvm::Value *sentinel)
+{
+    m_builder->CreateStore(sentinel, m_sentinelVar);
+    end();
 }

src/engine/internal/llvm/llvmcoroutine.h

@@ -24,16 +24,19 @@ class LLVMCoroutine
         void createSuspend();
         static llvm::Value *createResume(llvm::Module *module, llvm::IRBuilder<> *builder, llvm::Function *function, llvm::Value *coroHandle);
         void end();
+        void endWithSentinel(llvm::Value *sentinel);
 
     private:
         llvm::Module *m_module = nullptr;
         llvm::IRBuilder<> *m_builder = nullptr;
         llvm::Function *m_function = nullptr;
         llvm::Value *m_handle = nullptr;
         llvm::BasicBlock *m_suspendBlock = nullptr;
+        llvm::BasicBlock *m_finalSuspendBlock = nullptr;
         llvm::BasicBlock *m_cleanupBlock = nullptr;
         llvm::BasicBlock *m_freeMemRetBlock = nullptr;
         llvm::Value *m_didSuspendVar = nullptr;
+        llvm::Value *m_sentinelVar = nullptr;
 };
 
 } // namespace libscratchcpp