[NetKAT] Port the pointer-keyed unique table to PacketSetManager; polish flat-node ergonomics.

Three follow-ups to the flat-node work:

* PacketSetManager's unique table now stores stable pointers into the
  node arena instead of full DecisionNode copies, with transparent
  by-value lookup — the same design just applied to the transformer
  manager, and more impactful here since packet-set nodes (every Match
  and every And/Or/Not result) far outnumber transformer nodes.

* The transformer's NodeHash streams the canonical flat sequence
  through a single absl hash state (via an AbslHashValue adapter)
  instead of re-seeding absl::HashOf once per element.

* DecisionNode::Matches() iterates the "if" branches as
  (value, modifies-span) views, replacing the index-loop +
  MatchModifies(i) ceremony at six call sites.

With these, all five transformer benchmarks (including the new
Push/Pull read-path benchmark) run ahead of main: ReCompile* -6%,
FirstTimeCompile* -1 to -3%, PushAndPull -2% (min-of-4 interleaved
optimized runs).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: smolkaj

netkat/BUILD.bazel

@@ -132,6 +132,7 @@ cc_library(
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/hash",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -375,6 +376,7 @@ cc_library(
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/hash",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",

netkat/packet_set.cc

@@ -24,6 +24,7 @@
 #include "absl/algorithm/container.h"
 #include "absl/container/fixed_array.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/hash/hash.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -87,6 +88,14 @@ const PacketSetManager::DecisionNode& PacketSetManager::GetNodeOrDie(
   return nodes_[packet_set.node_index_];
 }
 
