LLVMCodeBuilder: Implement if statements

adazem009 · adazem009 · commit f26e8cf42fad · 2024-10-06T12:58:48.000+02:00
diff --git a/src/dev/engine/internal/llvmcodebuilder.cpp b/src/dev/engine/internal/llvmcodebuilder.cpp
@@ -38,6 +38,7 @@ std::shared_ptr<ExecutableCode> LLVMCodeBuilder::finalize()
 {
     size_t functionIndex = 0;
     llvm::Function *currentFunc = beginFunction(functionIndex);
+    std::vector<IfStatement> ifStatements;
 
     // Execute recorded steps
     for (const Step &step : m_steps) {
@@ -71,6 +72,68 @@ std::shared_ptr<ExecutableCode> LLVMCodeBuilder::finalize()
                 endFunction(currentFunc, functionIndex);
                 currentFunc = beginFunction(++functionIndex);
                 break;
+
+            case Step::Type::BeginIf: {
+                IfStatement statement;
+                statement.beforeIf = m_builder.GetInsertBlock();
+                statement.body = llvm::BasicBlock::Create(m_ctx, "", currentFunc);
+
+                // Convert last reg to bool
+                assert(step.args.size() == 1);
+                statement.condition = m_builder.CreateCall(resolve_value_toBool(), step.args[0]->value);
+
+                // Switch to body branch
+                m_builder.SetInsertPoint(statement.body);
+
+                ifStatements.push_back(statement);
+                break;
+            }
+
+            case Step::Type::BeginElse: {
+                assert(!ifStatements.empty());
+                IfStatement &statement = ifStatements.back();
+
+                // Jump to the branch after the if statement
+                assert(!statement.afterIf);
+                statement.afterIf = llvm::BasicBlock::Create(m_ctx, "", currentFunc);
+                m_builder.CreateBr(statement.afterIf);
+
+                // Create else branch
+                assert(!statement.elseBranch);
+                statement.elseBranch = llvm::BasicBlock::Create(m_ctx, "", currentFunc);
+
+                // Since there's an else branch, the conditional instruction should jump to it
+                m_builder.SetInsertPoint(statement.beforeIf);
+                m_builder.CreateCondBr(statement.condition, statement.body, statement.elseBranch);
+
+                // Switch to the else branch
+                m_builder.SetInsertPoint(statement.elseBranch);
+                break;
+            }
+
+            case Step::Type::EndIf: {
+                assert(!ifStatements.empty());
+                IfStatement &statement = ifStatements.back();
+
+                // Jump to the branch after the if statement
+                if (!statement.afterIf)
+                    statement.afterIf = llvm::BasicBlock::Create(m_ctx, "", currentFunc);
+
+                m_builder.CreateBr(statement.afterIf);
+
+                if (statement.elseBranch) {
+                } else {
+                    // If there wasn't an 'else' branch, create a conditional instruction which skips the if statement if false
+                    m_builder.SetInsertPoint(statement.beforeIf);
+                    m_builder.CreateCondBr(statement.condition, statement.body, statement.afterIf);
+                }
+
+                // Switch to the branch after the if statement
+                m_builder.SetInsertPoint(statement.afterIf);
+
+                ifStatements.pop_back();
+                break;
+            }
         }
     }
 
@@ -140,14 +203,21 @@ void LLVMCodeBuilder::addListContents(List *list)
 
 void LLVMCodeBuilder::beginIfStatement()
 {
+    Step step(Step::Type::BeginIf);
+    assert(!m_tmpRegs.empty());
+    step.args.push_back(m_tmpRegs.back());
+    m_tmpRegs.pop_back();
+    m_steps.push_back(step);
 }
 
 void LLVMCodeBuilder::beginElseBranch()
 {
+    m_steps.push_back(Step(Step::Type::BeginElse));
 }
 
 void LLVMCodeBuilder::endIf()
 {
+    m_steps.push_back(Step(Step::Type::EndIf));
 }
 
 void LLVMCodeBuilder::beginLoop()
@@ -289,3 +359,8 @@ llvm::FunctionCallee LLVMCodeBuilder::resolve_value_assign_special()
 {
     return resolveFunction("value_assign_special", llvm::FunctionType::get(m_builder.getVoidTy(), { m_valueDataType->getPointerTo(), m_builder.getInt32Ty() }, false));
 }
+
+llvm::FunctionCallee LLVMCodeBuilder::resolve_value_toBool()
+{
+    return resolveFunction("value_toBool", llvm::FunctionType::get(m_builder.getInt1Ty(), m_valueDataType->getPointerTo(), false));
+}
diff --git a/src/dev/engine/internal/llvmcodebuilder.h b/src/dev/engine/internal/llvmcodebuilder.h
@@ -49,7 +49,10 @@ class LLVMCodeBuilder : public ICodeBuilder
                 enum class Type
                 {
                     FunctionCall,
-                    Yield
+                    Yield,
+                    BeginIf,
+                    BeginElse,
+                    EndIf
                 };
 
                 Step(Type type) :
