Move unit test under AndroidTest so that it can be run on simulator

Author: cherylEnkidu

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/AutoIndexingExperiment.java

@@ -3,16 +3,50 @@
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testCollection;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.toDataMap;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitFor;
+import static com.google.firebase.firestore.testutil.TestUtil.doc;
+import static com.google.firebase.firestore.testutil.TestUtil.docMap;
+import static com.google.firebase.firestore.testutil.TestUtil.filter;
 import static com.google.firebase.firestore.testutil.TestUtil.map;
+import static com.google.firebase.firestore.testutil.TestUtil.query;
 import static org.junit.Assert.assertEquals;
 
+import android.content.Context;
+import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.android.gms.common.internal.Preconditions;
 import com.google.android.gms.tasks.Task;
+import com.google.firebase.database.collection.ImmutableSortedMap;
+import com.google.firebase.database.collection.ImmutableSortedSet;
+import com.google.firebase.firestore.auth.User;
+import com.google.firebase.firestore.core.Query;
+import com.google.firebase.firestore.core.View;
+import com.google.firebase.firestore.local.DocumentOverlayCache;
+import com.google.firebase.firestore.local.IndexManager;
+import com.google.firebase.firestore.local.LocalDocumentsView;
+import com.google.firebase.firestore.local.LocalSerializer;
+import com.google.firebase.firestore.local.LruGarbageCollector;
+import com.google.firebase.firestore.local.MutationQueue;
+import com.google.firebase.firestore.local.Persistence;
+import com.google.firebase.firestore.local.QueryContext;
+import com.google.firebase.firestore.local.QueryEngine;
+import com.google.firebase.firestore.local.RemoteDocumentCache;
+import com.google.firebase.firestore.local.SQLitePersistence;
+import com.google.firebase.firestore.model.DatabaseId;
+import com.google.firebase.firestore.model.Document;
+import com.google.firebase.firestore.model.DocumentKey;
+import com.google.firebase.firestore.model.DocumentSet;
+import com.google.firebase.firestore.model.FieldIndex;
+import com.google.firebase.firestore.model.MutableDocument;
+import com.google.firebase.firestore.remote.RemoteSerializer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -127,4 +161,211 @@ public void testIndexedWithNonIndexedTime() {
       }
     }
   }
