Port production usages of QueryCache

Firestore/Example/Tests/Local/FSTLRUGarbageCollectorTests.mm

@@ -25,12 +25,12 @@
 #import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
 #import "Firestore/Source/Local/FSTMutationQueue.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
-#import "Firestore/Source/Local/FSTQueryCache.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #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/query_cache.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 +42,7 @@
 using firebase::firestore::auth::User;
 using firebase::firestore::local::LruParams;
 using firebase::firestore::local::LruResults;
+using firebase::firestore::local::QueryCache;
 using firebase::firestore::local::RemoteDocumentCache;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::DocumentKeyHash;
@@ -58,7 +59,7 @@ @implementation FSTLRUGarbageCollectorTests {
   FSTObjectValue *_testValue;
   FSTObjectValue *_bigObjectValue;
   id<FSTPersistence> _persistence;
-  id<FSTQueryCache> _queryCache;
+  QueryCache *_queryCache;
   RemoteDocumentCache *_documentCache;
   id<FSTMutationQueue> _mutationQueue;
   id<FSTLRUDelegate> _lruDelegate;
@@ -151,7 +152,7 @@ - (FSTQueryData *)nextTestQuery {
 
 - (FSTQueryData *)addNextQueryInTransaction {
   FSTQueryData *queryData = [self nextTestQuery];
-  [_queryCache addQueryData:queryData];
+  _queryCache->AddTarget(queryData);
   return queryData;
 }
 
@@ -161,7 +162,7 @@ - (void)updateTargetInTransaction:(FSTQueryData *)queryData {
       [queryData queryDataByReplacingSnapshotVersion:queryData.snapshotVersion
                                          resumeToken:token
                                       sequenceNumber:_persistence.currentSequenceNumber];
-  [_queryCache updateQueryData:updated];
+  _queryCache->UpdateTarget(updated);
 }
 
 - (FSTQueryData *)addNextQuery {
@@ -195,11 +196,11 @@ - (void)markDocumentEligibleForGCInTransaction:(const DocumentKey &)docKey {
 }
 
 - (void)addDocument:(const DocumentKey &)docKey toTarget:(TargetId)targetId {
-  [_queryCache addMatchingKeys:DocumentKeySet{docKey} forTargetID:targetId];
+  _queryCache->AddMatchingKeys(DocumentKeySet{docKey}, targetId);
 }
 
 - (void)removeDocument:(const DocumentKey &)docKey fromTarget:(TargetId)targetId {
-  [_queryCache removeMatchingKeys:DocumentKeySet{docKey} forTargetID:targetId];
+  _queryCache->RemoveMatchingKeys(DocumentKeySet{docKey}, targetId);
 }
 
 /**
@@ -388,9 +389,9 @@ - (void)testRemoveQueriesUpThroughSequenceNumber {
   XCTAssertEqual(10, removed);
   // Make sure we removed the even targets with targetID <= 20.
   _persistence.run("verify remaining targets are > 20 or odd", [&]() {
-    [_queryCache enumerateTargetsUsingBlock:^(FSTQueryData *queryData, BOOL *stop) {
+    _queryCache->EnumerateTargets(^(FSTQueryData *queryData, BOOL *stop) {
       XCTAssertTrue(queryData.targetID > 20 || queryData.targetID % 2 == 1);
-    }];
+    });
   });
   [_persistence shutdown];
 }
@@ -463,7 +464,7 @@ - (void)testRemoveOrphanedDocuments {
   _persistence.run("verify", [&]() {
     for (const DocumentKey &key : toBeRemoved) {
       XCTAssertNil(_documentCache->Get(key));
-      XCTAssertFalse([_queryCache containsKey:key]);
+      XCTAssertFalse(_queryCache->Contains(key));
     }
     for (const DocumentKey &key : expectedRetained) {
       XCTAssertNotNil(_documentCache->Get(key), @"Missing document %s", key.ToString().c_str());
@@ -645,7 +646,7 @@ - (void)testRemoveTargetsThenGC {
     for (const DocumentKey &key : expectedRemoved) {
       XCTAssertNil(_documentCache->Get(key), @"Did not expect to find %s in document cache",
                    key.ToString().c_str());
-      XCTAssertFalse([_queryCache containsKey:key], @"Did not expect to find %s in queryCache",
+      XCTAssertFalse(_queryCache->Contains(key), @"Did not expect to find %s in queryCache",
                      key.ToString().c_str());
       [self expectSentinelRemoved:key];
     }

Firestore/Example/Tests/Local/FSTLevelDBMigrationsTests.mm

@@ -24,10 +24,10 @@
 #import "Firestore/Protos/objc/firestore/local/Target.pbobjc.h"
 #import "Firestore/Source/Local/FSTLevelDB.h"
 #import "Firestore/Source/Local/FSTLevelDBMutationQueue.h"
-#import "Firestore/Source/Local/FSTLevelDBQueryCache.h"
 
 #include "Firestore/core/src/firebase/firestore/local/leveldb_key.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_migrations.h"
+#include "Firestore/core/src/firebase/firestore/local/leveldb_query_cache.h"
 #include "Firestore/core/src/firebase/firestore/util/ordered_code.h"
 #include "Firestore/core/test/firebase/firestore/testutil/testutil.h"
 #include "absl/strings/match.h"
@@ -43,6 +43,7 @@
 using firebase::firestore::local::LevelDbMigrations;
 using firebase::firestore::local::LevelDbMutationKey;
 using firebase::firestore::local::LevelDbMutationQueueKey;
+using firebase::firestore::local::LevelDbQueryCache;
 using firebase::firestore::local::LevelDbQueryTargetKey;
 using firebase::firestore::local::LevelDbRemoteDocumentKey;
 using firebase::firestore::local::LevelDbTargetDocumentKey;
@@ -86,11 +87,11 @@ - (void)tearDown {
 }
 
 - (void)testAddsTargetGlobal {
-  FSTPBTargetGlobal *metadata = [FSTLevelDBQueryCache readTargetMetadataFromDB:_db.get()];
+  FSTPBTargetGlobal *metadata = LevelDbQueryCache::ReadMetadata(_db.get());
   XCTAssertNil(metadata, @"Not expecting metadata yet, we should have an empty db");
   LevelDbMigrations::RunMigrations(_db.get());
 
-  metadata = [FSTLevelDBQueryCache readTargetMetadataFromDB:_db.get()];
+  metadata = LevelDbQueryCache::ReadMetadata(_db.get());
   XCTAssertNotNil(metadata, @"Migrations should have added the metadata");
 }
 
@@ -166,7 +167,7 @@ - (void)testDropsTheQueryCache {
       ASSERT_FOUND(transaction, key);
     }
 
-    FSTPBTargetGlobal *metadata = [FSTLevelDBQueryCache readTargetMetadataFromDB:_db.get()];
+    FSTPBTargetGlobal *metadata = LevelDbQueryCache::ReadMetadata(_db.get());
     XCTAssertNotNil(metadata, @"Metadata should have been added");
     XCTAssertEqual(metadata.targetCount, 0);
   }
@@ -209,7 +210,7 @@ - (void)testAddsSentinelRows {
     LevelDbTransaction transaction(_db.get(), "Setup");
 
     // Set up target global
-    FSTPBTargetGlobal *metadata = [FSTLevelDBQueryCache readTargetMetadataFromDB:_db.get()];
+    FSTPBTargetGlobal *metadata = LevelDbQueryCache::ReadMetadata(_db.get());
     // Expect that documents missing a row will get the new number
     metadata.highestListenSequenceNumber = new_sequence_number;
     transaction.Put(LevelDbTargetGlobalKey::Key(), metadata);

Firestore/Example/Tests/Local/FSTLevelDBQueryCacheTests.mm

@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-#import "Firestore/Source/Local/FSTLevelDBQueryCache.h"
-
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Local/FSTLevelDB.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
@@ -55,7 +53,7 @@ @interface FSTLevelDBQueryCacheTests : FSTQueryCacheTests
 @implementation FSTLevelDBQueryCacheTests
 
 - (LevelDbQueryCache *)getCache:(id<FSTPersistence>)persistence {
-  return ((FSTLevelDBQueryCache *)([persistence queryCache])).cache;
+  return static_cast<LevelDbQueryCache *>(persistence.queryCache);
 }
 
 - (void)setUp {

Firestore/Example/Tests/Local/FSTLocalStoreTests.mm

@@ -22,7 +22,6 @@
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Local/FSTLocalWriteResult.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
-#import "Firestore/Source/Local/FSTQueryCache.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTDocumentSet.h"

Firestore/Example/Tests/Local/FSTMemoryQueryCacheTests.mm

@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-#import "Firestore/Source/Local/FSTMemoryQueryCache.h"
-
 #import "Firestore/Source/Local/FSTMemoryPersistence.h"
 
 #import "Firestore/Example/Tests/Local/FSTPersistenceTestHelpers.h"
@@ -37,7 +35,7 @@ - (void)setUp {
   [super setUp];
 
   self.persistence = [FSTPersistenceTestHelpers eagerGCMemoryPersistence];
-  self.queryCache = ((FSTMemoryQueryCache *)([self.persistence queryCache])).cache;
+  self.queryCache = self.persistence.queryCache;
 }
 
 - (void)tearDown {

Firestore/Source/Local/FSTLRUGarbageCollector.h

@@ -23,8 +23,6 @@
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/types.h"
 
-@protocol FSTQueryCache;
-
 @class FSTLRUGarbageCollector;
 
 extern const firebase::firestore::model::ListenSequenceNumber kFSTListenSequenceNumberInvalid;
@@ -108,7 +106,7 @@ struct LruResults {
 - (size_t)byteSize;
 
 /** Returns the number of targets and orphaned documents cached. */
-- (int32_t)sequenceNumberCount;
+- (size_t)sequenceNumberCount;
 
 /** Access to the underlying LRU Garbage collector instance. */
 @property(strong, nonatomic, readonly) FSTLRUGarbageCollector *gc;

Firestore/Source/Local/FSTLRUGarbageCollector.mm

@@ -22,7 +22,6 @@
 
 #import "Firestore/Source/Local/FSTMutationQueue.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
-#import "Firestore/Source/Local/FSTQueryCache.h"
 #include "Firestore/core/include/firebase/firestore/timestamp.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/util/log.h"
@@ -148,7 +147,7 @@ - (LruResults)runGCWithLiveTargets:(NSDictionary<NSNumber *, FSTQueryData *> *)l
 }
 
 - (int)queryCountForPercentile:(NSUInteger)percentile {
-  int totalCount = [_delegate sequenceNumberCount];
+  size_t totalCount = [_delegate sequenceNumberCount];
   int setSize = (int)((percentile / 100.0f) * totalCount);
   return setSize;
 }

Firestore/Source/Local/FSTLevelDB.mm

@@ -23,7 +23,6 @@
 #import "Firestore/Source/Core/FSTListenSequence.h"
 #import "Firestore/Source/Local/FSTLRUGarbageCollector.h"
 #import "Firestore/Source/Local/FSTLevelDBMutationQueue.h"
-#import "Firestore/Source/Local/FSTLevelDBQueryCache.h"
 #import "Firestore/Source/Local/FSTReferenceSet.h"
 #import "Firestore/Source/Remote/FSTSerializerBeta.h"
 
@@ -32,6 +31,7 @@
 #include "Firestore/core/src/firebase/firestore/core/database_info.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_key.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_migrations.h"
+#include "Firestore/core/src/firebase/firestore/local/leveldb_query_cache.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_remote_document_cache.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_transaction.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_util.h"
@@ -61,6 +61,7 @@
 using firebase::firestore::local::LevelDbDocumentTargetKey;
 using firebase::firestore::local::LevelDbMigrations;
 using firebase::firestore::local::LevelDbMutationKey;
+using firebase::firestore::local::LevelDbQueryCache;
 using firebase::firestore::local::LevelDbRemoteDocumentCache;
 using firebase::firestore::local::LevelDbTransaction;
 using firebase::firestore::local::LruParams;
@@ -87,6 +88,8 @@ - (size_t)byteSize;
 
 @property(nonatomic, assign, getter=isStarted) BOOL started;
 
+- (firebase::firestore::local::LevelDbQueryCache *)queryCache;
+
 @end
 
 /**
@@ -125,7 +128,7 @@ - (instancetype)initWithPersistence:(FSTLevelDB *)persistence lruParams:(LruPara
 }
 
 - (void)start {
-  ListenSequenceNumber highestSequenceNumber = _db.queryCache.highestListenSequenceNumber;
+  ListenSequenceNumber highestSequenceNumber = _db.queryCache->highest_listen_sequence_number();
   _listenSequence = [[FSTListenSequence alloc] initStartingAfter:highestSequenceNumber];
 }
 
@@ -156,7 +159,7 @@ - (void)removeTarget:(FSTQueryData *)queryData {
       [queryData queryDataByReplacingSnapshotVersion:queryData.snapshotVersion
                                          resumeToken:queryData.resumeToken
                                       sequenceNumber:[self currentSequenceNumber]];
-  [_db.queryCache updateQueryData:updated];
+  _db.queryCache->UpdateTarget(updated);
 }
 
 - (void)addReference:(const DocumentKey &)key {
@@ -195,29 +198,26 @@ - (BOOL)isPinned:(const DocumentKey &)docKey {
 }
 
 - (void)enumerateTargetsUsingBlock:(void (^)(FSTQueryData *queryData, BOOL *stop))block {
-  FSTLevelDBQueryCache *queryCache = _db.queryCache;
-  [queryCache enumerateTargetsUsingBlock:block];
+  _db.queryCache->EnumerateTargets(block);
 }
 
 - (void)enumerateMutationsUsingBlock:
     (void (^)(const DocumentKey &key, ListenSequenceNumber sequenceNumber, BOOL *stop))block {
-  FSTLevelDBQueryCache *queryCache = _db.queryCache;
-  [queryCache enumerateOrphanedDocumentsUsingBlock:block];
+  _db.queryCache->EnumerateOrphanedDocuments(block);
 }
 
 - (int)removeOrphanedDocumentsThroughSequenceNumber:(ListenSequenceNumber)upperBound {
-  FSTLevelDBQueryCache *queryCache = _db.queryCache;
   __block int count = 0;
-  [queryCache enumerateOrphanedDocumentsUsingBlock:^(
-                  const DocumentKey &docKey, ListenSequenceNumber sequenceNumber, BOOL *stop) {
-    if (sequenceNumber <= upperBound) {
-      if (![self isPinned:docKey]) {
-        count++;
-        self->_db.remoteDocumentCache->Remove(docKey);
-        [self removeSentinel:docKey];
-      }
-    }
-  }];
+  _db.queryCache->EnumerateOrphanedDocuments(
+      ^(const DocumentKey &docKey, ListenSequenceNumber sequenceNumber, BOOL *stop) {
+        if (sequenceNumber <= upperBound) {
+          if (![self isPinned:docKey]) {
+            count++;
+            self->_db.remoteDocumentCache->Remove(docKey);
+            [self removeSentinel:docKey];
+          }
+        }
+      });
   return count;
 }
 
@@ -227,12 +227,11 @@ - (void)removeSentinel:(const DocumentKey &)key {
 
 - (int)removeTargetsThroughSequenceNumber:(ListenSequenceNumber)sequenceNumber
                               liveQueries:(NSDictionary<NSNumber *, FSTQueryData *> *)liveQueries {
-  FSTLevelDBQueryCache *queryCache = _db.queryCache;
-  return [queryCache removeQueriesThroughSequenceNumber:sequenceNumber liveQueries:liveQueries];
+  return _db.queryCache->RemoveTargets(sequenceNumber, liveQueries);
 }
 
-- (int32_t)sequenceNumberCount {
-  __block int32_t totalCount = [_db.queryCache count];
+- (size_t)sequenceNumberCount {
+  __block size_t totalCount = _db.queryCache->size();
   [self enumerateMutationsUsingBlock:^(const DocumentKey &key, ListenSequenceNumber sequenceNumber,
                                        BOOL *stop) {
     totalCount++;
@@ -272,7 +271,7 @@ @implementation FSTLevelDB {
   std::unique_ptr<LevelDbRemoteDocumentCache> _documentCache;
   FSTTransactionRunner _transactionRunner;
   FSTLevelDBLRUDelegate *_referenceDelegate;
-  FSTLevelDBQueryCache *_queryCache;
+  std::unique_ptr<LevelDbQueryCache> _queryCache;
   std::set<std::string> _users;
 }
 
@@ -339,14 +338,14 @@ - (instancetype)initWithLevelDB:(std::unique_ptr<leveldb::DB>)db
     _ptr = std::move(db);
     _directory = std::move(directory);
     _serializer = serializer;
-    _queryCache = [[FSTLevelDBQueryCache alloc] initWithDB:self serializer:self.serializer];
+    _queryCache = absl::make_unique<LevelDbQueryCache>(self, _serializer);
     _documentCache = absl::make_unique<LevelDbRemoteDocumentCache>(self, _serializer);
     _referenceDelegate =
         [[FSTLevelDBLRUDelegate alloc] initWithPersistence:self lruParams:lruParams];
     _transactionRunner.SetBackingPersistence(self);
     _users = std::move(users);
     // TODO(gsoltis): set up a leveldb transaction for these operations.
-    [_queryCache start];
+    _queryCache->Start();
     [_referenceDelegate start];
   }
   return self;
@@ -462,8 +461,8 @@ - (LevelDbTransaction *)currentTransaction {
   return [FSTLevelDBMutationQueue mutationQueueWithUser:user db:self serializer:self.serializer];
 }
 
-- (id<FSTQueryCache>)queryCache {
-  return _queryCache;
+- (LevelDbQueryCache *)queryCache {
+  return _queryCache.get();
 }
 
 - (RemoteDocumentCache *)remoteDocumentCache {