Fix parametric ProcDefs in bytecode interpreter.

There are several problems being fixed here, all of which are verified by the added interpreter test case:
* Mapping of data by ProcDef* without the effective parametrics.
* Lack of knowledge of the effective parametrics in the ProcInitializer.
* The ProcInitializer being created based on evaluation of member initializers before their types were determined (wrong ordering in flatten_in_type_order).
* Missing storage of TypeReference constexpr values for generics.
* Incorrect comparison of TypeReference constants in interp_value.cc
* An auxiliary proc map in type_info.cc using unstable pointers to objects in a vector.

PiperOrigin-RevId: 930116263

Author: richmckeever

xls/dslx/bytecode/bytecode_interpreter.cc

@@ -357,6 +357,9 @@ absl::Status BytecodeInterpreter::Run(bool* progress_made) {
       const std::vector<Bytecode>& bytecodes = frame->bf()->bytecodes();
       const Bytecode& bytecode = bytecodes.at(frame->pc());
       VLOG(2) << "Bytecode: " << bytecode.ToString(file_table());
+      if (proc_id_.has_value()) {
+        VLOG(2) << "Proc: " << proc_id_->ToString();
+      }
       VLOG(2) << std::hex << "PC: " << frame->pc() << " : "
               << bytecode.ToString(file_table());
       VLOG(3) << absl::StreamFormat(" - stack depth %d [%s]", stack_.size(),

xls/dslx/bytecode/proc_hierarchy_interpreter.cc

@@ -398,7 +398,7 @@ absl::Status ProcHierarchyInterpreter::AddProcDefInstance(
                          ti->GetCanonicalProcInitializer(external_initializer));
     XLS_RETURN_IF_ERROR(AddProcDefInstance(
         external_initializer.GetProcInitializerOrDie().proc_def(),
-        external_initializer, canonical_initializer.constructor_type_info,
+        external_initializer, canonical_initializer.next_type_info,
         canonical_initializer.constructor_env, import_data, options));
   }
 

xls/dslx/bytecode/proc_hierarchy_interpreter_test.cc

@@ -1125,5 +1125,63 @@ impl CounterTest {
   EXPECT_EQ(result.result(), TestResult::kAllPassed);
 }
 
+TEST_F(ProcHierarchyInterpreterTest, ParametricProcDef) {
+  constexpr std::string_view kProgram = R"(
+#![feature(explicit_state_access)]
+#![feature(generics)]
+
+proc Counter<T: type> {
+  ch_out: chan<T> out,
+  i: T,
+}
+
+impl Counter {
+  fn new(ch_out: chan<T> out) -> Self {
+    Counter<T> { ch_out: ch_out, i: 0 }
+  }
+
+  fn next(self) {
+    let old_value = read(self.i);
+    send(token(), self.ch_out, old_value);
+    write(self.i, old_value + 1);
+  }
+}
+
+#[test]
+proc CounterTest {
+  terminator: chan<bool> out,
+  ch_in0: chan<u32> in,
+  ch_in1: chan<s16> in,
+}
+
+impl CounterTest {
+  fn new(terminator: chan<bool> out) -> Self {
+    let (counter_out0, counter_in0) = chan<u32>("counter0");
+    let (counter_out1, counter_in1) = chan<s16>("counter1");
+    Counter<u32>::new(counter_out0).spawn();
+    Counter<s16>::new(counter_out1).spawn();
+    CounterTest {
+      terminator: terminator,
+      ch_in0: counter_in0,
+      ch_in1: counter_in1
+    }
+  }
+
+  fn next(self) {
+    let (_, value0) = recv(token(), self.ch_in0);
+    let (_, value1) = recv(token(), self.ch_in1);
+    send_if(token(), self.terminator, value0 > 2 && value1 > 2, true);
+  }
+})";
+
+  XLS_ASSERT_OK_AND_ASSIGN(auto temp_file,
+                           TempFile::CreateWithContent(kProgram, "_test.x"));
+  ParseAndTestOptions options;
+  XLS_ASSERT_OK_AND_ASSIGN(
+      TestResultData result,
+      ParseAndTest(kProgram, "test", std::string{temp_file.path()}, options));
+  EXPECT_EQ(result.result(), TestResult::kAllPassed);
+}
+
 }  // namespace
 }  // namespace xls::dslx

xls/dslx/interp_value.cc

@@ -284,6 +284,11 @@ std::string InterpValue::ToStringInternal(bool humanize,
     case InterpValueTag::kProcInitializer:
       const auto& initializer =
           std::get<std::shared_ptr<ProcInitializer>>(payload_);
+      std::vector<std::string> parametric_strings;
+      parametric_strings.reserve(initializer->parametrics().size());
+      for (const InterpValue& parametric : initializer->parametrics()) {
+        parametric_strings.push_back(parametric.ToString());
+      }
       std::vector<std::string> member_strings;
       for (const InterpValue& member : initializer->members()) {
         member_strings.push_back(member.ToString());
@@ -292,8 +297,12 @@ std::string InterpValue::ToStringInternal(bool humanize,
         member_strings.push_back(absl::Substitute(
             "fwd $0: $1", param->identifier(), value.ToString()));
       }
-      return absl::Substitute("$0($1)", initializer->proc_def()->identifier(),
-                              absl::StrJoin(member_strings, ", "));
+      return absl::Substitute(
+          "$0$1($2)", initializer->proc_def()->identifier(),
+          parametric_strings.empty()
+              ? ""
+              : absl::StrCat("<", absl::StrJoin(parametric_strings, ", "), ">"),
+          absl::StrJoin(member_strings, ", "));
   }
   return "<INVALID>";
 }
