Add option to allow SDK create cache indexes automatically to improve query execution locally

Firestore/Example/Tests/Integration/API/FIRDatabaseTests.mm

@@ -1829,4 +1829,63 @@ - (void)testUnlimitedCacheSize {
   XCTAssertEqualObjects(result.data, data);
 }
 
+- (void)testGetValidPersistentCacheIndexManager {
+  [FIRApp configure];
+
+  FIRFirestore *db1 = [FIRFirestore firestoreForDatabase:@"PersistentCacheIndexManagerDB1"];
+  FIRFirestoreSettings *settings1 = [db1 settings];
+  [settings1 setCacheSettings:[[FIRPersistentCacheSettings alloc] init]];
+  [db1 setSettings:settings1];
+
+  XCTAssertNotNil(db1.persistentCacheIndexManager);
+
+  // Use persistent disk cache (default)
+  FIRFirestore *db2 = [FIRFirestore firestoreForDatabase:@"PersistentCacheIndexManagerDB2"];
+  XCTAssertNotNil(db2.persistentCacheIndexManager);
+
+  // Disable persistent disk cache
+  FIRFirestore *db3 = [FIRFirestore firestoreForDatabase:@"PersistentCacheIndexManagerDB3"];
+  FIRFirestoreSettings *settings3 = [db3 settings];
+  [settings3 setCacheSettings:[[FIRMemoryCacheSettings alloc] init]];
+  [db3 setSettings:settings3];
+
+  XCTAssertNil(db3.persistentCacheIndexManager);
+
+  // Disable persistent disk cache (deprecated)
+  FIRFirestore *db4 = [FIRFirestore firestoreForDatabase:@"PersistentCacheIndexManagerDB4"];
+  FIRFirestoreSettings *settings4 = [db4 settings];
+  settings4.persistenceEnabled = NO;
+  [db4 setSettings:settings4];
+  XCTAssertNil(db4.persistentCacheIndexManager);
+}
+
+- (void)testCanGetSameOrDifferentPersistentCacheIndexManager {
+  [FIRApp configure];
+  // Use persistent disk cache (explicit)
+  FIRFirestore *db1 = [FIRFirestore firestoreForDatabase:@"PersistentCacheIndexManagerDB5"];
+  FIRFirestoreSettings *settings1 = [db1 settings];
+  [settings1 setCacheSettings:[[FIRPersistentCacheSettings alloc] init]];
+  [db1 setSettings:settings1];
+  XCTAssertEqual(db1.persistentCacheIndexManager, db1.persistentCacheIndexManager);
+
+  // Use persistent disk cache (default)
+  FIRFirestore *db2 = [FIRFirestore firestoreForDatabase:@"PersistentCacheIndexManagerDB6"];
+  XCTAssertEqual(db2.persistentCacheIndexManager, db2.persistentCacheIndexManager);
+
+  XCTAssertNotEqual(db1.persistentCacheIndexManager, db2.persistentCacheIndexManager);
+
+  FIRFirestore *db3 = [FIRFirestore firestoreForDatabase:@"PersistentCacheIndexManagerDB7"];
+  FIRFirestoreSettings *settings3 = [db3 settings];
+  [settings3 setCacheSettings:[[FIRPersistentCacheSettings alloc] init]];
+  [db3 setSettings:settings3];
+  XCTAssertNotEqual(db1.persistentCacheIndexManager, db3.persistentCacheIndexManager);
+  XCTAssertNotEqual(db2.persistentCacheIndexManager, db3.persistentCacheIndexManager);
+
+  // Use persistent disk cache (default)
+  FIRFirestore *db4 = [FIRFirestore firestoreForDatabase:@"PersistentCacheIndexManagerDB8"];
+  XCTAssertNotEqual(db1.persistentCacheIndexManager, db4.persistentCacheIndexManager);
+  XCTAssertNotEqual(db2.persistentCacheIndexManager, db4.persistentCacheIndexManager);
+  XCTAssertNotEqual(db3.persistentCacheIndexManager, db4.persistentCacheIndexManager);
+}
+
 @end

