[XLS] When comparing TreeBitLocations, compare node ID rather than pointer

Do the same when caching BDD trees for each node in the `node_locations_` cache in the BDD query engine.

This is generally both a correctness & a safety improvement. In this specific case, it was found as a fix for nondeterminism! Due to false matches when memory locations were reused for new Nodes, the BDD query engine's variable caches could nondeterministically match with an (no-longer-used) variable created for a deleted Node. This was safe, but resulted in effective changes to the variable order, which could cause the BDD to saturate in different places nondeterministically.

PiperOrigin-RevId: 761556944

Author: ericastor

xls/ir/node.h

@@ -29,7 +29,6 @@
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #include "xls/common/casts.h"
 #include "xls/common/status/status_macros.h"
@@ -50,6 +49,35 @@ absl::Span<ChangeListener* const> GetChangeListeners(
 // Forward declaration to avoid circular dependency.
 class DfsVisitor;
 
+// A (non-owning) reference to a node; contains both the node's ID and the
+// pointer to the node, enabling efficient access & safe comparison.
+class NodeRef {
+ public:
+  explicit NodeRef(Node* node);
+
+  int64_t id() const { return id_; }
+  Node* node() const { return node_; }
+
+  // On dereference, behaves exactly like the underlying node.
+  Node& operator*() const { return *node_; }
+  Node* operator->() const { return node_; }
+
+  friend bool operator==(const NodeRef& a, const NodeRef& b) {
+    DCHECK(a.id() != b.id() || a.node() == b.node())
+        << "False node match due to reused ID";
+    return a.id() == b.id();
+  }
+
+  template <typename H>
+  friend H AbslHashValue(H h, const NodeRef& node) {
+    return h.combine(std::move(h), node.id());
+  }
+
+ private:
+  int64_t id_;
+  Node* node_;
+};
+
 // Abstract type for a node (representing an expression) in the high level IR.
 //
 // Node is subtyped and can be checked-converted via the As* methods below.
@@ -336,6 +364,9 @@ class Node {
   absl::InlinedVector<Node*, 2> users_;
 };
 
+inline NodeRef::NodeRef(Node* node)
+    : id_(node == nullptr ? -1 : node->id()), node_(node) {}
+
 inline std::ostream& operator<<(std::ostream& os, const Node& node) {
   os << node.ToString();
   return os;
@@ -344,10 +375,20 @@ inline std::ostream& operator<<(std::ostream& os, const Node* node) {
   os << (node == nullptr ? std::string("<nullptr Node*>") : node->ToString());
   return os;
 }
+inline std::ostream& operator<<(std::ostream& os, const NodeRef& node) {
+  os << node->ToString();
+  return os;
+}
 
-inline void NodeAppend(std::string* out, const Node* n) {
-  absl::StrAppend(out, n->ToString());
+inline bool operator==(const NodeRef& a, Node* b) {
+  if (b == nullptr) {
+    return a.node() == nullptr;
+  }
+  DCHECK(a.id() != b->id() || a.node() == b)
+      << "False node match due to reused ID";
+  return a.id() == b->id();
 }
+inline bool operator==(Node* a, const NodeRef& b) { return b == a; }
 
 }  // namespace xls
 

xls/passes/bdd_query_engine.cc

@@ -482,16 +482,17 @@ BddTree BddQueryEngine::ComputeInfo(
     Node* node, absl::Span<const BddTree* const> operand_infos) const {
   if (!ShouldEvaluate(node)) {
     VLOG(3) << "  node filtered out by generic ShouldEvaluate heuristic.";
-    return leaf_type_tree::Clone(GetVariablesFor(node));
+    return leaf_type_tree::Clone(GetVariablesFor(NodeRef(node)));
   }
   if (node_filter_.has_value() && !(*node_filter_)(node)) {
     VLOG(3) << "  node filtered out by configured filter.";
-    return leaf_type_tree::Clone(GetVariablesFor(node));
+    return leaf_type_tree::Clone(GetVariablesFor(NodeRef(node)));
   }
 
   VLOG(3) << "  computing BDD value...";
-  BddNodeEvaluator node_evaluator(
-      *evaluator_, [this](Node* node) { return GetVariablesFor(node); });
+  BddNodeEvaluator node_evaluator(*evaluator_, [this](Node* node) {
+    return GetVariablesFor(NodeRef(node));
+  });
   absl::flat_hash_set<Node*> injected_operands;
   injected_operands.reserve(node->operand_count());
   for (auto [operand, operand_info] :
@@ -1255,7 +1256,7 @@ bool BddQueryEngine::IsFullyKnown(
 }
 
 BddNodeIndex BddQueryEngine::GetVariableFor(TreeBitLocation location) const {
-  if (auto it = node_variables_.find(location.node());
+  if (auto it = node_variables_.find(location.node_ref());
       it != node_variables_.end()) {
     if (it->second->type() == location.node()->GetType()) {
       return std::get<BddNodeIndex>(
@@ -1271,7 +1272,7 @@ BddNodeIndex BddQueryEngine::GetVariableFor(TreeBitLocation location) const {
   CHECK(bit_variables_.emplace(location, result).second);
   return result;
 }
-BddTreeView BddQueryEngine::GetVariablesFor(Node* node) const {
+BddTreeView BddQueryEngine::GetVariablesFor(NodeRef node) const {
   if (auto it = node_variables_.find(node);
       it != node_variables_.end() && it->second->type() == node->GetType()) {
     // If a node has changed type (which can happen!), we need a new set of

xls/passes/bdd_query_engine.h

@@ -226,10 +226,11 @@ class BddQueryEngine
   mutable absl::flat_hash_map<TreeBitLocation, BddNodeIndex> bit_variables_;
   BddNodeIndex GetVariableFor(TreeBitLocation location) const;
 
-  // A map from nodes to BDD variables used to represent fully-unknown values;
-  // used to avoid creating new variables for the same node.
-  mutable absl::flat_hash_map<Node*, std::unique_ptr<BddTree>> node_variables_;
-  BddTreeView GetVariablesFor(Node* node) const;
+  // A map from node IDs to BDD variables used to represent fully-unknown
+  // values; used to avoid creating new variables for the same node.
+  mutable absl::flat_hash_map<NodeRef, std::unique_ptr<BddTree>>
+      node_variables_;
+  BddTreeView GetVariablesFor(NodeRef node) const;
 };
 
 }  // namespace xls

xls/passes/query_engine.h

@@ -50,7 +50,10 @@ class TreeBitLocation {
         bit_index_(bit_index),
         tree_index_(tree_index.begin(), tree_index.end()) {}
 
-  Node* node() const { return node_; }
+  // TODO: https://github.com/google/xls/issues/2235 - Replace node() with this.
+  const NodeRef& node_ref() const { return node_; }
+
+  Node* node() const { return node_.node(); }
 
   int64_t bit_index() const { return bit_index_; }
 
@@ -77,7 +80,7 @@ class TreeBitLocation {
   }
 
  private:
-  Node* node_;
+  NodeRef node_;
   int64_t bit_index_;
   std::vector<int64_t> tree_index_;
 };