add case statement support

3 files changed, 129 insertions, 0 deletions

diff --git a/node/Node.cc b/node/Node.cc
index 0dcf8ce..9325160 100644
--- a/node/Node.cc
+++ b/node/Node.cc
@@ -501,4 +501,45 @@ int ContinueStatementNode::Travel(Traveler& traveler) {
     traveler << TravelPart::END;
     return 0;
 }
+
+int SwitchStatementNode::Travel(Traveler& traveler) {
+    traveler << TravelPart::PREFIX;
+    traveler << TravelPart::NAME_BEGIN << "SwitchStatementNode"
+             << TravelPart::NAME_END;
+    traveler << TravelPart::DESCRPTION_BEGIN << "NumCases="
+             << CasePairs_.size() + (DefaultAction_ == nullptr ? 0 : 1)
+             << (DefaultAction_ == nullptr ? "" : " HasDefault")
+             << TravelPart::DESCRPTION_END;
+    traveler << TravelPart::CHILD_BEGIN;
+    CaseVariable_->Travel(traveler);
+    traveler << TravelPart::CHILD_END;
+    std::shared_ptr<BaseNode> PrevAction = nullptr;
+    std::vector<std::shared_ptr<BaseNode>> actions;
+    for (auto&& child : CasePairs_) {
+        std::shared_ptr<BaseNode> action = std::get<1>(child);
+        if (action != PrevAction) {
+            actions.push_back(action);
+            PrevAction = action;
+        }
+    }
+    auto lambda = [](Traveler& traveler, auto&& action) {
+        action->Travel(traveler);
+    };
+    NonterminalContainerTravel(
+        traveler,
+        [](Traveler& traveler, auto&& child) {
+            std::get<0>(child)->Travel(traveler);
+        },
+        CasePairs_);
+    if (DefaultAction_ == nullptr) {
+        TerminalContainerTravel(traveler, lambda, actions);
+    } else {
+        NonterminalContainerTravel(traveler, lambda, actions);
+        traveler << TravelPart::LAST_CHILD_BEGIN;
+        DefaultAction_->Travel(traveler);
+        traveler << TravelPart::LAST_CHILD_END;
+    }
+    traveler << TravelPart::END;
+    return 0;
+}
 } // namespace pcc
diff --git a/node/StatementNode.cc b/node/StatementNode.cc
index 07babe2..adce4f0 100644
--- a/node/StatementNode.cc
+++ b/node/StatementNode.cc
@@ -294,4 +294,74 @@ Value ContinueStatementNode::CodeGen() {
     context_->SetJumpingFlag(true);
     return nullptr;
 }
+
+SwitchStatementNode::SwitchStatementNode(
+    Context* context, std::shared_ptr<ExprNode> CaseVariable,
+    std::list<CasePair>&& CasePairs, std::shared_ptr<BaseNode> DefaultAction)
+    : BaseNode(context)
+    , CaseVariable_(CaseVariable)
+    , CasePairs_(std::move(CasePairs))
+    , DefaultAction_(DefaultAction) {}
+
+Value SwitchStatementNode::CodeGen() {
+    llvm::IRBuilder<>* builder = context_->GetBuilder();
+    Value SwitchValue = CaseVariable_->CodeGen();
+    std::shared_ptr<IntegerBaseType> type =
+        std::dynamic_pointer_cast<IntegerBaseType>(CaseVariable_->GetType());
+    if (type == nullptr) {
+        Log(LogLevel::PCC_ERROR,
+            "variable of switch statement variable should have a integer-like "
+            "type, not %s",
+            CaseVariable_->GetType()->GetCommonName());
+        return nullptr;
+    }
+    llvm::Function* body = builder->GetInsertBlock()->getParent();
+    llvm::BasicBlock* EndBlock =
+        llvm::BasicBlock::Create(GlobalLLVMContext, "done");
+    llvm::BasicBlock* DefaultBlock =
+        (DefaultAction_ == nullptr
+             ? EndBlock
+             : llvm::BasicBlock::Create(GlobalLLVMContext, "default"));
+    llvm::SwitchInst* inst =
+        builder->CreateSwitch(SwitchValue, DefaultBlock, CasePairs_.size() + 1);
+    std::shared_ptr<BaseNode> PrevAcion = nullptr;
+    llvm::BasicBlock* PrevBlock = nullptr;
+    for (auto&& CaseChild : CasePairs_) {
+        llvm::Value* CaseValue = std::get<0>(CaseChild)->CodeGen();
+        std::shared_ptr<IntegerBaseType> CaseType =
+            std::dynamic_pointer_cast<IntegerBaseType>(
+                std::get<0>(CaseChild)->GetType());
+        if (CaseType == nullptr) {
+            Log(LogLevel::PCC_ERROR,
+                "case of switch statement variable should have a "
+                "integer-like "
+                "type, not %s",
+                std::get<0>(CaseChild)->GetType()->GetCommonName());
+            return nullptr;
+        }
+        CaseValue = builder->CreateSExtOrTrunc(CaseValue, type->GetLLVMType());
+        llvm::ConstantInt* CaseConst = llvm::cast<llvm::ConstantInt>(CaseValue);
+        std::shared_ptr<BaseNode> CaseAction = std::get<1>(CaseChild);
+        if (CaseAction == PrevAcion) {
+            inst->addCase(CaseConst, PrevBlock);
+        } else {
+            llvm::BasicBlock* CaseBlock =
+                llvm::BasicBlock::Create(GlobalLLVMContext, "case", body);
+            inst->addCase(CaseConst, CaseBlock);
+            builder->SetInsertPoint(CaseBlock);
+            CaseAction->CodeGen();
+            builder->CreateBr(EndBlock);
+            PrevAcion = CaseAction, PrevBlock = CaseBlock;
+        }
+    }
+    if (DefaultAction_ != nullptr) {
+        body->getBasicBlockList().push_back(DefaultBlock);
+        builder->SetInsertPoint(DefaultBlock);
+        DefaultAction_->CodeGen();
+        builder->CreateBr(EndBlock);
+    }
+    body->getBasicBlockList().push_back(EndBlock);
+    builder->SetInsertPoint(EndBlock);
+    return nullptr;
+}
 } // namespace pcc
diff --git a/node/StatementNode.hh b/node/StatementNode.hh
index b729001..81cbb93 100644
--- a/node/StatementNode.hh
+++ b/node/StatementNode.hh
@@ -105,6 +105,24 @@ class ContinueStatementNode : public BaseNode {
     virtual Value CodeGen() override;
     virtual int Travel(Traveler& traveler) override;
 };
+
+using CasePair =
+    std::pair<std::shared_ptr<ExprNode>, std::shared_ptr<BaseNode>>;
+
+class SwitchStatementNode : public BaseNode {
+  public:
+    SwitchStatementNode(Context* context,
+                        std::shared_ptr<ExprNode> CaseVariable,
+                        std::list<CasePair>&& CasePairs,
+                        std::shared_ptr<BaseNode> DefaultAction);
+    virtual Value CodeGen() override;
+    virtual int Travel(Traveler& traveler) override;
+
+  protected:
+    std::shared_ptr<ExprNode> CaseVariable_;
+    std::list<CasePair> CasePairs_;
+    std::shared_ptr<BaseNode> DefaultAction_;
+};
 } // namespace pcc
 
 #endif