+
+  private Persistence persistence;
+  private RemoteDocumentCache remoteDocumentCache;
+  private MutationQueue mutationQueue;
+  private DocumentOverlayCache documentOverlayCache;
+
+  protected IndexManager indexManager;
+  protected QueryEngine queryEngine;
+
+  private @Nullable Boolean expectFullCollectionScan;
+
+  public static SQLitePersistence createSQLitePersistence() {
+    DatabaseId databaseId = DatabaseId.forProject("projectId");
+    LocalSerializer serializer = new LocalSerializer(new RemoteSerializer(databaseId));
+    Context context = ApplicationProvider.getApplicationContext();
+    SQLitePersistence persistence =
+        new SQLitePersistence(
+            context, "test", databaseId, serializer, LruGarbageCollector.Params.Default());
+    persistence.start();
+    return persistence;
+  }
+
+  @Before
+  public void setUp() {
+    expectFullCollectionScan = null;
+
+    persistence = createSQLitePersistence();
+
+    indexManager = persistence.getIndexManager(User.UNAUTHENTICATED);
+    mutationQueue = persistence.getMutationQueue(User.UNAUTHENTICATED, indexManager);
+    documentOverlayCache = persistence.getDocumentOverlayCache(User.UNAUTHENTICATED);
+    remoteDocumentCache = persistence.getRemoteDocumentCache();
+    queryEngine = new QueryEngine();
+
+    indexManager.start();
+    mutationQueue.start();
+
+    remoteDocumentCache.setIndexManager(indexManager);
+
+    LocalDocumentsView localDocuments =
+        new LocalDocumentsView(
+            remoteDocumentCache, mutationQueue, documentOverlayCache, indexManager) {
+          @Override
+          public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
+              com.google.firebase.firestore.core.Query query, FieldIndex.IndexOffset offset) {
+            assertEquals(
+                "Observed query execution mode did not match expectation",
+                expectFullCollectionScan,
+                FieldIndex.IndexOffset.NONE.equals(offset));
+            return super.getDocumentsMatchingQuery(query, offset);
+          }
+        };
+    queryEngine.initialize(localDocuments, indexManager, false);
+  }
+
+  /** Adds the provided documents to the remote document cache. */
+  protected void addDocument(MutableDocument... docs) {
+    persistence.runTransaction(
+        "addDocument",
+        () -> {
+          for (MutableDocument doc : docs) {
+            remoteDocumentCache.add(doc, doc.getVersion());
+          }
+        });
+  }
+
+  protected <T> T expectOptimizedCollectionScan(Callable<T> c) throws Exception {
+    try {
+      expectFullCollectionScan = false;
+      return c.call();
+    } finally {
+      expectFullCollectionScan = null;
+    }
+  }
+
+  private <T> T expectFullCollectionScan(Callable<T> c) throws Exception {
+    try {
+      expectFullCollectionScan = true;
+      return c.call();
+    } finally {
+      expectFullCollectionScan = null;
+    }
+  }
+
+  protected DocumentSet runQuery(
+      com.google.firebase.firestore.core.Query query, boolean autoIndexing, QueryContext counter) {
+    Preconditions.checkNotNull(
+        expectFullCollectionScan,
+        "Encountered runQuery() call not wrapped in expectOptimizedCollectionQuery()/expectFullCollectionQuery()");
+    ImmutableSortedMap<DocumentKey, Document> docs =
+        queryEngine.getDocumentsMatchingQueryTest(query, autoIndexing, counter);
+    View view =
+        new View(query, new ImmutableSortedSet<>(Collections.emptyList(), DocumentKey::compareTo));
+    View.DocumentChanges viewDocChanges = view.computeDocChanges(docs);
+    return view.applyChanges(viewDocChanges).getSnapshot().getDocuments();
+  }
+
+  private void createTestingDocument(
+      String basePath, int documentID, boolean isMatched, int numOfFields) {
+    Map<String, Object> fields = map("match", isMatched);
+
+    // Randomly generate the rest of fields
+    for (int i = 2; i <= numOfFields; i++) {
+      int valueIndex = (int) (Math.random() * values.size()) % values.size();
+      fields.put("field" + i, values.get(valueIndex));
+    }
+
+    MutableDocument doc = doc(basePath + "/" + documentID, 1, fields);
+    addDocument(doc);
+
+    indexManager.updateIndexEntries(docMap(doc));
+    indexManager.updateCollectionGroup(basePath, FieldIndex.IndexOffset.fromDocument(doc));
+  }
+
+  private void createTestingCollection(
+      String basePath, int totalSetCount, int portion /*0 - 10*/, int numOfFields /* 1 - 30*/) {
+    int documentCounter = 0;
+    for (int i = 1; i <= totalSetCount; i++) {
+      // Generate a random order list of 0 ... 9
+      ArrayList<Integer> indexes = new ArrayList<>();
+      for (int index = 0; index < 10; index++) {
+        indexes.add(index);
+      }
+      Collections.shuffle(indexes);
+
+      for (int match = 0; match < portion; match++) {
+        int currentID = documentCounter + indexes.get(match);
+        createTestingDocument(basePath, currentID, true, numOfFields);
+      }
+      for (int unmatch = portion; unmatch < 10; unmatch++) {
+        int currentID = documentCounter + indexes.get(unmatch);
+        createTestingDocument(basePath, currentID, false, numOfFields);
+      }
+      documentCounter += 10;
+    }
+  }
+
+  @Test
+  public void testCombinesIndexedWithNonIndexedResults() throws Exception {
+    // Every set contains 10 documents
+    final int numOfSet = 100;
+    // could overflow. Currently it is safe when numOfSet set to 1000 and running on macbook M1
+    long totalBeforeIndex = 0;
+    long totalAfterIndex = 0;
+    long totalDocumentCount = 0;
+    long totalResultCount = 0;
+
+    // Temperate heuristic
+    double without = 3.7;
+    double with = 4.9;
+
+    for (int totalSetCount = 1; totalSetCount <= numOfSet; totalSetCount *= 10) {
+      // portion stands for the percentage of documents matching query
+      for (int portion = 0; portion <= 10; portion++) {
+        for (int numOfFields = 1; numOfFields <= 31; numOfFields += 10) {
+          String basePath = "documentCount" + totalSetCount;
+          // Auto indexing
+          Query query = query(basePath).filter(filter("match", "==", true));
+          indexManager.createTargetIndices(query.toTarget());
+          createTestingCollection(basePath, totalSetCount, portion, numOfFields);
+
+          QueryContext counterWithoutIndex = new QueryContext();
+          long beforeAutoStart = System.nanoTime();
+          DocumentSet results =
+              expectFullCollectionScan(() -> runQuery(query, false, counterWithoutIndex));
+          long beforeAutoEnd = System.nanoTime();
+          long millisecondsBeforeAuto =
+              TimeUnit.MILLISECONDS.convert(
+                  (beforeAutoEnd - beforeAutoStart), TimeUnit.NANOSECONDS);
+          totalBeforeIndex += (beforeAutoEnd - beforeAutoStart);
+          totalDocumentCount += counterWithoutIndex.fullScanCount;
+          assertEquals(portion * totalSetCount, results.size());
+
+          QueryContext counterWithIndex = new QueryContext();
+          long autoStart = System.nanoTime();
+          results = expectOptimizedCollectionScan(() -> runQuery(query, true, counterWithIndex));
+          long autoEnd = System.nanoTime();
+          long millisecondsAfterAuto =
+              TimeUnit.MILLISECONDS.convert((autoEnd - autoStart), TimeUnit.NANOSECONDS);
+          totalAfterIndex += (autoEnd - autoStart);
+          assertEquals(portion * totalSetCount, results.size());
+          totalResultCount += results.size();
+          if (millisecondsBeforeAuto > millisecondsAfterAuto) {
+            System.out.println(
+                "Auto Indexing saves time when total of documents inside collection is "
+                    + totalSetCount * 10
+                    + ". The matching percentage is "
+                    + portion
+                    + "0%. And each document contains "
+                    + numOfFields
+                    + " fields.\n"
+                    + "Weight result for without auto indexing is "
+                    + without * counterWithoutIndex.fullScanCount
+                    + ". And weight result for auto indexing is "
+                    + with * results.size());
+          }
+        }
+      }
+    }
+
+    System.out.println(
+        "The time heuristic is "
+            + (totalBeforeIndex / totalDocumentCount)
+            + " before auto indexing");
+    System.out.println(
+        "The time heuristic is " + (totalAfterIndex / totalResultCount) + " after auto indexing");
+  }
 }

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

