From c306314426008143de70928738e32894d58d983d Mon Sep 17 00:00:00 2001
From: Lenny Truong <lenny@stanford.edu>
Date: Wed, 3 Jun 2020 14:32:50 -0700
Subject: [PATCH 1/7] Add support for If statement

---
 include/verilogAST.hpp             | 17 +++++++++++++++++
 include/verilogAST/transformer.hpp |  2 ++
 src/transformer.cpp                | 26 ++++++++++++++++++++++----
 src/verilogAST.cpp                 | 23 ++++++++++++++++++++++-
 tests/basic.cpp                    | 11 +++++++++++
 tests/common.cpp                   | 21 +++++++++++++++++----
 tests/transformer.cpp              | 17 ++++++++++++++---
 7 files changed, 105 insertions(+), 12 deletions(-)

diff --git a/include/verilogAST.hpp b/include/verilogAST.hpp
index 8184325..69bb6da 100644
--- a/include/verilogAST.hpp
+++ b/include/verilogAST.hpp
@@ -770,6 +770,23 @@ class Always : public StructuralStatement {
   ~Always(){};
 };
 
+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;
+
+  If(std::unique_ptr<Expression> cond,
+     std::vector<std::unique_ptr<BehavioralStatement>> true_body,
+     std::vector<std::unique_ptr<BehavioralStatement>> false_body)
+      : cond(std::move(cond)),
+        true_body(std::move(true_body)),
+        false_body(std::move(false_body)){};
+
+  std::string toString();
+  ~If(){};
+};
+
 class AbstractModule : public Node {};
 
 class Module : public AbstractModule {
diff --git a/include/verilogAST/transformer.hpp b/include/verilogAST/transformer.hpp
index 4f11900..2dd976b 100644
--- a/include/verilogAST/transformer.hpp
+++ b/include/verilogAST/transformer.hpp
@@ -62,6 +62,8 @@ class Transformer {
   virtual std::unique_ptr<BlockComment> visit(
       std::unique_ptr<BlockComment> node);
 
+  virtual std::unique_ptr<If> visit(std::unique_ptr<If> node);
+
   virtual std::unique_ptr<InlineVerilog> visit(
       std::unique_ptr<InlineVerilog> node);
 
diff --git a/src/transformer.cpp b/src/transformer.cpp
index 959990a..5b154eb 100644
--- a/src/transformer.cpp
+++ b/src/transformer.cpp
@@ -71,14 +71,12 @@ std::unique_ptr<Identifier> Transformer::visit(
   return node;
 }
 
-std::unique_ptr<Cast> Transformer::visit(
-    std::unique_ptr<Cast> node) {
+std::unique_ptr<Cast> Transformer::visit(std::unique_ptr<Cast> node) {
   node->expr = this->visit(std::move(node->expr));
   return node;
 }
 
-std::unique_ptr<Attribute> Transformer::visit(
-    std::unique_ptr<Attribute> node) {
+std::unique_ptr<Attribute> Transformer::visit(std::unique_ptr<Attribute> node) {
   node->value = this->visit(std::move(node->value));
   return node;
 }
@@ -199,6 +197,22 @@ std::unique_ptr<BlockComment> Transformer::visit(
   return node;
 }
 
+std::unique_ptr<If> Transformer::visit(std::unique_ptr<If> node) {
+  std::vector<std::unique_ptr<BehavioralStatement>> new_true_body;
+  for (auto&& item : node->true_body) {
+    new_true_body.push_back(this->visit(std::move(item)));
+  }
+  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);
+
+  return node;
+}
+
 std::unique_ptr<InlineVerilog> Transformer::visit(
     std::unique_ptr<InlineVerilog> node) {
   return node;
@@ -269,6 +283,10 @@ std::unique_ptr<BehavioralStatement> Transformer::visit(
     node.release();
     return this->visit(std::unique_ptr<BlockComment>(ptr));
   }
+  if (auto ptr = dynamic_cast<If*>(node.get())) {
+    node.release();
+    return this->visit(std::unique_ptr<If>(ptr));
+  }
   throw std::runtime_error("Unreachable");  // LCOV_EXCL_LINE
   return node;                              // LCOV_EXCL_LINE
 }
diff --git a/src/verilogAST.cpp b/src/verilogAST.cpp
index e6f6479..f33e7d6 100644
--- a/src/verilogAST.cpp
+++ b/src/verilogAST.cpp
@@ -123,7 +123,7 @@ std::string Identifier::toString() {
 }
 
 std::string Cast::toString() {
-    return std::to_string(this->width) + "'(" + this->expr->toString() + ")";
+  return std::to_string(this->width) + "'(" + this->expr->toString() + ")";
 }
 
 std::string Attribute::toString() {
@@ -422,6 +422,27 @@ std::string Always::toString() {
   return always_str;
 }
 
+std::string If::toString() {
+  std::string if_str = "";
+  if_str += "if (";
+  if_str += this->cond->toString();
+  if_str += ") begin\n";
+
+  for (auto &statement : this->true_body) {
+    if_str += "    " + statement->toString() + "\n";
+  }
+
+  if_str += "end";
+  if (this->false_body.size()) {
+    if_str += " else begin\n";
+    for (auto &statement : false_body) {
+      if_str += "    " + statement->toString() + "\n";
+    }
+    if_str += "end";
+  }
+  return if_str;
+}
+
 std::string File::toString() {
   std::string file_str = "";
 
diff --git a/tests/basic.cpp b/tests/basic.cpp
index e22faa0..a3f0515 100644
--- a/tests/basic.cpp
+++ b/tests/basic.cpp
@@ -442,6 +442,11 @@ TEST(BasicTests, TestAlways) {
       "a = b;\n"
       "b <= c;\n"
       "$display(\"b=%d, c=%d\", b, c);\n"
+      "if (b) begin\n"
+      "    e = f;\n"
+      "end else begin\n"
+      "    e = g;\n"
+      "end\n"
       "end\n";
   EXPECT_EQ(always.toString(), expected_str);
 }
@@ -452,6 +457,7 @@ TEST(BasicTests, TestAlwaysStar) {
       std::unique_ptr<vAST::NegEdge>, std::unique_ptr<vAST::Star>>>
       sensitivity_list;
   sensitivity_list.push_back(std::make_unique<vAST::Star>());
+
   vAST::Always always_star(std::move(sensitivity_list),
                            make_simple_always_body());
   std::string expected_str =
@@ -459,6 +465,11 @@ TEST(BasicTests, TestAlwaysStar) {
       "a = b;\n"
       "b <= c;\n"
       "$display(\"b=%d, c=%d\", b, c);\n"
+      "if (b) begin\n"
+      "    e = f;\n"
+      "end else begin\n"
+      "    e = g;\n"
+      "end\n"
       "end\n";
   EXPECT_EQ(always_star.toString(), expected_str);
 }
diff --git a/tests/common.cpp b/tests/common.cpp
index d1cded5..485123a 100644
--- a/tests/common.cpp
+++ b/tests/common.cpp
@@ -12,11 +12,11 @@ vAST::Parameters make_simple_params() {
 }
 
 std::unique_ptr<vAST::Connections> make_simple_connections() {
-  std::unique_ptr<vAST::Connections> connections = std::make_unique<vAST::Connections>();
+  std::unique_ptr<vAST::Connections> connections =
+      std::make_unique<vAST::Connections>();
   connections->insert("a", vAST::make_id("a"));
-  connections->insert(
-      "b",
-      std::make_unique<vAST::Index>(vAST::make_id("b"), vAST::make_num("0")));
+  connections->insert("b", std::make_unique<vAST::Index>(vAST::make_id("b"),
+                                                         vAST::make_num("0")));
   connections->insert(
       "c", std::make_unique<vAST::Slice>(
                vAST::make_id("c"), vAST::make_num("31"), vAST::make_num("0")));
@@ -67,5 +67,18 @@ make_simple_always_body() {
   args.push_back(std::make_unique<vAST::Identifier>("c"));
   body.push_back(std::make_unique<vAST::CallStmt>("$display", std::move(args)));
 
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> true_body;
+  true_body.push_back(std::make_unique<vAST::BlockingAssign>(
+      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::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)));
+
   return body;
 }
diff --git a/tests/transformer.cpp b/tests/transformer.cpp
index f7a03df..f95ad24 100644
--- a/tests/transformer.cpp
+++ b/tests/transformer.cpp
@@ -43,6 +43,8 @@ class AlwaysTransformer : public vAST::Transformer {
       return vAST::make_id("z");
     } else if (node->value == "b") {
       return vAST::make_id("y");
+    } else if (node->value == "e") {
+      return vAST::make_id("w");
     }
     return node;
   };
@@ -95,8 +97,7 @@ TEST(TransformerTests, TestXtoZ) {
               vAST::UnOp::INVERT)),
       std::make_unique<vAST::Replicate>(vAST::make_num("2"),
                                         vAST::make_id("x"))));
-  call_args.push_back(
-      std::make_unique<vAST::Cast>(5, vAST::make_id("x")));
+  call_args.push_back(std::make_unique<vAST::Cast>(5, vAST::make_id("x")));
   std::unique_ptr<vAST::Expression> expr =
       std::make_unique<vAST::CallExpr>("my_func", std::move(call_args));
   XtoZ transformer;
@@ -136,6 +137,11 @@ TEST(TransformerTests, TestAlways) {
       "z = y;\n"
       "y <= c;\n"
       "$display(\"b=%d, c=%d\", y, c);\n"
+      "if (b) begin\n"
+      "    w = f;\n"
+      "end else begin\n"
+      "    w = g;\n"
+      "end\n"
       "// Test comment\n"
       "/*\nTest comment\non multiple lines\n*/\n"
       "end\n";
@@ -200,7 +206,7 @@ TEST(TransformerTests, TestModule) {
       "    .param1(1)\n"
       ") other_module_inst (\n"
       "    .a(a),\n"
-      "    .b(b[0]),\n" 
+      "    .b(b[0]),\n"
       "    .c(g[31:0])\n"
       ");\n"
       "assign g = b;\n"
@@ -210,6 +216,11 @@ TEST(TransformerTests, TestModule) {
       "a = b;\n"
       "b <= g;\n"
       "$display(\"b=%d, c=%d\", b, g);\n"
+      "if (b) begin\n"
+      "    e = f;\n"
+      "end else begin\n"
+      "    e = g;\n"
+      "end\n"
       "end\n\n"
       "// Test comment\n"
       "/*\nTest comment\non multiple lines\n*/\n"

From 79dd538153231c8038d2d82bc7b40d23aa0c16df Mon Sep 17 00:00:00 2001
From: Lenny Truong <lenny@stanford.edu>
Date: Thu, 4 Jun 2020 10:37:06 -0700
Subject: [PATCH 2/7] Recurse into if statement cond

---
 src/transformer.cpp   | 2 ++
 tests/transformer.cpp | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/transformer.cpp b/src/transformer.cpp
index 5b154eb..4a1e202 100644
--- a/src/transformer.cpp
+++ b/src/transformer.cpp
@@ -198,6 +198,8 @@ std::unique_ptr<BlockComment> Transformer::visit(
 }
 
 std::unique_ptr<If> Transformer::visit(std::unique_ptr<If> node) {
+  node->cond = this->visit(std::move(node->cond));
+
   std::vector<std::unique_ptr<BehavioralStatement>> new_true_body;
   for (auto&& item : node->true_body) {
     new_true_body.push_back(this->visit(std::move(item)));
diff --git a/tests/transformer.cpp b/tests/transformer.cpp
index f95ad24..9d63dab 100644
--- a/tests/transformer.cpp
+++ b/tests/transformer.cpp
@@ -137,7 +137,7 @@ TEST(TransformerTests, TestAlways) {
       "z = y;\n"
       "y <= c;\n"
       "$display(\"b=%d, c=%d\", y, c);\n"
-      "if (b) begin\n"
+      "if (y) begin\n"
       "    w = f;\n"
       "end else begin\n"
       "    w = g;\n"

From 187e5fa952dd0a60527254fa65691afdb00b1726 Mon Sep 17 00:00:00 2001
From: Lenny Truong <lenny@stanford.edu>
Date: Thu, 4 Jun 2020 13:33:38 -0700
Subject: [PATCH 3/7] Update tab logic

---
 src/verilogAST.cpp | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/verilogAST.cpp b/src/verilogAST.cpp
index f33e7d6..3c7ca25 100644
--- a/src/verilogAST.cpp
+++ b/src/verilogAST.cpp
@@ -1,6 +1,7 @@
 #include "verilogAST.hpp"
 
 #include <regex>
+#include <sstream>
 #include <unordered_set>
 
 template <typename... Ts>
@@ -422,6 +423,18 @@ std::string Always::toString() {
   return always_str;
 }
 
+std::string add_tab(std::string block) {
+  // Indents each line by four spaces, adds an extra newline at the end
+  std::istringstream block_stream(block);
+  std::string new_block;
+  while (!block_stream.eof()) {
+    std::string line;
+    std::getline(block_stream, line);
+    new_block += "    " + line + "\n";
+  }
+  return new_block;
+}
+
 std::string If::toString() {
   std::string if_str = "";
   if_str += "if (";
@@ -429,14 +442,14 @@ std::string If::toString() {
   if_str += ") begin\n";
 
   for (auto &statement : this->true_body) {
-    if_str += "    " + statement->toString() + "\n";
+    if_str += add_tab(statement->toString());
   }
 
   if_str += "end";
   if (this->false_body.size()) {
     if_str += " else begin\n";
     for (auto &statement : false_body) {
-      if_str += "    " + statement->toString() + "\n";
+      if_str += add_tab(statement->toString());
     }
     if_str += "end";
   }

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 4/7] 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"

From 0b5147f13ea1a25aeb77468cb0b3d29a66b6c06b Mon Sep 17 00:00:00 2001
From: Lenny Truong <lenny@stanford.edu>
Date: Fri, 5 Jun 2020 10:02:10 -0700
Subject: [PATCH 5/7] Fixup always body type

---
 include/verilogAST.hpp                |  8 ++----
 include/verilogAST/assign_inliner.hpp |  6 ++++
 src/assign_inliner.cpp                | 41 +++++++++++++++++++++++++++
 src/transformer.cpp                   |  4 +--
 src/verilogAST.cpp                    |  4 +--
 tests/assign_inliner.cpp              | 31 +++++++++++---------
 tests/basic.cpp                       |  4 +--
 tests/common.cpp                      |  7 ++---
 tests/transformer.cpp                 |  5 ++--
 9 files changed, 74 insertions(+), 36 deletions(-)

diff --git a/include/verilogAST.hpp b/include/verilogAST.hpp
index ed22730..8d81da4 100644
--- a/include/verilogAST.hpp
+++ b/include/verilogAST.hpp
@@ -748,17 +748,13 @@ class Always : public StructuralStatement {
       std::variant<std::unique_ptr<Identifier>, std::unique_ptr<PosEdge>,
                    std::unique_ptr<NegEdge>, std::unique_ptr<Star>>>
       sensitivity_list;
-  std::vector<std::variant<std::unique_ptr<BehavioralStatement>,
-                           std::unique_ptr<Declaration>>>
-      body;
+  std::vector<std::unique_ptr<BehavioralStatement>> body;
 
   Always(std::vector<
              std::variant<std::unique_ptr<Identifier>, std::unique_ptr<PosEdge>,
                           std::unique_ptr<NegEdge>, std::unique_ptr<Star>>>
              sensitivity_list,
-         std::vector<std::variant<std::unique_ptr<BehavioralStatement>,
-                                  std::unique_ptr<Declaration>>>
-             body)
+         std::vector<std::unique_ptr<BehavioralStatement>> body)
       : body(std::move(body)) {
     if (sensitivity_list.empty()) {
       throw std::runtime_error(
diff --git a/include/verilogAST/assign_inliner.hpp b/include/verilogAST/assign_inliner.hpp
index b02ec4e..1e4e7e1 100644
--- a/include/verilogAST/assign_inliner.hpp
+++ b/include/verilogAST/assign_inliner.hpp
@@ -30,6 +30,8 @@ class AssignMapBuilder : public Transformer {
   virtual std::unique_ptr<Port> visit(std::unique_ptr<Port> node);
   virtual std::unique_ptr<ContinuousAssign> visit(
       std::unique_ptr<ContinuousAssign> node);
+  virtual std::unique_ptr<BlockingAssign> visit(
+      std::unique_ptr<BlockingAssign> node);
 };
 
 class WireReadCounter : public Transformer {
@@ -47,6 +49,8 @@ class WireReadCounter : public Transformer {
   // Skip target of assign (not read)
   virtual std::unique_ptr<ContinuousAssign> visit(
       std::unique_ptr<ContinuousAssign> node);
+  virtual std::unique_ptr<BlockingAssign> visit(
+      std::unique_ptr<BlockingAssign> node);
   // Skip declarations (not read)
   virtual std::unique_ptr<Declaration> visit(std::unique_ptr<Declaration> node);
 };
@@ -96,6 +100,8 @@ class AssignInliner : public Transformer {
   virtual std::unique_ptr<Expression> visit(std::unique_ptr<Expression> node);
   virtual std::unique_ptr<ContinuousAssign> visit(
       std::unique_ptr<ContinuousAssign> node);
+  virtual std::unique_ptr<BlockingAssign> visit(
+      std::unique_ptr<BlockingAssign> node);
   virtual std::unique_ptr<Wire> visit(std::unique_ptr<Wire> node);
   virtual std::unique_ptr<Module> visit(std::unique_ptr<Module> node);
 };
diff --git a/src/assign_inliner.cpp b/src/assign_inliner.cpp
index 78f32c6..2a0107c 100644
--- a/src/assign_inliner.cpp
+++ b/src/assign_inliner.cpp
@@ -35,6 +35,12 @@ std::unique_ptr<ContinuousAssign> WireReadCounter::visit(
   return node;
 }
 
+std::unique_ptr<BlockingAssign> WireReadCounter::visit(
+    std::unique_ptr<BlockingAssign> node) {
+  node->value = this->visit(std::move(node->value));
+  return node;
+}
+
 std::unique_ptr<Port> AssignMapBuilder::visit(std::unique_ptr<Port> node) {
   std::string port_str = std::visit(
       [](auto&& value) -> std::string {
@@ -58,6 +64,17 @@ std::unique_ptr<Port> AssignMapBuilder::visit(std::unique_ptr<Port> node) {
   return node;
 }
 
+std::unique_ptr<BlockingAssign> AssignMapBuilder::visit(
+    std::unique_ptr<BlockingAssign> node) {
+  node = Transformer::visit(std::move(node));
+  std::string key =
+      std::visit([](auto&& value) -> std::string { return value->toString(); },
+                 node->target);
+  this->assign_map[key] = node->value->clone();
+  this->assign_count[key]++;
+  return node;
+}
+
 std::unique_ptr<ContinuousAssign> AssignMapBuilder::visit(
     std::unique_ptr<ContinuousAssign> node) {
   node = Transformer::visit(std::move(node));
@@ -136,6 +153,30 @@ std::unique_ptr<ContinuousAssign> AssignInliner::visit(
   return node;
 }
 
+std::unique_ptr<BlockingAssign> AssignInliner::visit(
+    std::unique_ptr<BlockingAssign> node) {
+  node->value = this->visit(std::move(node->value));
+  std::string key =
+      std::visit([](auto&& value) -> std::string { return value->toString(); },
+                 node->target);
+  bool remove = false;
+  std::visit(
+      [&](auto&& value) {
+        if (auto ptr = dynamic_cast<Identifier*>(value.get())) {
+          if (this->can_inline(key) && this->non_input_ports.count(key) == 0) {
+            remove = true;
+          } else if (this->inlined_outputs.count(ptr->toString())) {
+            remove = true;
+          };
+        }
+      },
+      node->target);
+  if (remove) {
+    return std::unique_ptr<BlockingAssign>{};
+  }
+  return node;
+}
+
 std::vector<std::variant<std::unique_ptr<StructuralStatement>,
                          std::unique_ptr<Declaration>>>
 AssignInliner::do_inline(
diff --git a/src/transformer.cpp b/src/transformer.cpp
index 84c3648..9d13e1d 100644
--- a/src/transformer.cpp
+++ b/src/transformer.cpp
@@ -342,9 +342,7 @@ std::unique_ptr<Always> Transformer::visit(std::unique_ptr<Always> node) {
     new_sensitivity_list.push_back(this->visit(std::move(item)));
   }
   node->sensitivity_list = std::move(new_sensitivity_list);
-  std::vector<std::variant<std::unique_ptr<BehavioralStatement>,
-                           std::unique_ptr<Declaration>>>
-      new_body;
+  std::vector<std::unique_ptr<BehavioralStatement>> new_body;
   for (auto&& item : node->body) {
     new_body.push_back(this->visit(std::move(item)));
   }
diff --git a/src/verilogAST.cpp b/src/verilogAST.cpp
index 2ca3c53..3cc3b4c 100644
--- a/src/verilogAST.cpp
+++ b/src/verilogAST.cpp
@@ -414,9 +414,7 @@ std::string Always::toString() {
 
   // emit body
   for (auto &statement : body) {
-    always_str += variant_to_string<std::unique_ptr<BehavioralStatement>,
-                                    std::unique_ptr<Declaration>>(statement) +
-                  "\n";
+    always_str += statement->toString() + "\n";
   }
 
   always_str += "end\n";
diff --git a/tests/assign_inliner.cpp b/tests/assign_inliner.cpp
index 5d32d4b..1a5bc70 100644
--- a/tests/assign_inliner.cpp
+++ b/tests/assign_inliner.cpp
@@ -51,6 +51,10 @@ TEST(InlineAssignTests, TestBasic) {
       std::make_unique<vAST::Identifier>("o_vec"),
       std::make_unique<vAST::Identifier>("x_vec")));
 
+  std::vector<std::variant<std::unique_ptr<vAST::BehavioralStatement>,
+                           std::unique_ptr<vAST::Declaration>>>
+      always_body;
+
   std::unique_ptr<vAST::AbstractModule> module = std::make_unique<vAST::Module>(
       "test_module", std::move(ports), std::move(body));
 
@@ -540,14 +544,14 @@ TEST(InlineAssignTests, TestMultipleAssign) {
 TEST(InlineAssignTests, TestInstConn) {
   // Should not inline a wire that is assigned multiple times
   std::vector<std::unique_ptr<vAST::AbstractPort>> ports;
-  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("i"),
-                                               vAST::INPUT, vAST::WIRE));
-  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("a"),
-                                               vAST::INPUT, vAST::WIRE));
-  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("o"),
-                                               vAST::OUTPUT, vAST::WIRE));
-  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("b"),
-                                               vAST::OUTPUT, vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("i"), vAST::INPUT,
+                                               vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("a"), vAST::INPUT,
+                                               vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("o"), vAST::OUTPUT,
+                                               vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("b"), vAST::OUTPUT,
+                                               vAST::WIRE));
 
   std::vector<std::variant<std::unique_ptr<vAST::StructuralStatement>,
                            std::unique_ptr<vAST::Declaration>>>
@@ -568,7 +572,8 @@ TEST(InlineAssignTests, TestInstConn) {
       std::make_unique<vAST::Identifier>("a")));
 
   vAST::Parameters parameters;
-  std::unique_ptr<vAST::Connections> connections = std::make_unique<vAST::Connections>();
+  std::unique_ptr<vAST::Connections> connections =
+      std::make_unique<vAST::Connections>();
   connections->insert("c", vAST::make_id("a"));
   connections->insert("i", vAST::make_id("x"));
   connections->insert("o", vAST::make_id("y"));
@@ -626,12 +631,12 @@ TEST(InlineAssignTests, TestInstConn) {
 TEST(InlineAssignTests, TestNoInlineIndex) {
   std::vector<std::unique_ptr<vAST::AbstractPort>> ports;
   ports.push_back(std::make_unique<vAST::Port>(
-      std::make_unique<vAST::Vector>(vAST::make_id("i1"),
-                                     vAST::make_num("1"), vAST::make_num("0")),
+      std::make_unique<vAST::Vector>(vAST::make_id("i1"), vAST::make_num("1"),
+                                     vAST::make_num("0")),
       vAST::INPUT, vAST::WIRE));
   ports.push_back(std::make_unique<vAST::Port>(
-      std::make_unique<vAST::Vector>(vAST::make_id("i2"),
-                                     vAST::make_num("1"), vAST::make_num("0")),
+      std::make_unique<vAST::Vector>(vAST::make_id("i2"), vAST::make_num("1"),
+                                     vAST::make_num("0")),
       vAST::INPUT, vAST::WIRE));
   ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("o"), vAST::OUTPUT,
                                                vAST::WIRE));
diff --git a/tests/basic.cpp b/tests/basic.cpp
index c788a3d..49c2aee 100644
--- a/tests/basic.cpp
+++ b/tests/basic.cpp
@@ -487,9 +487,7 @@ TEST(BasicTests, TestAlwaysEmpty) {
       std::unique_ptr<vAST::Identifier>, std::unique_ptr<vAST::PosEdge>,
       std::unique_ptr<vAST::NegEdge>, std::unique_ptr<vAST::Star>>>
       sensitivity_list;
-  std::vector<std::variant<std::unique_ptr<vAST::BehavioralStatement>,
-                           std::unique_ptr<vAST::Declaration>>>
-      body;
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> body;
 
   ASSERT_THROW(
       vAST::Always always_empty(std::move(sensitivity_list), std::move(body)),
diff --git a/tests/common.cpp b/tests/common.cpp
index 5be5f64..d3b8888 100644
--- a/tests/common.cpp
+++ b/tests/common.cpp
@@ -49,12 +49,9 @@ make_simple_body() {
   return body;
 }
 
-std::vector<std::variant<std::unique_ptr<vAST::BehavioralStatement>,
-                         std::unique_ptr<vAST::Declaration>>>
+std::vector<std::unique_ptr<vAST::BehavioralStatement>>
 make_simple_always_body() {
-  std::vector<std::variant<std::unique_ptr<vAST::BehavioralStatement>,
-                           std::unique_ptr<vAST::Declaration>>>
-      body;
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> body;
   body.push_back(std::make_unique<vAST::BlockingAssign>(
       std::make_unique<vAST::Identifier>("a"),
       std::make_unique<vAST::Identifier>("b")));
diff --git a/tests/transformer.cpp b/tests/transformer.cpp
index 1de8616..3015f2e 100644
--- a/tests/transformer.cpp
+++ b/tests/transformer.cpp
@@ -122,9 +122,8 @@ TEST(TransformerTests, TestAlways) {
       std::make_unique<vAST::NegEdge>(std::make_unique<vAST::Identifier>("c")));
   sensitivity_list.push_back(std::make_unique<vAST::Star>());
 
-  std::vector<std::variant<std::unique_ptr<vAST::BehavioralStatement>,
-                           std::unique_ptr<vAST::Declaration>>>
-      always_body = make_simple_always_body();
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> always_body =
+      make_simple_always_body();
   always_body.push_back(
       std::make_unique<vAST::SingleLineComment>("Test comment"));
   always_body.push_back(

From 3d5f338679c1df701de2f5e3705956c262057cd5 Mon Sep 17 00:00:00 2001
From: Lenny Truong <lenny@stanford.edu>
Date: Fri, 5 Jun 2020 10:13:47 -0700
Subject: [PATCH 6/7] Refactor code to avoid duplication

---
 include/verilogAST/assign_inliner.hpp |  9 ++++
 src/assign_inliner.cpp                | 62 +++++++++++----------------
 2 files changed, 34 insertions(+), 37 deletions(-)

diff --git a/include/verilogAST/assign_inliner.hpp b/include/verilogAST/assign_inliner.hpp
index 1e4e7e1..f0435b4 100644
--- a/include/verilogAST/assign_inliner.hpp
+++ b/include/verilogAST/assign_inliner.hpp
@@ -14,6 +14,9 @@ class AssignMapBuilder : public Transformer {
   std::set<std::string> &output_ports;
   std::set<std::string> &input_ports;
 
+  template <typename T>
+  std::unique_ptr<T> process_assign(std::unique_ptr<T> node);
+
  public:
   AssignMapBuilder(
       std::map<std::string, int> &assign_count,
@@ -39,6 +42,9 @@ class WireReadCounter : public Transformer {
   //
   std::map<std::string, int> &read_count;
 
+  template <typename T>
+  std::unique_ptr<T> process_assign(std::unique_ptr<T> node);
+
  public:
   WireReadCounter(std::map<std::string, int> &read_count)
       : read_count(read_count){};
@@ -92,6 +98,9 @@ class AssignInliner : public Transformer {
 
   bool can_inline(std::string key);
 
+  template <typename T>
+  std::unique_ptr<T> process_assign(std::unique_ptr<T> node);
+
  public:
   AssignInliner() : wire_blacklist(){};
   explicit AssignInliner(std::set<std::string> wire_blacklist)
diff --git a/src/assign_inliner.cpp b/src/assign_inliner.cpp
index 2a0107c..f525196 100644
--- a/src/assign_inliner.cpp
+++ b/src/assign_inliner.cpp
@@ -29,16 +29,20 @@ std::unique_ptr<Declaration> WireReadCounter::visit(
   return node;
 }
 
-std::unique_ptr<ContinuousAssign> WireReadCounter::visit(
-    std::unique_ptr<ContinuousAssign> node) {
+template <typename T>
+std::unique_ptr<T> WireReadCounter::process_assign(std::unique_ptr<T> node) {
   node->value = this->visit(std::move(node->value));
   return node;
 }
 
+std::unique_ptr<ContinuousAssign> WireReadCounter::visit(
+    std::unique_ptr<ContinuousAssign> node) {
+  return this->process_assign(std::move(node));
+}
+
 std::unique_ptr<BlockingAssign> WireReadCounter::visit(
     std::unique_ptr<BlockingAssign> node) {
-  node->value = this->visit(std::move(node->value));
-  return node;
+  return this->process_assign(std::move(node));
 }
 
 std::unique_ptr<Port> AssignMapBuilder::visit(std::unique_ptr<Port> node) {
@@ -63,9 +67,8 @@ std::unique_ptr<Port> AssignMapBuilder::visit(std::unique_ptr<Port> node) {
   }
   return node;
 }
-
-std::unique_ptr<BlockingAssign> AssignMapBuilder::visit(
-    std::unique_ptr<BlockingAssign> node) {
+template <typename T>
+std::unique_ptr<T> AssignMapBuilder::process_assign(std::unique_ptr<T> node) {
   node = Transformer::visit(std::move(node));
   std::string key =
       std::visit([](auto&& value) -> std::string { return value->toString(); },
@@ -75,15 +78,14 @@ std::unique_ptr<BlockingAssign> AssignMapBuilder::visit(
   return node;
 }
 
+std::unique_ptr<BlockingAssign> AssignMapBuilder::visit(
+    std::unique_ptr<BlockingAssign> node) {
+  return this->process_assign(std::move(node));
+}
+
 std::unique_ptr<ContinuousAssign> AssignMapBuilder::visit(
     std::unique_ptr<ContinuousAssign> node) {
-  node = Transformer::visit(std::move(node));
-  std::string key =
-      std::visit([](auto&& value) -> std::string { return value->toString(); },
-                 node->target);
-  this->assign_map[key] = node->value->clone();
-  this->assign_count[key]++;
-  return node;
+  return this->process_assign(std::move(node));
 }
 
 bool AssignInliner::can_inline(std::string key) {
@@ -129,8 +131,8 @@ std::unique_ptr<Wire> AssignInliner::visit(std::unique_ptr<Wire> node) {
   return node;
 }
 
-std::unique_ptr<ContinuousAssign> AssignInliner::visit(
-    std::unique_ptr<ContinuousAssign> node) {
+template <typename T>
+std::unique_ptr<T> AssignInliner::process_assign(std::unique_ptr<T> node) {
   node->value = this->visit(std::move(node->value));
   std::string key =
       std::visit([](auto&& value) -> std::string { return value->toString(); },
@@ -148,33 +150,19 @@ std::unique_ptr<ContinuousAssign> AssignInliner::visit(
       },
       node->target);
   if (remove) {
-    return std::unique_ptr<ContinuousAssign>{};
+    return std::unique_ptr<T>{};
   }
   return node;
 }
 
+std::unique_ptr<ContinuousAssign> AssignInliner::visit(
+    std::unique_ptr<ContinuousAssign> node) {
+  return this->process_assign(std::move(node));
+}
+
 std::unique_ptr<BlockingAssign> AssignInliner::visit(
     std::unique_ptr<BlockingAssign> node) {
-  node->value = this->visit(std::move(node->value));
-  std::string key =
-      std::visit([](auto&& value) -> std::string { return value->toString(); },
-                 node->target);
-  bool remove = false;
-  std::visit(
-      [&](auto&& value) {
-        if (auto ptr = dynamic_cast<Identifier*>(value.get())) {
-          if (this->can_inline(key) && this->non_input_ports.count(key) == 0) {
-            remove = true;
-          } else if (this->inlined_outputs.count(ptr->toString())) {
-            remove = true;
-          };
-        }
-      },
-      node->target);
-  if (remove) {
-    return std::unique_ptr<BlockingAssign>{};
-  }
-  return node;
+  return this->process_assign(std::move(node));
 }
 
 std::vector<std::variant<std::unique_ptr<StructuralStatement>,

From d6af7b92592ccc11f6608b125c12a0f839249de6 Mon Sep 17 00:00:00 2001
From: Lenny Truong <lenny@stanford.edu>
Date: Fri, 5 Jun 2020 10:24:22 -0700
Subject: [PATCH 7/7] Add if statement inline test

---
 tests/assign_inliner.cpp | 84 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 81 insertions(+), 3 deletions(-)

diff --git a/tests/assign_inliner.cpp b/tests/assign_inliner.cpp
index 1a5bc70..63729c4 100644
--- a/tests/assign_inliner.cpp
+++ b/tests/assign_inliner.cpp
@@ -16,6 +16,8 @@ TEST(InlineAssignTests, TestBasic) {
                                                vAST::WIRE));
   ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("o"), vAST::OUTPUT,
                                                vAST::WIRE));
+  ports.push_back(std::make_unique<vAST::Port>(vAST::make_id("o1"),
+                                               vAST::OUTPUT, vAST::WIRE));
   ports.push_back(std::make_unique<vAST::Port>(
       std::make_unique<vAST::Vector>(vAST::make_id("o_vec"),
                                      vAST::make_num("1"), vAST::make_num("0")),
@@ -51,9 +53,49 @@ TEST(InlineAssignTests, TestBasic) {
       std::make_unique<vAST::Identifier>("o_vec"),
       std::make_unique<vAST::Identifier>("x_vec")));
 
-  std::vector<std::variant<std::unique_ptr<vAST::BehavioralStatement>,
-                           std::unique_ptr<vAST::Declaration>>>
-      always_body;
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> always_body;
+
+  std::vector<std::unique_ptr<vAST::BehavioralStatement>> true_body;
+  true_body.push_back(std::make_unique<vAST::BlockingAssign>(
+      std::make_unique<vAST::Identifier>("o1"),
+      std::make_unique<vAST::Identifier>("x")));
+
+  std::vector<
+      std::pair<std::unique_ptr<vAST::Expression>,
+                std::vector<std::unique_ptr<vAST::BehavioralStatement>>>>
+      else_ifs;
+  for (int i = 1; i < 3; i++) {
+    std::unique_ptr<vAST::Expression> cond = std::make_unique<vAST::BinaryOp>(
+        std::make_unique<vAST::Identifier>("x"), vAST::BinOp::EQ,
+        vAST::make_num(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>("o1"),
+        std::make_unique<vAST::BinaryOp>(
+            std::make_unique<vAST::Identifier>("x"), vAST::BinOp::ADD,
+            vAST::make_num(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>("o1"),
+      std::make_unique<vAST::BinaryOp>(std::make_unique<vAST::Identifier>("x"),
+                                       vAST::BinOp::ADD, vAST::make_num("3"))));
+
+  always_body.push_back(std::make_unique<vAST::If>(
+      std::make_unique<vAST::BinaryOp>(std::make_unique<vAST::Identifier>("x"),
+                                       vAST::BinOp::EQ, vAST::make_num("0")),
+      std::move(true_body), std::move(else_ifs), std::move(else_body)));
+
+  std::vector<std::variant<
+      std::unique_ptr<vAST::Identifier>, std::unique_ptr<vAST::PosEdge>,
+      std::unique_ptr<vAST::NegEdge>, std::unique_ptr<vAST::Star>>>
+      sensitivity_list;
+  sensitivity_list.push_back(std::make_unique<vAST::Star>());
+
+  body.push_back(std::make_unique<vAST::Always>(std::move(sensitivity_list),
+                                                std::move(always_body)));
 
   std::unique_ptr<vAST::AbstractModule> module = std::make_unique<vAST::Module>(
       "test_module", std::move(ports), std::move(body));
@@ -62,6 +104,7 @@ TEST(InlineAssignTests, TestBasic) {
       "module test_module (\n"
       "    input i,\n"
       "    output o,\n"
+      "    output o1,\n"
       "    output [1:0] o_vec\n"
       ");\n"
       "wire x;\n"
@@ -70,6 +113,17 @@ TEST(InlineAssignTests, TestBasic) {
       "assign o = x;\n"
       "assign x_vec = {i,i};\n"
       "assign o_vec = x_vec;\n"
+      "always @(*) begin\n"
+      "if (x == 0) begin\n"
+      "    o1 = x;\n"
+      "end else if (x == 1) begin\n"
+      "    o1 = x + 1;\n"
+      "end else if (x == 2) begin\n"
+      "    o1 = x + 2;\n"
+      "end else begin\n"
+      "    o1 = x + 3;\n"
+      "end\n"
+      "end\n\n"
       "endmodule\n";
 
   EXPECT_EQ(module->toString(), raw_str);
@@ -83,12 +137,24 @@ TEST(InlineAssignTests, TestBasic) {
       "module test_module (\n"
       "    input i,\n"
       "    output o,\n"
+      "    output o1,\n"
       "    output [1:0] o_vec\n"
       ");\n"
       "wire [1:0] x_vec;\n"
       "assign o = i;\n"
       "assign x_vec = {i,i};\n"
       "assign o_vec = x_vec;\n"
+      "always @(*) begin\n"
+      "if (i == 0) begin\n"
+      "    o1 = i;\n"
+      "end else if (i == 1) begin\n"
+      "    o1 = i + 1;\n"
+      "end else if (i == 2) begin\n"
+      "    o1 = i + 2;\n"
+      "end else begin\n"
+      "    o1 = i + 3;\n"
+      "end\n"
+      "end\n\n"
       "endmodule\n";
 
   blacklist = {"x_vec"};
@@ -100,10 +166,22 @@ TEST(InlineAssignTests, TestBasic) {
       "module test_module (\n"
       "    input i,\n"
       "    output o,\n"
+      "    output o1,\n"
       "    output [1:0] o_vec\n"
       ");\n"
       "assign o = i;\n"
       "assign o_vec = {i,i};\n"
+      "always @(*) begin\n"
+      "if (i == 0) begin\n"
+      "    o1 = i;\n"
+      "end else if (i == 1) begin\n"
+      "    o1 = i + 1;\n"
+      "end else if (i == 2) begin\n"
+      "    o1 = i + 2;\n"
+      "end else begin\n"
+      "    o1 = i + 3;\n"
+      "end\n"
+      "end\n\n"
       "endmodule\n";
 
   vAST::AssignInliner transformer;