[NetKAT] Add an NOT cache to packet set manager

.github/workflows/build.yml

@@ -37,10 +37,10 @@ jobs:
         uses: actions/cache/restore@v4
         with:
           path: "~/.cache/bazel"
-          key: bazel-${{ hashFiles('*.bazel', '*.bazelrc') }}-${{ github.ref_name }}
+          key: bazel-ubuntu-22.04-${{ hashFiles('*.bazel', '*.bazelrc') }}-${{ github.ref_name }}
           restore-keys: |
-            bazel-${{hashFiles('*.bazel', '*.bazelrc') }}
-            bazel-
+            bazel-ubuntu-22.04-${{ hashFiles('*.bazel', '*.bazelrc') }}
+            bazel-ubuntu-22.04-
 
       - name: Save start time
         uses: josStorer/get-current-time@v2
@@ -89,4 +89,4 @@ jobs:
         if: always() && (github.ref_name == 'main' || env.duration > 300)
         with:
           path: "~/.cache/bazel"
-          key: bazel-${{ hashFiles('*.bazel', '*.bazelrc') }}-${{ github.ref_name }}-${{ github.run_id }}
+          key: bazel-ubuntu-22.04-${{ hashFiles('*.bazel', '*.bazelrc') }}-${{ github.ref_name }}-${{ github.run_id }}

netkat/packet_set.cc

@@ -231,6 +231,11 @@ PacketSetHandle PacketSetManager::Not(PacketSetHandle negand) {
   if (IsEmptySet(negand)) return FullSet();
   if (IsFullSet(negand)) return EmptySet();
 
+  auto it = not_cache_.find(negand);
+  if (it != not_cache_.end()) {
+    return it->second;
+  }
+
   // Compute result the hard way.
   const DecisionNode& negand_node = GetNodeOrDie(negand);
   DecisionNode result_node{
@@ -248,7 +253,7 @@ PacketSetHandle PacketSetManager::Not(PacketSetHandle negand) {
         std::make_pair(value, negated_branch);
   }
 
-  return NodeToPacket(std::move(result_node));
+  return not_cache_[negand] = NodeToPacket(std::move(result_node));
 }
 
 PacketSetHandle PacketSetManager::And(PacketSetHandle left,

netkat/packet_set.h

@@ -320,6 +320,10 @@ class PacketSetManager {
                       PacketSetHandle>
       and_cache_;
 
+  // A memoization table for the `Not` operation.
+  // Maps an argument handle to its computed complement handle.
+  absl::flat_hash_map<PacketSetHandle, PacketSetHandle> not_cache_;
+
   // INVARIANT: All `DecisionNode` fields are interned by this manager.
   PacketFieldManager field_manager_;
 

netkat/packet_set_benchmark.cc

@@ -216,4 +216,49 @@ void BM_ReCompileAndWithHighOverlappingPredicate(benchmark::State& state) {
 }
 BENCHMARK(BM_ReCompileAndWithHighOverlappingPredicate);
 
+// Benchmarks the cost of applying NOT to the same packet set multiple times.
+// This heavily exercises the `not_cache_`.
+void BM_NotOnSamePacketSet(benchmark::State& state) {
+  PredicateProto proto = CreateFixedArbitraryPredicateProto();
+  PacketTransformerManager transformer;
+  PacketSetManager& manager = transformer.GetPacketSetManager();
+  PacketSetHandle handle = manager.Compile(proto);
+
+  for (auto s : state) {
+    PacketSetHandle negated = manager.Not(handle);
+    benchmark::DoNotOptimize(negated);
+  }
+}
+BENCHMARK(BM_NotOnSamePacketSet);
+
+// Benchmarks Or operations that share operands. Since Or is implemented via
+// And and Not, sharing operands means we repeatedly call Not on the same
+// BDD handles, exercising the `not_cache_`.
+void BM_OrOperationsSharingOperands(benchmark::State& state) {
+  PredicateProto proto1 = CreateFixedArbitraryPredicateProto(/*id_suffix=*/1);
+  PredicateProto google::protobuf =
+      CreateFixedArbitraryPredicateProto(/*id_suffix=*/2);
+  PredicateProto proto3 = CreateFixedArbitraryPredicateProto(/*id_suffix=*/3);
+
+  for (auto s : state) {
+    PacketTransformerManager transformer;
+    PacketSetManager& manager = transformer.GetPacketSetManager();
+    PacketSetHandle h1 = manager.Compile(proto1);
+    PacketSetHandle h2 = manager.Compile(google::protobuf);
+    PacketSetHandle h3 = manager.Compile(proto3);
+
+    // Or(h1, h2) calls Not(h1), Not(h2), And, Not
+    PacketSetHandle or1 = manager.Or(h1, h2);
+    // Or(h1, h3) calls Not(h1) [cache hit!], Not(h3), And, Not
+    PacketSetHandle or2 = manager.Or(h1, h3);
+    // Or(h2, h3) calls Not(h2) [cache hit!], Not(h3) [cache hit!], And, Not
+    PacketSetHandle or3 = manager.Or(h2, h3);
+
+    benchmark::DoNotOptimize(or1);
+    benchmark::DoNotOptimize(or2);
+    benchmark::DoNotOptimize(or3);
+  }
+}
+BENCHMARK(BM_OrOperationsSharingOperands);
+
 }  // namespace netkat