Allow more constructs in a generate loop.

Previously the generate loop contained a StatementBlock which limited what could be put inside of loops to statements. However, other constructs can be put in loops like always blocks, etc. This change expands what may be placed in the body of a GenerateLoop.

Author: meheff

xls/codegen/module_builder.cc

@@ -662,8 +662,8 @@ absl::Status ModuleBuilder::EmitArrayCopyAndUpdateViaGenerate(
       loop = module_->Add<GenerateLoop>(SourceInfo(), genvar_name, init, limit,
                                         std::move(label));
     } else {
-      loop = innermost_loop->body()->Add<GenerateLoop>(
-          SourceInfo(), genvar_name, init, limit, std::move(label));
+      loop = innermost_loop->Add<GenerateLoop>(SourceInfo(), genvar_name, init,
+                                               limit, std::move(label));
     }
     innermost_loop = loop;
     return loop;
@@ -750,8 +750,8 @@ absl::Status ModuleBuilder::EmitArrayCopyAndUpdateViaGenerate(
     assignment_section()->Add<ContinuousAssignment>(SourceInfo(), lhs_indexed,
                                                     assignment_rhs);
   } else {
-    innermost_loop->body()->Add<ContinuousAssignment>(SourceInfo(), lhs_indexed,
-                                                      assignment_rhs);
+    innermost_loop->Add<ContinuousAssignment>(SourceInfo(), lhs_indexed,
+                                              assignment_rhs);
   }
 
   return absl::OkStatus();

xls/codegen/vast/vast.cc

@@ -603,8 +603,7 @@ GenerateLoop::GenerateLoop(std::string_view genvar_name, Expression* init,
           loc, file->Make<GenvarDef>(loc, std::string(genvar_name)))),
       init_(init),
       limit_(limit),
-      label_(std::move(label)),
-      body_(file->Make<StatementBlock>(loc)) {}
+      label_(std::move(label)) {}
 
 std::string GenerateLoop::Emit(LineInfo* line_info) const {
   LineInfoStart(line_info, this);
@@ -618,13 +617,15 @@ std::string GenerateLoop::Emit(LineInfo* line_info) const {
       "for (genvar %s = %s; %s < %s; %s = %s + 1) begin%s", genvar_str,
       init_str, genvar_str, limit_str, genvar_str, genvar_str, label_suffix));
   LineInfoIncrease(line_info, 1);
