From f58eb196c3863144d5dc23b3d64179bbc0f8247f Mon Sep 17 00:00:00 2001
From: vla5924 <vla5924@gmail.com>
Date: Sat, 25 May 2024 12:10:54 +0300
Subject: [PATCH] ret terminator

---
 .../optree_to_llvmir/llvmir_generator.hpp     |  3 +++
 .../optree_to_llvmir/llvmir_generator.cpp     | 22 +++++++++++++++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/compiler/include/compiler/codegen/optree_to_llvmir/llvmir_generator.hpp b/compiler/include/compiler/codegen/optree_to_llvmir/llvmir_generator.hpp
index 2914c6e1..12bfd373 100644
--- a/compiler/include/compiler/codegen/optree_to_llvmir/llvmir_generator.hpp
+++ b/compiler/include/compiler/codegen/optree_to_llvmir/llvmir_generator.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <deque>
 #include <string>
 #include <string_view>
 #include <unordered_map>
@@ -34,11 +35,13 @@ class LLVMIRGenerator {
     std::unordered_map<llvm::Value *, llvm::Type *> typedValues;
     std::unordered_map<std::string, llvm::Value *> globalStrings;
     std::unordered_map<std::string_view, llvm::FunctionCallee> externalFunctions;
+    std::deque<llvm::BasicBlock *> basicBlocks;
 
     llvm::Value *findValue(const Value::Ptr &value) const;
     void saveValue(const Value::Ptr &value, llvm::Value *llvmValue);
     llvm::Type *convertType(const Type::Ptr &type);
     llvm::BasicBlock *createBlock();
+    void eraseDeadBlocks();
     llvm::Value *normalizePredicate(const Value::Ptr &cond);
     llvm::Value *getGlobalString(const std::string &str);
     llvm::FunctionCallee getExternalFunction(std::string_view name);
diff --git a/compiler/lib/codegen/optree_to_llvmir/llvmir_generator.cpp b/compiler/lib/codegen/optree_to_llvmir/llvmir_generator.cpp
index fd8fe63c..e7f560aa 100644
--- a/compiler/lib/codegen/optree_to_llvmir/llvmir_generator.cpp
+++ b/compiler/lib/codegen/optree_to_llvmir/llvmir_generator.cpp
@@ -8,6 +8,7 @@
 #include <system_error>
 #include <vector>
 
+#include <llvm/IR/CFG.h>
 #include <llvm/IR/Constants.h>
 #include <llvm/IR/DerivedTypes.h>
 #include <llvm/IR/Type.h>
@@ -83,7 +84,20 @@ llvm::Type *LLVMIRGenerator::convertType(const Type::Ptr &type) {
 }
 
 llvm::BasicBlock *LLVMIRGenerator::createBlock() {
-    return llvm::BasicBlock::Create(context, "bb", currentFunction);
+    auto *bb = llvm::BasicBlock::Create(context, "bb", currentFunction);
+    basicBlocks.push_back(bb);
+    return bb;
+}
+
+void LLVMIRGenerator::eraseDeadBlocks() {
+    for (auto *&bb : basicBlocks) {
+        if (bb == nullptr)
+            continue;
+        if (llvm::pred_empty(bb) && !bb->isEntryBlock()) {
+            bb->eraseFromParent();
+            bb = nullptr;
+        }
+    }
 }
 
 llvm::Value *LLVMIRGenerator::normalizePredicate(const Value::Ptr &cond) {
@@ -208,6 +222,8 @@ void LLVMIRGenerator::visit(const ReturnOp &op) {
         builder.CreateRetVoid();
     else
         builder.CreateRet(findValue(op.value()));
+    auto *bb = createBlock();
+    builder.SetInsertPoint(bb);
 }
 
 void LLVMIRGenerator::visit(const ConstantOp &op) {
@@ -347,6 +363,7 @@ void LLVMIRGenerator::visit(const IfOp &op) {
     auto *thenBlock = createBlock();
     builder.SetInsertPoint(thenBlock);
     visit(op.thenOp());
+    auto *newThenBlock = builder.GetInsertBlock();
     auto *elseBlock = createBlock();
     auto *nextBlock = elseBlock;
     if (auto elseOp = op.elseOp()) {
@@ -355,7 +372,7 @@ void LLVMIRGenerator::visit(const IfOp &op) {
         nextBlock = createBlock();
         builder.CreateBr(nextBlock);
     }
-    builder.SetInsertPoint(thenBlock);
+    builder.SetInsertPoint(newThenBlock);
     builder.CreateBr(nextBlock);
     builder.SetInsertPoint(prevBlock);
     builder.CreateCondBr(normalizePredicate(op.cond()), thenBlock, elseBlock);
@@ -431,6 +448,7 @@ void LLVMIRGenerator::visit(const PrintOp &op) {
 
 void LLVMIRGenerator::process(const Program &program) {
     visit(program.root);
+    eraseDeadBlocks();
 }
 
 std::string LLVMIRGenerator::dump() const {