Change node-dependency-analysis and some other dependencies of proc_state_narrowing_pass to use NodeMap & NodeSet

GEMINI.md

@@ -26,3 +26,10 @@ This file provides a context for the XLS project.
 **Documentation:**
 
 *   OSS docs are in `docs_src` and rendered with `mkdocs` at [https://google.github.io/xls/](https://google.github.io/xls/).
+
+**‼️ Agent Instructions ‼️**
+
+*   **NodeMap/NodeSet Usage:** New code should generally use `NodeMap` or
+    `NodeSet` instead of `absl::flat_hash_map` or `absl::flat_hash_set` when the
+    key is `Node*` and the value is not trivially copyable. Existing code should
+    not be modified unless specifically directed to do so.

xls/common/BUILD

@@ -640,6 +640,11 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "pointer_utils",
+    hdrs = ["pointer_utils.h"],
+)
+
 cc_test(
     name = "stopwatch_test",
     srcs = ["stopwatch_test.cc"],

xls/common/pointer_utils.h

@@ -0,0 +1,38 @@
+// Copyright 2025 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_COMMON_POINTER_UTILS_H_
+#define XLS_COMMON_POINTER_UTILS_H_
+
+#include <memory>
+
+namespace xls {
+
+namespace internal {
+using RawDeletePtr = void (*)(void*);
+}
+
+// type-erased version of unique ptr that keeps track of the appropriate
+// destructor. To use reinterpret_cast<T*>(x.get()).
+using TypeErasedUniquePtr = std::unique_ptr<void, internal::RawDeletePtr>;
+
+template <typename T, typename Deleter = std::default_delete<T>>
+TypeErasedUniquePtr EraseType(std::unique_ptr<T, Deleter> ptr) {
+  return TypeErasedUniquePtr(
+      ptr.release(), [](void* ptr) { Deleter()(reinterpret_cast<T*>(ptr)); });
+}
+
+}  // namespace xls
+
+#endif  // XLS_COMMON_POINTER_UTILS_H_

xls/data_structures/inline_bitmap.h

@@ -39,6 +39,10 @@ class BitmapView;
 // A bitmap that has 64-bits of inline storage by default.
 class InlineBitmap {
  public:
+  // How many bits are held in one word.
+  static constexpr int64_t kWordBits = 64;
+  // How many bytes are held in one word.
+  static constexpr int64_t kWordBytes = 8;
   // Constructs an InlineBitmap of width `bit_count` using the bits in
   // `word`. If `bit_count` is greater than 64, then all high bits are set to
   // `fill`.
@@ -345,9 +349,6 @@ class InlineBitmap {
   friend uint64_t GetWordBitsAtForTest(const InlineBitmap& ib,
                                        int64_t bit_offset);
 
-  static constexpr int64_t kWordBits = 64;
-  static constexpr int64_t kWordBytes = 8;
-
   // Gets the kWordBits bits following bit_offset with 'Get(bit_offset)' being
   // the LSB, Get(bit_offset + 1) being the next lsb etc.
   int64_t GetWordBitsAt(int64_t bit_offset) const;

xls/data_structures/transitive_closure.h

@@ -16,6 +16,7 @@
 #define XLS_DATA_STRUCTURES_TRANSITIVE_CLOSURE_H_
 
 #include <cstdint>
+#include <type_traits>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
@@ -25,6 +26,8 @@
 
 namespace xls {
 
+class Node;
+
 namespace internal {
 // Compute the transitive closure of a relation.
 template <typename Relation>
@@ -97,6 +100,28 @@ class DenseIdRelation {
   absl::Span<InlineBitmap> relation_;
 };
 
+template <typename NodeRelationBase>
+class NodeRelation {
+ public:
+  explicit NodeRelation(NodeRelationBase& relation) : relation_(relation) {}
+  template <typename F>
+  void ForEachKeyValue(F f) const {
+    for (auto& [j, from_j] : relation_) {
+      f(j, from_j);
+    }
+  }
+  bool Contains(const absl::flat_hash_set<Node*>& vs, Node* const& v) const {
+    return vs.contains(v);
+  }
+  void UnionInPlace(absl::flat_hash_set<Node*>& i,
+                    const absl::flat_hash_set<Node*>& k) const {
+    i.insert(k.begin(), k.end());
+  }
+
+ private:
+  NodeRelationBase& relation_;
+};
+
 }  // namespace internal
 
 template <typename V>
@@ -110,6 +135,19 @@ HashRelation<V> TransitiveClosure(HashRelation<V> v) {
   return v;
 }
 
+// Compute the transitive closure of a relation represented as an explicit
+// adjacency list using NodeMap.
+//
+// This has some terrible template stuff to ensure we don't actually need to
+// include NodeMap.
+template <typename NodeRelation>
+NodeRelation TransitiveClosure(NodeRelation v)
+  requires(NodeRelation::kIsNodeMap)
+{
+  internal::TransitiveClosure(internal::NodeRelation(v));
+  return v;
+}
+
 // TODO(allight): Using a more efficient bitmap format like croaring might give
 // a speedup here.
 using DenseIdRelation = absl::Span<InlineBitmap>;

xls/ir/BUILD

@@ -638,9 +638,11 @@ cc_library(
         "//xls/common:casts",
         "//xls/common:iterator_range",
         "//xls/common:math_util",
+        "//xls/common:pointer_utils",
         "//xls/common:visitor",
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
+        "//xls/data_structures:inline_bitmap",
         "//xls/data_structures:leaf_type_tree",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
@@ -2491,3 +2493,40 @@ cc_test(
         "@googletest//:gtest",
     ],
 )
+
+cc_library(
+    name = "node_map",
+    hdrs = ["node_map.h"],
+    deps = [
+        ":ir",
+        "//xls/common:pointer_utils",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:hash_container_defaults",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+    ],
+)
+
+cc_test(
+    name = "node_map_test",
+    srcs = ["node_map_test.cc"],
+    deps = [
+        ":benchmark_support",
+        ":bits",
+        ":function_builder",
+        ":ir",
+        ":ir_matcher",
+        ":ir_test_base",
+        ":node_map",
+        "//xls/common:xls_gunit_main",
+        "//xls/common/status:matchers",
+        "//xls/common/status:status_macros",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@google_benchmark//:benchmark",
+        "@googletest//:gtest",
+    ],
+)

xls/ir/node.cc

@@ -36,6 +36,7 @@
 #include "absl/strings/str_join.h"
 #include "absl/types/span.h"
 #include "xls/common/casts.h"
+#include "xls/common/pointer_utils.h"
 #include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
 #include "xls/ir/change_listener.h"
@@ -997,4 +998,28 @@ bool Node::OpIn(absl::Span<const Op> choices) const {
 
 Package* Node::package() const { return function_base()->package(); }
 
+std::optional<TypeErasedUniquePtr> Node::TakeUserData(int64_t idx) {
+  DCHECK(function_base()->package()->IsLiveUserDataId(idx)) << idx;
+  if (user_data_.size() <= idx) {
+    return std::nullopt;
+  }
+  std::optional<TypeErasedUniquePtr> data = std::move(user_data_[idx]);
+  user_data_[idx] = std::nullopt;
+  return data;
+}
+void* Node::GetUserData(int64_t idx) {
+  DCHECK(function_base()->package()->IsLiveUserDataId(idx)) << idx;
+  if (user_data_.size() <= idx) {
+    user_data_.resize(idx + 1);
+  }
+  const auto& v = user_data_[idx];
+  return v ? v->get() : nullptr;
+}
+void Node::SetUserData(int64_t idx, TypeErasedUniquePtr data) {
+  DCHECK(function_base()->package()->IsLiveUserDataId(idx)) << idx;
+  if (user_data_.size() <= idx) {
+    user_data_.resize(idx + 1);
+  }
+  user_data_[idx] = std::move(data);
+}
 }  // namespace xls

xls/ir/node.h

@@ -24,13 +24,15 @@
 #include <string>
 #include <string_view>
 #include <utility>
+#include <vector>
 
 #include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "xls/common/casts.h"
+#include "xls/common/pointer_utils.h"
 #include "xls/common/status/status_macros.h"
 #include "xls/ir/change_listener.h"
 #include "xls/ir/op.h"
@@ -316,6 +318,28 @@ class Node {
     absl::Format(&sink, "%s", node.GetName());
   }
 
+  // User-data access functions. Should not be directly used. Use NodeMap
+  // instead.
+  //
+  // Extreme care should be used when interacting with these functions and the
+  // Package ones since this is basically doing manual memory management.
+
+  // Get the pointer associated with this indexes user data or nullptr if
+  // never set. Use HasUserData to see if anything has ever been set.
+  //
+  // idx must be a value returned by Package::AllocateNodeUserData which has not
+  // had ReleaseNodeUserDataId called on it.
+  void* GetUserData(int64_t idx);
+  // Sets user data at idx to 'data'.
+  void SetUserData(int64_t idx, TypeErasedUniquePtr data);
+  // Removes user data at idx from the node. Returns std::nullopt if nothing has
+  // been set.
+  std::optional<TypeErasedUniquePtr> TakeUserData(int64_t idx);
+  // Checks if anything has ever been set at the given user data.
+  bool HasUserData(int64_t idx) {
+    return user_data_.size() > idx && user_data_[idx].has_value();
+  }
+
  protected:
   // FunctionBase needs to be a friend to access RemoveUser for deleting nodes
   // from the graph.
@@ -368,6 +392,15 @@ class Node {
 
   // Set of users sorted by node_id for stability.
   absl::InlinedVector<Node*, 2> users_;
+
+ private:
+  std::vector<std::optional<TypeErasedUniquePtr>> user_data_;
+
+  // Clear all user data.
+  void ClearUserData() { user_data_.clear(); };
+
+  // for ClearUserData
+  friend class Package;
 };
 
 inline NodeRef::NodeRef(Node* node)