@@ -48,14 +48,14 @@
  * in remoteDocumentCache or local mutations for the document). The view is computed by applying the
  * mutations in the MutationQueue to the RemoteDocumentCache.
  */
-class LocalDocumentsView {
+public class LocalDocumentsView {
 
   private final RemoteDocumentCache remoteDocumentCache;
   private final MutationQueue mutationQueue;
   private final DocumentOverlayCache documentOverlayCache;
   private final IndexManager indexManager;
 
-  LocalDocumentsView(
+  protected LocalDocumentsView(
       RemoteDocumentCache remoteDocumentCache,
       MutationQueue mutationQueue,
       DocumentOverlayCache documentOverlayCache,
@@ -276,7 +276,7 @@ ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
     }
   }
 
-  ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
+  protected ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       Query query, IndexOffset offset) {
     return getDocumentsMatchingQuery(query, offset, new QueryContext());
   }

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

@@ -94,7 +94,7 @@ private void setReferenceDelegate(ReferenceDelegate delegate) {
   }
 
   @Override
-  MutationQueue getMutationQueue(User user, IndexManager indexManager) {
+  public MutationQueue getMutationQueue(User user, IndexManager indexManager) {
     MemoryMutationQueue queue = mutationQueues.get(user);
     if (queue == null) {
       queue = new MemoryMutationQueue(this, user);
@@ -113,12 +113,12 @@ MemoryTargetCache getTargetCache() {
   }
 
   @Override
-  MemoryRemoteDocumentCache getRemoteDocumentCache() {
+  public MemoryRemoteDocumentCache getRemoteDocumentCache() {
     return remoteDocumentCache;
   }
 
   @Override
-  MemoryIndexManager getIndexManager(User user) {
+  public MemoryIndexManager getIndexManager(User user) {
     // We do not currently support indices for memory persistence, so we can return the same shared
     // instance of the memory index manager.
     return indexManager;
@@ -130,7 +130,7 @@ BundleCache getBundleCache() {
   }
 
   @Override
-  DocumentOverlayCache getDocumentOverlayCache(User user) {
+  public DocumentOverlayCache getDocumentOverlayCache(User user) {
     MemoryDocumentOverlayCache overlay = overlays.get(user);
     if (overlay == null) {
       overlay = new MemoryDocumentOverlayCache();
@@ -145,7 +145,7 @@ OverlayMigrationManager getOverlayMigrationManager() {
   }
 
   @Override
-  void runTransaction(String action, Runnable operation) {
+  public void runTransaction(String action, Runnable operation) {
     referenceDelegate.onTransactionStarted();
     try {
       operation.run();

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

@@ -24,7 +24,7 @@
 import java.util.List;
 
 /** A queue of mutations to apply to the remote store. */
-interface MutationQueue {
+public interface MutationQueue {
   /**
    * Starts the mutation queue, performing any initial reads that might be required to establish
    * invariants, etc.

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

@@ -78,22 +78,22 @@ public abstract class Persistence {
    * implementation to the extent possible (e.g. in the case of uid switching from
    * sally=>jack=>sally, sally's mutation queue will be preserved).
    */
-  abstract MutationQueue getMutationQueue(User user, IndexManager indexManager);
+  public abstract MutationQueue getMutationQueue(User user, IndexManager indexManager);
 
   /** Creates a TargetCache representing the persisted cache of queries. */
   abstract TargetCache getTargetCache();
 
   /** Creates a RemoteDocumentCache representing the persisted cache of remote documents. */
-  abstract RemoteDocumentCache getRemoteDocumentCache();
+  public abstract RemoteDocumentCache getRemoteDocumentCache();
 
   /** Creates an IndexManager that manages our persisted query indexes. */
-  abstract IndexManager getIndexManager(User user);
+  public abstract IndexManager getIndexManager(User user);
 
   /** Returns a BundleCache representing the persisted cache of loaded bundles. */
   abstract BundleCache getBundleCache();
 
   /** Returns a DocumentOverlayCache representing the documents that are mutated locally. */
-  abstract DocumentOverlayCache getDocumentOverlayCache(User user);
+  public abstract DocumentOverlayCache getDocumentOverlayCache(User user);
 
   /** Returns a OverlayMigrationManager that runs any pending data migration required by SDK. */
   abstract OverlayMigrationManager getOverlayMigrationManager();
@@ -106,7 +106,7 @@ public abstract class Persistence {
    * @param action A description of the action performed by this transaction, used for logging.
    * @param operation The operation to run inside a transaction.
    */
-  abstract void runTransaction(String action, Runnable operation);
+  public abstract void runTransaction(String action, Runnable operation);
 
   /**
    * Performs an operation inside a persistence transaction. Any reads or writes against persistence

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

@@ -15,7 +15,7 @@
 package com.google.firebase.firestore.local;
 
 public class QueryContext {
-  QueryContext() {}
+  public QueryContext() {}
 
-  int fullScanCount = 0;
+  public int fullScanCount = 0;
 }

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

@@ -31,7 +31,7 @@
  * meaning we can cache both Document instances (an actual document with data) as well as NoDocument
  * instances (indicating that the document is known to not exist).
  */
-interface RemoteDocumentCache {
+public interface RemoteDocumentCache {
 
   /** Sets the index manager to use for managing the collectionGroup index. */
   void setIndexManager(IndexManager indexManager);