Simplify ObjectValue (#4933)

Author: schmidt-sebastian

packages/firestore/src/exp/snapshot.ts

@@ -278,7 +278,7 @@ export class DocumentSnapshot<
       return this._converter.fromFirestore(snapshot, options);
     } else {
       return this._userDataWriter.convertValue(
-        this._document.data.toProto(),
+        this._document.data.value,
         options.serverTimestamps
       ) as T;
     }

packages/firestore/src/lite/snapshot.ts

@@ -174,9 +174,7 @@ export class DocumentSnapshot<T = DocumentData> {
       );
       return this._converter.fromFirestore(snapshot);
     } else {
-      return this._userDataWriter.convertValue(
-        this._document.data.toProto()
-      ) as T;
+      return this._userDataWriter.convertValue(this._document.data.value) as T;
     }
   }
 

packages/firestore/src/lite/user_data_writer.ts

@@ -92,7 +92,7 @@ export abstract class AbstractUserDataWriter {
     serverTimestampBehavior: ServerTimestampBehavior
   ): DocumentData {
     const result: DocumentData = {};
-    forEach(mapValue.fields || {}, (key, value) => {
+    forEach(mapValue.fields, (key, value) => {
       result[key] = this.convertValue(value, serverTimestampBehavior);
     });
     return result;

packages/firestore/src/local/local_store_impl.ts

@@ -16,7 +16,7 @@
  */
 
 import { User } from '../auth/user';
-import { BundledDocuments, NamedQuery, BundleConverter } from '../core/bundle';
+import { BundleConverter, BundledDocuments, NamedQuery } from '../core/bundle';
 import { newQueryForPath, Query, queryToTarget } from '../core/query';
 import { SnapshotVersion } from '../core/snapshot_version';
 import { canonifyTarget, Target, targetEquals } from '../core/target';
@@ -330,7 +330,7 @@ export function localStoreWriteLocally(
                 new PatchMutation(
                   mutation.key,
                   baseValue,
-                  extractFieldMask(baseValue.toProto().mapValue!),
+                  extractFieldMask(baseValue.value.mapValue),
                   Precondition.exists(true)
                 )
               );

packages/firestore/src/local/memory_persistence.ts

@@ -474,7 +474,7 @@ export class MemoryLruDelegate implements ReferenceDelegate, LruDelegate {
   documentSize(document: Document): number {
     let documentSize = document.key.toString().length;
     if (document.isFoundDocument()) {
-      documentSize += estimateByteSize(document.data.toProto());
+      documentSize += estimateByteSize(document.data.value);
     }
     return documentSize;
   }

packages/firestore/src/model/document.ts

@@ -331,7 +331,7 @@ export class MutableDocument implements Document {
   toString(): string {
     return (
       `Document(${this.key}, ${this.version}, ${JSON.stringify(
-        this.data.toProto()
+        this.data.value
       )}, ` +
       `{documentType: ${this.documentType}}), ` +
       `{documentState: ${this.documentState}})`

packages/firestore/src/model/object_value.ts

@@ -25,45 +25,21 @@ import { forEach } from '../util/obj';
 import { FieldMask } from './field_mask';
 import { FieldPath } from './path';
 import { isServerTimestamp } from './server_timestamps';
-import { TypeOrder } from './type_order';
-import { isMapValue, typeOrder, valueEquals } from './values';
+import { deepClone, isMapValue, valueEquals } from './values';
 
 export interface JsonObject<T> {
   [name: string]: T;
 }
-
-/**
- * An Overlay, which contains an update to apply. Can either be Value proto, a
- * map of Overlay values (to represent additional nesting at the given key) or
- * `null` (to represent field deletes).
- */
-type Overlay = Map<string, Overlay> | ProtoValue | null;
-
 /**
  * An ObjectValue represents a MapValue in the Firestore Proto and offers the
  * ability to add and remove fields (via the ObjectValueBuilder).
  */
 export class ObjectValue {
-  /**
-   * The immutable Value proto for this object. Local mutations are stored in
-   * `overlayMap` and only applied when `buildProto()` is invoked.
-   */
-  private partialValue: { mapValue: ProtoMapValue };
-
-  /**
-   * A nested map that contains the accumulated changes that haven't yet been
-   * applied to `partialValue`. Values can either be `Value` protos, Map<String,
-   * Object> values (to represent additional nesting) or `null` (to represent
-   * field deletes).
-   */
-  private overlayMap = new Map<string, Overlay>();
-
-  constructor(proto: { mapValue: ProtoMapValue }) {
+  constructor(readonly value: { mapValue: ProtoMapValue }) {
     debugAssert(
-      !isServerTimestamp(proto),
+      !isServerTimestamp(value),
       'ServerTimestamps should be converted to ServerTimestampValue'
     );
-    this.partialValue = proto;
   }
 
   static empty(): ObjectValue {
@@ -77,12 +53,19 @@ export class ObjectValue {
    * @returns The value at the path or null if the path is not set.
    */
   field(path: FieldPath): ProtoValue | null {
-    return ObjectValue.extractNestedValue(this.buildProto(), path);
-  }
-
-  /** Returns the full protobuf representation. */
-  toProto(): { mapValue: ProtoMapValue } {
-    return this.field(FieldPath.emptyPath()) as { mapValue: ProtoMapValue };
+    if (path.isEmpty()) {
+      return this.value;
+    } else {
+      let currentLevel: ProtoValue = this.value;
+      for (let i = 0; i < path.length - 1; ++i) {
+        currentLevel = (currentLevel.mapValue!.fields || {})[path.get(i)];
+        if (!isMapValue(currentLevel)) {
+          return null;
+        }
+      }
+      currentLevel = (currentLevel.mapValue!.fields! || {})[path.lastSegment()];
+      return currentLevel || null;
+    }
   }
 
   /**
@@ -96,7 +79,8 @@ export class ObjectValue {
       !path.isEmpty(),
       'Cannot set field for empty path on ObjectValue'
     );
-    this.setOverlay(path, value);
+    const fieldsMap = this.getFieldsMap(path.popLast());
+    fieldsMap[path.lastSegment()] = deepClone(value);
   }
 
   /**
@@ -105,13 +89,30 @@ export class ObjectValue {
    * @param data - A map of fields to values (or null for deletes).
    */
   setAll(data: Map<FieldPath, ProtoValue | null>): void {
-    data.forEach((value, fieldPath) => {
+    let parent = FieldPath.emptyPath();
+
+    let upserts: { [key: string]: ProtoValue } = {};
+    let deletes: string[] = [];
+
+    data.forEach((value, path) => {
+      if (!parent.isImmediateParentOf(path)) {
+        // Insert the accumulated changes at this parent location
+        const fieldsMap = this.getFieldsMap(parent);
+        this.applyChanges(fieldsMap, upserts, deletes);
+        upserts = {};
+        deletes = [];
+        parent = path.popLast();
+      }
+
       if (value) {
-        this.set(fieldPath, value);
+        upserts[path.lastSegment()] = deepClone(value);
       } else {
-        this.delete(fieldPath);
+        deletes.push(path.lastSegment());
       }
     });
+
+    const fieldsMap = this.getFieldsMap(parent);
+    this.applyChanges(fieldsMap, upserts, deletes);
   }
 
   /**
@@ -125,138 +126,58 @@ export class ObjectValue {
       !path.isEmpty(),
       'Cannot delete field for empty path on ObjectValue'
     );
-    this.setOverlay(path, null);
+    const nestedValue = this.field(path.popLast());
+    if (isMapValue(nestedValue) && nestedValue.mapValue.fields) {
+      delete nestedValue.mapValue.fields[path.lastSegment()];
+    }
   }
 
   isEqual(other: ObjectValue): boolean {
-    return valueEquals(this.buildProto(), other.buildProto());
+    return valueEquals(this.value, other.value);
   }
 
   /**
-   * Adds `value` to the overlay map at `path`. Creates nested map entries if
-   * needed.
+   * Returns the map that contains the leaf element of `path`. If the parent
+   * entry does not yet exist, or if it is not a map, a new map will be created.
    */
-  private setOverlay(path: FieldPath, value: ProtoValue | null): void {
-    let currentLevel = this.overlayMap;
+  private getFieldsMap(path: FieldPath): Record<string, ProtoValue> {
+    let current = this.value;
 
-    for (let i = 0; i < path.length - 1; ++i) {
-      const currentSegment = path.get(i);
-      let currentValue = currentLevel.get(currentSegment);
-
-      if (currentValue instanceof Map) {
-        // Re-use a previously created map
-        currentLevel = currentValue;
-      } else if (
-        currentValue &&
-        typeOrder(currentValue) === TypeOrder.ObjectValue
-      ) {
-        // Convert the existing Protobuf MapValue into a map
-        currentValue = new Map<string, Overlay>(
-          Object.entries(currentValue.mapValue!.fields || {})
-        );
-        currentLevel.set(currentSegment, currentValue);
-        currentLevel = currentValue;
-      } else {
-        // Create an empty map to represent the current nesting level
-        currentValue = new Map<string, Overlay>();
-        currentLevel.set(currentSegment, currentValue);
-        currentLevel = currentValue;
-      }
+    if (!current.mapValue!.fields) {
+      current.mapValue = { fields: {} };
     }
 
-    currentLevel.set(path.lastSegment(), value);
-  }
-
-  /**
-   * Applies any overlays from `currentOverlays` that exist at `currentPath`
-   * and returns the merged data at `currentPath` (or null if there were no
-   * changes).
-   *
-   * @param currentPath - The path at the current nesting level. Can be set to
-   * FieldValue.emptyPath() to represent the root.
-   * @param currentOverlays - The overlays at the current nesting level in the
-   * same format as `overlayMap`.
-   * @returns The merged data at `currentPath` or null if no modifications
-   * were applied.
-   */
-  private applyOverlay(
-    currentPath: FieldPath,
-    currentOverlays: Map<string, Overlay>
-  ): { mapValue: ProtoMapValue } | null {
-    let modified = false;
-
-    const existingValue = ObjectValue.extractNestedValue(
-      this.partialValue,
-      currentPath
-    );
-    const resultAtPath = isMapValue(existingValue)
-      ? // If there is already data at the current path, base our
-        // modifications on top of the existing data.
-        { ...existingValue.mapValue.fields }
-      : {};
-
-    currentOverlays.forEach((value, pathSegment) => {
-      if (value instanceof Map) {
-        const nested = this.applyOverlay(currentPath.child(pathSegment), value);
-        if (nested != null) {
-          resultAtPath[pathSegment] = nested;
-          modified = true;
-        }
-      } else if (value !== null) {
-        resultAtPath[pathSegment] = value;
-        modified = true;
-      } else if (resultAtPath.hasOwnProperty(pathSegment)) {
-        delete resultAtPath[pathSegment];
-        modified = true;
+    for (let i = 0; i < path.length; ++i) {
+      let next = current.mapValue!.fields![path.get(i)];
+      if (!isMapValue(next) || !next.mapValue.fields) {
+        next = { mapValue: { fields: {} } };
+        current.mapValue!.fields![path.get(i)] = next;
       }
-    });
+      current = next as { mapValue: ProtoMapValue };
+    }
 
-    return modified ? { mapValue: { fields: resultAtPath } } : null;
+    return current.mapValue!.fields!;
   }
 
   /**
-   * Builds the Protobuf that backs this ObjectValue.
-   *
-   * This method applies any outstanding modifications and memoizes the result.
-   * Further invocations are based on this memoized result.
+   * Modifies `fieldsMap` by adding, replacing or deleting the specified
+   * entries.
    */
-  private buildProto(): { mapValue: ProtoMapValue } {
-    const mergedResult = this.applyOverlay(
-      FieldPath.emptyPath(),
-      this.overlayMap
-    );
-    if (mergedResult != null) {
-      this.partialValue = mergedResult;
-      this.overlayMap.clear();
-    }
-    return this.partialValue;
-  }
-
-  private static extractNestedValue(
-    proto: ProtoValue,
-    path: FieldPath
-  ): ProtoValue | null {
-    if (path.isEmpty()) {
-      return proto;
-    } else {
-      let value: ProtoValue = proto;
-      for (let i = 0; i < path.length - 1; ++i) {
-        if (!value.mapValue!.fields) {
-          return null;
-        }
-        value = value.mapValue!.fields[path.get(i)];
-        if (!isMapValue(value)) {
-          return null;
-        }
-      }
-
-      value = (value.mapValue!.fields || {})[path.lastSegment()];
-      return value || null;
+  private applyChanges(
+    fieldsMap: Record<string, ProtoValue>,
+    inserts: { [key: string]: ProtoValue },
+    deletes: string[]
+  ): void {
+    forEach(inserts, (key, val) => (fieldsMap[key] = val));
+    for (const field of deletes) {
+      delete fieldsMap[field];
     }
   }
 
   clone(): ObjectValue {
-    return new ObjectValue(this.buildProto());
+    return new ObjectValue(
+      deepClone(this.value) as { mapValue: ProtoMapValue }
+    );
   }
 }
 
@@ -265,7 +186,7 @@ export class ObjectValue {
  */
 export function extractFieldMask(value: ProtoMapValue): FieldMask {
   const fields: FieldPath[] = [];
-  forEach(value!.fields || {}, (key, value) => {
+  forEach(value!.fields, (key, value) => {
     const currentPath = new FieldPath([key]);
     if (isMapValue(value)) {
       const nestedMask = extractFieldMask(value.mapValue!);