Simplify Document

firebase-firestore/ktx/src/test/java/com/google/firebase/firestore/testutil/TestUtil.java

@@ -64,21 +64,21 @@ public static SnapshotVersion version(long versionMicros) {
 
   public static Document doc(String key, long version, Map<String, Object> data) {
     return new Document(
-        key(key), version(version), Document.DocumentState.SYNCED, wrapObject(data));
+        key(key), version(version), wrapObject(data), Document.DocumentState.SYNCED);
   }
 
   public static Document doc(DocumentKey key, long version, Map<String, Object> data) {
-    return new Document(key, version(version), Document.DocumentState.SYNCED, wrapObject(data));
+    return new Document(key, version(version), wrapObject(data), Document.DocumentState.SYNCED);
   }
 
   public static Document doc(
       String key, long version, ObjectValue data, Document.DocumentState documentState) {
-    return new Document(key(key), version(version), documentState, data);
+    return new Document(key(key), version(version), data, documentState);
   }
 
   public static Document doc(
       String key, long version, Map<String, Object> data, Document.DocumentState documentState) {
-    return new Document(key(key), version(version), documentState, wrapObject(data));
+    return new Document(key(key), version(version), wrapObject(data), documentState);
   }
 
   public static DocumentSet docSet(Comparator<Document> comparator, Document... documents) {

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

@@ -55,12 +55,7 @@ com.google.firebase.firestore.proto.MaybeDocument encodeMaybeDocument(MaybeDocum
       builder.setHasCommittedMutations(noDocument.hasCommittedMutations());
     } else if (document instanceof Document) {
       Document existingDocument = (Document) document;
-      // Use the memoized encoded form if it exists.
-      if (existingDocument.getProto() != null) {
-        builder.setDocument(existingDocument.getProto());
-      } else {
-        builder.setDocument(encodeDocument(existingDocument));
-      }
+      builder.setDocument(encodeDocument(existingDocument));
       builder.setHasCommittedMutations(existingDocument.hasCommittedMutations());
     } else if (document instanceof UnknownDocument) {
       builder.setUnknownDocument(encodeUnknownDocument((UnknownDocument) document));
@@ -112,15 +107,15 @@ private com.google.firestore.v1.Document encodeDocument(Document document) {
   private Document decodeDocument(
       com.google.firestore.v1.Document document, boolean hasCommittedMutations) {
     DocumentKey key = rpcSerializer.decodeKey(document.getName());
+    ObjectValue value = rpcSerializer.decodeFields(document.getFieldsMap());
     SnapshotVersion version = rpcSerializer.decodeVersion(document.getUpdateTime());
     return new Document(
         key,
         version,
+        value,
         hasCommittedMutations
             ? Document.DocumentState.COMMITTED_MUTATIONS
-            : Document.DocumentState.SYNCED,
-        document,
-        rpcSerializer::decodeValue);
+            : Document.DocumentState.SYNCED);
   }
 
   /** Encodes a NoDocument value to the equivalent proto. */

firebase-firestore/src/main/java/com/google/firebase/firestore/model/Document.java

@@ -14,17 +14,11 @@
 
 package com.google.firebase.firestore.model;
 
-import static com.google.firebase.firestore.util.Assert.hardAssert;
-
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
-import com.google.common.base.Function;
 import com.google.firebase.firestore.model.value.FieldValue;
 import com.google.firebase.firestore.model.value.ObjectValue;
-import com.google.firestore.v1.Value;
 import java.util.Comparator;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * Represents a document in Firestore with a key, version, data and whether the data has local
@@ -51,101 +45,25 @@ public static Comparator<Document> keyComparator() {
   }
 
   private final DocumentState documentState;
-  private @Nullable final com.google.firestore.v1.Document proto;
-  private @Nullable final Function<Value, FieldValue> converter;
-  private @Nullable ObjectValue objectValue;
-
-  /** A cache for FieldValues that have already been deserialized in `getField()`. */
-  private @Nullable Map<FieldPath, FieldValue> fieldValueCache;
+  private ObjectValue objectValue;
 
   public Document(
       DocumentKey key,
       SnapshotVersion version,
-      DocumentState documentState,
-      ObjectValue objectValue) {
+      ObjectValue objectValue,
+      DocumentState documentState) {
     super(key, version);
     this.documentState = documentState;
     this.objectValue = objectValue;
-    this.proto = null;
-    this.converter = null;
-  }
-
-  public Document(
-      DocumentKey key,
-      SnapshotVersion version,
-      DocumentState documentState,
-      com.google.firestore.v1.Document proto,
-      Function<com.google.firestore.v1.Value, FieldValue> converter) {
-    super(key, version);
-    this.documentState = documentState;
-    this.proto = proto;
-    this.converter = converter;
-  }
-
-  /**
-   * Memoized serialized form of the document for optimization purposes (avoids repeated
-   * serialization). Might be null.
-   */
-  public @Nullable com.google.firestore.v1.Document getProto() {
-    return proto;
   }
 
   @NonNull
   public ObjectValue getData() {
-    if (objectValue == null) {
-      hardAssert(proto != null && converter != null, "Expected proto and converter to be non-null");
-
-      ObjectValue.Builder result = ObjectValue.newBuilder();
-      for (Map.Entry<String, com.google.firestore.v1.Value> entry :
-          proto.getFieldsMap().entrySet()) {
-        FieldPath path = FieldPath.fromSingleSegment(entry.getKey());
-        FieldValue value = converter.apply(entry.getValue());
-        result.set(path, value);
-      }
-      objectValue = result.build();
-
-      // Once objectValue is computed, values inside the fieldValueCache are no longer accessed.
-      fieldValueCache = null;
-    }
-
     return objectValue;
   }
 
   public @Nullable FieldValue getField(FieldPath path) {
-    if (objectValue != null) {
-      return objectValue.get(path);
-    } else {
-      hardAssert(proto != null && converter != null, "Expected proto and converter to be non-null");
-
-      Map<FieldPath, FieldValue> fieldValueCache = this.fieldValueCache;
-      if (fieldValueCache == null) {
-        // TODO(b/136090445): Remove the cache when `getField` is no longer called during Query
-        // ordering.
-        fieldValueCache = new ConcurrentHashMap<>();
-        this.fieldValueCache = fieldValueCache;
-      }
-
-      FieldValue fieldValue = fieldValueCache.get(path);
-      if (fieldValue == null) {
-        // Instead of deserializing the full Document proto, we only deserialize the value at
-        // the requested field path. This speeds up Query execution as query filters can discard
-        // documents based on a single field.
-        Value protoValue = proto.getFieldsMap().get(path.getFirstSegment());
-        for (int i = 1; protoValue != null && i < path.length(); ++i) {
-          if (protoValue.getValueTypeCase() != Value.ValueTypeCase.MAP_VALUE) {
-            return null;
-          }
-          protoValue = protoValue.getMapValue().getFieldsMap().get(path.getSegment(i));
-        }
-
-        if (protoValue != null) {
-          fieldValue = converter.apply(protoValue);
-          fieldValueCache.put(path, fieldValue);
-        }
-      }
-
-      return fieldValue;
-    }
+    return objectValue.get(path);
   }
 
   public @Nullable Object getFieldValue(FieldPath path) {
@@ -180,15 +98,15 @@ public boolean equals(Object o) {
     return getVersion().equals(document.getVersion())
         && getKey().equals(document.getKey())
         && documentState.equals(document.documentState)
-        && getData().equals(document.getData());
+        && objectValue.equals(document.objectValue);
   }
 
   @Override
   public int hashCode() {
-    // Note: We deliberately decided to omit `getData()` since its computation is expensive.
     int result = getKey().hashCode();
     result = 31 * result + getVersion().hashCode();
     result = 31 * result + documentState.hashCode();
+    result = 31 * result + objectValue.hashCode();
     return result;
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/PatchMutation.java

@@ -112,7 +112,7 @@ public MaybeDocument applyToRemoteDocument(
 
     SnapshotVersion version = mutationResult.getVersion();
     ObjectValue newData = patchDocument(maybeDoc);
-    return new Document(getKey(), version, Document.DocumentState.COMMITTED_MUTATIONS, newData);
+    return new Document(getKey(), version, newData, Document.DocumentState.COMMITTED_MUTATIONS);
   }
 
   @Nullable
@@ -127,7 +127,7 @@ public MaybeDocument applyToLocalView(
 
     SnapshotVersion version = getPostMutationVersion(maybeDoc);
     ObjectValue newData = patchDocument(maybeDoc);
-    return new Document(getKey(), version, Document.DocumentState.LOCAL_MUTATIONS, newData);
+    return new Document(getKey(), version, newData, Document.DocumentState.LOCAL_MUTATIONS);
   }
 
   @Nullable

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/SetMutation.java

@@ -72,7 +72,7 @@ public MaybeDocument applyToRemoteDocument(
     // accepted the mutation so the precondition must have held.
 
     SnapshotVersion version = mutationResult.getVersion();
-    return new Document(getKey(), version, Document.DocumentState.COMMITTED_MUTATIONS, value);
+    return new Document(getKey(), version, value, Document.DocumentState.COMMITTED_MUTATIONS);
   }
 
   @Nullable
@@ -86,7 +86,7 @@ public MaybeDocument applyToLocalView(
     }
 
     SnapshotVersion version = getPostMutationVersion(maybeDoc);
-    return new Document(getKey(), version, Document.DocumentState.LOCAL_MUTATIONS, value);
+    return new Document(getKey(), version, value, Document.DocumentState.LOCAL_MUTATIONS);
   }
 
   /** Returns the object value to use when setting the document. */

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/TransformMutation.java

@@ -101,7 +101,7 @@ public MaybeDocument applyToRemoteDocument(
         serverTransformResults(doc, mutationResult.getTransformResults());
     ObjectValue newData = transformObject(doc.getData(), transformResults);
     return new Document(
-        getKey(), mutationResult.getVersion(), Document.DocumentState.COMMITTED_MUTATIONS, newData);
+        getKey(), mutationResult.getVersion(), newData, Document.DocumentState.COMMITTED_MUTATIONS);
   }
 
   @Nullable
@@ -118,7 +118,7 @@ public MaybeDocument applyToLocalView(
     List<FieldValue> transformResults = localTransformResults(localWriteTime, maybeDoc, baseDoc);
     ObjectValue newData = transformObject(doc.getData(), transformResults);
     return new Document(
-        getKey(), doc.getVersion(), Document.DocumentState.LOCAL_MUTATIONS, newData);
+        getKey(), doc.getVersion(), newData, Document.DocumentState.LOCAL_MUTATIONS);
   }
 
   @Nullable

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/RemoteSerializer.java

@@ -405,11 +405,11 @@ private Document decodeFoundDocument(BatchGetDocumentsResponse response) {
         response.getResultCase().equals(ResultCase.FOUND),
         "Tried to deserialize a found document from a missing document.");
     DocumentKey key = decodeKey(response.getFound().getName());
+    ObjectValue value = decodeFields(response.getFound().getFieldsMap());
     SnapshotVersion version = decodeVersion(response.getFound().getUpdateTime());
     hardAssert(
         !version.equals(SnapshotVersion.NONE), "Got a document response with no snapshot version");
-    return new Document(
-        key, version, Document.DocumentState.SYNCED, response.getFound(), this::decodeValue);
+    return new Document(key, version, value, Document.DocumentState.SYNCED);
   }
 
   private NoDocument decodeMissingDocument(BatchGetDocumentsResponse response) {
@@ -1066,13 +1066,8 @@ public WatchChange decodeWatchChange(ListenResponse protoChange) {
         SnapshotVersion version = decodeVersion(docChange.getDocument().getUpdateTime());
         hardAssert(
             !version.equals(SnapshotVersion.NONE), "Got a document change without an update time");
-        Document document =
-            new Document(
-                key,
-                version,
-                Document.DocumentState.SYNCED,
-                docChange.getDocument(),
-                this::decodeValue);
+        ObjectValue data = decodeFields(docChange.getDocument().getFieldsMap());
+        Document document = new Document(key, version, data, Document.DocumentState.SYNCED);
         watchChange = new WatchChange.DocumentChange(added, removed, document.getKey(), document);
         break;
       case DOCUMENT_DELETE:

firebase-firestore/src/test/java/com/google/firebase/firestore/local/LruGarbageCollectorTestCase.java

@@ -578,7 +578,7 @@ public void testRemoveTargetsThenGC() {
         () -> {
           SnapshotVersion newVersion = version(3);
           Document doc =
-              new Document(middleDocToUpdate, newVersion, Document.DocumentState.SYNCED, testValue);
+              new Document(middleDocToUpdate, newVersion, testValue, Document.DocumentState.SYNCED);
           documentCache.add(doc, newVersion);
           updateTargetInTransaction(middleTarget);
         });

firebase-firestore/src/test/java/com/google/firebase/firestore/model/DocumentTest.java

@@ -39,7 +39,7 @@ public class DocumentTest {
   public void testInstantiation() {
     Document document =
         new Document(
-            key("messages/first"), version(1), Document.DocumentState.SYNCED, wrapObject("a", 1));
+            key("messages/first"), version(1), wrapObject("a", 1), Document.DocumentState.SYNCED);
 
     assertEquals(key("messages/first"), document.getKey());
     assertEquals(version(1), document.getVersion());
@@ -56,7 +56,7 @@ public void testExtractFields() {
             "owner",
             map("name", "Jonny", "title", "scallywag"));
     Document document =
-        new Document(key("rooms/eros"), version(1), Document.DocumentState.SYNCED, data);
+        new Document(key("rooms/eros"), version(1), data, Document.DocumentState.SYNCED);
 
     assertEquals("Discuss all the project related stuff", document.getFieldValue(field("desc")));
     assertEquals("scallywag", document.getFieldValue(field("owner.title")));

firebase-firestore/src/test/java/com/google/firebase/firestore/model/MutationTest.java

@@ -169,8 +169,8 @@ public void testAppliesLocalServerTimestampTransformsToDocuments() {
         new Document(
             key("collection/key"),
             version(0),
-            Document.DocumentState.LOCAL_MUTATIONS,
-            expectedData);
+            expectedData,
+            Document.DocumentState.LOCAL_MUTATIONS);
     assertEquals(expectedDoc, transformedDoc);
   }
 

firebase-firestore/src/testUtil/java/com/google/firebase/firestore/testutil/TestUtil.java

@@ -179,21 +179,21 @@ public static SnapshotVersion version(long versionMicros) {
 
   public static Document doc(String key, long version, Map<String, Object> data) {
     return new Document(
-        key(key), version(version), Document.DocumentState.SYNCED, wrapObject(data));
+        key(key), version(version), wrapObject(data), Document.DocumentState.SYNCED);
   }
 
   public static Document doc(DocumentKey key, long version, Map<String, Object> data) {
-    return new Document(key, version(version), Document.DocumentState.SYNCED, wrapObject(data));
+    return new Document(key, version(version), wrapObject(data), Document.DocumentState.SYNCED);
   }
 
   public static Document doc(
       String key, long version, ObjectValue data, Document.DocumentState documentState) {
-    return new Document(key(key), version(version), documentState, data);
+    return new Document(key(key), version(version), data, documentState);
   }
 
   public static Document doc(
       String key, long version, Map<String, Object> data, Document.DocumentState documentState) {
-    return new Document(key(key), version(version), documentState, wrapObject(data));
+    return new Document(key(key), version(version), wrapObject(data), documentState);
   }
 
   public static NoDocument deletedDoc(String key, long version) {