@@ -64,6 +67,15 @@ class LLVMCodeBuilder : public ICodeBuilder
                 size_t functionReturnRegIndex = 0;
         };
 
+        struct IfStatement
+        {
+                llvm::Value *condition = nullptr;
+                llvm::BasicBlock *beforeIf = nullptr;
+                llvm::BasicBlock *body = nullptr;
+                llvm::BasicBlock *elseBranch = nullptr;
+                llvm::BasicBlock *afterIf = nullptr;
+        };
+
         void initTypes();
         llvm::Function *beginFunction(size_t index);
         void endFunction(llvm::Function *func, size_t index);
@@ -76,6 +88,7 @@ class LLVMCodeBuilder : public ICodeBuilder
         llvm::FunctionCallee resolve_value_assign_bool();
         llvm::FunctionCallee resolve_value_assign_cstring();
         llvm::FunctionCallee resolve_value_assign_special();
+        llvm::FunctionCallee resolve_value_toBool();
 
         std::string m_id;
         llvm::LLVMContext m_ctx;
diff --git a/test/dev/llvm/llvmcodebuilder_test.cpp b/test/dev/llvm/llvmcodebuilder_test.cpp
@@ -8,6 +8,8 @@
 
 using namespace libscratchcpp;
 
+using ::testing::Return;
+
 class LLVMCodeBuilderTest : public testing::Test
 {
     public:
@@ -103,3 +105,172 @@ TEST_F(LLVMCodeBuilderTest, Yield)
     ASSERT_EQ(testing::internal::GetCapturedStdout(), expected2);
     ASSERT_TRUE(code->isFinished(ctx.get()));
 }
+
+TEST_F(LLVMCodeBuilderTest, IfStatement)
+{
+    // Without else branch (const condition)
+    m_builder->addConstValue("true");
+    m_builder->beginIfStatement();
+    m_builder->addFunctionCall("test_function_no_args", 0, false);
+    m_builder->endIf();
+
+    m_builder->addConstValue("false");
+    m_builder->beginIfStatement();
+    m_builder->addFunctionCall("test_function_no_args", 0, false);
+    m_builder->endIf();
+
+    // Without else branch (condition returned by function)
+    m_builder->addFunctionCall("test_function_no_args_ret", 0, true);
+    m_builder->addConstValue("no_args_output");
+    m_builder->addFunctionCall("test_equals", 2, true);
+    m_builder->beginIfStatement();
+    m_builder->addConstValue(0);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->endIf();
+
+    m_builder->addFunctionCall("test_function_no_args_ret", 0, true);
+    m_builder->addConstValue("");
+    m_builder->addFunctionCall("test_equals", 2, true);
+    m_builder->beginIfStatement();
+    m_builder->addConstValue(1);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->endIf();
+
+    // With else branch (const condition)
+    m_builder->addConstValue("true");
+    m_builder->beginIfStatement();
+    m_builder->addConstValue(2);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->beginElseBranch();
+    m_builder->addConstValue(3);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->endIf();
+
+    m_builder->addConstValue("false");
+    m_builder->beginIfStatement();
+    m_builder->addConstValue(4);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->beginElseBranch();
+    m_builder->addConstValue(5);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->endIf();
+
+    // With else branch (condition returned by function)
+    m_builder->addFunctionCall("test_function_no_args_ret", 0, true);
+    m_builder->addConstValue("no_args_output");
+    m_builder->addFunctionCall("test_equals", 2, true);
+    m_builder->beginIfStatement();
+    m_builder->addConstValue(6);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->beginElseBranch();
+    m_builder->addConstValue(7);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->endIf();
+
+    m_builder->addFunctionCall("test_function_no_args_ret", 0, true);
+    m_builder->addConstValue("");
+    m_builder->addFunctionCall("test_equals", 2, true);
+    m_builder->beginIfStatement();
+    m_builder->addConstValue(8);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->beginElseBranch();
+    m_builder->addConstValue(9);
+    m_builder->addFunctionCall("test_function_1_arg", 1, false);
+    m_builder->endIf();
+
+    // Nested 1
+    m_builder->addConstValue(true);
+    m_builder->beginIfStatement();
+    {
+        m_builder->addConstValue(false);
+        m_builder->beginIfStatement();
+        {
+            m_builder->addConstValue(0);
+            m_builder->addFunctionCall("test_function_1_arg", 1, false);
+        }
+        m_builder->beginElseBranch();
+        {
+            m_builder->addConstValue(1);
+            m_builder->addFunctionCall("test_function_1_arg", 1, false);
+
+            m_builder->addConstValue(false);
+            m_builder->beginIfStatement();
+            m_builder->beginElseBranch();
+            {
+                m_builder->addConstValue(2);
+                m_builder->addFunctionCall("test_function_1_arg", 1, false);
+            }
+            m_builder->endIf();
+        }
+        m_builder->endIf();
+    }
+    m_builder->beginElseBranch();
+    {
+        m_builder->addConstValue(true);
+        m_builder->beginIfStatement();
+        {
+            m_builder->addConstValue(3);
+            m_builder->addFunctionCall("test_function_1_arg", 1, false);
+        }
+        m_builder->beginElseBranch();
+        {
+            m_builder->addConstValue(4);
+            m_builder->addFunctionCall("test_function_1_arg", 1, false);
+        }
+        m_builder->endIf();
+    }
+    m_builder->endIf();
+
+    // Nested 2
+    m_builder->addConstValue(false);
+    m_builder->beginIfStatement();
+    {
+        m_builder->addConstValue(false);
+        m_builder->beginIfStatement();
+        {
+            m_builder->addConstValue(5);
+            m_builder->addFunctionCall("test_function_1_arg", 1, false);
+        }
+        m_builder->beginElseBranch();
+        {
+            m_builder->addConstValue(6);
+            m_builder->addFunctionCall("test_function_1_arg", 1, false);
+        }
+        m_builder->endIf();
+    }
+    m_builder->beginElseBranch();
+    {
+        m_builder->addConstValue(true);
+        m_builder->beginIfStatement();
+        {
+            m_builder->addConstValue(7);
+            m_builder->addFunctionCall("test_function_1_arg", 1, false);
+        }
+        m_builder->beginElseBranch();
+        m_builder->endIf();
+    }
+    m_builder->endIf();
+
+    auto code = m_builder->finalize();
+    auto ctx = code->createExecutionContext(&m_target);
+
+    static const std::string expected =
+        "no_args\n"
+        "no_args_ret\n"
+        "1_arg 0\n"
+        "no_args_ret\n"
+        "1_arg 2\n"
+        "1_arg 5\n"
+        "no_args_ret\n"
+        "1_arg 6\n"
+        "no_args_ret\n"
+        "1_arg 9\n"
+        "1_arg 1\n"
+        "1_arg 2\n"
+        "1_arg 7\n";
+
+    EXPECT_CALL(m_target, isStage).WillRepeatedly(Return(false));
+    testing::internal::CaptureStdout();
+    code->run(ctx.get());
+    ASSERT_EQ(testing::internal::GetCapturedStdout(), expected);
+}
diff --git a/test/dev/llvm/testfunctions.cpp b/test/dev/llvm/testfunctions.cpp
@@ -73,4 +73,9 @@ extern "C"
         std::cout << "3_args " << s1 << " " << s2 << " " << s3 << std::endl;
         value_assign_cstring(ret, "3_args_output");
     }
