From bddf270bce142294c140cdb612556f23ff60753e Mon Sep 17 00:00:00 2001
From: Lenny Truong <lenny@stanford.edu>
Date: Thu, 4 Jun 2020 14:10:23 -0700
Subject: [PATCH] Add else if logic

---
 include/verilogAST.hpp | 14 +++++++++++---
 src/transformer.cpp    | 23 ++++++++++++++++++-----
 src/verilogAST.cpp     | 13 +++++++++++--
 tests/basic.cpp        |  8 ++++++++
 tests/common.cpp       | 23 +++++++++++++++++++----
 tests/transformer.cpp  |  8 ++++++++
 6 files changed, 75 insertions(+), 14 deletions(-)

diff --git a/include/verilogAST.hpp b/include/verilogAST.hpp
index 69bb6da..ed22730 100644
--- a/include/verilogAST.hpp
+++ b/include/verilogAST.hpp
@@ -774,14 +774,22 @@ class If : public BehavioralStatement {
  public:
   std::unique_ptr<Expression> cond;
   std::vector<std::unique_ptr<BehavioralStatement>> true_body;
-  std::vector<std::unique_ptr<BehavioralStatement>> false_body;
+  // pair of cond + body for else if cases
+  std::vector<std::pair<std::unique_ptr<Expression>,
+                        std::vector<std::unique_ptr<BehavioralStatement>>>>
+      else_ifs;
+  std::vector<std::unique_ptr<BehavioralStatement>> else_body;
 
   If(std::unique_ptr<Expression> cond,
      std::vector<std::unique_ptr<BehavioralStatement>> true_body,
-     std::vector<std::unique_ptr<BehavioralStatement>> false_body)
+     std::vector<std::pair<std::unique_ptr<Expression>,
+                           std::vector<std::unique_ptr<BehavioralStatement>>>>
+         else_ifs,
+     std::vector<std::unique_ptr<BehavioralStatement>> else_body)
       : cond(std::move(cond)),
         true_body(std::move(true_body)),
-        false_body(std::move(false_body)){};
+        else_ifs(std::move(else_ifs)),
+        else_body(std::move(else_body)){};
 
   std::string toString();
   ~If(){};
diff --git a/src/transformer.cpp b/src/transformer.cpp
index 4a1e202..84c3648 100644
--- a/src/transformer.cpp
+++ b/src/transformer.cpp
@@ -206,11 +206,24 @@ std::unique_ptr<If> Transformer::visit(std::unique_ptr<If> node) {
   }
   node->true_body = std::move(new_true_body);
 
-  std::vector<std::unique_ptr<BehavioralStatement>> new_false_body;
-  for (auto&& item : node->false_body) {
-    new_false_body.push_back(this->visit(std::move(item)));
-  }
-  node->false_body = std::move(new_false_body);
+  std::vector<std::pair<std::unique_ptr<Expression>,
+                        std::vector<std::unique_ptr<BehavioralStatement>>>>
+      new_else_ifs;
+  for (auto&& item : node->else_ifs) {
+    std::vector<std::unique_ptr<BehavioralStatement>> new_body;
+    for (auto&& inner_statement : item.second) {
+      new_body.push_back(this->visit(std::move(inner_statement)));
+    }
+    new_else_ifs.push_back(
+        {this->visit(std::move(item.first)), std::move(new_body)});
+  }
+  node->else_ifs = std::move(new_else_ifs);
+
+  std::vector<std::unique_ptr<BehavioralStatement>> new_else_body;
+  for (auto&& item : node->else_body) {
+    new_else_body.push_back(this->visit(std::move(item)));
+  }
+  node->else_body = std::move(new_else_body);
 
   return node;
 }
diff --git a/src/verilogAST.cpp b/src/verilogAST.cpp
index 3c7ca25..2ca3c53 100644
--- a/src/verilogAST.cpp
+++ b/src/verilogAST.cpp
@@ -446,9 +446,18 @@ std::string If::toString() {
   }
 
   if_str += "end";
