Remove local references from reference delegate

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

@@ -558,7 +558,15 @@ public void releaseQuery(Query query) {
             queryCache.updateQueryData(queryData);
           }
 
-          localViewReferences.removeReferencesForId(queryData.getTargetId());
+          // References for documents sent via Watch are automatically removed when we delete a
+          // query's target data from the reference delegate. Since this does not remove references
+          // for locally mutated documents, we have to remove the target associations for these
+          // documents manually.
+          ImmutableSortedSet<DocumentKey> removedReferences =
+              localViewReferences.removeReferencesForId(queryData.getTargetId());
+          for (DocumentKey key : removedReferences) {
+            persistence.getReferenceDelegate().removeReference(key);
+          }
           persistence.getReferenceDelegate().removeTarget(queryData);
           targetIds.remove(queryData.getTargetId());
 
@@ -622,7 +630,7 @@ private Set<DocumentKey> releaseBatchResults(List<MutationBatchResult> batchResu
     ArrayList<MutationBatch> batches = new ArrayList<>(batchResults.size());
     // TODO: Call queryEngine.handleDocumentChange() as appropriate.
     for (MutationBatchResult batchResult : batchResults) {
-      applyBatchResult(batchResult);
+      applyWriteToRemoteDocuments(batchResult);
       batches.add(batchResult.getBatch());
     }
 
@@ -647,7 +655,7 @@ private Set<DocumentKey> removeMutationBatches(List<MutationBatch> batches) {
     return affectedDocs;
   }
 
-  private void applyBatchResult(MutationBatchResult batchResult) {
+  private void applyWriteToRemoteDocuments(MutationBatchResult batchResult) {
     MutationBatch batch = batchResult.getBatch();
     Set<DocumentKey> docKeys = batch.getKeys();
     for (DocumentKey docKey : docKeys) {

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

@@ -77,19 +77,27 @@ public void removeReferences(ImmutableSortedSet<DocumentKey> keys, int targetOrB
     }
   }
 
-  /** Clears all references with a given ID. Calls removeReference() for each key removed. */
-  public void removeReferencesForId(int targetId) {
+  /**
+   * Clears all references with a given ID. Calls removeReference() for each key removed.
+   *
+   * @return The keys of the documents that were removed.
+   */
+  public ImmutableSortedSet<DocumentKey> removeReferencesForId(int targetId) {
     DocumentKey emptyKey = DocumentKey.empty();
     DocumentReference startRef = new DocumentReference(emptyKey, targetId);
     Iterator<DocumentReference> it = referencesByTarget.iteratorFrom(startRef);
+    ImmutableSortedSet<DocumentKey> keys = DocumentKey.emptyKeySet();
     while (it.hasNext()) {
       DocumentReference ref = it.next();
       if (ref.getId() == targetId) {
+        keys = keys.insert(ref.getKey());
         removeReference(ref);
       } else {
         break;
       }
     }
+
+    return keys;
   }
 
   /** Clears all references for all IDs. */