+
+    void test_equals(Target *target, ValueData *ret, ValueData *a, ValueData *b)
+    {
+        value_assign_bool(ret, value_equals(a, b));
+    }
 }
diff --git a/test/dev/llvm/testfunctions.h b/test/dev/llvm/testfunctions.h
@@ -18,6 +18,8 @@ extern "C"
     void test_function_1_arg_ret(Target *target, ValueData *ret, const ValueData *arg1);
     void test_function_3_args(Target *target, const ValueData *arg1, const ValueData *arg2, const ValueData *arg3);
     void test_function_3_args_ret(Target *target, ValueData *ret, const ValueData *arg1, const ValueData *arg2, const ValueData *arg3);
+
+    void test_equals(Target *target, ValueData *ret, ValueData *a, ValueData *b);
 }
 
 } // namespace libscratchcpp

Original file line number	Diff line number	Diff line change
`@@ -73,4 +73,9 @@ extern "C"`
`73`	`73`	`std::cout << "3_args " << s1 << " " << s2 << " " << s3 << std::endl;`
`74`	`74`	`value_assign_cstring(ret, "3_args_output");`
`75`	`75`	`}`
	`76`	`+`
	`77`	`+ void test_equals(Target target, ValueData ret, ValueData a, ValueData b)`
	`78`	`+ {`
	`79`	`+ value_assign_bool(ret, value_equals(a, b));`
	`80`	`+ }`
`76`	`81`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,8 @@ extern "C"`
`18`	`18`	`void test_function_1_arg_ret(Target target, ValueData ret, const ValueData *arg1);`
`19`	`19`	`void test_function_3_args(Target target, const ValueData arg1, const ValueData arg2, const ValueData arg3);`
`20`	`20`	`void test_function_3_args_ret(Target target, ValueData ret, const ValueData arg1, const ValueData arg2, const ValueData *arg3);`
	`21`	`+`
	`22`	`+ void test_equals(Target target, ValueData ret, ValueData a, ValueData b);`
`21`	`23`	`}`
`22`	`24`
`23`	`25`	`} // namespace libscratchcpp`