Add Overlay class

Author: schmidt-sebastian

firebase-firestore/src/main/java/com/google/firebase/firestore/local/DocumentOverlayCache.java

@@ -18,6 +18,7 @@
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.mutation.Mutation;
+import com.google.firebase.firestore.model.mutation.Overlay;
 import java.util.Map;
 
 /**
@@ -35,7 +36,7 @@ public interface DocumentOverlayCache {
    * for that key.
    */
   @Nullable
-  Mutation getOverlay(DocumentKey key);
+  Overlay getOverlay(DocumentKey key);
 
   /**
    * Saves the given document key to mutation map to persistence as overlays. All overlays will have
@@ -52,6 +53,7 @@ public interface DocumentOverlayCache {
    * @param collection The collection path to get the overlays for.
    * @param sinceBatchId The minimum batch ID to filter by (exclusive). Only overlays that contain a
    *     change past `sinceBatchId` are returned.
+   * @return Mapping of each document key in the collection to its overlay.
    */
-  Map<DocumentKey, Mutation> getOverlays(ResourcePath collection, int sinceBatchId);
+  Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceBatchId);
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalDocumentsView.java

@@ -29,6 +29,7 @@
 import com.google.firebase.firestore.model.mutation.FieldMask;
 import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.MutationBatch;
+import com.google.firebase.firestore.model.mutation.Overlay;
 import com.google.firebase.firestore.model.mutation.PatchMutation;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -42,7 +43,6 @@
  * in remoteDocumentCache or local mutations for the document). The view is computed by applying the
  * mutations in the MutationQueue to the RemoteDocumentCache.
  */
