[DSLX] Support aliased state types and channel arrays in impl-style procs.

xls/dslx/interp_value_utils.cc

@@ -1005,7 +1005,9 @@ absl::StatusOr<std::pair<InterpValue, InterpValue>> CreateChannelReferencePair(
       lhs_elements.push_back(lhs_rhs.first);
       rhs_elements.push_back(lhs_rhs.second);
     }
-    int64_t array_id = (*channel_instance_allocator)();
+    int64_t array_id = channel_instance_allocator.has_value()
+                           ? (*channel_instance_allocator)()
+                           : 0;
     return std::make_pair(
         InterpValue::MakeChannelArray(ChannelDirection::kOut, array_id,
                                       definer.has_value() ? *definer : nullptr,

xls/dslx/interp_value_utils_test.cc

@@ -1241,6 +1241,29 @@ TEST(InterpValueHelpersTest, GetLeafChannelReferences) {
               testing::ElementsAre(ch0, ch1, ch2, ch3));
 }
 
+TEST(InterpValueHelpersTest,
+     CreateChannelReferencePairWithoutAllocatorOnArrayType) {
+  auto chan_type = std::make_unique<ChannelType>(BitsType::MakeU32(),
+                                                 ChannelDirection::kOut);
+  auto inner_array_type =
+      std::make_unique<ArrayType>(std::move(chan_type), TypeDim::CreateU32(3));
+  ArrayType outer_array_type(std::move(inner_array_type),
+                             TypeDim::CreateU32(2));
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      auto pair,
+      CreateChannelReferencePair(&outer_array_type,
+                                 /*channel_instance_allocator=*/std::nullopt));
+  EXPECT_TRUE(pair.first.IsChannelArray());
+  EXPECT_TRUE(pair.second.IsChannelArray());
+  EXPECT_EQ(pair.first.GetChannelArrayOrDie().direction(),
+            ChannelDirection::kOut);
+  EXPECT_EQ(pair.second.GetChannelArrayOrDie().direction(),
+            ChannelDirection::kIn);
+  EXPECT_EQ(GetLeafChannelReferences(pair.first).size(), 6);
+  EXPECT_EQ(GetLeafChannelReferences(pair.second).size(), 6);
+}
+
 TEST(InterpValueHelpersTest, ValueToInterpValueSumRejectsInvalidTag) {
   FileTable file_table;
   Module module("test", /*fs_path=*/std::nullopt, file_table);

xls/dslx/ir_convert/function_converter.cc

@@ -3920,6 +3920,32 @@ absl::Status FunctionConverter::DefineProcDefChannelOrArrayIfLocal(
       /*name_def=*/std::nullopt, value, &direction_type, channel_decls);
 }
 