@@ -456,6 +465,13 @@ bool InterpValue::ProcInitializer::Eq(const ProcInitializer& other) const {
   if (proc_def_ != other.proc_def_) {
     return false;
   }
+  CHECK_EQ(parametrics_.size(), other.parametrics_.size());
+  for (int i = 0; i < parametrics_.size(); i++) {
+    if (parametrics_[i] != other.parametrics_[i]) {
+      return false;
+    }
+  }
+
   if (members_.size() != other.members_.size()) {
     return false;
   }
@@ -473,6 +489,7 @@ bool InterpValue::ProcInitializer::Eq(const ProcInitializer& other) const {
       return false;
     }
   }
+
   return true;
 }
 
@@ -1223,6 +1240,14 @@ bool InterpValue::operator<(const InterpValue& rhs) const {
              rhs_inst.proc_def()->owner()->name();
     }
 
+    if (inst.parametrics().size() != rhs_inst.parametrics().size()) {
+      return inst.parametrics().size() < rhs_inst.parametrics().size();
+    }
+    for (int i = 0; i < inst.parametrics().size(); i++) {
+      if (!inst.parametrics()[i].Eq(rhs_inst.parametrics()[i])) {
+        return inst.parametrics()[i] < rhs_inst.parametrics()[i];
+      }
+    }
     if (inst.members().size() != rhs_inst.members().size()) {
       return inst.members().size() < rhs_inst.members().size();
     }
@@ -1250,6 +1275,13 @@ bool InterpValue::operator<(const InterpValue& rhs) const {
     return false;
   }
 
