Continue work on ReferenceSet

Firestore/Example/Firestore.xcodeproj/project.pbxproj

@@ -25,6 +25,7 @@
 /* Begin PBXBuildFile section */
 		020AFD89BB40E5175838BB76 /* local_serializer_test.cc in Sources */ = {isa = PBXBuildFile; fileRef = F8043813A5D16963EC02B182 /* local_serializer_test.cc */; };
 		0535C1B65DADAE1CE47FA3CA /* string_format_apple_test.mm in Sources */ = {isa = PBXBuildFile; fileRef = 9CFD366B783AE27B9E79EE7A /* string_format_apple_test.mm */; };
+		132E3483789344640A52F223 /* reference_set_test.cc in Sources */ = {isa = PBXBuildFile; fileRef = 132E32997D781B896672D30A /* reference_set_test.cc */; };
 		132E3E53179DE287D875F3F2 /* FSTLevelDBTransactionTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 132E36BB104830BD806351AC /* FSTLevelDBTransactionTests.mm */; };
 		132E3EE56C143B2C9ACB6187 /* FSTLevelDBBenchmarkTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 132E3BB3D5C42282B4ACFB20 /* FSTLevelDBBenchmarkTests.mm */; };
 		1CAA9012B25F975D445D5978 /* strerror_test.cc in Sources */ = {isa = PBXBuildFile; fileRef = 358C3B5FE573B1D60A4F7592 /* strerror_test.cc */; };
@@ -102,7 +103,6 @@
 		5492E0AC2021552D00B64F25 /* FSTMutationQueueTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0962021552C00B64F25 /* FSTMutationQueueTests.mm */; };
 		5492E0AD2021552D00B64F25 /* FSTMemoryMutationQueueTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0972021552C00B64F25 /* FSTMemoryMutationQueueTests.mm */; };
 		5492E0AE2021552D00B64F25 /* FSTLevelDBQueryCacheTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0982021552C00B64F25 /* FSTLevelDBQueryCacheTests.mm */; };
-		5492E0AF2021552D00B64F25 /* FSTReferenceSetTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E09A2021552C00B64F25 /* FSTReferenceSetTests.mm */; };
 		5492E0B12021552D00B64F25 /* FSTRemoteDocumentCacheTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E09C2021552D00B64F25 /* FSTRemoteDocumentCacheTests.mm */; };
 		5492E0B92021555100B64F25 /* FSTDocumentKeyTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0B22021555000B64F25 /* FSTDocumentKeyTests.mm */; };
 		5492E0BA2021555100B64F25 /* FSTDocumentSetTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0B32021555100B64F25 /* FSTDocumentSetTests.mm */; };
@@ -300,6 +300,7 @@
 		11984BA0A99D7A7ABA5B0D90 /* Pods-Firestore_Example_iOS-Firestore_SwiftTests_iOS.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Firestore_Example_iOS-Firestore_SwiftTests_iOS.release.xcconfig"; path = "Pods/Target Support Files/Pods-Firestore_Example_iOS-Firestore_SwiftTests_iOS/Pods-Firestore_Example_iOS-Firestore_SwiftTests_iOS.release.xcconfig"; sourceTree = "<group>"; };
 		1277F98C20D2DF0867496976 /* Pods-Firestore_IntegrationTests_iOS.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Firestore_IntegrationTests_iOS.debug.xcconfig"; path = "Pods/Target Support Files/Pods-Firestore_IntegrationTests_iOS/Pods-Firestore_IntegrationTests_iOS.debug.xcconfig"; sourceTree = "<group>"; };
 		12F4357299652983A615F886 /* LICENSE */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text; name = LICENSE; path = ../LICENSE; sourceTree = "<group>"; };
+		132E32997D781B896672D30A /* reference_set_test.cc */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.cpp; path = reference_set_test.cc; sourceTree = "<group>"; };
 		132E36BB104830BD806351AC /* FSTLevelDBTransactionTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTLevelDBTransactionTests.mm; sourceTree = "<group>"; };
 		132E3BB3D5C42282B4ACFB20 /* FSTLevelDBBenchmarkTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTLevelDBBenchmarkTests.mm; sourceTree = "<group>"; };
 		2A0CF41BA5AED6049B0BEB2C /* type_traits_apple_test.mm */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.cpp.objcpp; path = type_traits_apple_test.mm; sourceTree = "<group>"; };