+namespace {
+
+// Returns all leaf `ChannelReference` values from a `ChannelReference` or
+// `ChannelArray` in the same index order produced by
+// `ChannelScope::CreateAllArrayElementSuffixes()`.
+std::vector<InterpValue> GetChannelArrayLeafReferencesInScopeOrder(
+    const InterpValue& channel_or_array) {
+  std::vector<InterpValue> current = {channel_or_array};
+  while (!current.empty() && current.front().IsChannelArray()) {
+    size_t dim_size = current.front().GetChannelArrayOrDie().elements().size();
+    std::vector<InterpValue> next;
+    next.reserve(dim_size * current.size());
+    for (size_t element_index = 0; element_index < dim_size; ++element_index) {
+      for (const InterpValue& item : current) {
+        const auto& elements = item.GetChannelArrayOrDie().elements();
+        CHECK_LT(element_index, elements.size());
+        next.push_back(elements[element_index]);
+      }
+    }
+    current = std::move(next);
+  }
+  return current;
+}
+
+}  // namespace
+
 void FunctionConverter::PopulateChannelOrArrayIdToObject(
     const InterpValue& value, const ChannelOrArray& channel_or_array) {
   channel_or_array_id_to_object_[GetChannelOrArrayId(value)] = channel_or_array;
@@ -3928,7 +3954,8 @@ void FunctionConverter::PopulateChannelOrArrayIdToObject(
   }
 
   ChannelArray* array = std::get<ChannelArray*>(channel_or_array);
-  std::vector<InterpValue> leaves = GetLeafChannelReferences(value);
+  std::vector<InterpValue> leaves =
+      GetChannelArrayLeafReferencesInScopeOrder(value);
   if (leaves.size() == array->channels().size()) {
     for (int i = 0; i < leaves.size(); ++i) {
       channel_or_array_id_to_object_
@@ -4031,13 +4058,12 @@ absl::Status FunctionConverter::InitProcDefChannels(
     const ProcDef* proc_def,
     const InterpValue::ProcInitializer& canonical_initializer,
     const TypeInfo* constructor_ti) {
-  ParametricEnv env;
   proc_def_channel_scope_ = std::make_unique<ProcScopedChannelScope>(
       package_data_.conversion_info, import_data_, options_,
       absl::down_cast<ProcBuilder*>(function_builder_.get()),
       [this](const Attr* attr) { return GetChannelArrayForAttr(attr); });
   proc_def_channel_scope_->EnterFunctionContext(current_type_info_,
-                                                std::move(env));
+                                                GetParametricEnv());
   channel_scope_ = proc_def_channel_scope_.get();
 
   XLS_ASSIGN_OR_RETURN(const Type* type,
@@ -4143,7 +4169,7 @@ absl::Status FunctionConverter::ExpandProcDefChannelReference(
     std::vector<ChannelInterface*>& out) {
   if (channel_or_array_value.IsChannelArray()) {
     for (const InterpValue& elem :
-         channel_or_array_value.GetChannelArrayOrDie().elements()) {
+         GetChannelArrayLeafReferencesInScopeOrder(channel_or_array_value)) {
       XLS_RETURN_IF_ERROR(ExpandProcDefChannelReference(elem, out));
     }
     return absl::OkStatus();
@@ -4318,6 +4344,7 @@ absl::Status FunctionConverter::ConvertProcDef(
 
   SetNodeToIr((*next_fn)->params()[0]->name_def(), instance.get());
   ScopedTypeInfoSwap stis_next(this, next_ti);
+  proc_def_channel_scope_->EnterFunctionContext(next_ti, env);
   XLS_RETURN_IF_ERROR(Visit((*next_fn)->body()));
 
   XLS_ASSIGN_OR_RETURN(

xls/dslx/ir_convert/ir_converter_test.cc

@@ -8744,5 +8744,81 @@ fn sum_enumerated(values: u8[3]) -> u32 {
       ConvertOneFunctionForTest(kProgram, "sum_enumerated"));
   ExpectIr(converted);
 }
+
+TEST_F(IrConverterTest, ProcDefWith2DChannelArrayAndParametricBounds) {
+  constexpr std::string_view kModule = R"(
+#![feature(explicit_state_access)]
+#![feature(generics)]
+
+proc Node<ROW: u32, COL: u32> {
+  from_west: chan<u32> in,
+  to_east: chan<u32> out,
+}
+
+impl Node<ROW, COL> {
+  fn new(from_west: chan<u32> in, to_east: chan<u32> out) -> Self {
+    Self { from_west, to_east }
+  }
+  fn next(self) {
+    let (tok, val) = recv(join(), self.from_west);
+    send(tok, self.to_east, val);
+  }
+}
+
+proc Grid<ROWS: u32, COLS: u32> {
+  from_wests: chan<u32>[COLS + 1][ROWS] in,
+  to_easts: chan<u32>[COLS + 1][ROWS] out,
+}
+
+impl Grid<ROWS, COLS> {
+  fn new() -> Self {
+    let (to_easts, from_wests) = chan<u32>[COLS + 1][ROWS]("east_west");
+    unroll_for! (row, _): (u32, ()) in u32:0..ROWS {
+      unroll_for! (col, _): (u32, ()) in u32:0..COLS {
+        Node<row, col>::new(from_wests[row][col], to_easts[row][col + 1]).spawn();
+      }(());
+    }(());
+    Self { from_wests, to_easts }
+  }
+  fn next(self) {
+    const FIRST_COL: u32 = 0;
+    const LAST_COL = COLS;
+    unroll_for! (row, _): (u32, ()) in u32:0..ROWS {
+      send(join(), self.to_easts[row][FIRST_COL], row);
+      recv(join(), self.from_wests[row][LAST_COL]);
+    }(());
+  }
+}
+
+proc Top {}
+impl Top {
+  fn new() -> Self {
+    Grid<2, 2>::new().spawn();
+    Self {}
+  }
+  fn next(self) {}
+}
+)";
+
+  ConvertOptions options = kProcScopedChannelOptions;
+  XLS_ASSERT_OK_AND_ASSIGN(std::string converted,
+                           ConvertModuleForTest(kModule, options));
+  EXPECT_THAT(converted,
+              HasSubstr("proc_instantiation __test_module__Node__0_0_next_inst("
+                        "_east_west__0_0, _east_west__0_1, "
+                        "proc=__test_module__Node__0_0_next)"));
+  EXPECT_THAT(converted,
+              HasSubstr("proc_instantiation __test_module__Node__1_0_next_inst("
+                        "_east_west__0_1, _east_west__0_2, "
+                        "proc=__test_module__Node__1_0_next)"));
+  EXPECT_THAT(converted,
+              HasSubstr("proc_instantiation __test_module__Node__0_1_next_inst("
+                        "_east_west__1_0, _east_west__1_1, "
+                        "proc=__test_module__Node__0_1_next)"));
+  EXPECT_THAT(converted,
+              HasSubstr("proc_instantiation __test_module__Node__1_1_next_inst("
+                        "_east_west__1_1, _east_west__1_2, "
+                        "proc=__test_module__Node__1_1_next)"));
+}
 }  // namespace
 }  // namespace xls::dslx

xls/dslx/ir_convert/testdata/ir_converter_test_ProcDefTaking2DChannelArrayAndDealingOutSubArrays.ir

@@ -30,8 +30,8 @@ proc __test_module__MiddleProc_next<_ins__0_0: bits[32] in, _ins__1_0: bits[32]
   chan_interface _ins__1_0(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
   chan_interface _ins__0_1(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
   chan_interface _ins__1_1(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
-  proc_instantiation __test_module__InnerProc_next_inst(_ins__0_0, _ins__1_0, proc=__test_module__InnerProc_next)
-  proc_instantiation __test_module__InnerProc_next_inst_1(_ins__0_1, _ins__1_1, proc=__test_module__InnerProc_next)
+  proc_instantiation __test_module__InnerProc_next_inst(_ins__0_0, _ins__0_1, proc=__test_module__InnerProc_next)
+  proc_instantiation __test_module__InnerProc_next_inst_1(_ins__1_0, _ins__1_1, proc=__test_module__InnerProc_next)
   __token: token = literal(value=token, id=13)
   literal.14: bits[1] = literal(value=1, id=14)
   tuple.15: () = tuple(id=15)

xls/dslx/ir_convert/testdata/ir_converter_test_ProcDefWith2DChannelArrayDealingOutSubArrays.ir

@@ -104,8 +104,8 @@ proc __test_module__TopProc_next<>() {
   chan _reply__2__1_1(bits[32], id=23, kind=streaming, ops=send_receive, flow_control=ready_valid, strictness=proven_mutually_exclusive)
   chan_interface _reply__2__1_1(direction=send, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
   chan_interface _reply__2__1_1(direction=receive, kind=streaming, strictness=proven_mutually_exclusive, flow_control=ready_valid, flop_kind=none)
-  proc_instantiation __test_module__SomeProc_next_inst(_ins_outs__1__0_0, _ins_outs__1__1_0, _reply__1__0_0, _reply__1__1_0, proc=__test_module__SomeProc_next)
-  proc_instantiation __test_module__SomeProc_next_inst_1(_ins_outs__2__0_1, _ins_outs__2__1_1, _reply__2__0_1, _reply__2__1_1, proc=__test_module__SomeProc_next)
+  proc_instantiation __test_module__SomeProc_next_inst(_ins_outs__1__0_0, _ins_outs__1__0_1, _reply__1__0_0, _reply__1__0_1, proc=__test_module__SomeProc_next)
+  proc_instantiation __test_module__SomeProc_next_inst_1(_ins_outs__2__1_0, _ins_outs__2__1_1, _reply__2__1_0, _reply__2__1_1, proc=__test_module__SomeProc_next)
   i: bits[32] = literal(value=0, id=25)
   literal.28: bits[32] = literal(value=2, id=28)
   literal.48: bits[32] = literal(value=2, id=48)

xls/dslx/translators/dslx_to_verilog_main.cc

@@ -141,10 +141,11 @@ absl::Status RealMain(absl::Span<const std::string_view> paths) {
         ParseAndTypecheck(text, path, module_name, &import_data, {}));
 
     for (const auto& def : tm.module->GetTypeDefinitions()) {
-      // Ignore private type definitions.
+      // Ignore private type definitions and proc definitions.
       XLS_ASSIGN_OR_RETURN(const TypeInfo::TypeSource type_definition_source,
                            tm.type_info->ResolveTypeDefinition(def));
-      if (!TypeDefinitionSourceIsPublic(type_definition_source)) {
+      if (!TypeDefinitionSourceIsPublic(type_definition_source) ||
+          std::holds_alternative<ProcDef*>(type_definition_source.definition)) {
         continue;
       }
       AstNode* def_node = TypeDefinitionToAstNode(def);

xls/dslx/type_system_v2/typecheck_module_v2_proc_test.cc

@@ -1628,5 +1628,92 @@ impl Top {
   XLS_EXPECT_OK(TypecheckV2(kProgram, "main", &import_data));
 }
 
+TEST(TypecheckV2ProcTest, ProcDefWithArrayTypeAliasState) {
+  EXPECT_THAT(R"(
+#![feature(explicit_state_access)]
+
+pub const MEM_SIZE = u32:256;
+pub type MemWord = u32;
+type State = MemWord[MEM_SIZE];
+
+pub proc MemoryProc {
+    addr_in: chan<u8> in,
+    val_out: chan<MemWord> out,
+    state: State,
+}
+
+impl MemoryProc {
+    pub fn new(addr_in: chan<u8> in, val_out: chan<MemWord> out) -> Self {
+        Self { addr_in, val_out, state: zero!<State>() }
+    }
+
+    pub fn next(self) {
+        let state = read(self.state);
+        let (tok, addr) = recv(join(), self.addr_in);
+        let word = state[addr];
+        send(tok, self.val_out, word);
+        let next_state = update(state, addr, word + u32:1);
+        write(self.state, next_state);
+    }
+}
+)",
+              TypecheckSucceeds(::testing::_));
+}
+
+TEST(TypecheckV2ProcTest, ProcDefWith2DChannelArrayInConstructor) {
+  EXPECT_THAT(R"(
+#![feature(explicit_state_access)]
+#![feature(generics)]
+
+proc Node<ROW: u32, COL: u32> {
+    from_west: chan<u32> in,
+    to_east: chan<u32> out,
+}
+
+impl Node<ROW, COL> {
+    fn new(from_west: chan<u32> in, to_east: chan<u32> out) -> Self {
+        Self { from_west, to_east }
+    }
+    fn next(self) {
+        let (tok, val) = recv(join(), self.from_west);
+        send(tok, self.to_east, val);
+    }
+}
+
+proc Grid<ROWS: u32, COLS: u32> {
+    from_wests: chan<u32>[COLS + u32:1][ROWS] in,
+    to_easts: chan<u32>[COLS + u32:1][ROWS] out,
+}
+
+impl Grid<ROWS, COLS> {
+    fn new() -> Self {
+        let (to_easts, from_wests) = chan<u32>[COLS + u32:1][ROWS]("east_west");
+        unroll_for! (row, _): (u32, ()) in u32:0..ROWS {
+            unroll_for! (col, _): (u32, ()) in u32:0..COLS {
+                Node<row, col>::new(from_wests[row][col], to_easts[row][col + u32:1]).spawn();
+            }(());
+        }(());
+        Self { from_wests, to_easts }
+    }
+    fn next(self) {
+        const FIRST_COL = u32:0;
+        const LAST_COL = COLS;
+        unroll_for! (row, _): (u32, ()) in u32:0..ROWS {
+            send(join(), self.to_easts[row][FIRST_COL], row);
+            recv(join(), self.from_wests[row][LAST_COL]);
+        }(());
+    }
+}
+
+proc Top {}
+impl Top {
+    fn new() -> Self {
+        Grid<u32:2, u32:3>::new().spawn();
+        Self {}
+    }
+}
+)",
+              TypecheckSucceeds(::testing::_));
+}
 }  // namespace
 }  // namespace xls::dslx