From 791ffa01eb314557fb160c7c7ab92c93c3443a24 Mon Sep 17 00:00:00 2001
From: Ignacio Encinas <ignacio@iencinas.com>
Date: Tue, 19 Nov 2024 20:14:02 +0100
Subject: [PATCH] verilog: CST: add getters for NetVariableAssignment

---
 verilog/CST/statement.cc      | 13 +++++
 verilog/CST/statement.h       | 10 ++++
 verilog/CST/statement_test.cc | 93 +++++++++++++++++++++++++++++++++++
 3 files changed, 116 insertions(+)

diff --git a/verilog/CST/statement.cc b/verilog/CST/statement.cc
index c19cd681c..b9aa6d820 100644
--- a/verilog/CST/statement.cc
+++ b/verilog/CST/statement.cc
@@ -552,4 +552,17 @@ const verible::SyntaxTreeNode *GetAssignModifyLhs(
   return verible::GetSubtreeAsNode(assign_modify,
                                    NodeEnum::kAssignModifyStatement, 0);
 }
+
+const verible::SyntaxTreeNode *GetNetVariableAssignmentLhs(
+    const verible::SyntaxTreeNode &assignment) {
+  return verible::GetSubtreeAsNode(assignment, NodeEnum::kNetVariableAssignment,
+                                   0);
+}
+
+const verible::SyntaxTreeLeaf *GetNetVariableAssignmentOperator(
+    const verible::SyntaxTreeNode &assignment) {
+  return verible::GetSubtreeAsLeaf(assignment, NodeEnum::kNetVariableAssignment,
+                                   1);
+}
+
 }  // namespace verilog
diff --git a/verilog/CST/statement.h b/verilog/CST/statement.h
index 4d3fffe85..367d8297d 100644
--- a/verilog/CST/statement.h
+++ b/verilog/CST/statement.h
@@ -237,6 +237,16 @@ const verible::SyntaxTreeNode *GetAssignModifyRhs(
 const verible::SyntaxTreeNode *GetAssignModifyLhs(
     const verible::SyntaxTreeNode &assign_modify);
 
+// Return the left hand side (Lhs) from a NetVariableAssignment
+// Example: get 'x' from 'x = y'
+const verible::SyntaxTreeNode *GetNetVariableAssignmentLhs(
+    const verible::SyntaxTreeNode &assignment);
+
+// Return the operator from a NetVariableAssignment
+// Example: get '=' from 'x = y'
+const verible::SyntaxTreeLeaf *GetNetVariableAssignmentOperator(
+    const verible::SyntaxTreeNode &assignment);
+
 }  // namespace verilog
 
 #endif  // VERIBLE_VERILOG_CST_STATEMENT_H_
diff --git a/verilog/CST/statement_test.cc b/verilog/CST/statement_test.cc
index 215498a4b..63f2f1038 100644
--- a/verilog/CST/statement_test.cc
+++ b/verilog/CST/statement_test.cc
@@ -1745,5 +1745,98 @@ TEST(GetAssignModifyRhs, Various) {
   }
 }
 
+TEST(GetNetVariableAssignmentLhs, Various) {
+  constexpr int kTag = 1;  // value doesn't matter
+  const SyntaxTreeSearchTestCase kTestCases[] = {
+      {""},
+      {"module m;\nendmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n",
+       {kTag, "k"},
+       " =  1;\nend\n",
+       "endmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n",
+       "k &= 1;\nend\n"
+       "endmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n"
+       "k |= 1;\nend\n"
+       "endmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_ff begin\n"
+       "k <= 1;\nend\n"
+       "endmodule\n"},
+  };
+  for (const auto &test : kTestCases) {
+    TestVerilogSyntaxRangeMatches(
+        __FUNCTION__, test, [](const TextStructureView &text_structure) {
+          const auto &root = text_structure.SyntaxTree();
+          const auto &net_var_assignments = SearchSyntaxTree(
+              *ABSL_DIE_IF_NULL(root), NodekNetVariableAssignment());
+
+          std::vector<TreeSearchMatch> left_hand_sides;
+          for (const auto &assignment : net_var_assignments) {
+            const auto *lhs = GetNetVariableAssignmentLhs(
+                verible::SymbolCastToNode(*assignment.match));
+            left_hand_sides.emplace_back(
+                TreeSearchMatch{lhs, {/* ignored context */}});
+          }
+          return left_hand_sides;
+        });
+  }
+}
+
+TEST(GetNetVariableAssignmentOperator, Various) {
+  constexpr int kTag = 1;  // value doesn't matter
+  const SyntaxTreeSearchTestCase kTestCases[] = {
+      {""},
+      {"module m;\nendmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n",
+       "k ",
+       {kTag, "="},
+       " 1;\nend\n",
+       "endmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n",
+       "k &= 1;\nend\n"
+       "endmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n"
+       "k |= 1;\nend\n"
+       "endmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_ff begin\n"
+       "k <= 1;\nend\n"
+       "endmodule\n"},
+  };
+  for (const auto &test : kTestCases) {
+    TestVerilogSyntaxRangeMatches(
+        __FUNCTION__, test, [](const TextStructureView &text_structure) {
+          const auto &root = text_structure.SyntaxTree();
+          const auto &net_var_assignments = SearchSyntaxTree(
+              *ABSL_DIE_IF_NULL(root), NodekNetVariableAssignment());
+
+          std::vector<TreeSearchMatch> left_hand_sides;
+          for (const auto &assignment : net_var_assignments) {
+            const auto *lhs = GetNetVariableAssignmentOperator(
+                verible::SymbolCastToNode(*assignment.match));
+            left_hand_sides.emplace_back(
+                TreeSearchMatch{lhs, {/* ignored context */}});
+          }
+          return left_hand_sides;
+        });
+  }
+}
+
 }  // namespace
 }  // namespace verilog