Introduce ReferenceDelegates (#1222)

* Bump sequence number on resume token refresh

* Style

* Fix comment formatting

* Add FSTReferenceDelegate definition and documentation

* Add methods to return nil for delegates, wire up inMemoryPins

* Add hook for removing a reference

* Start work on reference delegates

* Fix up tests to support adding documents at a sequence number

* Implement removing references

* Remove from target when dropped from local view

* Fix warning

* Add hooks for removal from mutation queue

* Add hooks for limbo document updates

* Style

* Drop commented-out code

* Fixup after merging master

* Drop sequence number plumbing

* Style

* Drop errant semicolon

10 files changed, 109 insertions, 10 deletions

diff --git a/Firestore/Source/Local/FSTLevelDB.mm b/Firestore/Source/Local/FSTLevelDB.mm
index fae85e7..bc2f2eb 100644
--- a/Firestore/Source/Local/FSTLevelDB.mm
+++ b/Firestore/Source/Local/FSTLevelDB.mm
@@ -243,6 +243,10 @@ using leveldb::WriteOptions;
   _ptr.reset();
 }
 
+- (_Nullable id<FSTReferenceDelegate>)referenceDelegate {
+  return nil;
+}
+
 #pragma mark - Error and Status
 
 + (nullable NSError *)errorWithStatus:(Status)status description:(NSString *)description, ... {
diff --git a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
index 75c3cf6..2c9f68d 100644
--- a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
+++ b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
@@ -511,6 +511,7 @@ using leveldb::WriteOptions;
                                              documentKey:mutation.key
                                                  batchID:batchID];
       _db.currentTransaction->Delete(key);
+      [_db.referenceDelegate removeMutationReference:mutation.key];
       [garbageCollector addPotentialGarbageKey:mutation.key];
     }
   }
diff --git a/Firestore/Source/Local/FSTLevelDBQueryCache.mm b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
index 97b218d..68b6f98 100644
--- a/Firestore/Source/Local/FSTLevelDBQueryCache.mm
+++ b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
@@ -294,7 +294,8 @@ using firebase::firestore::model::DocumentKeySet;
         [FSTLevelDBTargetDocumentKey keyWithTargetID:targetID documentKey:key], emptyBuffer);
     self->_db.currentTransaction->Put(
         [FSTLevelDBDocumentTargetKey keyWithDocumentKey:key targetID:targetID], emptyBuffer);
-  }
+    [self->_db.referenceDelegate addReference:key target:targetID];
+  };
 }
 
 - (void)removeMatchingKeys:(const DocumentKeySet &)keys forTargetID:(FSTTargetID)targetID {
@@ -303,6 +304,7 @@ using firebase::firestore::model::DocumentKeySet;
         [FSTLevelDBTargetDocumentKey keyWithTargetID:targetID documentKey:key]);
     self->_db.currentTransaction->Delete(
         [FSTLevelDBDocumentTargetKey keyWithDocumentKey:key targetID:targetID]);
+    [self->_db.referenceDelegate removeReference:key target:targetID];
     [self.garbageCollector addPotentialGarbageKey:key];
   }
 }
diff --git a/Firestore/Source/Local/FSTLocalStore.mm b/Firestore/Source/Local/FSTLocalStore.mm
index 29d0395..0d6a785 100644
--- a/Firestore/Source/Local/FSTLocalStore.mm
+++ b/Firestore/Source/Local/FSTLocalStore.mm
@@ -114,6 +114,7 @@ NS_ASSUME_NONNULL_BEGIN
     _localDocuments = [FSTLocalDocumentsView viewWithRemoteDocumentCache:_remoteDocumentCache
                                                            mutationQueue:_mutationQueue];
     _localViewReferences = [[FSTReferenceSet alloc] init];
+    [_persistence.referenceDelegate addInMemoryPins:_localViewReferences];
 
     _garbageCollector = garbageCollector;
     [_garbageCollector addGarbageSource:_queryCache];
