Fix overlay patch mutation bug (#10579)

* Fix overlay patch mutation bug

* Add change log

* Add include

* Feedback

Firestore/CHANGELOG.md

@@ -1,3 +1,6 @@
+# Unreleased
+- [fixed] Fix an issue that stops some performance optimization being applied.
+
 # 10.3.0
 - [feature] Add MultiDb support.
 - [fixed] Fix App crashed when there are nested data structures inside IN

Firestore/core/src/local/local_documents_view.cc

@@ -266,6 +266,9 @@ model::OverlayedDocumentMap LocalDocumentsView::ComputeViews(
           {doc->key(), overlay_it->second.mutation().field_mask()});
       overlay_it->second.mutation().ApplyToLocalView(*doc, absl::nullopt,
                                                      Timestamp::Now());
+    } else {  // No overlay for this document
+      // Using empty mask to indicate there is no overlay for the document.
+      mutated_fields.emplace(doc->key(), FieldMask{});
     }
   }
 

Firestore/core/src/model/overlayed_document.h

@@ -26,8 +26,10 @@ namespace firebase {
 namespace firestore {
 namespace model {
 
-/** Represents a local view (overlay) of a document, and the fields that are
- * locally mutated. */
+/**
+ * Represents a local view (overlay) of a document, and the fields that are
+ * locally mutated.
+ */
 class OverlayedDocument {
  public:
   OverlayedDocument(model::Document document,
@@ -44,6 +46,13 @@ class OverlayedDocument {
     return std::move(document_);
   }
 
+  /**
+   * The fields that are locally mutated by patch mutations.
+   *
+   * If the overlayed document is from set or delete mutations, returns
+   * `nullopt`. If there is no overlay (mutation) for the document, returns
+   * empty `FieldMask`.
+   */
   const absl::optional<model::FieldMask>& mutated_fields() const {
     return mutated_fields_;
   }

Firestore/core/test/unit/local/counting_query_engine.cc

@@ -60,6 +60,7 @@ void CountingQueryEngine::ResetCounts() {
   overlays_read_by_key_ = 0;
   overlays_read_by_collection_ = 0;
   overlays_read_by_collection_group_ = 0;
+  overlay_types_.clear();
 }
 
 // MARK: - WrappedMutationQueue
@@ -198,7 +199,13 @@ model::MutableDocumentMap WrappedRemoteDocumentCache::GetAll(
 absl::optional<model::Overlay> WrappedDocumentOverlayCache::GetOverlay(
     const model::DocumentKey& key) const {
   ++query_engine_->overlays_read_by_key_;
-  return subject_->GetOverlay(key);
+  auto result = subject_->GetOverlay(key);
+  if (result.has_value()) {
+    query_engine_->overlay_types_.emplace(key,
+                                          result.value().mutation().type());
+  }
+
+  return result;
 }
 
 void WrappedDocumentOverlayCache::SaveOverlays(
@@ -214,6 +221,10 @@ OverlayByDocumentKeyMap WrappedDocumentOverlayCache::GetOverlays(
     const model::ResourcePath& collection, int since_batch_id) const {
   auto result = subject_->GetOverlays(collection, since_batch_id);
   query_engine_->overlays_read_by_collection_ += result.size();
+  for (const auto& r : result) {
+    query_engine_->overlay_types_.emplace(r.first, r.second.mutation().type());
+  }
+
   return result;
 }
 
@@ -223,6 +234,10 @@ OverlayByDocumentKeyMap WrappedDocumentOverlayCache::GetOverlays(
     std::size_t count) const {
   auto result = subject_->GetOverlays(collection_group, since_batch_id, count);
   query_engine_->overlays_read_by_collection_group_ += result.size();
+  for (const auto& r : result) {
+    query_engine_->overlay_types_.emplace(r.first, r.second.mutation().type());
+  }
+
   return result;
 }
 

Firestore/core/test/unit/local/counting_query_engine.h

@@ -19,6 +19,7 @@
 
 #include <memory>
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -26,6 +27,7 @@
 #include "Firestore/core/src/local/mutation_queue.h"
 #include "Firestore/core/src/local/query_engine.h"
 #include "Firestore/core/src/local/remote_document_cache.h"
+#include "Firestore/core/src/model/document_key.h"
 #include "Firestore/core/src/model/model_fwd.h"
 #include "Firestore/core/src/util/hard_assert.h"
 
@@ -90,6 +92,13 @@ class CountingQueryEngine : public QueryEngine {
     return overlays_read_by_key_;
   }
 
+  std::unordered_map<model::DocumentKey,
+                     model::Mutation::Type,
+                     model::DocumentKeyHash>
+  overlay_types() const {
+    return overlay_types_;
+  }
+
  private:
   friend class WrappedDocumentOverlayCache;
   friend class WrappedMutationQueue;
@@ -107,6 +116,10 @@ class CountingQueryEngine : public QueryEngine {
   size_t overlays_read_by_key_ = 0;
   size_t overlays_read_by_collection_ = 0;
   size_t overlays_read_by_collection_group_ = 0;
+  std::unordered_map<model::DocumentKey,
+                     model::Mutation::Type,
+                     model::DocumentKeyHash>
+      overlay_types_;
 };
 
 /** A MutationQueue that counts document reads. */

Firestore/core/test/unit/local/leveldb_local_store_test.cc

@@ -45,6 +45,7 @@ using testutil::Key;
 using testutil::MakeFieldIndex;
 using testutil::Map;
 using testutil::OrderBy;
+using testutil::OverlayTypeMap;
 using testutil::SetMutation;
 using testutil::UpdateRemoteEvent;
 using testutil::Vector;
@@ -238,6 +239,10 @@ TEST_F(LevelDbLocalStoreTest, UsesPartiallyIndexedOverlaysWhenAvailable) {
       testutil::Query("coll").AddingFilter(Filter("matches", "==", true));
   ExecuteQuery(query);
   FSTAssertOverlaysRead(/* byKey= */ 1, /* byCollection= */ 1);
+  FSTAssertOverlayTypes(
+      OverlayTypeMap({{Key("coll/a"), model::Mutation::Type::Set},
+                      {Key("coll/b"), model::Mutation::Type::Set}}));
+
   FSTAssertQueryReturned("coll/a", "coll/b");
 }
 
@@ -265,6 +270,9 @@ TEST_F(LevelDbLocalStoreTest, DoesNotUseLimitWhenIndexIsOutdated) {
   // query without limit.
   FSTAssertRemoteDocumentsRead(/* byKey= */ 5, /* byCollection= */ 0);
   FSTAssertOverlaysRead(/* byKey= */ 5, /* byCollection= */ 1);
+  FSTAssertOverlayTypes(
+      OverlayTypeMap({{Key("coll/b"), model::Mutation::Type::Delete}}));
+
   FSTAssertQueryReturned("coll/a", "coll/c");
 }
 
@@ -289,6 +297,8 @@ TEST_F(LevelDbLocalStoreTest, UsesIndexForLimitQueryWhenIndexIsUpdated) {
   ExecuteQuery(query);
   FSTAssertRemoteDocumentsRead(/* byKey= */ 2, /* byCollection= */ 0);
   FSTAssertOverlaysRead(/* byKey= */ 2, /* byCollection= */ 0);
+  FSTAssertOverlayTypes(OverlayTypeMap({}));
+
   FSTAssertQueryReturned("coll/a", "coll/c");
 }
 
@@ -306,6 +316,9 @@ TEST_F(LevelDbLocalStoreTest, IndexesServerTimestamps) {
 
   ExecuteQuery(query);
   FSTAssertOverlaysRead(/* byKey= */ 1, /* byCollection= */ 0);
+  FSTAssertOverlayTypes(
+      OverlayTypeMap({{Key("coll/a"), model::Mutation::Type::Set}}));
+
   FSTAssertQueryReturned("coll/a");
 }
 

Firestore/core/test/unit/local/local_store_test.cc

@@ -16,6 +16,7 @@
 
 #include "Firestore/core/test/unit/local/local_store_test.h"
 
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -31,8 +32,10 @@
 #include "Firestore/core/src/local/target_data.h"
 #include "Firestore/core/src/model/delete_mutation.h"
 #include "Firestore/core/src/model/document.h"
+#include "Firestore/core/src/model/document_key.h"
 #include "Firestore/core/src/model/field_index.h"
 #include "Firestore/core/src/model/mutable_document.h"
+#include "Firestore/core/src/model/mutation.h"
 #include "Firestore/core/src/model/mutation_batch_result.h"
 #include "Firestore/core/src/model/patch_mutation.h"
 #include "Firestore/core/src/model/set_mutation.h"
@@ -87,6 +90,7 @@ using testutil::DeletedDoc;
 using testutil::Doc;
 using testutil::Key;
 using testutil::Map;
+using testutil::OverlayTypeMap;
 using testutil::Query;
 using testutil::UnknownDoc;
 using testutil::UpdateRemoteEvent;
@@ -906,6 +910,8 @@ TEST_P(LocalStoreTest, ReadsAllDocumentsForInitialCollectionQueries) {
 
   FSTAssertRemoteDocumentsRead(/* by_key= */ 0, /* by_query= */ 2);
   FSTAssertOverlaysRead(/* by_key= */ 0, /* by_query= */ 1);
+  FSTAssertOverlayTypes(
+      OverlayTypeMap({{Key("foo/bonk"), model::Mutation::Type::Set}}));
 }
 
 TEST_P(LocalStoreTest, PersistsResumeTokens) {
@@ -1663,6 +1669,21 @@ TEST_P(LocalStoreTest, MultipleFieldPatchesOnLocalDocs) {
       Doc("foo/bar", 0, Map("likes", 1, "stars", 2)).SetHasLocalMutations());
 }
 
+TEST_P(LocalStoreTest, PatchMutationLeadsToPatchOverlay) {
+  AllocateQuery(Query("foo"));
+  ApplyRemoteEvent(UpdateRemoteEvent(Doc("foo/baz", 10, Map("a", 1)), {2}, {}));
+  ApplyRemoteEvent(UpdateRemoteEvent(Doc("foo/bar", 20, Map()), {2}, {}));
+  WriteMutation(testutil::PatchMutation("foo/baz", Map("b", 2)));
+
+  ResetPersistenceStats();
+
+  ExecuteQuery(Query("foo"));
+  FSTAssertRemoteDocumentsRead(0, 2);
+  FSTAssertOverlaysRead(0, 1);
+  FSTAssertOverlayTypes(
+      OverlayTypeMap({{Key("foo/baz"), model::Mutation::Type::Patch}}));
+}
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase

Firestore/core/test/unit/local/local_store_test.h

@@ -214,9 +214,15 @@ class LocalStoreTest : public LocalStoreTestBase,
 #define FSTAssertOverlaysRead(by_key, by_query)                        \
   do {                                                                 \
     ASSERT_EQ(query_engine_.overlays_read_by_key(), (by_key))          \
-        << "Mutations read (by key)";                                  \
+        << "Overlays read (by key)";                                   \
     ASSERT_EQ(query_engine_.overlays_read_by_collection(), (by_query)) \
-        << "Mutations read (by query)";                                \
+        << "Overlays read (by query)";                                 \
+  } while (0)
+
+#define FSTAssertOverlayTypes(expected_types)                \
+  do {                                                       \
+    ASSERT_EQ(query_engine_.overlay_types(), expected_types) \
+        << "Overlay types read";                             \
   } while (0)
 
 /**

Firestore/core/test/unit/testutil/testutil.h

@@ -20,6 +20,7 @@
 #include <algorithm>
 #include <memory>
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -30,6 +31,7 @@
 #include "Firestore/core/src/model/document_set.h"
 #include "Firestore/core/src/model/field_index.h"
 #include "Firestore/core/src/model/model_fwd.h"
+#include "Firestore/core/src/model/mutation.h"
 #include "Firestore/core/src/model/precondition.h"
 #include "Firestore/core/src/model/value_util.h"
 #include "Firestore/core/src/nanopb/byte_string.h"
@@ -158,6 +160,10 @@ nanopb::Message<google_firestore_v1_Value> Value(
 nanopb::Message<google_firestore_v1_Value> Value(
     const model::ObjectValue& value);
 
+using OverlayTypeMap = std::unordered_map<model::DocumentKey,
+                                          model::Mutation::Type,
+                                          model::DocumentKeyHash>;
+
 namespace details {
 
 /**