From 0617c13978051a7155872d1c5ea9c32369136dcb Mon Sep 17 00:00:00 2001
From: Hadi Ravanbakhsh <ravanbakhsh@google.com>
Date: Fri, 7 Feb 2025 13:14:49 -0800
Subject: [PATCH] Store recursion analysis results per policy.

Protos of same type could get customized through different policies, which could lead to different recursive structure.

PiperOrigin-RevId: 724449290
---
 .../internal/domains/protobuf_domain_impl.h     | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/fuzztest/internal/domains/protobuf_domain_impl.h b/fuzztest/internal/domains/protobuf_domain_impl.h
index dc00eb80..bf20201e 100644
--- a/fuzztest/internal/domains/protobuf_domain_impl.h
+++ b/fuzztest/internal/domains/protobuf_domain_impl.h
@@ -18,6 +18,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <functional>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <string>
@@ -233,7 +234,10 @@ class ProtoPolicy {
  public:
   ProtoPolicy()
       : optional_policies_({{/*filter=*/IncludeAll<FieldDescriptor>(),
-                             /*value=*/OptionalPolicy::kWithNull}}) {}
+                             /*value=*/OptionalPolicy::kWithNull}}) {
+    absl::BitGen gen;
+    id_ = absl::Uniform(gen, 0, std::numeric_limits<int64_t>::max());
+  }
 
   void SetOptionalPolicy(OptionalPolicy optional_policy) {
     SetOptionalPolicy(IncludeAll<FieldDescriptor>(), optional_policy);
@@ -314,7 +318,11 @@ class ProtoPolicy {
     return max;
   }
 
+  int64_t id() const { return id_; }
+
  private:
+  int64_t id_;
+
   template <typename T>
   struct FilterToValue {
     Filter filter;
@@ -1707,12 +1715,13 @@ class ProtobufDomainUntypedImpl
   bool IsFieldFinitelyRecursive(const FieldDescriptor* field) {
     if (!field->message_type()) return false;
     ABSL_CONST_INIT static absl::Mutex mutex(absl::kConstInit);
-    static absl::NoDestructor<absl::flat_hash_map<const FieldDescriptor*, bool>>
+    static absl::NoDestructor<
+        absl::flat_hash_map<std::pair<int64_t, const FieldDescriptor*>, bool>>
         cache ABSL_GUARDED_BY(mutex);
     bool can_use_cache = IsCustomizedRecursivelyOnly();
     if (can_use_cache) {
       absl::MutexLock l(&mutex);
-      auto it = cache->find(field);
+      auto it = cache->find({policy_.id(), field});
       if (it != cache->end()) return it->second;
     }
     absl::flat_hash_set<decltype(field->message_type())> parents;
@@ -1720,7 +1729,7 @@ class ProtobufDomainUntypedImpl
                                    RecursionType::kFinitelyRecursive);
     if (can_use_cache) {
       absl::MutexLock l(&mutex);
-      cache->insert({field, result});
+      cache->insert({{policy_.id(), field}, result});
     }
     return result;
   }