@@ -389,7 +390,6 @@
 		5492E0972021552C00B64F25 /* FSTMemoryMutationQueueTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTMemoryMutationQueueTests.mm; sourceTree = "<group>"; };
 		5492E0982021552C00B64F25 /* FSTLevelDBQueryCacheTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTLevelDBQueryCacheTests.mm; sourceTree = "<group>"; };
 		5492E0992021552C00B64F25 /* FSTPersistenceTestHelpers.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = FSTPersistenceTestHelpers.h; sourceTree = "<group>"; };
-		5492E09A2021552C00B64F25 /* FSTReferenceSetTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTReferenceSetTests.mm; sourceTree = "<group>"; };
 		5492E09C2021552D00B64F25 /* FSTRemoteDocumentCacheTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTRemoteDocumentCacheTests.mm; sourceTree = "<group>"; };
 		5492E0B22021555000B64F25 /* FSTDocumentKeyTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTDocumentKeyTests.mm; sourceTree = "<group>"; };
 		5492E0B32021555100B64F25 /* FSTDocumentSetTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTDocumentSetTests.mm; sourceTree = "<group>"; };
@@ -768,6 +768,7 @@
 				54995F6E205B6E12004EFFA0 /* leveldb_key_test.cc */,
 				332485C4DCC6BA0DBB5E31B7 /* leveldb_util_test.cc */,
 				F8043813A5D16963EC02B182 /* local_serializer_test.cc */,
+				132E32997D781B896672D30A /* reference_set_test.cc */,
 			);
 			path = local;
 			sourceTree = "<group>";
@@ -1112,7 +1113,6 @@
 				5492E08D2021552B00B64F25 /* FSTPersistenceTestHelpers.mm */,
 				5492E0952021552C00B64F25 /* FSTQueryCacheTests.h */,
 				5492E0892021552A00B64F25 /* FSTQueryCacheTests.mm */,
-				5492E09A2021552C00B64F25 /* FSTReferenceSetTests.mm */,
 				5492E0852021552A00B64F25 /* FSTRemoteDocumentCacheTests.h */,
 				5492E09C2021552D00B64F25 /* FSTRemoteDocumentCacheTests.mm */,
 			);
@@ -1907,7 +1907,6 @@
 				5492E0A22021552D00B64F25 /* FSTQueryCacheTests.mm in Sources */,
 				5492E064202154B900B64F25 /* FSTQueryListenerTests.mm in Sources */,
 				5492E068202154B900B64F25 /* FSTQueryTests.mm in Sources */,
-				5492E0AF2021552D00B64F25 /* FSTReferenceSetTests.mm in Sources */,
 				5492E0B12021552D00B64F25 /* FSTRemoteDocumentCacheTests.mm in Sources */,
 				5492E0C92021557E00B64F25 /* FSTRemoteEventTests.mm in Sources */,
 				5492E0C72021557E00B64F25 /* FSTSerializerBetaTests.mm in Sources */,
@@ -2001,6 +2000,7 @@
 				C80B10E79CDD7EF7843C321E /* type_traits_apple_test.mm in Sources */,
 				ABC1D7DE2023A05300BA84F0 /* user_test.cc in Sources */,
 				544129DE21C2DDC800EFB9CC /* write.pb.cc in Sources */,
