LLVMCodeBuilder: Implement list item index

adazem009 · adazem009 · commit b05996c9dc18 · 2024-11-21T14:24:02.000+01:00
diff --git a/src/dev/engine/internal/icodebuilder.h b/src/dev/engine/internal/icodebuilder.h
@@ -24,6 +24,7 @@ class ICodeBuilder
         virtual void addVariableValue(Variable *variable) = 0;
         virtual void addListContents(List *list) = 0;
         virtual void addListItem(List *list) = 0;
+        virtual void addListItemIndex(List *list) = 0;
         virtual void addListSize(List *list) = 0;
 
         virtual void createAdd() = 0;
diff --git a/src/dev/engine/internal/llvm/llvmcodebuilder.cpp b/src/dev/engine/internal/llvm/llvmcodebuilder.cpp
@@ -590,6 +590,14 @@ std::shared_ptr<ExecutableCode> LLVMCodeBuilder::finalize()
                 break;
             }
 
+            case LLVMInstruction::Type::GetListItemIndex: {
+                assert(step.args.size() == 1);
+                const auto &arg = step.args[0];
+                const LLVMListPtr &listPtr = m_listPtrs[step.workList];
+                step.functionReturnReg->value = getListItemIndex(listPtr, arg.second, func);
+                break;
+            }
+
             case LLVMInstruction::Type::Yield:
                 if (!m_warp) {
                     freeHeap();
@@ -929,6 +937,13 @@ void LLVMCodeBuilder::addListItem(List *list)
     m_instructions.push_back(ins);
 }
 