Firestore/Example/Tests/Integration/API/FIRIndexingTests.mm

@@ -18,6 +18,10 @@
 
 #import <XCTest/XCTest.h>
 
+#import "Firestore/Source/API/FIRFirestore+Internal.h"
+#import "Firestore/Source/API/FIRPersistentCacheIndexManager+Internal.h"
+
+#import "Firestore/Example/Tests/Util/FSTHelpers.h"
 #import "Firestore/Example/Tests/Util/FSTIntegrationTestCase.h"
 
 @interface FIRIndexingTests : FSTIntegrationTestCase
@@ -29,15 +33,15 @@ @implementation FIRIndexingTests
 - (void)setUp {
   [super setUp];
   self.db = [self firestore];
-  XCTestExpectation* exp = [self expectationWithDescription:@"clear persistence"];
-  [self.db clearPersistenceWithCompletion:^(NSError*) {
+  XCTestExpectation *exp = [self expectationWithDescription:@"clear persistence"];
+  [self.db clearPersistenceWithCompletion:^(NSError *) {
     [exp fulfill];
   }];
   [self awaitExpectation:exp];
 }
 
 - (void)testCanConfigureIndexes {
-  NSString* json = @"{\n"
+  NSString *json = @"{\n"
                     "\t\"indexes\": [{\n"
                     "\t\t\t\"collectionGroup\": \"restaurants\",\n"
                     "\t\t\t\"queryScope\": \"COLLECTION\",\n"
@@ -64,22 +68,22 @@ - (void)testCanConfigureIndexes {
                     "}";
 
   [self.db setIndexConfigurationFromJSON:json
-                              completion:^(NSError* error) {
+                              completion:^(NSError *error) {
                                 XCTAssertNil(error);
                               }];
 }
 
 - (void)testBadJsonDoesNotCrashClient {
   [self.db setIndexConfigurationFromJSON:@"{,"
-                              completion:^(NSError* error) {
+                              completion:^(NSError *error) {
                                 XCTAssertNotNil(error);
                                 XCTAssertEqualObjects(error.domain, FIRFirestoreErrorDomain);
                                 XCTAssertEqual(error.code, FIRFirestoreErrorCodeInvalidArgument);
                               }];
 }
 
 - (void)testBadIndexDoesNotCrashClient {
-  NSString* json = @"{\n"
+  NSString *json = @"{\n"
                     "\t\"indexes\": [{\n"
                     "\t\t\"collectionGroup\": \"restaurants\",\n"
                     "\t\t\"queryScope\": \"COLLECTION\",\n"
@@ -92,11 +96,114 @@ - (void)testBadIndexDoesNotCrashClient {
                     "}";
 
   [self.db setIndexConfigurationFromJSON:json
-                              completion:^(NSError* error) {
+                              completion:^(NSError *error) {
                                 XCTAssertNotNil(error);
                                 XCTAssertEqualObjects(error.domain, FIRFirestoreErrorDomain);
                                 XCTAssertEqual(error.code, FIRFirestoreErrorCodeInvalidArgument);
                               }];
 }
 
+/**
+ * After Auto Index Creation is enabled, through public API there is no way to see the indexes
+ * sitting inside SDK. So this test only checks the API of auto index creation.
+ */
+- (void)testAutoIndexCreationSetSuccessfully {
+  // Use persistent disk cache (explict)
+  FIRFirestoreSettings *settings = [self.db settings];
+  [settings setCacheSettings:[[FIRPersistentCacheSettings alloc] init]];
+  [self.db setSettings:settings];
+
+  FIRCollectionReference *coll = [self collectionRef];
+  NSDictionary *testDocs = @{
+    @"a" : @{@"match" : @YES},
+    @"b" : @{@"match" : @NO},
+    @"c" : @{@"match" : @NO},
+  };
+  [self writeAllDocuments:testDocs toCollection:coll];
+
+  FIRQuery *query = [coll queryWhereField:@"match" isEqualTo:@YES];
+
+  [query getDocumentsWithSource:FIRFirestoreSourceCache
+                     completion:^(FIRQuerySnapshot *results, NSError *error) {
+                       XCTAssertNil(error);
+                       XCTAssertEqual(results.count, 1);
+                     }];
+
+  XCTAssertNoThrow([self.db.persistentCacheIndexManager enableIndexAutoCreation]);
+  [query getDocumentsWithSource:FIRFirestoreSourceCache
+                     completion:^(FIRQuerySnapshot *results, NSError *error) {
+                       XCTAssertNil(error);
+                       XCTAssertEqual(results.count, 1);
+                     }];
+
+  XCTAssertNoThrow([self.db.persistentCacheIndexManager disableIndexAutoCreation]);
+  [query getDocumentsWithSource:FIRFirestoreSourceCache
+                     completion:^(FIRQuerySnapshot *results, NSError *error) {
+                       XCTAssertNil(error);
+                       XCTAssertEqual(results.count, 1);
+                     }];
+
+  XCTAssertNoThrow([self.db.persistentCacheIndexManager deleteAllIndexes]);
+  [query getDocumentsWithSource:FIRFirestoreSourceCache
+                     completion:^(FIRQuerySnapshot *results, NSError *error) {
+                       XCTAssertNil(error);
+                       XCTAssertEqual(results.count, 1);
+                     }];
+}
+
+- (void)testAutoIndexCreationSetSuccessfullyUsingDefault {
+  // Use persistent disk cache (default)
+  FIRCollectionReference *coll = [self collectionRef];
+  NSDictionary *testDocs = @{
+    @"a" : @{@"match" : @YES},
+    @"b" : @{@"match" : @NO},
+    @"c" : @{@"match" : @NO},
+  };
+  [self writeAllDocuments:testDocs toCollection:coll];
+
+  FIRQuery *query = [coll queryWhereField:@"match" isEqualTo:@YES];
+
+  [query getDocumentsWithSource:FIRFirestoreSourceCache
+                     completion:^(FIRQuerySnapshot *results, NSError *error) {
+                       XCTAssertNil(error);
+                       XCTAssertEqual(results.count, 1);
+                     }];
+
+  XCTAssertNoThrow([self.db.persistentCacheIndexManager enableIndexAutoCreation]);
+  [query getDocumentsWithSource:FIRFirestoreSourceCache
+                     completion:^(FIRQuerySnapshot *results, NSError *error) {
+                       XCTAssertNil(error);
+                       XCTAssertEqual(results.count, 1);
+                     }];
+
+  XCTAssertNoThrow([self.db.persistentCacheIndexManager disableIndexAutoCreation]);
+  [query getDocumentsWithSource:FIRFirestoreSourceCache
+                     completion:^(FIRQuerySnapshot *results, NSError *error) {
+                       XCTAssertNil(error);
+                       XCTAssertEqual(results.count, 1);
+                     }];
+
+  XCTAssertNoThrow([self.db.persistentCacheIndexManager deleteAllIndexes]);
+  [query getDocumentsWithSource:FIRFirestoreSourceCache
+                     completion:^(FIRQuerySnapshot *results, NSError *error) {
+                       XCTAssertNil(error);
+                       XCTAssertEqual(results.count, 1);
+                     }];
+}
+
+- (void)testAutoIndexCreationAfterFailsTermination {
+  [self terminateFirestore:self.db];
+
+  FSTAssertThrows([self.db.persistentCacheIndexManager enableIndexAutoCreation],
+                  @"The client has already been terminated.");
+
+  FSTAssertThrows([self.db.persistentCacheIndexManager disableIndexAutoCreation],
+                  @"The client has already been terminated.");
+
+  FSTAssertThrows([self.db.persistentCacheIndexManager deleteAllIndexes],
+                  @"The client has already been terminated.");
+}
+
+// TODO(b/296100693) Add testing hooks to verify indexes are created as expected.
+
 @end

Firestore/Source/API/FIRFirestore+Internal.h

@@ -25,6 +25,7 @@
 
 @class FIRApp;
 @class FSTUserDataReader;
+@class FIRPersistentCacheIndexManager;
 
 namespace firebase {
 namespace firestore {
@@ -78,6 +79,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (const std::shared_ptr<firebase::firestore::util::AsyncQueue> &)workerQueue;
 
+// TODO(csi): make this function public
+@property(nonatomic, readonly) FIRPersistentCacheIndexManager *persistentCacheIndexManager;
+
 @property(nonatomic, assign, readonly) std::shared_ptr<api::Firestore> wrapped;
 
 @property(nonatomic, assign, readonly) const model::DatabaseId &databaseID;

Firestore/Source/API/FIRFirestore.mm

@@ -21,6 +21,7 @@
 #include <utility>
 
 #import "FIRFirestoreSettings+Internal.h"
+#import "FIRPersistentCacheIndexManager+Internal.h"
 #import "FIRTransactionOptions+Internal.h"
 #import "FIRTransactionOptions.h"
 
@@ -106,6 +107,7 @@ @implementation FIRFirestore {
   std::shared_ptr<Firestore> _firestore;
   FIRFirestoreSettings *_settings;
   __weak id<FSTFirestoreInstanceRegistry> _registry;
+  FIRPersistentCacheIndexManager *_indexManager;
 }
 
 + (void)initialize {
@@ -533,6 +535,17 @@ @implementation FIRFirestore (Internal)
   return _firestore->worker_queue();
 }
 
+- (FIRPersistentCacheIndexManager *)persistentCacheIndexManager {
+  if (!_indexManager) {
+    auto index_manager = _firestore->persistent_cache_index_manager();
+    if (index_manager) {
+      _indexManager = [[FIRPersistentCacheIndexManager alloc]
+          initWithPersistentCacheIndexManager:index_manager];
+    }
+  }
+  return _indexManager;
+}
+
 - (const DatabaseId &)databaseID {
   return _firestore->database_id();
 }

Firestore/Source/API/FIRPersistentCacheIndexManager+Internal.h

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <Foundation/Foundation.h>
+
+#include <memory>
+
+#include "Firestore/core/src/api/persistent_cache_index_manager.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+// TODO(sum/avg) move the contents of this category to
+// ../Public/FirebaseFirestore/FIRPersistentCacheIndexManager.h
+/**
+ * A PersistentCacheIndexManager which you can config persistent cache indexes used for
+ * local query execution.
+ */
+NS_SWIFT_NAME(PersistentCacheIndexManager)
+@interface FIRPersistentCacheIndexManager : NSObject
+
+/** :nodoc: */
+- (instancetype)init
+    __attribute__((unavailable("FIRPersistentCacheIndexManager cannot be created directly.")));
+
+/**
+ * Enables SDK to create persistent cache indexes automatically for local query execution when SDK
+ * believes cache indexes can help improves performance.
+ *
+ * This feature is disabled by default.
+ */
+- (void)enableIndexAutoCreation NS_SWIFT_NAME(enableIndexAutoCreation());
+
+/**
+ * Stops creating persistent cache indexes automatically for local query execution. The indexes
+ * which have been created by calling `enableIndexAutoCreation` still take effect.
+ */
+- (void)disableIndexAutoCreation NS_SWIFT_NAME(disableIndexAutoCreation());
+
+/**
+ * Removes all persistent cache indexes. Please note this function will also deletes indexes
+ * generated by [[FIRFirestore firestore] setIndexConfigurationFromJSON] and [[FIRFirestore
+ * firestore] setIndexConfigurationFromStream], which are deprecated.
+ */
+- (void)deleteAllIndexes NS_SWIFT_NAME(deleteAllIndexes());
+
+@end
+
+@interface FIRPersistentCacheIndexManager (/* Init */)
+
+- (instancetype)initWithPersistentCacheIndexManager:
+    (std::shared_ptr<const firebase::firestore::api::PersistentCacheIndexManager>)indexManager
+    NS_DESIGNATED_INITIALIZER;
+
+@end
+
+NS_ASSUME_NONNULL_END