+				132E3483789344640A52F223 /* reference_set_test.cc in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};

Firestore/Example/Tests/Local/FSTLRUGarbageCollectorTests.mm

@@ -31,6 +31,7 @@
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Util/FSTClasses.h"
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
+#include "Firestore/core/src/firebase/firestore/local/reference_set.h"
 #include "Firestore/core/src/firebase/firestore/local/remote_document_cache.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 #include "Firestore/core/src/firebase/firestore/model/precondition.h"
@@ -42,6 +43,7 @@
 using firebase::firestore::auth::User;
 using firebase::firestore::local::LruParams;
 using firebase::firestore::local::LruResults;
+using firebase::firestore::local::ReferenceSet;
 using firebase::firestore::local::RemoteDocumentCache;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::DocumentKeyHash;
@@ -65,6 +67,7 @@ @implementation FSTLRUGarbageCollectorTests {
   FSTLRUGarbageCollector *_gc;
   ListenSequenceNumber _initialSequenceNumber;
   User _user;
+  ReferenceSet _additionalReferences;
 }
 
 - (void)setUp {
@@ -85,6 +88,7 @@ - (BOOL)isTestBaseClass {
 - (void)newTestResourcesWithLruParams:(LruParams)lruParams {
   HARD_ASSERT(_persistence == nil, "Persistence already created");
   _persistence = [self newPersistenceWithLruParams:lruParams];
+  [_persistence.referenceDelegate addInMemoryPins:&_additionalReferences];
   _queryCache = [_persistence queryCache];
   _documentCache = [_persistence remoteDocumentCache];
   _mutationQueue = [_persistence mutationQueueForUser:_user];

Firestore/Example/Tests/Local/FSTLevelDBMutationQueueTests.mm

@@ -28,6 +28,7 @@
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_key.h"
+#include "Firestore/core/src/firebase/firestore/local/reference_set.h"
 #include "Firestore/core/src/firebase/firestore/util/ordered_code.h"
 #include "absl/strings/string_view.h"
 #include "leveldb/db.h"
@@ -36,6 +37,7 @@
 
 using firebase::firestore::auth::User;
 using firebase::firestore::local::LevelDbMutationKey;
+using firebase::firestore::local::ReferenceSet;
 using firebase::firestore::model::BatchId;
 using firebase::firestore::util::OrderedCode;
 using leveldb::DB;
@@ -68,11 +70,13 @@ @interface FSTLevelDBMutationQueueTests : FSTMutationQueueTests
 
 @implementation FSTLevelDBMutationQueueTests {
   FSTLevelDB *_db;
+  ReferenceSet _additionalReferences;
 }
 
 - (void)setUp {
   [super setUp];
   _db = [FSTPersistenceTestHelpers levelDBPersistence];
+  [_db.referenceDelegate addInMemoryPins:&_additionalReferences];
   self.mutationQueue = [_db mutationQueueForUser:User("user")];
   self.persistence = _db;
 

Firestore/Example/Tests/Local/FSTLevelDBQueryCacheTests.mm

@@ -26,11 +26,13 @@
 #import "Firestore/Example/Tests/Local/FSTQueryCacheTests.h"
 
 #include "Firestore/core/include/firebase/firestore/timestamp.h"
+#include "Firestore/core/src/firebase/firestore/local/reference_set.h"
 #include "Firestore/core/src/firebase/firestore/model/database_id.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 using firebase::Timestamp;
+using firebase::firestore::local::ReferenceSet;
 using firebase::firestore::model::DatabaseId;
 using firebase::firestore::model::ListenSequenceNumber;
 using firebase::firestore::model::ResourcePath;
@@ -48,13 +50,16 @@ @interface FSTLevelDBQueryCacheTests : FSTQueryCacheTests
  * FSTQueryCacheTests. This class is merely responsible for setting up and tearing down the
  * @a queryCache.
  */
