Create shims for undefined procedures

Author: adazem009

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

@@ -44,7 +44,7 @@ LLVMInstruction *Procedures::buildCallProcedure(LLVMInstruction *ins)
     llvm::FunctionType *type = m_utils.scriptFunctionType(ins->procedurePrototype);
     std::vector<llvm::Value *> args;
 
-    m_utils.compilerCtx()->addUsedProcedure(ins->procedurePrototype, name, type);
+    m_utils.compilerCtx()->addUsedProcedure(ins->procedurePrototype, name);
 
     llvm::FunctionType *funcType = m_utils.scriptFunctionType(nullptr);
     int passArgCount = funcType->getNumParams();

src/engine/internal/llvm/llvmcompilercontext.cpp

@@ -7,6 +7,7 @@
 #include <llvm/IR/Verifier.h>
 
 #include <scratchcpp/target.h>
+#include <scratchcpp/blockprototype.h>
 #include <iostream>
 
 #include "llvmcompilercontext.h"
@@ -70,9 +71,9 @@ void LLVMCompilerContext::addDefinedProcedure(BlockPrototype *prototype)
     m_definedProcedures.insert(prototype);
 }
 
-void LLVMCompilerContext::addUsedProcedure(BlockPrototype *prototype, const std::string &functionName, llvm::FunctionType *functionType)
+void LLVMCompilerContext::addUsedProcedure(BlockPrototype *prototype, const std::string &functionName)
 {
-    m_usedProcedures[prototype] = { functionName, functionType };
+    m_usedProcedures[prototype] = functionName;
 }
 
 function_id_t LLVMCompilerContext::getNextFunctionId()
@@ -99,6 +100,9 @@ void LLVMCompilerContext::initJit()
     std::cout << "==============" << std::endl << std::endl;
 #endif
 
+    // Define shims for missing procedures
+    createProcedureShims();
+
     // Optimize
     optimize(llvm::OptimizationLevel::O3);
 
@@ -199,6 +203,27 @@ void LLVMCompilerContext::createTargetMachine()
     m_targetMachine = std::unique_ptr<llvm::TargetMachine>(target->createTargetMachine(targetTriple, cpu, features, opt, llvm::Reloc::PIC_));
 }
 
+void LLVMCompilerContext::createProcedureShims()
+{
+    llvm::IRBuilder<> builder(*m_llvmCtx);
+
+    for (const auto &[prototype, name] : m_usedProcedures) {
+        if (m_definedProcedures.find(prototype) == m_definedProcedures.cend()) {
+            std::cout << "warning: procedure \"" << prototype->procCode() << "\" is not defined" << std::endl;
+
+            // We need to define shims for undefined procedures (the JIT compiler crashes without them)
+            llvm::Function *func = m_module->getFunction(name); // since the function is used, it's already declared
+            llvm::BasicBlock *entry = llvm::BasicBlock::Create(*m_llvmCtx, "entry", func);
+            llvm::PointerType *pointerType = llvm::PointerType::get(*m_llvmCtx, 0);
+            llvm::Constant *nullPointer = llvm::ConstantPointerNull::get(pointerType);
+            builder.SetInsertPoint(entry);
+            builder.CreateRet(nullPointer);
+
+            verifyFunction(func);
+        }
+    }
+}
+
 void LLVMCompilerContext::optimize(llvm::OptimizationLevel optLevel)
 {
     llvm::PassBuilder passBuilder(m_targetMachine.get());

src/engine/internal/llvm/llvmcompilercontext.h

@@ -37,7 +37,7 @@ class LLVMCompilerContext : public CompilerContext
         const std::unordered_map<function_id_t, LLVMExecutableCode *> &codeMap() const;
 
         void addDefinedProcedure(BlockPrototype *prototype);
-        void addUsedProcedure(BlockPrototype *prototype, const std::string &functionName, llvm::FunctionType *functionType);
+        void addUsedProcedure(BlockPrototype *prototype, const std::string &functionName);
 
         function_id_t getNextFunctionId();
 
@@ -70,6 +70,7 @@ class LLVMCompilerContext : public CompilerContext
 
         void initTarget();
         void createTargetMachine();
+        void createProcedureShims();
         void optimize(llvm::OptimizationLevel optLevel);
 
         llvm::Function *createCoroResumeFunction();
@@ -98,7 +99,7 @@ class LLVMCompilerContext : public CompilerContext
         llvm::Type *m_functionIdType = nullptr;
 
         std::unordered_set<BlockPrototype *> m_definedProcedures;
-        std::unordered_map<BlockPrototype *, std::pair<std::string, llvm::FunctionType *>> m_usedProcedures; // function name, function type
+        std::unordered_map<BlockPrototype *, std::string> m_usedProcedures;
 };
 
 } // namespace libscratchcpp

test/llvm/llvmcodebuilder_test.cpp

@@ -3522,6 +3522,36 @@ TEST_F(LLVMCodeBuilderTest, Procedures)
     ASSERT_TRUE(code->isFinished(ctx.get()));
 }
 
+TEST_F(LLVMCodeBuilderTest, UndefinedProcedure)
+{
+    Sprite sprite;
+
+    BlockPrototype prototype;
+    prototype.setProcCode("procedure 1 %s %s %b");
+    prototype.setArgumentNames({ "any type 1", "any type 2", "bool" });
+    prototype.setArgumentIds({ "a", "b", "c" });
+    prototype.setWarp(false);
+
+    LLVMCodeBuilder *builder = m_utils.createBuilder(&sprite, false);
+    builder->createProcedureCall(&prototype, { builder->addConstValue("test"), builder->addConstValue(true), builder->addConstValue(false) });
+
+    CompilerValue *v = builder->addConstValue("test");
+    builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String }, { v });
+
+    std::string expected = "test\n";
+
+    auto code = builder->build();
+    Script script(&sprite, nullptr, nullptr);
+    script.setCode(code);
+    Thread thread(&sprite, nullptr, &script);
+    auto ctx = code->createExecutionContext(&thread);
+
+    testing::internal::CaptureStdout();
+    code->run(ctx.get());
+    ASSERT_EQ(testing::internal::GetCapturedStdout(), expected);
+    ASSERT_TRUE(code->isFinished(ctx.get()));
+}
+
 TEST_F(LLVMCodeBuilderTest, HatPredicates)
 {
     Sprite sprite;