-  if (this->false_body.size()) {
+
+  for (auto &entry : this->else_ifs) {
+    if_str += " else if (" + entry.first->toString() + ") begin\n";
+    for (auto &statement : entry.second) {
+      if_str += add_tab(statement->toString());
+    }
+    if_str += "end";
+  }
+
+  if (this->else_body.size()) {
     if_str += " else begin\n";
-    for (auto &statement : false_body) {
+    for (auto &statement : else_body) {
       if_str += add_tab(statement->toString());
     }
     if_str += "end";
diff --git a/tests/basic.cpp b/tests/basic.cpp
index a3f0515..c788a3d 100644
--- a/tests/basic.cpp
+++ b/tests/basic.cpp
@@ -444,6 +444,10 @@ TEST(BasicTests, TestAlways) {
       "$display(\"b=%d, c=%d\", b, c);\n"
       "if (b) begin\n"
       "    e = f;\n"
+      "end else if (x0) begin\n"
+      "    e = g0;\n"
+      "end else if (x1) begin\n"
+      "    e = g1;\n"
       "end else begin\n"
       "    e = g;\n"
       "end\n"
@@ -467,6 +471,10 @@ TEST(BasicTests, TestAlwaysStar) {
       "$display(\"b=%d, c=%d\", b, c);\n"
       "if (b) begin\n"
       "    e = f;\n"
+      "end else if (x0) begin\n"
+      "    e = g0;\n"
+      "end else if (x1) begin\n"
+      "    e = g1;\n"
       "end else begin\n"
       "    e = g;\n"
       "end\n"
diff --git a/tests/common.cpp b/tests/common.cpp
index 485123a..5be5f64 100644
--- a/tests/common.cpp
+++ b/tests/common.cpp
@@ -72,13 +72,28 @@ make_simple_always_body() {
       std::make_unique<vAST::Identifier>("e"),
       std::make_unique<vAST::Identifier>("f")));
 
-  std::vector<std::unique_ptr<vAST::BehavioralStatement>> false_body;
-  false_body.push_back(std::make_unique<vAST::BlockingAssign>(
+  std::vector<
+      std::pair<std::unique_ptr<vAST::Expression>,
+                std::vector<std::unique_ptr<vAST::BehavioralStatement>>>>
+      else_ifs;
+  for (int i = 0; i < 2; i++) {
+    std::unique_ptr<vAST::Expression> cond =
+        vAST::make_id("x" + std::to_string(i));
+    std::vector<std::unique_ptr<vAST::BehavioralStatement>> body;
+    body.push_back(std::make_unique<vAST::BlockingAssign>(
+        std::make_unique<vAST::Identifier>("e"),
+        std::make_unique<vAST::Identifier>("g" + std::to_string(i))));
+    else_ifs.push_back({std::move(cond), std::move(body)});
+  }
+
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> else_body;
+  else_body.push_back(std::make_unique<vAST::BlockingAssign>(
       std::make_unique<vAST::Identifier>("e"),
       std::make_unique<vAST::Identifier>("g")));
 
-  body.push_back(std::make_unique<vAST::If>(
-      vAST::make_id("b"), std::move(true_body), std::move(false_body)));
+  body.push_back(
+      std::make_unique<vAST::If>(vAST::make_id("b"), std::move(true_body),
+                                 std::move(else_ifs), std::move(else_body)));
 
   return body;
 }
diff --git a/tests/transformer.cpp b/tests/transformer.cpp
index 9d63dab..1de8616 100644
--- a/tests/transformer.cpp
+++ b/tests/transformer.cpp
@@ -139,6 +139,10 @@ TEST(TransformerTests, TestAlways) {
       "$display(\"b=%d, c=%d\", y, c);\n"
       "if (y) begin\n"
       "    w = f;\n"
+      "end else if (x0) begin\n"
+      "    w = g0;\n"
+      "end else if (x1) begin\n"
+      "    w = g1;\n"
       "end else begin\n"
       "    w = g;\n"
       "end\n"
@@ -218,6 +222,10 @@ TEST(TransformerTests, TestModule) {
       "$display(\"b=%d, c=%d\", b, g);\n"
       "if (b) begin\n"
       "    e = f;\n"
+      "end else if (x0) begin\n"
+      "    e = g0;\n"
+      "end else if (x1) begin\n"
+      "    e = g1;\n"
       "end else begin\n"
       "    e = g;\n"
       "end\n"