-@implementation FSTLevelDBQueryCacheTests
+@implementation FSTLevelDBQueryCacheTests {
+  ReferenceSet _additionalReferences;
+}
 
 - (void)setUp {
   [super setUp];
 
   self.persistence = [FSTPersistenceTestHelpers levelDBPersistence];
   self.queryCache = [self.persistence queryCache];
+  [self.persistence.referenceDelegate addInMemoryPins:&_additionalReferences];
 }
 
 - (void)tearDown {

Firestore/Example/Tests/Local/FSTMemoryMutationQueueTests.mm

@@ -22,8 +22,10 @@
 #import "Firestore/Example/Tests/Local/FSTPersistenceTestHelpers.h"
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
+#include "Firestore/core/src/firebase/firestore/local/reference_set.h"
 
 using firebase::firestore::auth::User;
+using firebase::firestore::local::ReferenceSet;
 
 @interface FSTMemoryMutationQueueTests : FSTMutationQueueTests
 @end
@@ -32,12 +34,15 @@ @interface FSTMemoryMutationQueueTests : FSTMutationQueueTests
  * The tests for FSTMemoryMutationQueue are performed on the FSTMutationQueue protocol in
  * FSTMutationQueueTests. This class is merely responsible for setting up the @a mutationQueue.
  */
-@implementation FSTMemoryMutationQueueTests
+@implementation FSTMemoryMutationQueueTests {
+  ReferenceSet _additionalReferences;
+}
 
 - (void)setUp {
   [super setUp];
 
   self.persistence = [FSTPersistenceTestHelpers eagerGCMemoryPersistence];
+  [self.persistence.referenceDelegate addInMemoryPins:&_additionalReferences];
   self.mutationQueue = [self.persistence mutationQueueForUser:User("user")];
 }
 

Firestore/Example/Tests/Local/FSTMemoryQueryCacheTests.mm

@@ -21,6 +21,10 @@
 #import "Firestore/Example/Tests/Local/FSTPersistenceTestHelpers.h"
 #import "Firestore/Example/Tests/Local/FSTQueryCacheTests.h"
 
+#include "Firestore/core/src/firebase/firestore/local/reference_set.h"
+
+using firebase::firestore::local::ReferenceSet;
+
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTMemoryQueryCacheTests : FSTQueryCacheTests
@@ -31,13 +35,16 @@ @interface FSTMemoryQueryCacheTests : FSTQueryCacheTests
  * FSTQueryCacheTests. This class is merely responsible for setting up and tearing down the
  * @a queryCache.
  */
-@implementation FSTMemoryQueryCacheTests
+@implementation FSTMemoryQueryCacheTests {
+  ReferenceSet _additionalReferences;
+}
 
 - (void)setUp {
   [super setUp];
 
   self.persistence = [FSTPersistenceTestHelpers eagerGCMemoryPersistence];
   self.queryCache = [self.persistence queryCache];
+  [self.persistence.referenceDelegate addInMemoryPins:&_additionalReferences];
 }
 
 - (void)tearDown {

Firestore/Example/Tests/Local/FSTQueryCacheTests.mm

@@ -26,6 +26,7 @@
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
 #import "Firestore/third_party/Immutable/Tests/FSTImmutableSortedSet+Testing.h"
 
+#include "Firestore/core/src/firebase/firestore/local/reference_set.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/test/firebase/firestore/testutil/testutil.h"
 

Firestore/Source/Core/FSTSyncEngine.mm

@@ -30,14 +30,14 @@
 #import "Firestore/Source/Local/FSTLocalViewChanges.h"
 #import "Firestore/Source/Local/FSTLocalWriteResult.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
-#import "Firestore/Source/Local/FSTReferenceSet.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTDocumentSet.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Remote/FSTRemoteEvent.h"
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/core/target_id_generator.h"
+#include "Firestore/core/src/firebase/firestore/local/reference_set.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/document_map.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
@@ -47,6 +47,7 @@
 using firebase::firestore::auth::HashUser;
 using firebase::firestore::auth::User;
 using firebase::firestore::core::TargetIdGenerator;
+using firebase::firestore::local::ReferenceSet;
 using firebase::firestore::model::BatchId;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::DocumentKeySet;
@@ -157,9 +158,6 @@ @interface FSTSyncEngine ()
 @property(nonatomic, strong, readonly)
     NSMutableDictionary<NSNumber *, FSTQueryView *> *queryViewsByTarget;
 
-/** Used to track any documents that are currently in limbo. */
-@property(nonatomic, strong, readonly) FSTReferenceSet *limboDocumentRefs;
-
 @end
 
 @implementation FSTSyncEngine {
@@ -183,6 +181,9 @@ @implementation FSTSyncEngine {
   std::map<TargetId, LimboResolution> _limboResolutionsByTarget;
 
   User _currentUser;
+
+  /** Used to track any documents that are currently in limbo. */
+  ReferenceSet _limboDocumentRefs;
 }
 
 - (instancetype)initWithLocalStore:(FSTLocalStore *)localStore
@@ -195,7 +196,6 @@ - (instancetype)initWithLocalStore:(FSTLocalStore *)localStore
     _queryViewsByQuery = [NSMutableDictionary dictionary];
     _queryViewsByTarget = [NSMutableDictionary dictionary];
 
-    _limboDocumentRefs = [[FSTReferenceSet alloc] init];
     _targetIdGenerator = TargetIdGenerator::SyncEngineTargetIdGenerator();
     _currentUser = initialUser;
   }
@@ -468,10 +468,10 @@ - (void)removeAndCleanupQuery:(FSTQueryView *)queryView {
   [self.queryViewsByQuery removeObjectForKey:queryView.query];
   [self.queryViewsByTarget removeObjectForKey:@(queryView.targetID)];
 
-  DocumentKeySet limboKeys = [self.limboDocumentRefs referencedKeysForID:queryView.targetID];
-  [self.limboDocumentRefs removeReferencesForID:queryView.targetID];
+  DocumentKeySet limboKeys = _limboDocumentRefs.ReferencedKeys(queryView.targetID);
+  _limboDocumentRefs.RemoveReferences(queryView.targetID);
   for (const DocumentKey &key : limboKeys) {
-    if (![self.limboDocumentRefs containsKey:key]) {
+    if (!_limboDocumentRefs.ContainsKey(key)) {
       // We removed the last reference for this key.
       [self removeLimboTargetForKey:key];
     }
@@ -531,14 +531,14 @@ - (void)updateTrackedLimboDocumentsWithChanges:(NSArray<FSTLimboDocumentChange *
   for (FSTLimboDocumentChange *limboChange in limboChanges) {
     switch (limboChange.type) {
       case FSTLimboDocumentChangeTypeAdded:
-        [self.limboDocumentRefs addReferenceToKey:limboChange.key forID:targetID];
+        _limboDocumentRefs.AddReference(limboChange.key, targetID);
         [self trackLimboChange:limboChange];
         break;
 
       case FSTLimboDocumentChangeTypeRemoved:
         LOG_DEBUG("Document no longer in limbo: %s", limboChange.key.ToString());
-        [self.limboDocumentRefs removeReferenceToKey:limboChange.key forID:targetID];
-        if (![self.limboDocumentRefs containsKey:limboChange.key]) {
+        _limboDocumentRefs.RemoveReference(limboChange.key, targetID);
+        if (!_limboDocumentRefs.ContainsKey(limboChange.key)) {
           // We removed the last reference for this key
           [self removeLimboTargetForKey:limboChange.key];
         }