Fix ProcId generation for multiple non-leaf spawns of same proc.

Previously, a spawn graph like this would fail to compile with an ID collision:

A -> B -> C
  -> B -> C

PiperOrigin-RevId: 660557462

xls/dslx/ir_convert/extract_conversion_order.cc

@@ -69,6 +69,20 @@ std::string ConversionRecordsToString(
 
 }  // namespace
 
+// -- class ProcIdFactory
+
+ProcId ProcIdFactory::CreateProcId(const ProcId& parent, Proc* spawnee,
+                                   bool count_as_new_instance) {
+  std::vector<std::pair<Proc*, int>> new_stack = parent.proc_instance_stack;
+  int& instance_count =
+      instance_counts_[std::make_pair(parent, spawnee->identifier())];
+  new_stack.push_back(std::make_pair(spawnee, instance_count));
+  if (count_as_new_instance) {
+    ++instance_count;
+  }
+  return ProcId{.proc_instance_stack = std::move(new_stack)};
+}
+
 // -- class Callee
 
 /* static */ absl::StatusOr<Callee> Callee::Make(
@@ -295,15 +309,13 @@ class InvocationVisitor : public ExprVisitor {
     if (maybe_proc.has_value()) {
       XLS_RET_CHECK(proc_id_.has_value()) << "Functions cannot spawn procs.";
 
-      std::vector<Proc*> proc_stack = proc_id_.value().proc_stack;
-      proc_stack.push_back(maybe_proc.value());
-      proc_id = ProcId{proc_stack, proc_instances_[proc_stack]};
-      // Only increment on next so that Config & Next have the same ID.
-      // This assumes that we call a Proc's config & next calls in that order,
-      // which is indeed the case.
-      if (callee_info->callee->tag() == FunctionTag::kProcNext) {
-        proc_instances_[proc_stack]++;
-      }
+      // Only count `next` as a new instance, so that `config` and `next` have
+      // the same ID. This assumes that we call a proc's `config` and `next` in
+      // that order, which is indeed the case.
+      const bool count_as_new_instance =
+          callee_info->callee->tag() == FunctionTag::kProcNext;
+      proc_id = proc_id_factory_.CreateProcId(*proc_id_, maybe_proc.value(),
+                                              count_as_new_instance);
     }
 
     // See if there are parametric bindings to use in the callee for this
@@ -438,7 +450,8 @@ class InvocationVisitor : public ExprVisitor {
     Module* module = (*info)->module;
     XLS_ASSIGN_OR_RETURN(Function * f,
                          module->GetMemberOrError<Function>(colon_ref->attr()));
-    return CalleeInfo{module, f, (*info)->type_info};
+    return CalleeInfo{
+        .module = module, .callee = f, .type_info = (*info)->type_info};
   }
 
   // Helper for invocations of NameRef callees.
@@ -488,14 +501,15 @@ class InvocationVisitor : public ExprVisitor {
                                  identifier);
     }
 
-    return CalleeInfo{this_m, f, callee_type_info};
+    return CalleeInfo{
+        .module = this_m, .callee = f, .type_info = callee_type_info};
   }
 
   Module* module_;
   TypeInfo* type_info_;
   const ParametricEnv& bindings_;
   std::optional<ProcId> proc_id_;
-  absl::flat_hash_map<std::vector<Proc*>, int> proc_instances_;
+  ProcIdFactory proc_id_factory_;
 
   // Built up list of callee records discovered during traversal.
   std::vector<Callee> callees_;
@@ -717,13 +731,14 @@ static absl::StatusOr<std::vector<ConversionRecord>> GetOrderForProc(
 
   // The next function of a proc is the entry function when converting a proc to
   // IR.
-  XLS_RETURN_IF_ERROR(AddToReady(&p->next(),
-                                 /*invocation=*/nullptr, p->owner(), type_info,
-                                 ParametricEnv(), &ready, ProcId{{p}, 0},
-                                 is_top));
+  XLS_RETURN_IF_ERROR(
+      AddToReady(&p->next(),
+                 /*invocation=*/nullptr, p->owner(), type_info, ParametricEnv(),
+                 &ready, ProcId{.proc_instance_stack = {{p, 0}}}, is_top));
   XLS_RETURN_IF_ERROR(AddToReady(&p->config(),
                                  /*invocation=*/nullptr, p->owner(), type_info,
-                                 ParametricEnv(), &ready, ProcId{{p}, 0}));
+                                 ParametricEnv(), &ready,
+                                 ProcId{.proc_instance_stack = {{p, 0}}}));
 
   // Constants and "member" vars are assigned and defined in Procs' "config'"
   // functions, so we need to execute those before their "next" functions.