+  if (IsTypeReference()) {
+    if (!rhs.IsTypeReference()) {
+      return false;
+    }
+    return Lt(rhs).value().IsTrue();
+  }
+
   if (IsTuple()) {
     if (rhs.IsArray()) {
       return true;

xls/dslx/interp_value.h

@@ -165,28 +165,32 @@ class InterpValue {
    public:
     ProcInitializer(
         const ProcDef* proc_def, const AstNode* definer,
-        std::vector<InterpValue> members,
+        std::vector<InterpValue> parametrics, std::vector<InterpValue> members,
         std::vector<std::pair<const Param*, InterpValue>> forwarded_values)
         : proc_def_(proc_def),
           definer_(definer),
+          parametrics_(std::move(parametrics)),
           members_(std::move(members)),
           forwarded_values_(std::move(forwarded_values)) {
+      CHECK_EQ(parametrics_.size(), proc_def->parametric_bindings().size());
       InitMemberMap();
     }
 
     ProcInitializer(const ProcInitializer& other)
-        : ProcInitializer(other.proc_def_, other.definer_, other.members_,
-                          std::move(other.forwarded_values_)) {}
+        : ProcInitializer(other.proc_def_, other.definer_, other.parametrics_,
+                          other.members_, other.forwarded_values_) {}
 
     ProcInitializer(ProcInitializer&& other)
         : ProcInitializer(other.proc_def_, other.definer_,
                           std::move(other.members_),
+                          std::move(other.parametrics_),
                           std::move(other.forwarded_values_)) {
       other.member_map_.clear();
     }
 
     ProcInitializer& operator=(const ProcInitializer& other) {
       proc_def_ = other.proc_def_;
+      parametrics_ = other.parametrics_;
       members_ = other.members_;
       forwarded_values_ = other.forwarded_values_;
       definer_ = other.definer_;
@@ -200,6 +204,7 @@ class InterpValue {
       }
       proc_def_ = other.proc_def_;
       members_ = std::move(other.members_);
+      parametrics_ = std::move(other.parametrics_);
       forwarded_values_ = std::move(other.forwarded_values_);
       definer_ = other.definer_;
       InitMemberMap();
@@ -209,6 +214,7 @@ class InterpValue {
 
     const ProcDef* proc_def() const { return proc_def_; }
     const AstNode* definer() const { return definer_; }
+    const std::vector<InterpValue>& parametrics() const { return parametrics_; }
     const std::vector<InterpValue>& members() const { return members_; }
     const std::vector<std::pair<const Param*, InterpValue>>& forwarded_values()
         const {
@@ -232,6 +238,7 @@ class InterpValue {
 
     const ProcDef* proc_def_;
     const AstNode* definer_;
+    std::vector<InterpValue> parametrics_;
     std::vector<InterpValue> members_;
     std::vector<std::pair<const Param*, InterpValue>> forwarded_values_;
 
@@ -298,12 +305,12 @@ class InterpValue {
 
   static InterpValue MakeProcInitializer(
       const ProcDef* proc_def, const AstNode* definer,
-      std::vector<InterpValue> members,
+      std::vector<InterpValue> parametrics, std::vector<InterpValue> members,
       std::vector<std::pair<const Param*, InterpValue>> forwarded_values) {
-    return InterpValue(
-        InterpValueTag::kProcInitializer,
-        std::make_shared<ProcInitializer>(proc_def, definer, std::move(members),
-                                          std::move(forwarded_values)));
+    return InterpValue(InterpValueTag::kProcInitializer,
+                       std::make_shared<ProcInitializer>(
+                           proc_def, definer, std::move(parametrics),
+                           std::move(members), std::move(forwarded_values)));
   }
 
   static absl::StatusOr<InterpValue> MakeArray(

xls/dslx/type_system/type_info.cc

@@ -248,6 +248,7 @@ absl::Status TypeInfo::NoteProcConstructorInvocation(
   XLS_RET_CHECK(it != root->decorated_canonical_proc_initializer_.end());
   external_to_canonical_proc_initializer_.emplace(
       std::move(external_proc_initializer), std::move(canonical_initializer));
+
   return absl::OkStatus();
 }
 
@@ -277,9 +278,11 @@ TypeInfo::GetCanonicalProcInitializer(const InterpValue& external_initializer) {
   const auto it =
       external_to_canonical_proc_initializer_.find(external_initializer);
   XLS_RET_CHECK(it != external_to_canonical_proc_initializer_.end());
+  TypeInfo* root = GetRoot();
   const auto decorated_it =
-      decorated_canonical_proc_initializer_.find(it->second);
-  XLS_RET_CHECK(decorated_it != decorated_canonical_proc_initializer_.end());
+      root->decorated_canonical_proc_initializer_.find(it->second);
+  XLS_RET_CHECK(decorated_it !=
+                root->decorated_canonical_proc_initializer_.end());
   return *decorated_it->second;
 }
 

xls/dslx/type_system/type_info.h

@@ -18,6 +18,7 @@
 #define XLS_DSLX_TYPE_SYSTEM_TYPE_INFO_H_
 
 #include <cstdint>
+#include <list>
 #include <memory>
 #include <optional>
 #include <string>
@@ -609,7 +610,7 @@ class TypeInfo {
   absl::flat_hash_map<const Proc*, std::vector<SpawnData>> spawns_;
 
   // Initializers for each callee proc.
-  absl::node_hash_map<const ProcDef*, std::vector<ProcInitializerWithTypeInfo>>
+  absl::node_hash_map<const ProcDef*, std::list<ProcInitializerWithTypeInfo>>
       proc_def_initializers_by_callee_proc_;
 
   // The same objects in `proc_def_initializers_by_callee_proc`, but with the

xls/dslx/type_system_v2/BUILD

@@ -252,6 +252,7 @@ cc_library(
         "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

xls/dslx/type_system_v2/constant_collector.cc

@@ -757,6 +757,7 @@ class Visitor : public AstNodeVisitorWithDefault {
         invocation, table_.GetParametricEnv(parametric_context_),
         InterpValue::MakeProcInitializer(
             internal_initializer.proc_def(), internal_initializer.definer(),
+            internal_initializer.parametrics(),
             std::move(external_instance_args),
             std::move(external_instance_forwarded_values))));
     return absl::OkStatus();
@@ -931,9 +932,18 @@ class Visitor : public AstNodeVisitorWithDefault {
       }
     }
 
+    std::vector<InterpValue> parametrics;
+    parametrics.reserve(type.struct_def_base().parametric_bindings().size());
+    for (const ParametricBinding* binding :
+         type.struct_def_base().parametric_bindings()) {
+      XLS_ASSIGN_OR_RETURN(InterpValue value,
+                           ti_->GetConstExpr(binding->name_def()));
+      parametrics.push_back(value);
+    }
+
     auto inst = InterpValue::MakeProcInitializer(
         &absl::down_cast<const ProcDef&>(type.struct_def_base()), node,
-        member_values, std::move(forwarded_values));
+        parametrics, member_values, std::move(forwarded_values));
 
     ti_->NoteCanonicalProcInitializer(
         node, table_.GetParametricEnv(parametric_context_), inst);

xls/dslx/type_system_v2/flatten_in_type_order.cc

@@ -76,11 +76,11 @@ class Flattener : public AstNodeVisitorWithDefault {
   }
 
   absl::Status HandleStructInstance(const StructInstance* node) override {
-    nodes_.push_back(node);
     XLS_RETURN_IF_ERROR(node->struct_ref()->Accept(this));
     for (const auto& [_, member] : node->GetUnorderedMembers()) {
       XLS_RETURN_IF_ERROR(member->Accept(this));
     }
+    nodes_.push_back(node);
     return absl::OkStatus();
   }