Firestore: QoL improvements for converters

packages/firestore/src/api/database.ts

@@ -127,6 +127,7 @@ import {
   NextFn,
   PartialObserver
 } from './observer';
+import { NestedPartialWithFieldValue, WithFieldValue } from '../lite/reference';
 
 /**
  * A persistence provider for either memory-only or IndexedDB persistence.
@@ -440,21 +441,24 @@ export class Transaction implements PublicTransaction, Compat<ExpTransaction> {
 
   set<T>(
     documentRef: DocumentReference<T>,
-    data: Partial<T>,
+    data: NestedPartialWithFieldValue<T>,
     options: PublicSetOptions
   ): Transaction;
-  set<T>(documentRef: DocumentReference<T>, data: T): Transaction;
+  set<T>(
+    documentRef: DocumentReference<T>,
+    data: WithFieldValue<T>
+  ): Transaction;
   set<T>(
     documentRef: PublicDocumentReference<T>,
-    data: T | Partial<T>,
+    data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
     options?: PublicSetOptions
   ): Transaction {
     const ref = castReference(documentRef);
     if (options) {
       validateSetOptions('Transaction.set', options);
-      this._delegate.set(ref, data, options);
+      this._delegate.set(ref, data as NestedPartialWithFieldValue<T>, options);
     } else {
-      this._delegate.set(ref, data);
+      this._delegate.set(ref, data as WithFieldValue<T>);
     }
     return this;
   }
@@ -501,21 +505,24 @@ export class WriteBatch implements PublicWriteBatch, Compat<ExpWriteBatch> {
   constructor(readonly _delegate: ExpWriteBatch) {}
   set<T>(
     documentRef: DocumentReference<T>,
-    data: Partial<T>,
+    data: NestedPartialWithFieldValue<T>,
     options: PublicSetOptions
   ): WriteBatch;
-  set<T>(documentRef: DocumentReference<T>, data: T): WriteBatch;
+  set<T>(
+    documentRef: DocumentReference<T>,
+    data: WithFieldValue<T>
+  ): WriteBatch;
   set<T>(
     documentRef: PublicDocumentReference<T>,
-    data: T | Partial<T>,
+    data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
     options?: PublicSetOptions
   ): WriteBatch {
     const ref = castReference(documentRef);
     if (options) {
       validateSetOptions('WriteBatch.set', options);
-      this._delegate.set(ref, data, options);
+      this._delegate.set(ref, data as NestedPartialWithFieldValue<T>, options);
     } else {
-      this._delegate.set(ref, data);
+      this._delegate.set(ref, data as WithFieldValue<T>);
     }
     return this;
   }
@@ -597,19 +604,19 @@ class FirestoreDataConverter<U>
     );
   }
 
-  toFirestore(modelObject: U): PublicDocumentData;
+  toFirestore(modelObject: WithFieldValue<U>): PublicDocumentData;
   toFirestore(
-    modelObject: Partial<U>,
+    modelObject: NestedPartialWithFieldValue<U>,
     options: PublicSetOptions
   ): PublicDocumentData;
   toFirestore(
-    modelObject: U | Partial<U>,
+    modelObject: WithFieldValue<U> | NestedPartialWithFieldValue<U>,
     options?: PublicSetOptions
   ): PublicDocumentData {
     if (!options) {
       return this._delegate.toFirestore(modelObject as U);
     } else {
-      return this._delegate.toFirestore(modelObject, options);
+      return this._delegate.toFirestore(modelObject as Partial<U>, options);
     }
   }
 
@@ -733,7 +740,15 @@ export class DocumentReference<T = PublicDocumentData>
   set(value: T | Partial<T>, options?: PublicSetOptions): Promise<void> {
     options = validateSetOptions('DocumentReference.set', options);
     try {
-      return setDoc(this._delegate, value, options);
+      if (options) {
+        return setDoc(
+          this._delegate,
+          value as NestedPartialWithFieldValue<T>,
+          options
+        );
+      } else {
+        return setDoc(this._delegate, value as WithFieldValue<T>);
+      }
     } catch (e) {
       throw replaceFunctionName(e, 'setDoc()', 'DocumentReference.set()');
     }

packages/firestore/src/exp/reference_impl.ts

@@ -42,9 +42,12 @@ import {
   CollectionReference,
   doc,
   DocumentReference,
+  NestedPartialWithFieldValue,
   Query,
   SetOptions,
-  UpdateData
+  TypedUpdateData,
+  UpdateData,
+  WithFieldValue
 } from '../lite/reference';
 import { applyFirestoreDataConverter } from '../lite/reference_impl';
 import {
@@ -243,7 +246,7 @@ export function getDocsFromServer<T>(
  */
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: T
+  data: WithFieldValue<T>
 ): Promise<void>;
 /**
  * Writes to the document referred to by the specified `DocumentReference`. If
@@ -258,12 +261,12 @@ export function setDoc<T>(
  */
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: Partial<T>,
+  data: NestedPartialWithFieldValue<T>,
   options: SetOptions
 ): Promise<void>;
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: T,
+  data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
   options?: SetOptions
 ): Promise<void> {
   reference = cast<DocumentReference<T>>(reference, DocumentReference);
@@ -300,9 +303,9 @@ export function setDoc<T>(
  * @returns A Promise resolved once the data has been successfully written
  * to the backend (note that it won't resolve while you're offline).
  */
-export function updateDoc(
-  reference: DocumentReference<unknown>,
-  data: UpdateData
+export function updateDoc<T>(
+  reference: DocumentReference<T>,
+  data: TypedUpdateData<T>
 ): Promise<void>;
 /**
  * Updates fields in the document referred to by the specified
@@ -325,9 +328,9 @@ export function updateDoc(
   value: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void>;
-export function updateDoc(
+export function updateDoc<T>(
   reference: DocumentReference<unknown>,
-  fieldOrUpdateData: string | FieldPath | UpdateData,
+  fieldOrUpdateData: string | FieldPath | TypedUpdateData<T>,
   value?: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void> {
@@ -393,7 +396,7 @@ export function deleteDoc(
  */
 export function addDoc<T>(
   reference: CollectionReference<T>,
-  data: T
+  data: WithFieldValue<T>
 ): Promise<DocumentReference<T>> {
   const firestore = cast(reference.firestore, Firestore);
 

packages/firestore/src/exp/snapshot.ts

@@ -18,7 +18,14 @@
 import { newQueryComparator } from '../core/query';
 import { ChangeType, ViewSnapshot } from '../core/view_snapshot';
 import { FieldPath } from '../lite/field_path';
-import { DocumentData, Query, queryEqual, SetOptions } from '../lite/reference';
+import {
+  DocumentData,
+  NestedPartialWithFieldValue,
+  Query,
+  queryEqual,
+  SetOptions,
+  WithFieldValue
+} from '../lite/reference';
 import {
   DocumentSnapshot as LiteDocumentSnapshot,
   fieldPathFromArgument,
@@ -84,15 +91,18 @@ export interface FirestoreDataConverter<T>
    * Firestore database). To use `set()` with `merge` and `mergeFields`,
    * `toFirestore()` must be defined with `Partial<T>`.
    */
-  toFirestore(modelObject: T): DocumentData;
+  toFirestore(modelObject: WithFieldValue<T>): DocumentData;
 
   /**
    * Called by the Firestore SDK to convert a custom model object of type `T`
    * into a plain JavaScript object (suitable for writing directly to the
    * Firestore database). Used with {@link (setDoc:1)}, {@link (WriteBatch.set:1)}
    * and {@link (Transaction.set:1)} with `merge:true` or `mergeFields`.
    */
-  toFirestore(modelObject: Partial<T>, options: SetOptions): DocumentData;
+  toFirestore(
+    modelObject: NestedPartialWithFieldValue<T>,
+    options: SetOptions
+  ): DocumentData;
 
   /**
    * Called by the Firestore SDK to convert Firestore data into an object of

packages/firestore/src/lite/reference.ts

@@ -36,6 +36,7 @@ import { AutoId } from '../util/misc';
 
 import { Firestore } from './database';
 import { FieldPath } from './field_path';
+import { FieldValue } from './field_value';
 import { FirestoreDataConverter } from './snapshot';
 
 /**
@@ -48,6 +49,21 @@ export interface DocumentData {
   [field: string]: any;
 }
 
+type Primitive = string | number | boolean | bigint | undefined | null;
+// eslint-disable-next-line @typescript-eslint/ban-types
+type Builtin = Primitive | Function;
+
+/** Like Partial but recursive */
+export type NestedPartialWithFieldValue<T> = T extends Builtin
+  ? T
+  : T extends Map<infer K, infer V>
+  ? Map<NestedPartialWithFieldValue<K>, NestedPartialWithFieldValue<V>>
+  : T extends {}
+  ? { [K in keyof T]?: NestedPartialWithFieldValue<T[K]> | FieldValue }
+  : Partial<T>;
+
+export type WithFieldValue<T> = { [P in keyof T]: T[P] | FieldValue };
+
 /**
  * Update data (for use with {@link @firebase/firestore/lite#(updateDoc:1)}) consists of field paths (e.g.
  * 'foo' or 'foo.baz') mapped to values. Fields that contain dots reference
@@ -58,6 +74,39 @@ export interface UpdateData {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   [fieldPath: string]: any;
 }
+// Represents an update object to Firestore document data, which can contain either fields like {a: 2}
+// or dot-separated paths such as {"a.b" : 2} (which updates the nested property "b" in map field "a").
+export type TypedUpdateData<T> = T extends Builtin
+  ? T
+  : T extends Map<infer K, infer V>
+  ? Map<TypedUpdateData<K>, TypedUpdateData<V>>
+  : T extends {}
+  ? { [K in keyof T]?: TypedUpdateData<T[K]> | FieldValue } &
+      NestedUpdateFields<T>
+  : Partial<T>;
+
+// For each field (e.g. "bar"), calculate its nested keys (e.g. {"bar.baz": T1, "bar.quax": T2}), and then
+// intersect them together to make one giant map containing all possible keys (all marked as optional).
+type NestedUpdateFields<T extends Record<string, any>> = UnionToIntersection<
+  {
+    [K in keyof T & string]: T[K] extends Record<string, any> // Only allow nesting for map values
+      ? AddPrefixToKeys<K, TypedUpdateData<T[K]>> // Recurse into map and add "bar." in front of every key
+      : never;
+  }[keyof T & string]
+>;
+
+// Return a new map where every key is prepended with Prefix + dot.
+type AddPrefixToKeys<Prefix extends string, T extends Record<string, any>> =
+  // Remap K => Prefix.K. See https://www.typescriptlang.org/docs/handbook/2/mapped-types.html#key-remapping-via-as
+  { [K in keyof T & string as `${Prefix}.${K}`]+?: T[K] };
+
+// This takes union type U = T1 | T2 | ... and returns a intersected type (T1 & T2 & ...)
+type UnionToIntersection<U> =
+  // Works because "multiple candidates for the same type variable in contra-variant positions causes an intersection type to be inferred"
+  // https://www.typescriptlang.org/docs/handbook/advanced-types.html#type-inference-in-conditional-types
+  (U extends any ? (k: U) => void : never) extends (k: infer I) => void
+    ? I
+    : never;
 
 /**
  * An options object that configures the behavior of {@link @firebase/firestore/lite#(setDoc:1)}, {@link

packages/firestore/src/lite/reference_impl.ts

@@ -41,9 +41,12 @@ import {
   CollectionReference,
   doc,
   DocumentReference,
+  NestedPartialWithFieldValue,
   Query,
   SetOptions,
-  UpdateData
+  TypedUpdateData,
+  UpdateData,
+  WithFieldValue
 } from './reference';
 import {
   DocumentSnapshot,
@@ -197,7 +200,7 @@ export function getDocs<T>(query: Query<T>): Promise<QuerySnapshot<T>> {
  */
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: T
+  data: WithFieldValue<T>
 ): Promise<void>;
 /**
  * Writes to the document referred to by the specified `DocumentReference`. If
@@ -217,12 +220,12 @@ export function setDoc<T>(
  */
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: Partial<T>,
+  data: NestedPartialWithFieldValue<T>,
   options: SetOptions
 ): Promise<void>;
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: T,
+  data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
   options?: SetOptions
 ): Promise<void> {
   reference = cast<DocumentReference<T>>(reference, DocumentReference);
@@ -264,9 +267,9 @@ export function setDoc<T>(
  * @returns A Promise resolved once the data has been successfully written
  * to the backend.
  */
-export function updateDoc(
-  reference: DocumentReference<unknown>,
-  data: UpdateData
+export function updateDoc<T>(
+  reference: DocumentReference<T>,
+  data: TypedUpdateData<T>
 ): Promise<void>;
 /**
  * Updates fields in the document referred to by the specified
@@ -294,9 +297,9 @@ export function updateDoc(
   value: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void>;
-export function updateDoc(
+export function updateDoc<T>(
   reference: DocumentReference<unknown>,
-  fieldOrUpdateData: string | FieldPath | UpdateData,
+  fieldOrUpdateData: string | FieldPath | TypedUpdateData<T>,
   value?: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void> {

packages/firestore/src/lite/snapshot.ts

@@ -27,9 +27,11 @@ import { FieldPath } from './field_path';
 import {
   DocumentData,
   DocumentReference,
+  NestedPartialWithFieldValue,
   Query,
   queryEqual,
-  SetOptions
+  SetOptions,
+  WithFieldValue
 } from './reference';
 import {
   fieldPathFromDotSeparatedString,
@@ -83,15 +85,18 @@ export interface FirestoreDataConverter<T> {
    * Firestore database). Used with {@link @firebase/firestore/lite#(setDoc:1)}, {@link @firebase/firestore/lite#(WriteBatch.set:1)}
    * and {@link @firebase/firestore/lite#(Transaction.set:1)}.
    */
-  toFirestore(modelObject: T): DocumentData;
+  toFirestore(modelObject: WithFieldValue<T>): DocumentData;
 
   /**
    * Called by the Firestore SDK to convert a custom model object of type `T`
    * into a plain Javascript object (suitable for writing directly to the
    * Firestore database). Used with {@link @firebase/firestore/lite#(setDoc:1)}, {@link @firebase/firestore/lite#(WriteBatch.set:1)}
    * and {@link @firebase/firestore/lite#(Transaction.set:1)} with `merge:true` or `mergeFields`.
    */
-  toFirestore(modelObject: Partial<T>, options: SetOptions): DocumentData;
+  toFirestore(
+    modelObject: NestedPartialWithFieldValue<T>,
+    options: SetOptions
+  ): DocumentData;
 
   /**
    * Called by the Firestore SDK to convert Firestore data into an object of