xls/dslx/ir_convert/extract_conversion_order.h

@@ -21,9 +21,12 @@
 #include <variant>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
 #include "xls/dslx/frontend/ast.h"
 #include "xls/dslx/frontend/proc.h"
@@ -40,39 +43,72 @@ namespace xls::dslx {
 struct ProcId {
   // Contains the "spawn chain": the series of Procs through which this Proc was
   // spawned, with the oldest/"root" proc as element 0.  Contains the current
-  // proc, as well.
-  std::vector<Proc*> proc_stack;
-
-  // The index of this Proc in the proc stack. If a Proc is spawned > 1 inside
-  // the same Proc config function, this index differentiates the spawnees.
-  // Each unique proc stack will have its count start at 0 - in other words,
-  // the sequence:
-  // spawn foo()(c0)
-  // spawn bar()(c1)
-  // spawn foo()(c2)
-  // Would result in IDs of:
-  // foo:0, bar:0, and foo:1, respectively.
-  int instance;
+  // proc, as well. The second element of each pair is a zero-based spawn index
+  // of that same proc by the spawning proc. For example, with a spawn chain
+  // like:
+  //     A |-> D |-> B -> C
+  //       |-> B -> C
+  //       |-> B -> C
+  //       |-> E |-> B -> C
+  //
+  // the `proc_instance_stack` for each `C` would look like:
+  //    [{A, 0}, {D, 0}, {B, 0}, {C, 0}]
+  //    [{A, 0}, {B, 0}, {C, 0}]
+  //    [{A, 0}, {B, 1}, {C, 0}]
+  //    [{A, 0}, {E, 0}, {B, 0}, {C, 0}]
+  std::vector<std::pair<Proc*, int>> proc_instance_stack;
 
   std::string ToString() const {
-    return absl::StrCat(absl::StrJoin(proc_stack, "->",
-                                      [](std::string* out, const Proc* p) {
-                                        out->append(p->identifier());
-                                      }),
-                        ":", instance);
-    return "";
+    if (proc_instance_stack.empty()) {
+      return "";
+    }
+    // The first proc in a chain never needs an instance count. Leaving it out
+    // specifically when the chain length is more than 1 gets us the historical
+    // output in most cases (where there was only an instance count at the end
+    // of the chain).
+    CHECK_EQ(proc_instance_stack[0].second, 0);
+    const bool omit_first_instance_count = proc_instance_stack.size() > 1;
+    std::string part_with_instance_counts = absl::StrJoin(
+        proc_instance_stack.begin() + (omit_first_instance_count ? 1 : 0),
+        proc_instance_stack.end(), "->",
+        [](std::string* out, const std::pair<Proc*, int> p) {
+          absl::StrAppendFormat(out, "%s:%d", p.first->identifier(), p.second);
+        });
+    return omit_first_instance_count
+               ? absl::StrCat(proc_instance_stack[0].first->identifier(), "->",
+                              part_with_instance_counts)
+               : part_with_instance_counts;
   }
 
   bool operator==(const ProcId& other) const {
-    return proc_stack == other.proc_stack && instance == other.instance;
+    return proc_instance_stack == other.proc_instance_stack;
   }
 
   template <typename H>
   friend H AbslHashValue(H h, const ProcId& pid) {
-    return H::combine(std::move(h), pid.proc_stack, pid.instance);
+    return H::combine(std::move(h), pid.proc_instance_stack);
   }
 };
 
+// An object that deals out `ProcId` instances.
+class ProcIdFactory {
+ public:
+  // Creates a `ProcId` representing the given `spawnee` spawned by the given
+  // `parent` context. If `count_as_new_instance` is true, then subsequent calls
+  // with the same `parent` and `spawnee` will get a new instance count value.
+  // Otherwise, subsequent calls will get an equivalent `ProcId` to the one
+  // returned by this call.
+  ProcId CreateProcId(const ProcId& parent, Proc* spawnee,
+                      bool count_as_new_instance = true);
+
+ private:
+  // Maps each `parent` and `spawnee` identifier passed to `CreateProcId` to the
+  // number of instances of that pairing, i.e., the number of times that
+  // `parent` and `spawnee` have been passed in with `true` for
+  // `count_as_new_instance`.
+  absl::flat_hash_map<std::pair<ProcId, std::string>, int> instance_counts_;
+};
+
 // Describes a callee function in the conversion order (see
 // ConversionRecord::callees).
 class Callee {