-// TODO: Turn this into the UnifiedDocumentCache / whatever.
 class LocalDocumentsView {
 
   private final RemoteDocumentCache remoteDocumentCache;
@@ -88,14 +88,14 @@ IndexManager getIndexManager() {
    *     for it.
    */
   Document getDocument(DocumentKey key) {
-    Mutation overlay = documentOverlayCache.getOverlay(key);
+    Overlay overlay = documentOverlayCache.getOverlay(key);
     // Only read from remote document cache if overlay is a patch.
     MutableDocument document =
-        (overlay == null || overlay instanceof PatchMutation)
+        (overlay == null || overlay.getMutation() instanceof PatchMutation)
             ? remoteDocumentCache.get(key)
             : MutableDocument.newInvalidDocument(key);
     if (overlay != null) {
-      overlay.applyToLocalView(document, null, Timestamp.now());
+      overlay.getMutation().applyToLocalView(document, null, Timestamp.now());
     }
 
     return document;
@@ -125,18 +125,18 @@ ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
     ImmutableSortedMap<DocumentKey, Document> results = emptyDocumentMap();
     Map<DocumentKey, MutableDocument> recalculateDocuments = new HashMap<>();
     for (Map.Entry<DocumentKey, MutableDocument> entry : docs.entrySet()) {
-      Mutation overlay = documentOverlayCache.getOverlay(entry.getKey());
+      Overlay overlay = documentOverlayCache.getOverlay(entry.getKey());
       // Recalculate an overlay if the document's existence state is changed due to a remote
       // event *and* the overlay is a PatchMutation. This is because document existence state
       // can change if some patch mutation's preconditions are met.
       // NOTE: we recalculate when `overlay` is null as well, because there might be a patch
       // mutation whose precondition does not match before the change (hence overlay==null),
       // but would now match.
       if (existenceStateChanged.contains(entry.getKey())
-          && (overlay == null || overlay instanceof PatchMutation)) {
+          && (overlay == null || overlay.getMutation() instanceof PatchMutation)) {
         recalculateDocuments.put(entry.getKey(), docs.get(entry.getKey()));
       } else if (overlay != null) {
-        overlay.applyToLocalView(entry.getValue(), null, Timestamp.now());
+        overlay.getMutation().applyToLocalView(entry.getValue(), null, Timestamp.now());
       }
     }
 
@@ -260,11 +260,11 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
       Query query, IndexOffset offset) {
     Map<DocumentKey, MutableDocument> remoteDocuments =
         remoteDocumentCache.getAll(query.getPath(), offset);
-    Map<DocumentKey, Mutation> overlays = documentOverlayCache.getOverlays(query.getPath(), -1);
+    Map<DocumentKey, Overlay> overlays = documentOverlayCache.getOverlays(query.getPath(), -1);
 
     // As documents might match the query because of their overlay we need to include documents
     // for all overlays in the initial document set.
-    for (Map.Entry<DocumentKey, Mutation> entry : overlays.entrySet()) {
+    for (Map.Entry<DocumentKey, Overlay> entry : overlays.entrySet()) {
       if (!remoteDocuments.containsKey(entry.getKey())) {
         remoteDocuments.put(entry.getKey(), MutableDocument.newInvalidDocument(entry.getKey()));
       }
@@ -273,9 +273,9 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
     // Apply the overlays and match against the query.
     ImmutableSortedMap<DocumentKey, Document> results = emptyDocumentMap();
     for (Map.Entry<DocumentKey, MutableDocument> docEntry : remoteDocuments.entrySet()) {
-      Mutation overlay = overlays.get(docEntry.getKey());
+      Overlay overlay = overlays.get(docEntry.getKey());
       if (overlay != null) {
-        overlay.applyToLocalView(docEntry.getValue(), null, Timestamp.now());
+        overlay.getMutation().applyToLocalView(docEntry.getValue(), null, Timestamp.now());
       }
       // Finally, insert the documents that still match the query
       if (query.matches(docEntry.getValue())) {

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryDocumentOverlayCache.java

@@ -14,11 +14,11 @@
 
 package com.google.firebase.firestore.local;
 
-import android.util.Pair;
 import androidx.annotation.Nullable;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.mutation.Mutation;
+import com.google.firebase.firestore.model.mutation.Overlay;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -28,17 +28,13 @@
 public class MemoryDocumentOverlayCache implements DocumentOverlayCache {
   // A map sorted by DocumentKey, whose value is a pair of the largest batch id for the overlay
   // and the overlay itself.
-  private final TreeMap<DocumentKey, Pair<Integer, Mutation>> overlays = new TreeMap<>();
+  private final TreeMap<DocumentKey, Overlay> overlays = new TreeMap<>();
   private final Map<Integer, Set<DocumentKey>> overlayByBatchId = new HashMap<>();
 
   @Nullable
   @Override
-  public Mutation getOverlay(DocumentKey key) {
-    Pair<Integer, Mutation> overlay = overlays.get(key);
-    if (overlay != null) {
-      return overlay.second;
-    }
-    return null;
+  public Overlay getOverlay(DocumentKey key) {
+    return overlays.get(key);
   }
 
   private void saveOverlay(int largestBatchId, @Nullable Mutation mutation) {
@@ -47,12 +43,12 @@ private void saveOverlay(int largestBatchId, @Nullable Mutation mutation) {
     }
 
     // Remove the association of the overlay to its batch id.
-    Pair<Integer, Mutation> existing = this.overlays.get(mutation.getKey());
+    Overlay existing = this.overlays.get(mutation.getKey());
     if (existing != null) {
-      overlayByBatchId.get(existing.first).remove(mutation.getKey());
+      overlayByBatchId.get(existing.getLargestBatchId()).remove(mutation.getKey());
     }
 
-    overlays.put(mutation.getKey(), new Pair<>(largestBatchId, mutation));
+    overlays.put(mutation.getKey(), Overlay.create(largestBatchId, mutation));
 
     // Create the associate of this overlay to the given largestBatchId.
     if (overlayByBatchId.get(largestBatchId) == null) {
@@ -80,15 +76,15 @@ public void removeOverlaysForBatchId(int batchId) {
   }
 
   @Override
-  public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection, int sinceBatchId) {
-    Map<DocumentKey, Mutation> result = new HashMap<>();
+  public Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceBatchId) {
+    Map<DocumentKey, Overlay> result = new HashMap<>();
 
     int immediateChildrenPathLength = collection.length() + 1;
     DocumentKey prefix = DocumentKey.fromPath(collection.append(""));
-    Map<DocumentKey, Pair<Integer, Mutation>> view = overlays.tailMap(prefix);
+    Map<DocumentKey, Overlay> view = overlays.tailMap(prefix);
 
-    for (Map.Entry<DocumentKey, Pair<Integer, Mutation>> entry : view.entrySet()) {
-      DocumentKey key = entry.getKey();
+    for (Overlay overlay : view.values()) {
+      DocumentKey key = overlay.getKey();
       if (!collection.isPrefixOf(key.getPath())) {
         break;
       }
@@ -97,9 +93,8 @@ public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection, int since
         continue;
       }
 
-      Pair<Integer, Mutation> batchIdToOverlay = entry.getValue();
-      if (batchIdToOverlay.first > sinceBatchId) {
-        result.put(entry.getKey(), batchIdToOverlay.second);
+      if (overlay.getLargestBatchId() > sinceBatchId) {
+        result.put(overlay.getKey(), overlay);
       }
     }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteDocumentOverlayCache.java

@@ -21,6 +21,7 @@
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.mutation.Mutation;
+import com.google.firebase.firestore.model.mutation.Overlay;
 import com.google.firestore.v1.Write;
 import com.google.protobuf.InvalidProtocolBufferException;
 import java.util.HashMap;
@@ -39,23 +40,14 @@ public SQLiteDocumentOverlayCache(SQLitePersistence db, LocalSerializer serializ
 
   @Nullable
   @Override
-  public Mutation getOverlay(DocumentKey key) {
+  public Overlay getOverlay(DocumentKey key) {
     String collectionPath = EncodedPath.encode(key.getPath().popLast());
     String documentId = key.getPath().getLastSegment();
     return db.query(
-            "SELECT overlay_mutation FROM document_overlays WHERE uid = ? AND collection_path = ? AND document_id = ?")
+            "SELECT overlay_mutation, largest_batch_id FROM document_overlays "
+                + "WHERE uid = ? AND collection_path = ? AND document_id = ?")
         .binding(uid, collectionPath, documentId)
-        .firstValue(
-            row -> {
-              if (row == null) return null;
-
-              try {
-                Write mutation = Write.parseFrom(row.getBlob(0));
-                return serializer.decodeMutation(mutation);
-              } catch (InvalidProtocolBufferException e) {
-                throw fail("Overlay failed to parse: %s", e);
-              }
-            });
+        .firstValue(this::decodeOverlay);
   }
 
   private void saveOverlay(int largestBatchId, DocumentKey key, @Nullable Mutation mutation) {
@@ -90,28 +82,31 @@ public void removeOverlaysForBatchId(int batchId) {
   }
 
   @Override
-  public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection, int sinceBatchId) {
+  public Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceBatchId) {
     String collectionPath = EncodedPath.encode(collection);
 
-    Map<DocumentKey, Mutation> result = new HashMap<>();
-
+    Map<DocumentKey, Overlay> result = new HashMap<>();
     db.query(
-            "SELECT document_id, overlay_mutation FROM document_overlays "
+            "SELECT overlay_mutation, largest_batch_id FROM document_overlays "
                 + "WHERE uid = ? AND collection_path = ? AND largest_batch_id > ?")
         .binding(uid, collectionPath, sinceBatchId)
         .forEach(
             row -> {
-              try {
-                String documentId = row.getString(0);
-                Write write = Write.parseFrom(row.getBlob(1));
-                Mutation mutation = serializer.decodeMutation(write);
-
-                result.put(DocumentKey.fromPath(collection.append(documentId)), mutation);
-              } catch (InvalidProtocolBufferException e) {
-                throw fail("Overlay failed to parse: %s", e);
-              }
+              Overlay overlay = decodeOverlay(row);
+              result.put(overlay.getKey(), overlay);
             });
 
     return result;
   }
+
+  private Overlay decodeOverlay(android.database.Cursor row) {
+    try {
+      Write write = Write.parseFrom(row.getBlob(0));
+      int largestBatchId = row.getInt(1);
+      Mutation mutation = serializer.decodeMutation(write);
+      return Overlay.create(largestBatchId, mutation);
+    } catch (InvalidProtocolBufferException e) {
+      throw fail("Overlay failed to parse: %s", e);
+    }
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteSchema.java

@@ -297,7 +297,7 @@ private void removeMutationBatch(String uid, int batchId) {
     mutationDeleter.bindString(1, uid);
     mutationDeleter.bindLong(2, batchId);
     int deleted = mutationDeleter.executeUpdateDelete();
-    hardAssert(deleted != 0, "Mutatiohn batch (%s, %d) did not exist", uid, batchId);
+    hardAssert(deleted != 0, "Mutation batch (%s, %d) did not exist", uid, batchId);
 
     // Delete all index entries for this batch
     db.execSQL(

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/Overlay.java

@@ -0,0 +1,39 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.firestore.model.mutation;
+
+import com.google.auto.value.AutoValue;
+import com.google.firebase.firestore.model.DocumentKey;
+
+/**
+ * Representation of an overlay computed by Firestore.
+ *
+ * <p>Holds information about a mutation and the largest batch id in Firestore when the mutation was
+ * created.
+ */
+@AutoValue
+public abstract class Overlay {
+  public static Overlay create(int largestBatchId, Mutation mutation) {
+    return new AutoValue_Overlay(largestBatchId, mutation);
+  }
+
+  public abstract int getLargestBatchId();
+
+  public abstract Mutation getMutation();
+
+  public DocumentKey getKey() {
+    return getMutation().getKey();
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/util/Util.java

@@ -392,11 +392,11 @@ public static <K, V> Map<K, V> firstNEntries(Map<K, V> data, int n, Comparator<V
     if (data.size() <= n) {
       return data;
     } else {
-      List<Map.Entry<K, V>> sortedVlaues = new ArrayList<>(data.entrySet());
-      Collections.sort(sortedVlaues, (l, r) -> comp.compare(l.getValue(), r.getValue()));
+      List<Map.Entry<K, V>> sortedValues = new ArrayList<>(data.entrySet());
+      Collections.sort(sortedValues, (l, r) -> comp.compare(l.getValue(), r.getValue()));
       Map<K, V> result = new HashMap<>();
       for (int i = 0; i < n; ++i) {
-        result.put(sortedVlaues.get(i).getKey(), sortedVlaues.get(i).getValue());
+        result.put(sortedValues.get(i).getKey(), sortedValues.get(i).getValue());
       }
       return result;
     }