Implement ReversibleFlatMap.

PiperOrigin-RevId: 966760746
diff --git a/doc/domains-reference.md b/doc/domains-reference.md
index e945809..6d8e56d 100644
--- a/doc/domains-reference.md
+++ b/doc/domains-reference.md
@@ -674,7 +674,7 @@
 Note: The [note for `Map`](#map) about aggregate types with nested C-style
 arrays also applies to `ReversibleMap`.
 
-### FlatMap
+### FlatMap {#flat-map}
 
 Sometimes we need to fuzz parameters that are dependent on each other. Think of
 a property function that takes a string, and valid index into that string. This
@@ -722,11 +722,54 @@
 Note: Domains defined using `FlatMap()` don't support
 [initial seeds](fuzz-test-macro.md#initial-seeds). If you need a
 [seeded domain](#seeded-domains)—a domain skewed toward certain values—consider
-seeding the input domains passed to `FlatMap()`.
+seeding the input domains passed to `FlatMap()`. Otherwise, if you need full
+support for seeds, consider using [`ReversibleFlatMap()`](#reversible-flat-map).
 
 Note: The [note for `Map`](#map) about aggregate types with nested C-style
 arrays also applies to `FlatMap`.
 
+### ReversibleFlatMap {#reversible-flat-map}
+
+The `ReversibleFlatMap()` domain combinator is similar to
+[`FlatMap()`](#flat-map) but is able to support
+[initial seeds](fuzz-test-macro.md#initial-seeds). It is strongly suggested that
+readers read the documentation of [`FlatMap()`](#flat-map) first. In the same
+way as `FlatMap()`, `ReversibleFlatMap()` takes as input a domain *factory*
+function, but it takes as second argument a function that would pick any value
+that was generated from that domain and map it back to the original value.
+
+We can revisit the `FlatMap()` example for a vector of *equal sized strings*:
+
+```c++
+auto AnyVectorOfFixedLengthStrings(int size) {
+  return VectorOf(Arbitrary<std::string>().WithSize(size));
+}
+std::optional<std::tuple<int>> SizeOfStringsIfIdentical(
+    const std::vector<std::string>& input) {
+  if (input.empty()) {
+    // Note that we are returning 0, but any value in the original (0, 10) range
+    // would work!
+    return std::tuple<int>(0);
+  }
+  const int size = input[0].size();
+  for (const std::string& s : input) {
+    if (s.size() != size) return std::nullopt;
+  }
+  return std::tuple<int>(size);
+}
+auto AnyVectorOfEqualSizedStrings() {
+  return ReversibleFlatMap(
+          AnyVectorOfFixedLengthStrings,
+          SizeOfStringsIfIdentical,
+          /*size=*/InRange(0, 10));
+}
+```
+
+IMPORTANT: The mapping function `f` and the inverse mapping function `g` must
+satisfy the following property: If `g` maps a value `y` to a tuple of inner
+values `x` (i.e., `g(y)` is not `std::nullopt`), then `y` must be a valid value
+from the domain `f(x)`. That is, `y` is in `f(g(y))`.
+
 ### Filter
 
 The `Filter` domain combinator takes a domain and a predicate and returns a new
diff --git a/domain_tests/map_filter_combinator_test.cc b/domain_tests/map_filter_combinator_test.cc
index 82c80f9..353b123 100644
--- a/domain_tests/map_filter_combinator_test.cc
+++ b/domain_tests/map_filter_combinator_test.cc
@@ -477,5 +477,154 @@
               testing::HasSubstr("Traversal budget exceeded"));
 }
 
+TEST(ReversibleFlatMap, WorksWithSameCorpusType) {
+  auto domain = ReversibleFlatMap(
+      [](int a) { return Just(~a); },
+      [](int a) { return std::optional(std::tuple(~a)); }, Arbitrary<int>());
+  absl::BitGen bitgen;
+  Value value(domain, bitgen);
+  // Corpus value is a tuple: (output_corpus, input_corpus...)
+  EXPECT_EQ(value.user_value, ~std::get<1>(value.corpus_value));
+}
+
+TEST(ReversibleFlatMap, AcceptsMultipleInnerDomains) {
+  auto domain = ReversibleFlatMap(
+      [](int len, char c) { return StringOf(Just(c)).WithSize(len); },
+      [](const std::string& s) -> std::optional<std::tuple<int, char>> {
+        if (s.empty()) return std::nullopt;
+        return std::tuple(s.size(), s[0]);
+      },
+      InRange(2, 4), ElementOf({'A', 'B'}));
+
+  absl::BitGen bitgen;
+  Set<std::string> values;
+  while (values.size() < 6) {
+    values.insert(Value(domain, bitgen).user_value);
+  }
+  EXPECT_THAT(values,
+              UnorderedElementsAre("AA", "AAA", "AAAA", "BB", "BBB", "BBBB"));
+}
+
+TEST(ReversibleFlatMap, WorksWithSeeds) {
+  // This is the core feature that standard FlatMap lacks.
+  auto domain =
+      ReversibleFlatMap([](int len) { return AsciiString().WithSize(len); },
+                        [](const std::string& s) {
+                          return std::optional(std::tuple<int>(s.size()));
+                        },
+                        InRange(2, 5))
+          .WithSeeds({"ABC", "WXYZ"});
+
+  EXPECT_THAT(GenerateInitialValues(domain, 20), Contains("ABC"));
+  EXPECT_THAT(GenerateInitialValues(domain, 20), Contains("WXYZ"));
+}
+
+TEST(ReversibleFlatMap, FromValueReturnsNulloptWhenInverseOrConstraintsFail) {
+  auto domain = ReversibleFlatMap(
+      [](int len) { return StringOf(ElementOf({'A', 'B'})).WithSize(len); },
+      [](const std::string& s) -> std::optional<std::tuple<int>> {
+        if (s.empty()) return std::nullopt;
+        return std::tuple(s.size());
+      },
+      InRange(2, 4));
+
+  // inv_mapper itself returns nullopt
+  EXPECT_EQ(domain.FromValue(""), std::nullopt);
+  // inv_mapper succeeds, but `int` 1 is rejected by `InRange(2, 4)`
+  EXPECT_EQ(domain.FromValue("A"), std::nullopt);
+  // inv_mapper succeeds, length is valid, but 'C' is rejected by
+  // `ElementOf({'A', 'B'})`
+  EXPECT_EQ(domain.FromValue("CCC"), std::nullopt);
+}
+
+TEST(ReversibleFlatMap, SerializationRoundTrip) {
+  auto domain =
+      ReversibleFlatMap([](int len) { return AsciiString().WithSize(len); },
+                        [](const std::string& s) {
+                          return std::optional(std::tuple<int>(s.size()));
+                        },
+                        InRange(0, 10));
+  absl::BitGen bitgen;
+  Value value(domain, bitgen);
+  auto serialized = domain.SerializeCorpus(value.corpus_value);
+  EXPECT_EQ(domain.ParseCorpus(serialized), value.corpus_value);
+}
+
+TEST(ReversibleFlatMap, ParseCorpusRejectsInvalidInputValues) {
+  absl::BitGen bitgen;
+
+  auto domain_a = ReversibleFlatMap(
+      [](int a) { return Just(a); },
+      [](int a) { return std::optional(std::tuple(a)); }, InRange(0, 9));
+  auto domain_b = ReversibleFlatMap(
+      [](int a) { return Just(a); },
+      [](int a) { return std::optional(std::tuple(a)); }, InRange(10, 19));
+
+  Value value(domain_a, bitgen);
+  auto serialized = domain_a.SerializeCorpus(value.corpus_value);
+
+  // domain_b should fail to parse domain_a's serialized corpus because the
+  // input corpus value (0-9) is out of bounds for domain_b (10-19).
+  EXPECT_EQ(domain_b.ParseCorpus(serialized), std::nullopt);
+}
+
+TEST(ReversibleFlatMap, ValidationRejectsInvalidValue) {
+  absl::BitGen bitgen;
+
+  auto domain_a = ReversibleFlatMap(
+      [](int a) { return Just(~a); },
+      [](int a) { return std::optional(std::tuple(~a)); }, InRange(0, 9));
+  auto domain_b = ReversibleFlatMap(
+      [](int a) { return Just(~a); },
+      [](int a) { return std::optional(std::tuple(~a)); }, InRange(10, 19));
+
+  Value value_a(domain_a, bitgen);
+  Value value_b(domain_b, bitgen);
+
+  ASSERT_OK(domain_a.ValidateCorpusValue(value_a.corpus_value));
+  ASSERT_OK(domain_b.ValidateCorpusValue(value_b.corpus_value));
+
+  EXPECT_THAT(
+      domain_a.ValidateCorpusValue(value_b.corpus_value),
+      IsInvalid(testing::MatchesRegex(
+          R"(Invalid value for ReversibleFlatMap\(\)-ed domain >> The value .+ is not InRange\(0, 9\))")));
+  EXPECT_THAT(
+      domain_b.ValidateCorpusValue(value_a.corpus_value),
+      IsInvalid(testing::MatchesRegex(
+          R"(Invalid value for ReversibleFlatMap\(\)-ed domain >> The value .+ is not InRange\(10, 19\))")));
+}
+
+TEST(ReversibleFlatMap, FlatMapperWorksWithMoveOnlyTypes) {
+  auto domain = ReversibleFlatMap(
+      [](std::unique_ptr<int> n) -> Domain<int> {
+        return n == nullptr ? Just(0) : Just(*n);
+      },
+      [](int n) -> std::optional<std::tuple<std::unique_ptr<int>>> {
+        return std::tuple(std::make_unique<int>(n));
+      },
+      UniquePtrOf(Just(1)));
+  EXPECT_THAT(MutateUntilFoundN(domain, /*n=*/2), UnorderedElementsAre(0, 1));
+}
+
+TEST(ReversibleFlatMap, MutationAcceptsShrinkingOutputDomains) {
+  auto domain =
+      ReversibleFlatMap([](int len) { return AsciiString().WithMaxSize(len); },
+                        [](const std::string& s) {
+                          return std::optional(std::tuple<int>(s.size()));
+                        },
+                        InRange(0, 10));
+  absl::BitGen bitgen;
+  std::optional<Value<decltype(domain)>> value;
+  // Generate something shrinkable
+  while (!value.has_value() || value->user_value.empty()) {
+    value = Value(domain, bitgen);
+  }
+  auto mutated = value->corpus_value;
+  while (!domain.GetValue(mutated).empty()) {
+    domain.Mutate(mutated, bitgen, {}, /*only_shrink=*/true);
+  }
+  EXPECT_THAT(domain.GetValue(mutated), IsEmpty());
+}
+
 }  // namespace
 }  // namespace fuzztest
diff --git a/fuzztest/domain_core.h b/fuzztest/domain_core.h
index 4da59f3..b7672e1 100644
--- a/fuzztest/domain_core.h
+++ b/fuzztest/domain_core.h
@@ -832,6 +832,39 @@
                                                      std::move(inner)...);
 }
 
+// ReversibleFlatMap(flat_mapper, inv_mapper, inner...) combinator creates a
+// domain that dynamically generates an output domain based on values from
+// `inner...`, while supporting `.WithSeeds()` via an inverse mapper.
+// Importantly, if `inv_mapper` maps `y` to some tuple `x` (i.e.,
+// `inv_mapper(y)` is not `std::nullopt`), then `y` must be a valid value from
+// the domain `flat_mapper(x)`. That is, `y` is in `flat_mapper(inv_mapper(y))`.
+
+//
+// To ensure this property, `inv_mapper` may need to return `std::nullopt`.
+//
+// Example:
+//   // Generate domain of two equal-sized strings.
+//   ReversibleFlatMap(
+//     [](int size) {
+//       return PairOf(Arbitrary<std::string>().WithSize(size),
+//                     Arbitrary<std::string>().WithSize(size)); },
+//     [](std::pair<std::string, std::string> pair) {
+//       const auto& [s1, s2] = pair;
+//       if (s1.size() != s2.size()) return std::nullopt;
+//       return std::optional(std::tuple<int>(s1.size()));
+//     },
+//     InRange(0, 10));
+//
+// The return type of `inv_mapper` should be `std::optional<std::tuple<T...>>`,
+// where `T...` are the input types of `flat_mapper`.
+template <int&... ExplicitArgumentBarrier, typename FlatMapper,
+          typename InvMapper, typename... Inner>
+auto ReversibleFlatMap(FlatMapper flat_mapper, InvMapper inv_mapper,
+                       Inner... inner) {
+  return internal::ReversibleFlatMapImpl<FlatMapper, InvMapper, Inner...>(
+      std::move(flat_mapper), std::move(inv_mapper), std::move(inner)...);
+}
+
 // VectorOf(inner) combinator creates a `std::vector` domain with elements of
 // the `inner` domain.
 //
diff --git a/fuzztest/internal/domains/flat_map_impl.h b/fuzztest/internal/domains/flat_map_impl.h
index 6368c99..37f0cb9 100644
--- a/fuzztest/internal/domains/flat_map_impl.h
+++ b/fuzztest/internal/domains/flat_map_impl.h
@@ -24,6 +24,7 @@
 #include "absl/random/distributions.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "./fuzztest/internal/domains/domain_base.h"
 #include "./fuzztest/internal/domains/serialization_helpers.h"
@@ -44,10 +45,20 @@
 using FlatMapOutputDomain = std::decay_t<
     std::invoke_result_t<FlatMapper, value_type_t<InputDomain>...>>;
 
-template <typename FlatMapper, typename... InputDomain>
-class FlatMapImpl
+// Base class for FlatMapImpl and ReversibleFlatMapImpl.
+// Meant to be used as:
+// ```c++
+// template <typename FlatMapper, typename... InputDomain>
+// class ConcreteFlatMapImpl
+//     : public FlatMapImplBase<ConcreteFlatMapImpl<FlatMapper, InputDomain...>,
+//                              FlatMapper, InputDomain...> {
+//   ...
+// };
+// ```
+template <typename Derived, typename FlatMapper, typename... InputDomain>
+class FlatMapImplBase
     : public domain_implementor::DomainBase<
-          FlatMapImpl<FlatMapper, InputDomain...>,
+          Derived,
           // The user value is the user value of the output domain.
           value_type_t<FlatMapOutputDomain<FlatMapper, InputDomain...>>,
           // The corpus value is a tuple where the first element is the corpus
@@ -57,11 +68,11 @@
               corpus_type_t<FlatMapOutputDomain<FlatMapper, InputDomain...>>,
               corpus_type_t<InputDomain>...>> {
  public:
-  using typename FlatMapImpl::DomainBase::corpus_type;
-  using typename FlatMapImpl::DomainBase::value_type;
+  using typename FlatMapImplBase::DomainBase::corpus_type;
+  using typename FlatMapImplBase::DomainBase::value_type;
 
-  FlatMapImpl() = default;
-  explicit FlatMapImpl(FlatMapper flat_mapper, InputDomain... input_domains)
+  FlatMapImplBase() = default;
+  explicit FlatMapImplBase(FlatMapper flat_mapper, InputDomain... input_domains)
       : flat_mapper_(std::move(flat_mapper)),
         input_domains_(std::move(input_domains)...) {}
 
@@ -108,12 +119,6 @@
     return GetOutputDomain(v).GetValue(std::get<0>(v));
   }
 
-  std::optional<corpus_type> FromValue(const value_type&) const {
-    // We cannot infer the input corpus from the output value, or even determine
-    // from which output domain the output value came.
-    return std::nullopt;
-  }
-
   auto GetPrinter() const {
     return FlatMappedPrinter<FlatMapper, InputDomain...>{flat_mapper_,
                                                          input_domains_};
@@ -153,7 +158,12 @@
         .ValidateCorpusValue(std::get<0>(corpus_value));
   }
 
- private:
+ protected:
+  const std::tuple<InputDomain...>& input_domains() const {
+    return input_domains_;
+  }
+  static constexpr size_t kNumInputValues = sizeof...(InputDomain);
+
   // Returns the output domain for a `tuple` with or without the output value
   // as the leading element, and with the input values as the last
   // `kNumInputValues` elements.
@@ -189,18 +199,110 @@
                 std::get<I>(input_domains_)
                     .ValidateCorpusValue(std::get<kOffset + I>(tuple));
             input_values_validity =
-                Prefix(s, "Invalid value for FlatMap()-ed domain");
+                Prefix(s, absl::StrCat("Invalid value for ", Derived::GetName(),
+                                       "()-ed domain"));
           }(),
           ...);
       return input_values_validity;
     });
   }
 
