View changes for held write acks

firebase-firestore/src/main/java/com/google/firebase/firestore/DocumentChange.java

@@ -142,7 +142,11 @@ static List<DocumentChange> changesFromSnapshot(
       for (DocumentViewChange change : snapshot.getChanges()) {
         Document document = change.getDocument();
         QueryDocumentSnapshot documentSnapshot =
-            QueryDocumentSnapshot.fromDocument(firestore, document, snapshot.isFromCache());
+            QueryDocumentSnapshot.fromDocument(
+                firestore,
+                document,
+                snapshot.isFromCache(),
+                snapshot.getMutatedKeys().contains(document.getKey()));
         hardAssert(
             change.getType() == DocumentViewChange.Type.ADDED,
             "Invalid added event for first snapshot");
@@ -163,7 +167,11 @@ static List<DocumentChange> changesFromSnapshot(
         }
         Document document = change.getDocument();
         QueryDocumentSnapshot documentSnapshot =
-            QueryDocumentSnapshot.fromDocument(firestore, document, snapshot.isFromCache());
+            QueryDocumentSnapshot.fromDocument(
+                firestore,
+                document,
+                snapshot.isFromCache(),
+                snapshot.getMutatedKeys().contains(document.getKey()));
         int oldIndex, newIndex;
         Type type = getType(change);
         if (type != Type.ADDED) {

firebase-firestore/src/main/java/com/google/firebase/firestore/DocumentReference.java

@@ -317,8 +317,12 @@ public Task<DocumentSnapshot> get(Source source) {
           .getDocumentFromLocalCache(key)
           .continueWith(
               Executors.DIRECT_EXECUTOR,
-              (Task<Document> doc) ->
-                  new DocumentSnapshot(firestore, key, doc.getResult(), /*isFromCache=*/ true));
+              (Task<Document> task) -> {
+                Document doc = task.getResult();
+                boolean hasPendingWrites = doc != null && doc.hasLocalMutations();
+                return new DocumentSnapshot(
+                    firestore, key, doc, /*isFromCache=*/ true, hasPendingWrites);
+              });
     } else {
       return getViaSnapshotListener(source);
     }
@@ -523,11 +527,16 @@ private ListenerRegistration addSnapshotListenerInternal(
                 Document document = snapshot.getDocuments().getDocument(key);
                 DocumentSnapshot documentSnapshot;
                 if (document != null) {
+                  boolean hasPendingWrites = snapshot.getMutatedKeys().contains(document.getKey());
                   documentSnapshot =
-                      DocumentSnapshot.fromDocument(firestore, document, snapshot.isFromCache());
+                      DocumentSnapshot.fromDocument(
+                          firestore, document, snapshot.isFromCache(), hasPendingWrites);
                 } else {
+                  // We don't raise `hasPendingWrites` for deleted documents.
+                  boolean hasPendingWrites = false;
                   documentSnapshot =
-                      DocumentSnapshot.fromNoDocument(firestore, key, snapshot.isFromCache());
+                      DocumentSnapshot.fromNoDocument(
+                          firestore, key, snapshot.isFromCache(), hasPendingWrites);
                 }
                 listener.onEvent(documentSnapshot, null);
               } else {

firebase-firestore/src/main/java/com/google/firebase/firestore/DocumentSnapshot.java

@@ -90,22 +90,25 @@ public enum ServerTimestampBehavior {
   private final SnapshotMetadata metadata;
 
   DocumentSnapshot(
-      FirebaseFirestore firestore, DocumentKey key, @Nullable Document doc, boolean isFromCache) {
+      FirebaseFirestore firestore,
+      DocumentKey key,
+      @Nullable Document doc,
+      boolean isFromCache,
+      boolean hasPendingWrites) {
     this.firestore = checkNotNull(firestore);
     this.key = checkNotNull(key);
     this.doc = doc;
-    boolean hasPendingWrites = this.doc != null && this.doc.hasLocalMutations();
     this.metadata = new SnapshotMetadata(hasPendingWrites, isFromCache);
   }
 
   static DocumentSnapshot fromDocument(
-      FirebaseFirestore firestore, Document doc, boolean fromCache) {
-    return new DocumentSnapshot(firestore, doc.getKey(), doc, fromCache);
+      FirebaseFirestore firestore, Document doc, boolean fromCache, boolean hasPendingWrites) {
+    return new DocumentSnapshot(firestore, doc.getKey(), doc, fromCache, hasPendingWrites);
   }
 
   static DocumentSnapshot fromNoDocument(
-      FirebaseFirestore firestore, DocumentKey key, boolean fromCache) {
-    return new DocumentSnapshot(firestore, key, null, fromCache);
+      FirebaseFirestore firestore, DocumentKey key, boolean fromCache, boolean hasPendingWrites) {
+    return new DocumentSnapshot(firestore, key, null, fromCache, hasPendingWrites);
   }
 
   /** @return The id of the document. */

firebase-firestore/src/main/java/com/google/firebase/firestore/QueryDocumentSnapshot.java

@@ -41,13 +41,17 @@
 public class QueryDocumentSnapshot extends DocumentSnapshot {
 
   private QueryDocumentSnapshot(
-      FirebaseFirestore firestore, DocumentKey key, @Nullable Document doc, boolean isFromCache) {
-    super(firestore, key, doc, isFromCache);
+      FirebaseFirestore firestore,
+      DocumentKey key,
+      @Nullable Document doc,
+      boolean isFromCache,
+      boolean hasPendingWrites) {
+    super(firestore, key, doc, isFromCache, hasPendingWrites);
   }
 
   static QueryDocumentSnapshot fromDocument(
-      FirebaseFirestore firestore, Document doc, boolean fromCache) {
-    return new QueryDocumentSnapshot(firestore, doc.getKey(), doc, fromCache);
+      FirebaseFirestore firestore, Document doc, boolean fromCache, boolean hasPendingWrites) {
+    return new QueryDocumentSnapshot(firestore, doc.getKey(), doc, fromCache, hasPendingWrites);
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/QuerySnapshot.java

@@ -194,7 +194,11 @@ public <T> List<T> toObjects(
   }
 
   private QueryDocumentSnapshot convertDocument(Document document) {
-    return QueryDocumentSnapshot.fromDocument(firestore, document, snapshot.isFromCache());
+    return QueryDocumentSnapshot.fromDocument(
+        firestore,
+        document,
+        snapshot.isFromCache(),
+        snapshot.getMutatedKeys().contains(document.getKey()));
   }
 
   @Override

firebase-firestore/src/main/java/com/google/firebase/firestore/Transaction.java

@@ -242,10 +242,10 @@ private Task<DocumentSnapshot> getAsync(DocumentReference documentRef) {
               MaybeDocument doc = docs.get(0);
               if (doc instanceof Document) {
                 return DocumentSnapshot.fromDocument(
-                    firestore, (Document) doc, /*fromCache=*/ false);
+                    firestore, (Document) doc, /*fromCache=*/ false, /*hasPendingWrites=*/ false);
               } else if (doc instanceof NoDocument) {
                 return DocumentSnapshot.fromNoDocument(
-                    firestore, doc.getKey(), /*fromCache=*/ false);
+                    firestore, doc.getKey(), /*fromCache=*/ false, /*hasPendingWrites=*/ false);
               } else {
                 throw fail(
                     "BatchGetDocumentsRequest returned unexpected document type: "

firebase-firestore/src/main/java/com/google/firebase/firestore/core/QueryListener.java

@@ -81,7 +81,7 @@ public void onViewSnapshot(ViewSnapshot newSnapshot) {
               newSnapshot.getOldDocuments(),
               documentChanges,
               newSnapshot.isFromCache(),
-              newSnapshot.hasPendingWrites(),
+              newSnapshot.getMutatedKeys(),
               newSnapshot.didSyncStateChange());
     }
 
@@ -160,7 +160,7 @@ private void raiseInitialEvent(ViewSnapshot snapshot) {
             DocumentSet.emptySet(snapshot.getQuery().comparator()),
             QueryListener.getInitialViewChanges(snapshot),
             snapshot.isFromCache(),
-            snapshot.hasPendingWrites(),
+            snapshot.getMutatedKeys(),
             /*didSyncStateChange=*/ true);
     raisedInitialEvent = true;
     listener.onEvent(snapshot, null);

firebase-firestore/src/main/java/com/google/firebase/firestore/core/SyncEngine.java

@@ -364,7 +364,9 @@ public void handleRejectedListen(int targetId, Status error) {
       // Ideally, we would have a method in the local store to purge a document. However, it would
       // be tricky to keep all of the local store's invariants with another method.
       Map<DocumentKey, MaybeDocument> documentUpdates =
-          Collections.singletonMap(limboKey, new NoDocument(limboKey, SnapshotVersion.NONE));
+          Collections.singletonMap(
+              limboKey,
+              new NoDocument(limboKey, SnapshotVersion.NONE, /*hasCommittedMutations=*/ false));
       Set<DocumentKey> limboDocuments = Collections.singleton(limboKey);
       RemoteEvent event =
           new RemoteEvent(

firebase-firestore/src/main/java/com/google/firebase/firestore/core/View.java

@@ -169,49 +169,66 @@ public <D extends MaybeDocument> DocumentChanges computeDocChanges(
         }
       }
 
-      if (newDoc != null) {
-        newDocumentSet = newDocumentSet.add(newDoc);
-        if (newDoc.hasLocalMutations()) {
-          newMutatedKeys = newMutatedKeys.insert(newDoc.getKey());
-        } else {
-          newMutatedKeys = newMutatedKeys.remove(newDoc.getKey());
-        }
-      } else {
-        newDocumentSet = newDocumentSet.remove(key);
-        newMutatedKeys = newMutatedKeys.remove(key);
-      }
+      boolean oldDocHadPendingMutations =
+          oldDoc != null && this.mutatedKeys.contains(oldDoc.getKey());
+
+      // We only consider committed mutations for documents that were mutated during the lifetime of
+      // the view.
+      boolean newDocHasPendingMutations =
+          newDoc != null
+              && (newDoc.hasLocalMutations()
+                  || (this.mutatedKeys.contains(newDoc.getKey())
+                      && newDoc.hasCommittedMutations()));
+
+      boolean changeApplied = false;
+
       // Calculate change
       if (oldDoc != null && newDoc != null) {
         boolean docsEqual = oldDoc.getData().equals(newDoc.getData());
-        if (!docsEqual || oldDoc.hasLocalMutations() != newDoc.hasLocalMutations()) {
-          // only report a change if document actually changed.
-          if (docsEqual) {
-            changeSet.addChange(DocumentViewChange.create(Type.METADATA, newDoc));
-          } else {
+        if (!docsEqual) {
+          if (!shouldWaitForSyncedDocument(oldDoc, newDoc)) {
             changeSet.addChange(DocumentViewChange.create(Type.MODIFIED, newDoc));
+            changeApplied = true;
           }
-
           if (lastDocInLimit != null && query.comparator().compare(newDoc, lastDocInLimit) > 0) {
             // This doc moved from inside the limit to after the limit. That means there may be some
             // doc in the local cache that's actually less than this one.
             needsRefill = true;
           }
+        } else if (oldDocHadPendingMutations != newDocHasPendingMutations) {
+          changeSet.addChange(DocumentViewChange.create(Type.METADATA, newDoc));
+          changeApplied = true;
         }
       } else if (oldDoc == null && newDoc != null) {
         changeSet.addChange(DocumentViewChange.create(Type.ADDED, newDoc));
+        changeApplied = true;
       } else if (oldDoc != null && newDoc == null) {
         changeSet.addChange(DocumentViewChange.create(Type.REMOVED, oldDoc));
+        changeApplied = true;
         if (lastDocInLimit != null) {
           // A doc was removed from a full limit query. We'll need to requery from the local cache
           // to see if we know about some other doc that should be in the results.
           needsRefill = true;
         }
       }
+
+      if (changeApplied) {
+        if (newDoc != null) {
+          newDocumentSet = newDocumentSet.add(newDoc);
+          if (newDoc.hasLocalMutations()) {
+            newMutatedKeys = newMutatedKeys.insert(newDoc.getKey());
+          } else {
+            newMutatedKeys = newMutatedKeys.remove(newDoc.getKey());
+          }
+        } else {
+          newDocumentSet = newDocumentSet.remove(key);
+          newMutatedKeys = newMutatedKeys.remove(key);
+        }
+      }
     }
 
     if (query.hasLimit()) {
-      // TODO: Make QuerySnapshot size be constant time.
-      while (newDocumentSet.size() > query.getLimit()) {
+      for (long i = newDocumentSet.size() - this.query.getLimit(); i > 0; --i) {
         Document oldDoc = newDocumentSet.getLastDocument();
         newDocumentSet = newDocumentSet.remove(oldDoc.getKey());
         newMutatedKeys = newMutatedKeys.remove(oldDoc.getKey());
@@ -226,6 +243,18 @@ public <D extends MaybeDocument> DocumentChanges computeDocChanges(
     return new DocumentChanges(newDocumentSet, changeSet, newMutatedKeys, needsRefill);
   }
 
+  private boolean shouldWaitForSyncedDocument(Document oldDoc, Document newDoc) {
+    // We suppress the initial change event for documents that were modified as part of a write
+    // acknowledgment (e.g. when the value of a server transform is applied) as Watch will send us
+    // the same document again. By suppressing the event, we only raise two user visible events (one
+    // with `hasPendingWrites` and the final state of the document) instead of three (one with
+    // `hasPendingWrites`, the modified document with `hasPendingWrites` and the final state of the
+    // document).
+    return (oldDoc.hasLocalMutations()
+        && newDoc.hasCommittedMutations()
+        && !newDoc.hasLocalMutations());
+  }
+
   /**
    * Updates the view with the given ViewDocumentChanges and updates limbo docs and sync state from
    * the given (optional) target change.
@@ -273,15 +302,14 @@ public ViewChange applyChanges(DocumentChanges docChanges, TargetChange targetCh
     ViewSnapshot snapshot = null;
     if (viewChanges.size() != 0 || syncStatedChanged) {
       boolean fromCache = newSyncState == SyncState.LOCAL;
-      boolean hasPendingWrites = !docChanges.mutatedKeys.isEmpty();
       snapshot =
           new ViewSnapshot(
               query,
               docChanges.documentSet,
               oldDocumentSet,
               viewChanges,
               fromCache,
-              hasPendingWrites,
+              docChanges.mutatedKeys,
               syncStatedChanged);
     }
     return new ViewChange(snapshot, limboDocumentChanges);

firebase-firestore/src/main/java/com/google/firebase/firestore/core/ViewSnapshot.java

@@ -14,6 +14,8 @@
 
 package com.google.firebase.firestore.core;
 
+import com.google.firebase.database.collection.ImmutableSortedSet;
+import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.DocumentSet;
 import java.util.List;
 
@@ -31,7 +33,7 @@ public enum SyncState {
   private final DocumentSet oldDocuments;
   private final List<DocumentViewChange> changes;
   private final boolean isFromCache;
-  private final boolean hasPendingWrites;
+  private final ImmutableSortedSet<DocumentKey> mutatedKeys;
   private final boolean didSyncStateChange;
 
   public ViewSnapshot(
@@ -40,14 +42,14 @@ public ViewSnapshot(
       DocumentSet oldDocuments,
       List<DocumentViewChange> changes,
       boolean isFromCache,
-      boolean hasPendingWrites,
+      ImmutableSortedSet<DocumentKey> mutatedKeys,
       boolean didSyncStateChange) {
     this.query = query;
     this.documents = documents;
     this.oldDocuments = oldDocuments;
     this.changes = changes;
     this.isFromCache = isFromCache;
-    this.hasPendingWrites = hasPendingWrites;
+    this.mutatedKeys = mutatedKeys;
     this.didSyncStateChange = didSyncStateChange;
   }
 
@@ -72,7 +74,11 @@ public boolean isFromCache() {
   }
 
   public boolean hasPendingWrites() {
-    return hasPendingWrites;
+    return !mutatedKeys.isEmpty();
+  }
+
+  public ImmutableSortedSet<DocumentKey> getMutatedKeys() {
+    return mutatedKeys;
   }
 
   public boolean didSyncStateChange() {
@@ -93,15 +99,15 @@ public boolean equals(Object o) {
     if (isFromCache != that.isFromCache) {
       return false;
     }
-    if (hasPendingWrites != that.hasPendingWrites) {
-      return false;
-    }
     if (didSyncStateChange != that.didSyncStateChange) {
       return false;
     }
     if (!query.equals(that.query)) {
       return false;
     }
+    if (!mutatedKeys.equals(that.mutatedKeys)) {
+      return false;
+    }
     if (!documents.equals(that.documents)) {
       return false;
     }
@@ -117,8 +123,8 @@ public int hashCode() {
     result = 31 * result + documents.hashCode();
     result = 31 * result + oldDocuments.hashCode();
     result = 31 * result + changes.hashCode();
+    result = 31 * result + mutatedKeys.hashCode();
     result = 31 * result + (isFromCache ? 1 : 0);
-    result = 31 * result + (hasPendingWrites ? 1 : 0);
     result = 31 * result + (didSyncStateChange ? 1 : 0);
     return result;
   }
@@ -135,8 +141,8 @@ public String toString() {
         + changes
         + ", isFromCache="
         + isFromCache
-        + ", hasPendingWrites="
-        + hasPendingWrites
+        + ", mutatedKeys="
+        + mutatedKeys.size()
         + ", didSyncStateChange="
         + didSyncStateChange
         + ")";