verilog: CST: add getters for AssignModifyStatements

verilog/CST/statement.cc

@@ -535,4 +535,21 @@ const verible::SyntaxTreeNode *GetIfHeaderExpression(
   return verible::GetSubtreeAsNode(*paren_group, NodeEnum::kParenGroup, 1);
 }
 
+const verible::SyntaxTreeLeaf *GetAssignModifyOperator(
+    const verible::SyntaxTreeNode &assign_modify) {
+  return verible::GetSubtreeAsLeaf(assign_modify,
+                                   NodeEnum::kAssignModifyStatement, 1);
+}
+
+const verible::SyntaxTreeNode *GetAssignModifyRhs(
+    const verible::SyntaxTreeNode &assign_modify) {
+  return verible::GetSubtreeAsNode(assign_modify,
+                                   NodeEnum::kAssignModifyStatement, 2);
+}
+
+const verible::SyntaxTreeNode *GetAssignModifyLhs(
+    const verible::SyntaxTreeNode &assign_modify) {
+  return verible::GetSubtreeAsNode(assign_modify,
+                                   NodeEnum::kAssignModifyStatement, 0);
+}
 }  // namespace verilog

verilog/CST/statement.h

@@ -222,6 +222,21 @@ const verible::SyntaxTreeNode *GetIfClauseHeader(
 const verible::SyntaxTreeNode *GetIfHeaderExpression(
     const verible::SyntaxTreeNode &if_header);
 
+// Return the operator from an AssignModifyStatement
+// Example: get '&' from 'x &= y'
+const verible::SyntaxTreeLeaf *GetAssignModifyOperator(
+    const verible::SyntaxTreeNode &assign_modify);
+
+// Return the right hand side (Rhs) from an AssignModifyStatement
+// Example: get 'y' from 'x &= y'
+const verible::SyntaxTreeNode *GetAssignModifyRhs(
+    const verible::SyntaxTreeNode &assign_modify);
+
+// Return the left hand side (Lhs) from an AssignModifyStatement
+// Example: get 'x' from 'x &= y'
+const verible::SyntaxTreeNode *GetAssignModifyLhs(
+    const verible::SyntaxTreeNode &assign_modify);
+
 }  // namespace verilog
 
 #endif  // VERIBLE_VERILOG_CST_STATEMENT_H_

verilog/CST/statement_test.cc

@@ -1585,5 +1585,165 @@ TEST(GetIfClauseExpressionTest, Various) {
   }
 }
 
+TEST(GetAssignModifyOperator, 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 ",
+       {kTag, "|="},
+       " 1;\nend\n"
+       "endmodule\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_ff @(posedge clk) begin "
+       "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 &assign_modify_statements = SearchSyntaxTree(
+              *ABSL_DIE_IF_NULL(root), NodekAssignModifyStatement());
+
+          std::vector<TreeSearchMatch> operators;
+          for (const auto &block : assign_modify_statements) {
+            const auto *operator_ = GetAssignModifyOperator(
+                verible::SymbolCastToNode(*block.match));
+            operators.emplace_back(
+                TreeSearchMatch{operator_, {/* ignored context */}});
+          }
+          return operators;
+        });
+  }
+}
+
+TEST(GetAssignModifyLhs, 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 ",
+       {kTag, "k"},
+       " &=  1;\nend\n"
+       "endmodule\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",
+       {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"
+       "always_comb begin\n"
+       "reg 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 &assign_modify_statements = SearchSyntaxTree(
+              *ABSL_DIE_IF_NULL(root), NodekAssignModifyStatement());
+
+          std::vector<TreeSearchMatch> left_hand_sides;
+          for (const auto &block : assign_modify_statements) {
+            const auto *lhs =
+                GetAssignModifyLhs(verible::SymbolCastToNode(*block.match));
+            left_hand_sides.emplace_back(
+                TreeSearchMatch{lhs, {/* ignored context */}});
+          }
+          return left_hand_sides;
+        });
+  }
+}
+
+TEST(GetAssignModifyRhs, 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 "
+       "k &= ",
+       {kTag, "1"},
+       ";\nend\nendmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n"
+       "k |= ",
+       {kTag, "(1 + 2)"},
+       ";\nend\nendmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n"
+       "k += ",
+       {kTag, "1"},
+       ";\nend\nendmodule\n"},
+      {"module m;\n"
+       "reg k;\n"
+       "always_comb begin\n"
+       "k = 1;\nend\n"
+       "endmodule\n"},
+      {"module m;\n"
+       "always_comb begin\n"
+       "reg 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 &assign_modify_statements = SearchSyntaxTree(
+              *ABSL_DIE_IF_NULL(root), NodekAssignModifyStatement());
+
+          std::vector<TreeSearchMatch> right_hand_sides;
+          for (const auto &block : assign_modify_statements) {
+            const auto *rhs =
+                GetAssignModifyRhs(verible::SymbolCastToNode(*block.match));
+            right_hand_sides.emplace_back(
+                TreeSearchMatch{rhs, {/* ignored context */}});
+          }
+          return right_hand_sides;
+        });
+  }
+}
+
 }  // namespace
 }  // namespace verilog