-  static constexpr size_t kNumInputValues = sizeof...(InputDomain);
+ private:
   FlatMapper flat_mapper_;
   std::tuple<InputDomain...> input_domains_;
 };
 
+template <typename FlatMapper, typename... InputDomain>
+class FlatMapImpl
+    : public FlatMapImplBase<FlatMapImpl<FlatMapper, InputDomain...>,
+                             FlatMapper, InputDomain...> {
+ public:
+  using typename FlatMapImpl::FlatMapImplBase::corpus_type;
+  using typename FlatMapImpl::FlatMapImplBase::value_type;
+
+  using FlatMapImpl::FlatMapImplBase::FlatMapImplBase;
+
+  static constexpr absl::string_view GetName() { return "FlatMap"; }
+
+  std::optional<corpus_type> FromValue(const value_type&) const {
+    // We cannot infer the input corpus from the output value, or even determine
+    // from which output domain the output value came.
+    return std::nullopt;
+  }
+};
+
+// -----------------------------------------------------------------------------
+// ReversibleFlatMap Implementation
+// -----------------------------------------------------------------------------
+
+template <typename FlatMapper, typename InvMapper, typename... InputDomain>
+class ReversibleFlatMapImpl
+    : public FlatMapImplBase<
+          ReversibleFlatMapImpl<FlatMapper, InvMapper, InputDomain...>,
+          FlatMapper, InputDomain...> {
+ public:
+  using typename ReversibleFlatMapImpl::FlatMapImplBase::corpus_type;
+  using typename ReversibleFlatMapImpl::FlatMapImplBase::value_type;
+
+  static_assert(
+      std::is_invocable_v<InvMapper, const value_type&> &&
+          std::is_same_v<
+              std::invoke_result_t<InvMapper, const value_type&>,
+              std::optional<std::tuple<value_type_t<InputDomain>...>>>,
+      "ReversibleFlatMap must have an inverse mapper that takes the output "
+      "value and returns an optional of a tuple of the input values.");
+
+  ReversibleFlatMapImpl() = default;
+  explicit ReversibleFlatMapImpl(FlatMapper flat_mapper, InvMapper inv_mapper,
+                                 InputDomain... input_domains)
+      : ReversibleFlatMapImpl::FlatMapImplBase(std::move(flat_mapper),
+                                               std::move(input_domains)...),
+        inv_mapper_(std::move(inv_mapper)) {}
+
+  static constexpr absl::string_view GetName() { return "ReversibleFlatMap"; }
+
+  std::optional<corpus_type> FromValue(const value_type& v) const {
+    // 1. Recover the input values using the user-provided inverse mapper.
+    auto input_values_opt = std::invoke(inv_mapper_, v);
+    if (!input_values_opt.has_value()) return std::nullopt;
+
+    // 2. Map input values into input corpus values.
+    auto input_corpus_opt =
+        ApplyIndex<ReversibleFlatMapImpl::FlatMapImplBase::kNumInputValues>(
+            [&](auto... I)
+                -> std::optional<std::tuple<corpus_type_t<InputDomain>...>> {
+              auto inner_corpus_vals =
+                  std::tuple{std::get<I>(this->input_domains())
+                                 .FromValue(std::get<I>(*input_values_opt))...};
+              bool has_nullopt =
+                  (!std::get<I>(inner_corpus_vals).has_value() || ...);
+              if (has_nullopt) return std::nullopt;
+              return std::tuple{*std::move(std::get<I>(inner_corpus_vals))...};
+            });
+    if (!input_corpus_opt.has_value()) return std::nullopt;
+
+    if (!this->ValidateInputValues(*input_corpus_opt).ok()) return std::nullopt;
+
+    // 3. Re-instantiate the dynamically generated output domain.
+    auto output_domain = this->GetOutputDomain(*input_corpus_opt);
+
+    // 4. Map the output value into the output corpus value.
+    auto output_corpus_opt = output_domain.FromValue(v);
+    if (!output_corpus_opt.has_value()) return std::nullopt;
+
+    if (!output_domain.ValidateCorpusValue(*output_corpus_opt).ok()) {
+      return std::nullopt;
+    }
+    // 5. Assemble the final corpus tuple (output corpus followed by input
+    // corpus).
+    return std::tuple_cat(std::make_tuple(*std::move(output_corpus_opt)),
+                          *std::move(input_corpus_opt));
+  }
+
+ private:
+  InvMapper inv_mapper_;
+};
+
 }  // namespace fuzztest::internal
 
 #endif  // FUZZTEST_FUZZTEST_INTERNAL_DOMAINS_FLAT_MAP_IMPL_H_