@@ -263,6 +264,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (FSTMaybeDocumentDictionary *)applyRemoteEvent:(FSTRemoteEvent *)remoteEvent {
   return self.persistence.run("Apply remote event", [&]() -> FSTMaybeDocumentDictionary * {
+    // TODO(gsoltis): move the sequence number into the reference delegate.
     FSTListenSequenceNumber sequenceNumber = [self.listenSequence next];
     id<FSTQueryCache> queryCache = self.queryCache;
 
@@ -309,6 +311,7 @@ NS_ASSUME_NONNULL_BEGIN
 
     // TODO(klimt): This could probably be an NSMutableDictionary.
     DocumentKeySet changedDocKeys;
+    const DocumentKeySet &limboDocuments = remoteEvent.limboDocumentChanges;
     for (const auto &kv : remoteEvent.documentUpdates) {
       const DocumentKey &key = kv.first;
       FSTMaybeDocument *doc = kv.second;
@@ -331,6 +334,9 @@ NS_ASSUME_NONNULL_BEGIN
       // The document might be garbage because it was unreferenced by everything.
       // Make sure to mark it as garbage if it is...
       [self.garbageCollector addPotentialGarbageKey:key];
+      if (limboDocuments.contains(key)) {
+        [self.persistence.referenceDelegate limboDocumentUpdated:key];
+      }
     }
 
     // HACK: The only reason we allow omitting snapshot version is so we can synthesize remote
@@ -365,6 +371,9 @@ NS_ASSUME_NONNULL_BEGIN
       FSTQueryData *queryData = [self.queryCache queryDataForQuery:view.query];
       FSTAssert(queryData, @"Local view changes contain unallocated query.");
       FSTTargetID targetID = queryData.targetID;
+      for (const DocumentKey &key : view.removedKeys) {
+        [self->_persistence.referenceDelegate removeReference:key target:targetID];
+      }
       [localViewReferences addReferencesToKeys:view.addedKeys forID:targetID];
       [localViewReferences removeReferencesToKeys:view.removedKeys forID:targetID];
     }
@@ -415,6 +424,7 @@ NS_ASSUME_NONNULL_BEGIN
     if (self.garbageCollector.isEager) {
       [self.queryCache removeQueryData:queryData];
     }
+    [self.persistence.referenceDelegate removeTarget:queryData];
     [self.targetIDs removeObjectForKey:@(queryData.targetID)];
 
     // If this was the last watch target, then we won't get any more watch snapshots, so we should
diff --git a/Firestore/Source/Local/FSTMemoryMutationQueue.h b/Firestore/Source/Local/FSTMemoryMutationQueue.h
index f0786cc..fd46a6e 100644
--- a/Firestore/Source/Local/FSTMemoryMutationQueue.h
+++ b/Firestore/Source/Local/FSTMemoryMutationQueue.h
@@ -16,6 +16,7 @@
 
 #import <Foundation/Foundation.h>
 
+#import "Firestore/Source/Local/FSTMemoryPersistence.h"
 #import "Firestore/Source/Local/FSTMutationQueue.h"
 
 @protocol FSTGarbageCollector;
@@ -24,7 +25,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTMemoryMutationQueue : NSObject <FSTMutationQueue>
 
-+ (instancetype)mutationQueue;
+- (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence NS_DESIGNATED_INITIALIZER;
+
+- (instancetype)init NS_UNAVAILABLE;
 
 /** The garbage collector to notify about potential garbage keys. */
 @property(nonatomic, weak, readwrite, nullable) id<FSTGarbageCollector> garbageCollector;
diff --git a/Firestore/Source/Local/FSTMemoryMutationQueue.mm b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
index b0c8760..e05ee64 100644
--- a/Firestore/Source/Local/FSTMemoryMutationQueue.mm
+++ b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
@@ -18,6 +18,7 @@
 
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Local/FSTDocumentReference.h"
+#import "Firestore/Source/Local/FSTMemoryPersistence.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Util/FSTAssert.h"
@@ -73,14 +74,13 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
 
 @end
 
-@implementation FSTMemoryMutationQueue
-
-+ (instancetype)mutationQueue {
-  return [[FSTMemoryMutationQueue alloc] init];
+@implementation FSTMemoryMutationQueue {
+  FSTMemoryPersistence *_persistence;
 }
 
-- (instancetype)init {
+- (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence {
   if (self = [super init]) {
+    _persistence = persistence;
     _queue = [NSMutableArray array];
     _batchesByDocumentKey =
         [FSTImmutableSortedSet setWithComparator:FSTDocumentReferenceComparatorByKey];
@@ -348,6 +348,7 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
     for (FSTMutation *mutation in batch.mutations) {
       const DocumentKey &key = mutation.key;
       [garbageCollector addPotentialGarbageKey:key];
+      [_persistence.referenceDelegate removeMutationReference:key];
 
       FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:batchID];
       references = [references setByRemovingObject:reference];
diff --git a/Firestore/Source/Local/FSTMemoryPersistence.mm b/Firestore/Source/Local/FSTMemoryPersistence.mm
index 8d74881..3466f3e 100644
--- a/Firestore/Source/Local/FSTMemoryPersistence.mm
+++ b/Firestore/Source/Local/FSTMemoryPersistence.mm
@@ -59,7 +59,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (instancetype)init {
   if (self = [super init]) {
-    _queryCache = [[FSTMemoryQueryCache alloc] init];
+    _queryCache = [[FSTMemoryQueryCache alloc] initWithPersistence:self];
     _remoteDocumentCache = [[FSTMemoryRemoteDocumentCache alloc] init];
   }
   return self;
@@ -78,6 +78,10 @@ NS_ASSUME_NONNULL_BEGIN
   self.started = NO;
 }
 
+- (_Nullable id<FSTReferenceDelegate>)referenceDelegate {
+  return nil;
+}
+
 - (const FSTTransactionRunner &)run {
   return _transactionRunner;
 }
@@ -85,7 +89,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (id<FSTMutationQueue>)mutationQueueForUser:(const User &)user {
   id<FSTMutationQueue> queue = _mutationQueues[user];
   if (!queue) {
-    queue = [FSTMemoryMutationQueue mutationQueue];
+    queue = [[FSTMemoryMutationQueue alloc] initWithPersistence:self];
     _mutationQueues[user] = queue;
   }
   return queue;
diff --git a/Firestore/Source/Local/FSTMemoryQueryCache.h b/Firestore/Source/Local/FSTMemoryQueryCache.h
index 98f0277..126ce59 100644
--- a/Firestore/Source/Local/FSTMemoryQueryCache.h
+++ b/Firestore/Source/Local/FSTMemoryQueryCache.h
@@ -20,11 +20,18 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+@class FSTMemoryPersistence;
+
 /**
  * An implementation of the FSTQueryCache protocol that merely keeps queries in memory, suitable
  * for online only clients with persistence disabled.
  */
 @interface FSTMemoryQueryCache : NSObject <FSTQueryCache>
+
+- (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence NS_DESIGNATED_INITIALIZER;
+
+- (instancetype)init NS_UNAVAILABLE;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Local/FSTMemoryQueryCache.mm b/Firestore/Source/Local/FSTMemoryQueryCache.mm
index 991faa9..2eba4f6 100644
--- a/Firestore/Source/Local/FSTMemoryQueryCache.mm
+++ b/Firestore/Source/Local/FSTMemoryQueryCache.mm
@@ -19,6 +19,7 @@
 #include <utility>
 
 #import "Firestore/Source/Core/FSTQuery.h"
+#import "Firestore/Source/Local/FSTMemoryPersistence.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Local/FSTReferenceSet.h"
 
@@ -27,6 +28,7 @@
 
 using firebase::firestore::model::SnapshotVersion;
 using firebase::firestore::model::DocumentKeySet;
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -46,12 +48,14 @@ NS_ASSUME_NONNULL_BEGIN
 @end
 
 @implementation FSTMemoryQueryCache {
+  FSTMemoryPersistence *_persistence;
   /** The last received snapshot version. */
   SnapshotVersion _lastRemoteSnapshotVersion;
 }
 
-- (instancetype)init {
+- (instancetype)initWithPersistence:(FSTMemoryPersistence *)persistence {
   if (self = [super init]) {
+    _persistence = persistence;
     _queries = [NSMutableDictionary dictionary];
     _references = [[FSTReferenceSet alloc] init];
     _lastRemoteSnapshotVersion = SnapshotVersion::None();
@@ -119,10 +123,16 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)addMatchingKeys:(const DocumentKeySet &)keys forTargetID:(FSTTargetID)targetID {
   [self.references addReferencesToKeys:keys forID:targetID];
+  for (const DocumentKey &key : keys) {
+    [_persistence.referenceDelegate addReference:key target:targetID];
+  }
 }
 
 - (void)removeMatchingKeys:(const DocumentKeySet &)keys forTargetID:(FSTTargetID)targetID {
   [self.references removeReferencesToKeys:keys forID:targetID];
+  for (const DocumentKey &key : keys) {
+    [_persistence.referenceDelegate removeReference:key target:targetID];
+  }
 }
 
 - (void)removeMatchingKeysForTargetID:(FSTTargetID)targetID {
diff --git a/Firestore/Source/Local/FSTPersistence.h b/Firestore/Source/Local/FSTPersistence.h
index 2294ef1..417ff3f 100644
--- a/Firestore/Source/Local/FSTPersistence.h
+++ b/Firestore/Source/Local/FSTPersistence.h
@@ -16,12 +16,16 @@
 
 #import <Foundation/Foundation.h>
 
+#import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 
+@class FSTDocumentKey;
 @protocol FSTMutationQueue;
 @protocol FSTQueryCache;
+@class FSTQueryData;
 @protocol FSTRemoteDocumentCache;
+@class FSTReferenceSet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -56,6 +60,7 @@ NS_ASSUME_NONNULL_BEGIN
  * of its reads and writes in order to avoid relying on being able to read back uncommitted writes.
  */
 struct FSTTransactionRunner;
+@protocol FSTReferenceDelegate;
 @protocol FSTPersistence <NSObject>
 
 /**
@@ -87,6 +92,12 @@ struct FSTTransactionRunner;
 
 @property(nonatomic, readonly, assign) const FSTTransactionRunner &run;
 
+/**
+ * This property provides access to hooks around the document reference lifecycle. It is initially
+ * nullable while being implemented, but the goal is to eventually have it be non-nil.
+ */
+@property(nonatomic, readonly, strong) _Nullable id<FSTReferenceDelegate> referenceDelegate;
+
 @end
 
 @protocol FSTTransactional
@@ -97,6 +108,52 @@ struct FSTTransactionRunner;
 
 @end
 
+/**
+ * An FSTReferenceDelegate instance handles all of the hooks into the document-reference lifecycle.
+ * This includes being added to a target, being removed from a target, being subject to mutation,
+ * and being mutated by the user.
+ *
+ * Different implementations may do different things with each of these events. Not every
+ * implementation needs to do something with every lifecycle hook.
+ *
+ * Implementations that care about sequence numbers are responsible for generating them and making
+ * them available.
+ */
+@protocol FSTReferenceDelegate
+
+/**
+ * Registers an FSTReferenceSet of documents that should be considered 'referenced' and not eligible
+ * for removal during garbage collection.
+ */
+- (void)addInMemoryPins:(FSTReferenceSet *)set;
+
+/**
+ * Notify the delegate that a target was removed.
+ */
+- (void)removeTarget:(FSTQueryData *)queryData;
+
+/**
+ * Notify the delegate that the given document was added to the given target.
+ */
+- (void)addReference:(FSTDocumentKey *)key target:(FSTTargetID)targetID;
+
+/**
+ * Notify the delegate that the given document was removed from the given target.
+ */
+- (void)removeReference:(FSTDocumentKey *)key target:(FSTTargetID)targetID;
+
+/**
+ * Notify the delegate that a document is no longer being mutated by the user.
+ */
+- (void)removeMutationReference:(FSTDocumentKey *)key;
+
+/**
+ * Notify the delegate that a limbo document was updated.
+ */
+- (void)limboDocumentUpdated:(FSTDocumentKey *)key;
+
+@end
+
 struct FSTTransactionRunner {
 // Intentionally disable nullability checking for this function. We cannot properly annotate
 // the function because this function can handle both pointer and non-pointer types. It is an error