+size_t PacketSetManager::NodeHash::operator()(const DecisionNode* node) const {
+  return absl::HashOf(*node);
+}
+
+size_t PacketSetManager::NodeHash::operator()(const DecisionNode& node) const {
+  return absl::HashOf(node);
+}
+
 PacketSetHandle PacketSetManager::NodeToPacket(DecisionNode&& node) {
   if (node.branch_by_field_value.empty()) return node.default_branch;
 
@@ -109,16 +118,19 @@ PacketSetHandle PacketSetManager::NodeToPacket(DecisionNode&& node) {
   }
 #endif
 
-  auto [it, inserted] =
-      packet_by_node_.try_emplace(node, PacketSetHandle(nodes_.size()));
-  if (inserted) {
-    nodes_.push_back(std::move(node));
-    LOG_IF(DFATAL, nodes_.size() > SentinelNodeIndex::kMinSentinel)
-        << "Internal invariant violated: Proper and sentinel node indices must "
-           "be disjoint. This indicates that we allocated more nodes than are "
-           "supported (> 2^32 - 2).";
+  // Look up the node by value via the transparent `NodeHash`/`NodeEq`
+  // functors; only new nodes get stored (exactly once, in `nodes_`).
+  if (auto it = packet_by_node_.find(node); it != packet_by_node_.end()) {
+    return it->second;
   }
-  return it->second;
+  PacketSetHandle packet(nodes_.size());
+  nodes_.push_back(std::move(node));
+  packet_by_node_.insert({&nodes_[packet.node_index_], packet});
+  LOG_IF(DFATAL, nodes_.size() > SentinelNodeIndex::kMinSentinel)
+      << "Internal invariant violated: Proper and sentinel node indices must "
+         "be disjoint. This indicates that we allocated more nodes than are "
+         "supported (> 2^32 - 2).";
+  return packet;
 }
 
 bool PacketSetManager::Contains(PacketSetHandle packet_set,
@@ -483,14 +495,19 @@ absl::Status PacketSetManager::CheckInternalInvariants() const {
   // Invariant: Proper and sentinel node indices are disjoint.
   RET_CHECK(nodes_.size() <= SentinelNodeIndex::kMinSentinel);
 
-  // Invariant: `packet_by_node_[n] = s` iff `nodes_[s.node_index_] == n`.
-  for (const auto& [node, packet] : packet_by_node_) {
+  // Invariant: `packet_by_node_[p] = s` iff `p == &nodes_[s.node_index_]`.
+  for (const auto& [node_ptr, packet] : packet_by_node_) {
     RET_CHECK(packet.node_index_ < nodes_.size());
-    RET_CHECK(nodes_[packet.node_index_] == node);
+    RET_CHECK(node_ptr == &nodes_[packet.node_index_]);
   }
   for (int i = 0; i < nodes_.size(); ++i) {
     const DecisionNode& node = nodes_[i];
-    auto it = packet_by_node_.find(node);
+    // Look up both by pointer and by value (exercising the transparent
+    // functors used by `NodeToPacket`).
+    auto it = packet_by_node_.find(&node);
+    RET_CHECK(it != packet_by_node_.end());
+    RET_CHECK(it->second == PacketSetHandle(i));
+    it = packet_by_node_.find(node);
     RET_CHECK(it != packet_by_node_.end());
     RET_CHECK(it->second == PacketSetHandle(i));
   }

netkat/packet_set.h

@@ -356,11 +356,38 @@ class PacketSetManager {
   // `And`, `Or`, `Not`). The class also avoids expensive relocations.
   PagedStableVector<DecisionNode, kPageSize> nodes_;
 
+  // Transparent hash and equality functors for the unique table
+  // (`packet_by_node_`), which is keyed by stable `DecisionNode*` pointers
+  // into `nodes_` (so each node is stored only once). Lookups work directly
+  // with a not-yet-interned `DecisionNode` value. Both functors are
+  // stateless: keys are pointers, and the pages holding the nodes are stable
+  // across moves of the manager.
+  struct NodeHash {
+    using is_transparent = void;
+    size_t operator()(const DecisionNode* node) const;
+    size_t operator()(const DecisionNode& node) const;
+  };
+  struct NodeEq {
+    using is_transparent = void;
+    bool operator()(const DecisionNode* a, const DecisionNode* b) const {
+      return a == b || *a == *b;
+    }
+    bool operator()(const DecisionNode* a, const DecisionNode& b) const {
+      return *a == b;
+    }
+    bool operator()(const DecisionNode& a, const DecisionNode* b) const {
+      return a == *b;
+    }
+  };
+
   // A so called "unique table" to ensure each node is only added to `nodes_`
   // once, and thus has a unique `PacketSetHandle::node_index`.
+  // Keyed by pointers into `nodes_` (stable, see `PagedStableVector`), so
+  // nodes are not stored twice.
   //
-  // INVARIANT: `packet_by_node_[n] = s` iff `nodes_[s.node_index_] == n`.
-  absl::flat_hash_map<DecisionNode, PacketSetHandle> packet_by_node_;
+  // INVARIANT: `packet_by_node_[p] = s` iff `p == &nodes_[s.node_index_]`.
+  absl::flat_hash_map<const DecisionNode*, PacketSetHandle, NodeHash, NodeEq>
+      packet_by_node_;
 
   // A map of a given `PredicateProto` to a `PacketSetHandle`.
   //

netkat/packet_transformer.cc

@@ -180,35 +180,14 @@ PacketTransformerManager::DecisionNode PacketTransformerManager::Flatten(
   return node;
 }
 
-template <typename NodeOrBuilder>
-size_t PacketTransformerManager::NodeHash::HashOfFlatSequence(
-    const NodeOrBuilder& node) {
-  size_t num_matches = 0;
-  size_t num_modifies = 0;
-  size_t hash = absl::HashOf(node.field, node.default_branch);
-  ForEachFlatEntry(
-      node,
-      [&](int match_value, uint32_t end_offset) {
-        hash = absl::HashOf(hash, match_value, end_offset);
-        ++num_matches;
-        return true;
-      },
-      [&](int modify_value, PacketTransformerHandle branch) {
-        hash = absl::HashOf(hash, modify_value, branch);
-        ++num_modifies;
-        return true;
-      });
-  return absl::HashOf(hash, num_matches, num_modifies);
-}
-
 size_t PacketTransformerManager::NodeHash::operator()(
     const DecisionNode* node) const {
-  return HashOfFlatSequence(*node);
+  return absl::HashOf(FlatSequenceView<DecisionNode>{*node});
 }
 
 size_t PacketTransformerManager::NodeHash::operator()(
     const DecisionNodeBuilder& builder) const {
-  return HashOfFlatSequence(builder);
+  return absl::HashOf(FlatSequenceView<DecisionNodeBuilder>{builder});
 }
 
 bool PacketTransformerManager::NodeEq::operator()(
@@ -240,10 +219,10 @@ PacketTransformerManager::ToBuilder(const DecisionNode& node) {
       .field = node.field,
       .default_branch = node.default_branch,
   };
-  for (size_t i = 0; i < node.matches.size(); ++i) {
+  for (const DecisionNode::Match& match : node.Matches()) {
     absl::btree_map<int, PacketTransformerHandle>& modify_branch_by_value =
-        builder.modify_branch_by_field_match[node.matches[i].first];
-    for (const DecisionNode::ModifyEntry& entry : node.MatchModifies(i)) {
+        builder.modify_branch_by_field_match[match.value];
+    for (const DecisionNode::ModifyEntry& entry : match.modifies) {
       modify_branch_by_value.insert(modify_branch_by_value.end(), entry);
     }
   }
@@ -950,8 +929,8 @@ PacketSetHandle PacketTransformerManager::GetAllPossibleOutputPackets(
   // Implements the `b_B` in the `fwd` function in section C.3 Push and Pull
   // in KATch: A Fast Symbolic Verifier for NetKAT.
   absl::flat_hash_set<int> branch_modify_values;
-  for (size_t i = 0; i < node.matches.size(); ++i) {
-    for (const auto& [modify_value, branch] : node.MatchModifies(i)) {
+  for (const DecisionNode::Match& match : node.Matches()) {
+    for (const auto& [modify_value, branch] : match.modifies) {
       branch_modify_values.insert(modify_value);
       add_to_output_by_field_value(modify_value,
                                    GetAllPossibleOutputPackets(branch));
@@ -1035,13 +1014,13 @@ PacketTransformerManager::GetAllInputPacketsThatProduceAnyOutput(
   // Implements the `b_A` in the `bwd` function in section C.3 Push and Pull
   // in KATch: A Fast Symbolic Verifier for NetKAT.
   absl::flat_hash_map<int, PacketSetHandle> branch_by_field_value_map;
-  for (size_t i = 0; i < node.matches.size(); ++i) {
+  for (const DecisionNode::Match& match : node.Matches()) {
     PacketSetHandle union_of_branches;
-    for (const auto& [modify_value, branch] : node.MatchModifies(i)) {
+    for (const auto& [modify_value, branch] : match.modifies) {
       union_of_branches = packet_set_manager_.Or(
           union_of_branches, GetAllInputPacketsThatProduceAnyOutput(branch));
     }
-    branch_by_field_value_map[node.matches[i].first] = union_of_branches;
+    branch_by_field_value_map[match.value] = union_of_branches;
   }
 
   // Case 3: Input packets that do not get matched on an explicit branch, but
@@ -1110,10 +1089,9 @@ std::string PacketTransformerManager::ToString(const DecisionNode& node) const {
       absl::CEscape(
           packet_set_manager_.field_manager_.GetFieldName(node.field)));
 
-  for (size_t i = 0; i < node.matches.size(); ++i) {
-    absl::StrAppendFormat(&result, "  %s == %d:\n", field,
-                          node.matches[i].first);
-    pretty_print_map(field, node.MatchModifies(i));
+  for (const DecisionNode::Match& match : node.Matches()) {
+    absl::StrAppendFormat(&result, "  %s == %d:\n", field, match.value);
+    pretty_print_map(field, match.modifies);
   }
   absl::StrAppendFormat(&result, "  %s == *:\n", field);
   pretty_print_map(field, node.DefaultModifies());
@@ -1160,10 +1138,9 @@ std::string PacketTransformerManager::ToString(
         "%v:'%s'", node.field,
         absl::CEscape(
             packet_set_manager_.field_manager_.GetFieldName(node.field)));
-    for (size_t i = 0; i < node.matches.size(); ++i) {
-      absl::StrAppendFormat(&result, "  %s == %d:\n", field,
-                            node.matches[i].first);
-      pretty_print_map(field, node.MatchModifies(i));
+    for (const DecisionNode::Match& match : node.Matches()) {
+      absl::StrAppendFormat(&result, "  %s == %d:\n", field, match.value);
+      pretty_print_map(field, match.modifies);
     }
     absl::StrAppendFormat(&result, "  %s == *:\n", field);
     pretty_print_map(field, node.DefaultModifies());
@@ -1216,10 +1193,9 @@ std::string PacketTransformerManager::ToDot(
         packet_set_manager_.field_manager_.GetFieldName(node.field);
     absl::StrAppendFormat(&result, "  %d [label=\"%s\"]\n",
                           transformer.node_index_, field);
-    for (size_t i = 0; i < node.matches.size(); ++i) {
-      int value = node.matches[i].first;
-      absl::Span<const DecisionNode::ModifyEntry> modify_map =
-          node.MatchModifies(i);
+    for (const DecisionNode::Match& match : node.Matches()) {
+      int value = match.value;
+      absl::Span<const DecisionNode::ModifyEntry> modify_map = match.modifies;
       if (modify_map.empty()) {
         absl::StrAppendFormat(&result, "  %d -> %d [label=\"%s==%s\"]\n",
                               transformer.node_index_, SentinelNodeIndex::kDeny,
@@ -1309,8 +1285,8 @@ absl::Status PacketTransformerManager::CheckInternalInvariants() const {
           }
           return true;
         };
-    for (size_t j = 0; j < node.matches.size(); ++j) {
-      RET_CHECK(is_strictly_sorted_by_value(node.MatchModifies(j)))
+    for (const DecisionNode::Match& match : node.Matches()) {
+      RET_CHECK(is_strictly_sorted_by_value(match.modifies))
           << ":\n"
           << ToString(node);
     }
@@ -1324,12 +1300,11 @@ absl::Status PacketTransformerManager::CheckInternalInvariants() const {
         << ":\n"
         << ToString(node);
 
-    for (size_t j = 0; j < node.matches.size(); ++j) {
-      int match_value = node.matches[j].first;
-      for (const auto& [modify_value, branch] : node.MatchModifies(j)) {
+    for (const DecisionNode::Match& match : node.Matches()) {
+      for (const auto& [modify_value, branch] : match.modifies) {
         // Invariant: Modify branches are not Deny unless `modify_value ==
-        // match_value`.
-        RET_CHECK(!IsDeny(branch) || modify_value == match_value)
+        // match.value`.
+        RET_CHECK(!IsDeny(branch) || modify_value == match.value)
             << ":\n"
             << ToString(node);
 

netkat/packet_transformer.h

@@ -376,6 +376,41 @@ class PacketTransformerManager {
                                  matches[i].second - begin);
     }
 
+    // A single "if (field == value)" branch: the match value together with
+    // its ModifyEntry range.
+    struct Match {
+      int value;
+      absl::Span<const ModifyEntry> modifies;
+    };
+
+    // Iterates the "if" branches as `Match` views, in order of strictly
+    // increasing match value. Allows range-for loops over the branches
+    // without manual index bookkeeping.
+    class MatchIterator {
+     public:
+      MatchIterator(const DecisionNode* node, size_t index)
+          : node_(node), index_(index) {}
+      Match operator*() const {
+        return {node_->matches[index_].first, node_->MatchModifies(index_)};
+      }
+      MatchIterator& operator++() {
+        ++index_;
+        return *this;
+      }
+      friend bool operator==(const MatchIterator& a,
+                             const MatchIterator& b) = default;
+
+     private:
+      const DecisionNode* node_;
+      size_t index_;
+    };
+    struct MatchRange {
+      const DecisionNode* node;
+      MatchIterator begin() const { return {node, 0}; }
+      MatchIterator end() const { return {node, node->matches.size()}; }
+    };
+    MatchRange Matches() const { return {this}; }
+
     // The ModifyEntry range of the "else" branch.
     absl::Span<const ModifyEntry> DefaultModifies() const {
       uint32_t begin = matches.empty() ? 0 : matches.back().second;
@@ -457,11 +492,35 @@ class PacketTransformerManager {
     size_t operator()(const DecisionNodeBuilder& builder) const;
 
    private:
-    // Shared implementation of both overloads: hashes the canonical flat
-    // element sequence (via `ForEachFlatEntry`), so flat nodes and builders
-    // with the same logical content hash identically.
+    // Adapter implementing both overloads: hashes the canonical flat element
+    // sequence (via `ForEachFlatEntry`) in a single streaming pass, so flat
+    // nodes and builders with the same logical content hash identically.
     template <typename NodeOrBuilder>
-    static size_t HashOfFlatSequence(const NodeOrBuilder& node);
+    struct FlatSequenceView {
+      const NodeOrBuilder& node;
+
+      // Hashing, see https://abseil.io/docs/cpp/guides/hash.
+      template <typename H>
+      friend H AbslHashValue(H h, const FlatSequenceView& view) {
+        size_t num_matches = 0;
+        size_t num_modifies = 0;
+        h = H::combine(std::move(h), view.node.field,
+                       view.node.default_branch);
+        ForEachFlatEntry(
+            view.node,
+            [&](int match_value, uint32_t end_offset) {
+              h = H::combine(std::move(h), match_value, end_offset);
+              ++num_matches;
+              return true;
+            },
+            [&](int modify_value, PacketTransformerHandle branch) {
+              h = H::combine(std::move(h), modify_value, branch);
+              ++num_modifies;
+              return true;
+            });
+        return H::combine(std::move(h), num_matches, num_modifies);
+      }
+    };
   };
   struct NodeEq {
     using is_transparent = void;

netkat/packet_transformer_test.cc

@@ -968,8 +968,8 @@ class PacketTransformerManagerTestPeer {
                                                                       value);
     };
     // Case 1: Output from explicit match+modify branches.
-    for (size_t i = 0; i < node.matches.size(); ++i) {
-      for (const auto& [modify_value, branch] : node.MatchModifies(i)) {
+    for (const auto& match : node.Matches()) {
+      for (const auto& [modify_value, branch] : match.modifies) {
         add_to_output(
             and_fn(match_fn(field, modify_value),
                    GetAllPossibleOutputPacketsReferenceImplementation(branch)));