[Proc testing] First iteration of TransformTestFunctions transformer.

xls/dslx/frontend/BUILD

@@ -834,6 +834,22 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "test_function_transformer",
+    srcs = ["test_function_transformer.cc"],
+    hdrs = ["test_function_transformer.h"],
+    deps = [
+        ":ast",
+        ":ast_cloner",
+        ":ast_node_visitor_with_default",
+        ":module",
+        "//xls/common:attribute_data",
+        "//xls/common/status:status_macros",
+        "//xls/dslx:import_data",
+        "@abseil-cpp//absl/status",
+    ],
+)
+
 cc_test(
     name = "fuzz_domain_rewriter_test",
     srcs = ["fuzz_domain_rewriter_test.cc"],

xls/dslx/frontend/test_function_transformer.cc

@@ -0,0 +1,113 @@
+// Copyright 2026 The XLS Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "xls/dslx/frontend/test_function_transformer.h"
+
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "xls/common/attribute_data.h"
+#include "xls/common/status/status_macros.h"
+#include "xls/dslx/frontend/ast.h"
+#include "xls/dslx/frontend/ast_cloner.h"
+#include "xls/dslx/frontend/ast_node_visitor_with_default.h"
+#include "xls/dslx/frontend/module.h"
+
+namespace xls {
+namespace dslx {
+
+namespace {
+class SpawnFinder : public AstNodeVisitorWithDefault {
+ public:
+  // Legacy spawn statement, not the new spawn attribute.
+  absl::Status HandleSpawn(const Spawn* node) override {
+    has_spawn_ = true;
+    return DefaultHandler(node);
+  }
+
+  absl::Status HandleInvocation(const Invocation* node) override {
+    // Trait-based spawn
+    if (const Attr* attr = dynamic_cast<const Attr*>(node->callee());
+        attr != nullptr) {
+      has_spawn_ = attr->attr() == "spawn";
+    }
+    return absl::OkStatus();
+  }
+
+  absl::Status HandleFunction(const Function* node) override {
+    has_spawn_ = false;
+    if (GetAttribute(node, AttributeKind::kTest).has_value()) {
+      // Skip functions that aren't test functions.
+      return absl::OkStatus();
+    }
+    // Recurse into children to see if there are any spawns.
+    return DefaultHandler(node);
+  }
+  absl::Status DefaultHandler(const AstNode* node) override {
+    for (const auto& child : node->GetChildren(false)) {
+      XLS_RETURN_IF_ERROR(child->Accept(this));
+    }
+    return absl::OkStatus();
+  }
+
+  bool has_spawn() const { return has_spawn_; }
+
+ private:
+  bool has_spawn_ = false;
+};
+
+absl::StatusOr<TestProc*> TransmuteToTestProc(const Function* fn,
+                                              Module* module) {
+  // TODO(davidplass): Implement this.
+  return nullptr;
+}
+
+}  // namespace
+
+absl::StatusOr<std::unique_ptr<Module>> TransformTestFunctions(
+    const Module& module) {
+  // Cannot reserve space in the vector because we don't know how many functions
+  // will be transformed.
+  std::vector<const AstNode*> test_functions_with_spawn;
+
+  // 1. Find all test functions with spawns.
+  for (auto* func : module.GetFunctions()) {
+    SpawnFinder spawn_finder;
+    XLS_RETURN_IF_ERROR(func->Accept(&spawn_finder));
+    if (spawn_finder.has_spawn()) {
+      test_functions_with_spawn.push_back(func);
+    }
+  }
+
+  // 2. Clone the module, removing test functions with spawns.
+  XLS_ASSIGN_OR_RETURN(
+      std::unique_ptr<Module> cloned_module,
+      CloneModuleRemovingMembers(module, test_functions_with_spawn));
+
+  // 3. For each test function with a spawn, transform it into a TestProc and
+  // add to the cloned module.
+  for (auto* node : test_functions_with_spawn) {
+    const Function* func = dynamic_cast<const Function*>(node);
+    XLS_ASSIGN_OR_RETURN(TestProc * test_proc,
+                         TransmuteToTestProc(func, cloned_module.get()));
+    // TODO(davidplass): Add collision error handler using the node's location
+    XLS_RETURN_IF_ERROR(
+        cloned_module->AddTop(test_proc, /*make_collision_error=*/nullptr));
+  }
+  return cloned_module;
+}
+}  // namespace dslx
+}  // namespace xls

xls/dslx/frontend/test_function_transformer.h

@@ -0,0 +1,34 @@
+// Copyright 2026 The XLS Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef XLS_DSLX_FRONTEND_TEST_FUNCTION_TRANSFORMER_H_
+#define XLS_DSLX_FRONTEND_TEST_FUNCTION_TRANSFORMER_H_
+
+#include <memory>
+
+#include "absl/status/status.h"
+#include "xls/dslx/frontend/module.h"
+
+namespace xls {
+namespace dslx {
+
+// Transforms test functions that spawn procs into TestProcs in the cloned
+// module. The original functions will be removed from the cloned module.
+absl::StatusOr<std::unique_ptr<Module>> TransformTestFunctions(
+    const Module& module);
+
+}  // namespace dslx
+}  // namespace xls
+
+#endif  // XLS_DSLX_FRONTEND_TEST_FUNCTION_TRANSFORMER_H_

xls/dslx/ir_convert/ir_converter_legacy_test.cc

@@ -3205,7 +3205,7 @@ TEST_F(IrConverterLegacyTest, InvalidChannelDecl) {
   auto import_data = CreateImportDataForTest();
   EXPECT_THAT(
       ConvertOneFunctionForTest(program, "main", import_data, kNoPosOptions),
-      StatusIs(absl::StatusCode::kInternal,
+      StatusIs(absl::StatusCode::kInvalidArgument,
                HasSubstr("Channels can only be declared in")));
 }
 
@@ -3225,8 +3225,8 @@ fn main() {
   auto import_data = CreateImportDataForTest();
   EXPECT_THAT(
       ConvertOneFunctionForTest(program, "main", import_data, kNoPosOptions),
-      StatusIs(absl::StatusCode::kUnimplemented,
-               HasSubstr("Functions cannot spawn procs.")));
+      StatusIs(absl::StatusCode::kInvalidArgument,
+               HasSubstr("Cannot spawn outside")));
 }
 
 TEST_F(IrConverterLegacyTest, InvalidSpawnProcScopedChannel) {
@@ -3244,8 +3244,8 @@ fn main() {
   auto import_data = CreateImportDataForTest();
   EXPECT_THAT(ConvertOneFunctionForTest(program, "main", import_data,
                                         kProcScopedChannelOptions),
-              StatusIs(absl::StatusCode::kUnimplemented,
-                       HasSubstr("Functions cannot spawn procs.")));
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Cannot spawn outside")));
 }
 
 TEST_F(IrConverterLegacyTest, InvalidSpawnInNextProcScoped) {

xls/dslx/ir_convert/ir_converter_test.cc

@@ -3232,8 +3232,9 @@ TEST_F(IrConverterTest, InvalidChannelDecl) {
   })";
 
   auto import_data = CreateImportDataForTest();
+  // This is now a typecheck violation instead of IR Converter check.
   EXPECT_THAT(ConvertOneFunctionForTest(program, "main", import_data),
-              StatusIs(absl::StatusCode::kInternal,
+              StatusIs(absl::StatusCode::kInvalidArgument,
                        HasSubstr("Channels can only be declared in")));
 }
 
@@ -3252,8 +3253,8 @@ fn main() {
 
   auto import_data = CreateImportDataForTest();
   EXPECT_THAT(ConvertOneFunctionForTest(program, "main", import_data),
-              StatusIs(absl::StatusCode::kUnimplemented,
-                       HasSubstr("Functions cannot spawn procs.")));
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Cannot spawn outside a")));
 }
 
 TEST_F(IrConverterTest, InvalidSpawnProcScopedChannel) {
@@ -3270,8 +3271,8 @@ fn main() {
 )";
   auto import_data = CreateImportDataForTest();
   EXPECT_THAT(ConvertOneFunctionForTest(program, "main", import_data),
-              StatusIs(absl::StatusCode::kUnimplemented,
-                       HasSubstr("Functions cannot spawn procs.")));
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Cannot spawn outside")));
 }
 
 TEST_F(IrConverterTest, InvalidSpawnInNextProcScoped) {

xls/dslx/type_system_v2/BUILD

@@ -223,6 +223,7 @@ cc_library(
         ":type_system_tracer",
         ":unify_type_annotations",
         ":validate_concrete_type",
+        "//xls/common:attribute_data",
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "//xls/dslx:constexpr_evaluator",

xls/dslx/type_system_v2/constant_collector.cc

@@ -805,9 +805,10 @@ class Visitor : public AstNodeVisitorWithDefault {
     XLS_ASSIGN_OR_RETURN(
         std::optional<const StructDefBase*> proc_def,
         GetContainingStructOrProcDef(invocation, import_data_));
-    XLS_RET_CHECK(proc_def.has_value());
-    ti_->AddProcDefSpawn(absl::down_cast<const ProcDef*>(*proc_def),
-                         external_initializer);
+    if (proc_def.has_value()) {
+      ti_->AddProcDefSpawn(absl::down_cast<const ProcDef*>(*proc_def),
+                           external_initializer);
+    }
     return absl::OkStatus();
   }
 

xls/dslx/type_system_v2/inference_table_converter_impl.cc

@@ -42,6 +42,7 @@
 #include "absl/strings/str_join.h"
 #include "absl/strings/substitute.h"
 #include "absl/types/span.h"
+#include "xls/common/attribute_data.h"
 #include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
 #include "xls/dslx/constexpr_evaluator.h"
@@ -534,6 +535,39 @@ class InferenceTableConverterImpl : public InferenceTableConverter,
           file_table_);
     }
 
+    // Only allow calling (new-style) spawn from inside a proc (legacy or
+    // impl-based), or (new!) a test function.
+    XLS_ASSIGN_OR_RETURN(bool is_proc_def_spawn,
+                         IsProcDefSpawnFunction(function));
+    if (is_proc_def_spawn) {
+      // If the spawn is in a proc, OK.
+      // If the spawn is in a test function, OK.
+      // Otherwise, error.
+      bool valid_spawn = false;
+      if (caller.has_value()) {
+        // Legacy procs can call new-style spawn.
+        valid_spawn = (*caller)->IsInProc();
+
+        XLS_ASSIGN_OR_RETURN(
+            std::optional<const StructDefBase*> containing_struct_or_proc,
+            GetContainingStructOrProcDef(invocation, import_data_));
+        if (containing_struct_or_proc.has_value() &&
+            (*containing_struct_or_proc)->kind() == AstNodeKind::kProcDef) {
+          // impl-based procs can call spawn
+          valid_spawn = true;
+        } else {
+          // Test functions can call spawn
+          valid_spawn |=
+              GetAttribute(*caller, AttributeKind::kTest).has_value();
+        }
+      }
+      if (!valid_spawn) {
+        return TypeInferenceErrorStatus(invocation->span(), nullptr,
+                                        "Cannot spawn outside a proc or test.",
+                                        file_table_);
+      }
+    }
+
     // Come up with the actual args by merging the possible target object
     // (`some_struct` in the case of `some_struct.foo(args)`), with the vector
     // of explicit actual args.