-  for (Statement* statement : body_->statements()) {
-    std::string pre_emit = statement->PreEmit(line_info);
+  for (const ModuleMember& member : members_) {
+    std::string pre_emit =
+        absl::visit([=](auto* m) { return m->PreEmit(line_info); }, member);
     if (!pre_emit.empty()) {
       lines.push_back(Indent(pre_emit, kDefaultIndentSpaces));
       LineInfoIncrease(line_info, 1);
     }
-    lines.push_back(Indent(statement->Emit(line_info), kDefaultIndentSpaces));
+    lines.push_back(Indent(
+        absl::visit([=](auto* m) { return m->Emit(line_info); }, member)));
     LineInfoIncrease(line_info, 1);
   }
   lines.push_back("end");

xls/codegen/vast/vast.h

@@ -731,38 +731,6 @@ class StatementBlock final : public VastNode {
   std::vector<Statement*> statements_;
 };
 
-// Represents a generate loop construct. Example:
-// ```verilog
-// for (genvar i = 0; i < 32; i = i + 1) begin : gen_loop
-//   assign output[i] = input[i];
-// end
-// ```
-//
-class GenerateLoop final : public Statement {
- public:
-  GenerateLoop(std::string_view genvar_name, Expression* init,
-               Expression* limit, std::optional<std::string> label,
-               VerilogFile* file, const SourceInfo& loc);
-
-  LogicRef* genvar() const { return genvar_; }
-  Expression* init() const { return init_; }
-  Expression* limit() const { return limit_; }
-  const std::optional<std::string>& label() const { return label_; }
-  StatementBlock* body() const { return body_; }
-  absl::Span<Statement* const> statements() const {
-    return body_->statements();
-  }
-
-  std::string Emit(LineInfo* line_info) const final;
-
- private:
-  LogicRef* genvar_;
-  Expression* init_;
-  Expression* limit_;
-  std::optional<std::string> label_;
-  StatementBlock* body_;
-};
-
 // Similar to statement block,  but for use if `ifdef `else `endif blocks (no
 // "begin" or "end").
 class MacroStatementBlock final : public VastNode {
@@ -2268,6 +2236,7 @@ class VerilogFunctionCall final : public Expression {
 
 class ModuleSection;
 class ModuleConditionalDirective;
+class GenerateLoop;
 
 // Represents a member of a module.
 using ModuleMember =
@@ -2291,6 +2260,38 @@ using ModuleMember =
                  // at elaboration time
                  GenerateLoop*>;
 
+// Represents a generate loop construct. Example:
+// ```verilog
+// for (genvar i = 0; i < 32; i = i + 1) begin : gen_loop
+//   assign output[i] = input[i];
+// end
+// ```
+//
+class GenerateLoop final : public Statement {
+ public:
+  GenerateLoop(std::string_view genvar_name, Expression* init,
+               Expression* limit, std::optional<std::string> label,
+               VerilogFile* file, const SourceInfo& loc);
+
+  LogicRef* genvar() const { return genvar_; }
+  Expression* init() const { return init_; }
+  Expression* limit() const { return limit_; }
+  const std::optional<std::string>& label() const { return label_; }
+
+  template <typename T, typename... Args>
+  T* Add(const SourceInfo& loc, Args&&... args);
+  void AddMember(ModuleMember member) { members_.push_back(member); }
+
+  std::string Emit(LineInfo* line_info) const final;
+
+ private:
+  LogicRef* genvar_;
+  Expression* init_;
+  Expression* limit_;
+  std::optional<std::string> label_;
+  std::vector<ModuleMember> members_;
+};
+
 // A ModuleSection is a container of ModuleMembers used to organize the contents
 // of a module. A Module contains a single top-level ModuleSection which may
 // contain other ModuleSections. ModuleSections enables modules to be
@@ -2993,6 +2994,12 @@ inline T* VerilogPackageSection::Add(const SourceInfo& loc, Args&&... args) {
   return ptr;
 }
 
+template <typename T, typename... Args>
+inline T* GenerateLoop::Add(const SourceInfo& loc, Args&&... args) {
+  T* ptr = file()->Make<T>(loc, std::forward<Args>(args)...);
+  AddMember(ptr);
+  return ptr;
+}
 }  // namespace verilog
 }  // namespace xls
 

xls/codegen/vast/vast_test.cc

@@ -21,14 +21,14 @@
 #include <utility>
 #include <vector>
 
-#include "gmock/gmock.h"
-#include "gtest/gtest.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/types/span.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
 #include "xls/common/status/matchers.h"
 #include "xls/ir/bits.h"
 #include "xls/ir/fileno.h"
@@ -2834,8 +2834,11 @@ TEST_P(VastTest, SimpleGenerateLoop) {
       m->Add<GenerateLoop>(si, std::string("i"), f.PlainLiteral(0, si),
                            f.PlainLiteral(32, si), "gen_loop");
   LogicRef* i = generate_loop->genvar();
-  generate_loop->body()->Add<ContinuousAssignment>(
-      si, f.Make<Index>(si, output, i), f.Make<Index>(si, input, i));
+  generate_loop->Add<ContinuousAssignment>(si, f.Make<Index>(si, output, i),
+                                           f.Make<Index>(si, input, i));
+  generate_loop->Add<BlankLine>(si);
+  generate_loop->Add<Comment>(si, "This is a comment.");
+  generate_loop->Add<InlineVerilogStatement>(si, "inline_verilog_statement;");
 
   LineInfo line_info;
   EXPECT_EQ(m->Emit(&line_info),
@@ -2845,6 +2848,9 @@ TEST_P(VastTest, SimpleGenerateLoop) {
 );
   for (genvar i = 0; i < 32; i = i + 1) begin : gen_loop
     assign my_output[i] = my_input[i];
+
+    // This is a comment.
+    inline_verilog_statement;
   end
 endmodule)");
 }
@@ -2864,11 +2870,11 @@ TEST_P(VastTest, NestedGenerateLoop) {
       m->Add<GenerateLoop>(si, std::string("i"), f.PlainLiteral(0, si),
                            f.PlainLiteral(4, si), "outer");
   LogicRef* i = outer_loop->genvar();
-  auto* inner_loop = outer_loop->body()->Add<GenerateLoop>(
-      si, std::string("j"), f.PlainLiteral(0, si), f.PlainLiteral(3, si),
-      "inner");
+  auto* inner_loop =
+      outer_loop->Add<GenerateLoop>(si, std::string("j"), f.PlainLiteral(0, si),
+                                    f.PlainLiteral(3, si), "inner");
   LogicRef* j = inner_loop->genvar();
-  inner_loop->body()->Add<ContinuousAssignment>(
+  inner_loop->Add<ContinuousAssignment>(
       si, f.Make<Index>(si, f.Make<Index>(si, dst, i), j),
       f.Make<Index>(si, f.Make<Index>(si, src, i), j));
 

xls/public/c_api_symbols.txt

@@ -312,7 +312,16 @@ xls_vast_data_type_width
 xls_vast_data_type_width_as_int64
 xls_vast_def_get_data_type
 xls_vast_def_get_name
-xls_vast_generate_loop_get_body
+xls_vast_generate_loop_add_always_comb
+xls_vast_generate_loop_add_always_ff
+xls_vast_generate_loop_add_blank_line
+xls_vast_generate_loop_add_comment
+xls_vast_generate_loop_add_continuous_assignment
+xls_vast_generate_loop_add_generate_loop
+xls_vast_generate_loop_add_inline_verilog_statement
+xls_vast_generate_loop_add_instantiation
+xls_vast_generate_loop_add_localparam
+xls_vast_generate_loop_add_localparam_with_def
 xls_vast_generate_loop_get_genvar
 xls_vast_index_as_expression
 xls_vast_index_as_indexable_expression
@@ -328,8 +337,6 @@ xls_vast_slice_as_expression
 xls_vast_statement_block_add_blocking_assignment
 xls_vast_statement_block_add_case
 xls_vast_statement_block_add_conditional
-xls_vast_statement_block_add_continuous_assignment
-xls_vast_statement_block_add_generate_loop
 xls_vast_statement_block_add_nonblocking_assignment
 xls_vast_verilog_file_add_include
 xls_vast_verilog_file_add_module
@@ -346,6 +353,7 @@ xls_vast_verilog_file_make_def
 xls_vast_verilog_file_make_extern_package_type
 xls_vast_verilog_file_make_index
 xls_vast_verilog_file_make_index_i64
+xls_vast_verilog_file_make_inline_verilog_statement
 xls_vast_verilog_file_make_instantiation
 xls_vast_verilog_file_make_int_def
 xls_vast_verilog_file_make_int_type