[Explicit State Access] Return the type as the initialized state in init

Currently it returns the same type as Proc Next field which is an empty tuple.

PiperOrigin-RevId: 862298516

Author: NL02

xls/dslx/frontend/parser.cc

@@ -3522,12 +3522,19 @@ absl::StatusOr<ModuleMember> Parser::ParseProcLike(const Pos& start_pos,
 
   // Just as with proc member decls, we need the init fn to have its own return
   // type, to avoid parent/child relationship violations.
-  XLS_ASSIGN_OR_RETURN(auto* init_return_type,
-                       CloneNode(proc_like_body.next->return_type(),
-                                 &PreserveTypeDefinitionsReplacer));
+  TypeAnnotation* type_to_clone;
+  if (module_->attributes().contains(ModuleAttribute::kExplicitStateAccess)) {
+    XLS_RET_CHECK(!proc_like_body.next->params().empty())
+        << "Proc 'next' function unexpectedly had no parameters.";
+    type_to_clone = proc_like_body.next->params()[0]->type_annotation();
+  } else {
+    type_to_clone = proc_like_body.next->return_type();
+  }
+  XLS_ASSIGN_OR_RETURN(
+      TypeAnnotation * init_return_type,
+      CloneNode(type_to_clone, &PreserveTypeDefinitionsReplacer));
   init_return_type->SetParentage();
-  proc_like_body.init->set_return_type(
-      down_cast<TypeAnnotation*>(init_return_type));
+  proc_like_body.init->set_return_type(init_return_type);
   proc_like_body.init->SetParentage();
 
   auto* proc_like = module_->Make<T>(span, body_span, name_def,

xls/dslx/frontend/parser_test.cc

@@ -3495,6 +3495,26 @@ proc simple {
   EXPECT_EQ(next.return_type()->ToString(), "()");
 }
 
+TEST_F(ParserTest, ExplicitStateAccessProcInitReturnsSameTypeAsSelf) {
+  constexpr std::string_view kProgram = R"(#![feature(explicit_state_access)]
+proc simple {
+    config() {
+        ()
+    }
+    init {
+        u32:0
+    }
+    next(state: u32) {
+    }
+})";
+  XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Module> module, Parse(kProgram));
+  XLS_ASSERT_OK_AND_ASSIGN(Proc * proc,
+                           module->GetMemberOrError<Proc>("simple"));
+  const Function& init = proc->init();
+  ASSERT_NE(init.return_type(), nullptr);
+  EXPECT_EQ(init.return_type()->ToString(), "u32");
+}
+
 // Verifies that we can walk backwards through a tree. In this case, from the
 // terminal node to the defining expr.
 TEST(ParserBackrefTest, CanFindDefiner) {