+void LLVMCodeBuilder::addListItemIndex(List *list)
+{
+    LLVMInstruction &ins = createOp(LLVMInstruction::Type::GetListItemIndex, Compiler::StaticType::Number, Compiler::StaticType::Unknown, 1);
+    ins.workList = list;
+    m_listPtrs[list] = LLVMListPtr();
+}
+
 void LLVMCodeBuilder::addListSize(List *list)
 {
     LLVMInstruction &ins = createOp(LLVMInstruction::Type::GetListSize, Compiler::StaticType::Number, {}, 0);
@@ -1752,6 +1767,56 @@ llvm::Value *LLVMCodeBuilder::getListItem(const LLVMListPtr &listPtr, llvm::Valu
     return m_builder.CreateGEP(m_valueDataType, m_builder.CreateLoad(m_valueDataType->getPointerTo(), listPtr.dataPtr), index);
 }
 
+llvm::Value *LLVMCodeBuilder::getListItemIndex(const LLVMListPtr &listPtr, LLVMRegisterPtr item, llvm::Function *func)
+{
+    llvm::Value *size = m_builder.CreateLoad(m_builder.getInt64Ty(), listPtr.sizePtr);
+    llvm::BasicBlock *condBlock = llvm::BasicBlock::Create(m_ctx, "", func);
+    llvm::BasicBlock *bodyBlock = llvm::BasicBlock::Create(m_ctx, "", func);
+    llvm::BasicBlock *cmpIfBlock = llvm::BasicBlock::Create(m_ctx, "", func);
+    llvm::BasicBlock *cmpElseBlock = llvm::BasicBlock::Create(m_ctx, "", func);
+    llvm::BasicBlock *notFoundBlock = llvm::BasicBlock::Create(m_ctx, "", func);
+    llvm::BasicBlock *nextBlock = llvm::BasicBlock::Create(m_ctx, "", func);
+
+    // index = 0
+    llvm::Value *index = m_builder.CreateAlloca(m_builder.getInt64Ty());
+    m_builder.CreateStore(m_builder.getInt64(0), index);
+    m_builder.CreateBr(condBlock);
+
+    // while (index < size)
+    m_builder.SetInsertPoint(condBlock);
+    llvm::Value *cond = m_builder.CreateICmpULT(m_builder.CreateLoad(m_builder.getInt64Ty(), index), size);
+    m_builder.CreateCondBr(cond, bodyBlock, notFoundBlock);
+
+    // if (list[index] == item)
+    m_builder.SetInsertPoint(bodyBlock);
+    LLVMRegisterPtr currentItem = std::make_shared<LLVMRegister>(listPtr.type);
+    currentItem->isRawValue = false;
+    currentItem->value = getListItem(listPtr, m_builder.CreateLoad(m_builder.getInt64Ty(), index), func);
+    llvm::Value *cmp = createComparison(currentItem, item, Comparison::EQ);
+    m_builder.CreateCondBr(cmp, cmpIfBlock, cmpElseBlock);
+
+    // goto nextBlock
+    m_builder.SetInsertPoint(cmpIfBlock);
+    m_builder.CreateBr(nextBlock);
+
+    // else index++
+    m_builder.SetInsertPoint(cmpElseBlock);
+    m_builder.CreateStore(m_builder.CreateAdd(m_builder.CreateLoad(m_builder.getInt64Ty(), index), m_builder.getInt64(1)), index);
+    m_builder.CreateBr(condBlock);
+
+    // notFoundBlock:
+    // index = -1
+    // goto nextBlock
+    m_builder.SetInsertPoint(notFoundBlock);
+    m_builder.CreateStore(llvm::ConstantInt::get(llvm::Type::getInt64Ty(m_ctx), -1, true), index);
+    m_builder.CreateBr(nextBlock);
+
+    // nextBlock:
+    m_builder.SetInsertPoint(nextBlock);
+
+    return m_builder.CreateSIToFP(m_builder.CreateLoad(m_builder.getInt64Ty(), index), m_builder.getDoubleTy());
+}
+
 llvm::Value *LLVMCodeBuilder::createValue(LLVMRegisterPtr reg)
 {
     if (reg->isConstValue) {
diff --git a/src/dev/engine/internal/llvm/llvmcodebuilder.h b/src/dev/engine/internal/llvm/llvmcodebuilder.h
@@ -30,6 +30,7 @@ class LLVMCodeBuilder : public ICodeBuilder
         void addVariableValue(Variable *variable) override;
         void addListContents(List *list) override;
         void addListItem(List *list) override;
+        void addListItemIndex(List *list) override;
         void addListSize(List *list) override;
 
         void createAdd() override;
@@ -123,6 +124,7 @@ class LLVMCodeBuilder : public ICodeBuilder
         void createValueCopy(llvm::Value *source, llvm::Value *target);
         void copyStructField(llvm::Value *source, llvm::Value *target, int index, llvm::StructType *structType, llvm::Type *fieldType);
         llvm::Value *getListItem(const LLVMListPtr &listPtr, llvm::Value *index, llvm::Function *func);
+        llvm::Value *getListItemIndex(const LLVMListPtr &listPtr, LLVMRegisterPtr item, llvm::Function *func);
         llvm::Value *createValue(LLVMRegisterPtr reg);
         llvm::Value *createComparison(LLVMRegisterPtr arg1, LLVMRegisterPtr arg2, Comparison type);
 
diff --git a/src/dev/engine/internal/llvm/llvminstruction.h b/src/dev/engine/internal/llvm/llvminstruction.h
@@ -49,6 +49,7 @@ struct LLVMInstruction
             ListReplace,
             GetListItem,
             GetListSize,
+            GetListItemIndex,
             Yield,
             BeginIf,
             BeginElse,
diff --git a/test/dev/llvm/llvmcodebuilder_test.cpp b/test/dev/llvm/llvmcodebuilder_test.cpp
@@ -2304,6 +2304,97 @@ TEST_F(LLVMCodeBuilderTest, GetListSize)
     ASSERT_EQ(localList->toString(), "Lorem ipsum dolor sit");
 }
 
+TEST_F(LLVMCodeBuilderTest, GetListItemIndex)
+{
+    EngineMock engine;
+    Stage stage;
+    Sprite sprite;
+    sprite.setEngine(&engine);
+    EXPECT_CALL(engine, stage()).WillRepeatedly(Return(&stage));
+
+    auto globalList = std::make_shared<List>("", "");
+    stage.addList(globalList);
+
+    auto localList = std::make_shared<List>("", "");
+    sprite.addList(localList);
+
+    globalList->append(1);
+    globalList->append(2);
+    globalList->append(3);
+
+    localList->append("Lorem");
+    localList->append("ipsum");
+    localList->append("dolor");
+    localList->append("sit");
+
+    createBuilder(&sprite, true);
+
+    m_builder->addConstValue(2);
+    m_builder->addListItemIndex(globalList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue(1);
+    m_builder->addListItemIndex(globalList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue(0);
+    m_builder->addListItemIndex(globalList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue(1);
+    m_builder->addConstValue("test");
+    m_builder->createListReplace(globalList.get());
+
+    m_builder->addConstValue(2);
+    m_builder->addListItemIndex(globalList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue(1);
+    m_builder->addListItemIndex(globalList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue("test");
+    m_builder->addListItemIndex(globalList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue("abc");
+    m_builder->addListItemIndex(globalList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue("doLor");
+    m_builder->addListItemIndex(localList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue(true);
+    m_builder->addListItemIndex(localList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    m_builder->addConstValue("site");
+    m_builder->addListItemIndex(localList.get());
+    m_builder->addFunctionCall("test_print_string", Compiler::StaticType::Void, { Compiler::StaticType::String });
+
+    static const std::string expected =
+        "1\n"
+        "0\n"
+        "-1\n"
+        "-1\n"
+        "0\n"
+        "1\n"
+        "-1\n"
+        "2\n"
+        "-1\n"
+        "-1\n";
+
+    auto code = m_builder->finalize();
+    auto ctx = code->createExecutionContext(&sprite);
+    testing::internal::CaptureStdout();
+    code->run(ctx.get());
+    ASSERT_EQ(testing::internal::GetCapturedStdout(), expected);
+
+    ASSERT_EQ(globalList->toString(), "1 test 3");
+    ASSERT_EQ(localList->toString(), "Lorem ipsum dolor sit");
+}
+
 TEST_F(LLVMCodeBuilderTest, Yield)
 {
     auto build = [this]() {
diff --git a/test/mocks/codebuildermock.h b/test/mocks/codebuildermock.h
@@ -14,6 +14,7 @@ class CodeBuilderMock : public ICodeBuilder
         MOCK_METHOD(void, addVariableValue, (Variable *), (override));
         MOCK_METHOD(void, addListContents, (List *), (override));
         MOCK_METHOD(void, addListItem, (List *), (override));
+        MOCK_METHOD(void, addListItemIndex, (List *), (override));
         MOCK_METHOD(void, addListSize, (List *), (override));
 
         MOCK_METHOD(void, createAdd, (), (override));
