Adding mergeFields support (#1141)

8 files changed, 172 insertions, 10 deletions

diff --git a/Firestore/Source/API/FIRDocumentReference.mm b/Firestore/Source/API/FIRDocumentReference.mm
index c2fc546..da67a5b 100644
--- a/Firestore/Source/API/FIRDocumentReference.mm
+++ b/Firestore/Source/API/FIRDocumentReference.mm
@@ -125,6 +125,11 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)setData:(NSDictionary<NSString *, id> *)documentData
+    mergeFields:(NSArray<id> *)mergeFields {
+  return [self setData:documentData mergeFields:mergeFields completion:nil];
+}
+
+- (void)setData:(NSDictionary<NSString *, id> *)documentData
      completion:(nullable void (^)(NSError *_Nullable error))completion {
   return [self setData:documentData merge:NO completion:completion];
 }
@@ -132,8 +137,19 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)setData:(NSDictionary<NSString *, id> *)documentData
           merge:(BOOL)merge
      completion:(nullable void (^)(NSError *_Nullable error))completion {
-  FSTParsedSetData *parsed = merge ? [self.firestore.dataConverter parsedMergeData:documentData]
-                                   : [self.firestore.dataConverter parsedSetData:documentData];
+  FSTParsedSetData *parsed =
+      merge ? [self.firestore.dataConverter parsedMergeData:documentData fieldMask:nil]
+            : [self.firestore.dataConverter parsedSetData:documentData];
+  return [self.firestore.client
+      writeMutations:[parsed mutationsWithKey:self.key precondition:Precondition::None()]
+          completion:completion];
+}
+
+- (void)setData:(NSDictionary<NSString *, id> *)documentData
+    mergeFields:(NSArray<id> *)mergeFields
+     completion:(nullable void (^)(NSError *_Nullable error))completion {
+  FSTParsedSetData *parsed =
+      [self.firestore.dataConverter parsedMergeData:documentData fieldMask:mergeFields];
   return [self.firestore.client
       writeMutations:[parsed mutationsWithKey:self.key precondition:Precondition::None()]
           completion:completion];
diff --git a/Firestore/Source/API/FIRTransaction.mm b/Firestore/Source/API/FIRTransaction.mm
index 668a359..b5bdefa 100644
--- a/Firestore/Source/API/FIRTransaction.mm
+++ b/Firestore/Source/API/FIRTransaction.mm
@@ -68,8 +68,19 @@ NS_ASSUME_NONNULL_BEGIN
                 forDocument:(FIRDocumentReference *)document
                       merge:(BOOL)merge {
   [self validateReference:document];
-  FSTParsedSetData *parsed = merge ? [self.firestore.dataConverter parsedMergeData:data]
-                                   : [self.firestore.dataConverter parsedSetData:data];
+  FSTParsedSetData *parsed = merge
+                                 ? [self.firestore.dataConverter parsedMergeData:data fieldMask:nil]
+                                 : [self.firestore.dataConverter parsedSetData:data];
+  [self.internalTransaction setData:parsed forDocument:document.key];
+  return self;
+}
+
+- (FIRTransaction *)setData:(NSDictionary<NSString *, id> *)data
+                forDocument:(FIRDocumentReference *)document
+                mergeFields:(NSArray<id> *)mergeFields {
+  [self validateReference:document];
+  FSTParsedSetData *parsed =
+      [self.firestore.dataConverter parsedMergeData:data fieldMask:mergeFields];
   [self.internalTransaction setData:parsed forDocument:document.key];
   return self;
 }
diff --git a/Firestore/Source/API/FIRWriteBatch.mm b/Firestore/Source/API/FIRWriteBatch.mm
index 1185dae..366c708 100644
--- a/Firestore/Source/API/FIRWriteBatch.mm
+++ b/Firestore/Source/API/FIRWriteBatch.mm
@@ -70,8 +70,21 @@ NS_ASSUME_NONNULL_BEGIN
                      merge:(BOOL)merge {
   [self verifyNotCommitted];
   [self validateReference:document];
-  FSTParsedSetData *parsed = merge ? [self.firestore.dataConverter parsedMergeData:data]
-                                   : [self.firestore.dataConverter parsedSetData:data];
+  FSTParsedSetData *parsed = merge
+                                 ? [self.firestore.dataConverter parsedMergeData:data fieldMask:nil]
+                                 : [self.firestore.dataConverter parsedSetData:data];
+  [self.mutations
+      addObjectsFromArray:[parsed mutationsWithKey:document.key precondition:Precondition::None()]];
+  return self;
+}
+
+- (FIRWriteBatch *)setData:(NSDictionary<NSString *, id> *)data
+               forDocument:(FIRDocumentReference *)document
+               mergeFields:(NSArray<id> *)mergeFields {
+  [self verifyNotCommitted];
+  [self validateReference:document];
+  FSTParsedSetData *parsed =
+      [self.firestore.dataConverter parsedMergeData:data fieldMask:mergeFields];
   [self.mutations
       addObjectsFromArray:[parsed mutationsWithKey:document.key precondition:Precondition::None()]];
   return self;
diff --git a/Firestore/Source/API/FSTUserDataConverter.h b/Firestore/Source/API/FSTUserDataConverter.h
index a3d8a2d..27a5f09 100644
--- a/Firestore/Source/API/FSTUserDataConverter.h
+++ b/Firestore/Source/API/FSTUserDataConverter.h
@@ -129,7 +129,7 @@ typedef id _Nullable (^FSTPreConverterBlock)(id _Nullable);
 - (FSTParsedSetData *)parsedSetData:(id)input;
 
 /** Parse document data from a setData call with `merge:YES`. */
-- (FSTParsedSetData *)parsedMergeData:(id)input;
+- (FSTParsedSetData *)parsedMergeData:(id)input fieldMask:(nullable NSArray<id> *)fieldMask;
 
 /** Parse update data from an updateData call. */
 - (FSTParsedUpdateData *)parsedUpdateData:(id)input;
diff --git a/Firestore/Source/API/FSTUserDataConverter.mm b/Firestore/Source/API/FSTUserDataConverter.mm
index 2794398..6d01c75 100644
--- a/Firestore/Source/API/FSTUserDataConverter.mm
+++ b/Firestore/Source/API/FSTUserDataConverter.mm
@@ -412,7 +412,7 @@ typedef NS_ENUM(NSInteger, FSTUserDataSource) {
   return self;
 }
 
-- (FSTParsedSetData *)parsedMergeData:(id)input {
+- (FSTParsedSetData *)parsedMergeData:(id)input fieldMask:(nullable NSArray<id> *)fieldMask {
   // NOTE: The public API is typed as NSDictionary but we type 'input' as 'id' since we can't trust
   // Obj-C to verify the type for us.
   if (![input isKindOfClass:[NSDictionary class]]) {
@@ -424,9 +424,45 @@ typedef NS_ENUM(NSInteger, FSTUserDataSource) {
                                     path:absl::make_unique<FieldPath>(FieldPath::EmptyPath())];
   FSTObjectValue *updateData = (FSTObjectValue *)[self parseData:input context:context];
 
+  FieldMask convertedFieldMask;
+  std::vector<FieldTransform> convertedFieldTransform;
+
+  if (fieldMask) {
+    __block std::vector<FieldPath> fieldMaskPaths{};
+    [fieldMask enumerateObjectsUsingBlock:^(id fieldPath, NSUInteger idx, BOOL *stop) {
+      FieldPath path{};
+
+      if ([fieldPath isKindOfClass:[NSString class]]) {
+        path = [FIRFieldPath pathWithDotSeparatedString:fieldPath].internalValue;
+      } else if ([fieldPath isKindOfClass:[FIRFieldPath class]]) {
+        path = ((FIRFieldPath *)fieldPath).internalValue;
+      } else {
+        FSTThrowInvalidArgument(
+            @"All elements in mergeFields: must be NSStrings or FIRFieldPaths.");
+      }
+
+      if ([updateData valueForPath:path] == nil) {
+        FSTThrowInvalidArgument(
+            @"Field '%s' is specified in your field mask but missing from your input data.",
+            path.CanonicalString().c_str());
+      }
+
+      fieldMaskPaths.push_back(path);
+    }];
+    convertedFieldMask = FieldMask(fieldMaskPaths);
+    std::copy_if(context.fieldTransforms->begin(), context.fieldTransforms->end(),
+                 std::back_inserter(convertedFieldTransform),
+                 [&](const FieldTransform &fieldTransform) {
+                   return convertedFieldMask.covers(fieldTransform.path());
+                 });
+  } else {
+    convertedFieldMask = FieldMask{*context.fieldMask};
+    convertedFieldTransform = *context.fieldTransforms;
+  }
+
   return [[FSTParsedSetData alloc] initWithData:updateData
-                                      fieldMask:FieldMask{*context.fieldMask}
-                                fieldTransforms:*context.fieldTransforms];
+                                      fieldMask:convertedFieldMask
+                                fieldTransforms:convertedFieldTransform];
 }
 
 - (FSTParsedSetData *)parsedSetData:(id)input {
diff --git a/Firestore/Source/Public/FIRDocumentReference.h b/Firestore/Source/Public/FIRDocumentReference.h
index 4aa8c45..7baa30a 100644
--- a/Firestore/Source/Public/FIRDocumentReference.h
+++ b/Firestore/Source/Public/FIRDocumentReference.h
@@ -92,6 +92,23 @@ NS_SWIFT_NAME(DocumentReference)
 - (void)setData:(NSDictionary<NSString *, id> *)documentData merge:(BOOL)merge;
 
 /**
+ * Writes to the document referred to by `document` and only replace the fields
+ * specified under `mergeFields`. Any field that is not specified in `mergeFields`
+ * is ignored and remains untouched. If the document doesn't yet exist,
+ * this method creates it and then sets the data.
+ *
+ * It is an error to include a field in `mergeFields` that does not have a corresponding
+ * value in the `data` dictionary.
+ *
+ * @param documentData An `NSDictionary` containing the fields that make up the document
+ * to be written.
+ * @param mergeFields An `NSArray` that contains a list of `NSString` or `FIRFieldPath` elements
+ *     specifying which fields to merge. Fields can contain dots to reference nested fields within
+ *     the document.
+ */
+- (void)setData:(NSDictionary<NSString *, id> *)documentData mergeFields:(NSArray<id> *)mergeFields;
+
+/**
  * Overwrites the document referred to by this `FIRDocumentReference`. If no document exists, it
  * is created. If a document already exists, it is overwritten.
  *
@@ -121,6 +138,28 @@ NS_SWIFT_NAME(DocumentReference)
      completion:(nullable void (^)(NSError *_Nullable error))completion;
 
 /**
+ * Writes to the document referred to by `document` and only replace the fields
+ * specified under `mergeFields`. Any field that is not specified in `mergeFields`
+ * is ignored and remains untouched. If the document doesn't yet exist,
+ * this method creates it and then sets the data.
+ *
+ * It is an error to include a field in `mergeFields` that does not have a corresponding
+ * value in the `data` dictionary.
+ *
+ * @param documentData An `NSDictionary` containing the fields that make up the document
+ * to be written.
+ * @param mergeFields An `NSArray` that contains a list of `NSString` or `FIRFieldPath` elements
+ *     specifying which fields to merge. Fields can contain dots to reference nested fields within
+ *     the document.
+ * @param completion A block to execute once the document has been successfully written to the
+ *     server. This block will not be called while the client is offline, though local
+ *     changes will be visible immediately.
+ */
+- (void)setData:(NSDictionary<NSString *, id> *)documentData
+    mergeFields:(NSArray<id> *)mergeFields
+     completion:(nullable void (^)(NSError *_Nullable error))completion;
+
+/**
  * Updates fields in the document referred to by this `FIRDocumentReference`.
  * If the document does not exist, the update fails (specify a completion block to be notified).
  *
diff --git a/Firestore/Source/Public/FIRTransaction.h b/Firestore/Source/Public/FIRTransaction.h
index 2fa4430..e53414d 100644
--- a/Firestore/Source/Public/FIRTransaction.h
+++ b/Firestore/Source/Public/FIRTransaction.h
@@ -65,6 +65,30 @@ NS_SWIFT_NAME(Transaction)
 // clang-format on
 
 /**
+ * Writes to the document referred to by `document` and only replace the fields
+ * specified under `mergeFields`. Any field that is not specified in `mergeFields`
+ * is ignored and remains untouched. If the document doesn't yet exist,
+ * this method creates it and then sets the data.
+ *
+ * It is an error to include a field in `mergeFields` that does not have a corresponding
+ * value in the `data` dictionary.
+ *
+ * @param data An `NSDictionary` containing the fields that make up the document
+ * to be written.
+ * @param document A reference to the document whose data should be overwritten.
+ * @param mergeFields An `NSArray` that contains a list of `NSString` or `FIRFieldPath` elements
+ *     specifying which fields to merge. Fields can contain dots to reference nested fields within
+ *     the document.
+ * @return This `FIRTransaction` instance. Used for chaining method calls.
+ */
+// clang-format off
+- (FIRTransaction *)setData:(NSDictionary<NSString *, id> *)data
+                forDocument:(FIRDocumentReference *)document
+                mergeFields:(NSArray<id> *)mergeFields
+    NS_SWIFT_NAME(setData(_:forDocument:mergeFields:));
+// clang-format on
+
+/**
  * Updates fields in the document referred to by `document`.
  * If the document does not exist, the transaction will fail.
  *
diff --git a/Firestore/Source/Public/FIRWriteBatch.h b/Firestore/Source/Public/FIRWriteBatch.h
index 1568723..22d1b16 100644
--- a/Firestore/Source/Public/FIRWriteBatch.h
+++ b/Firestore/Source/Public/FIRWriteBatch.h
@@ -68,6 +68,29 @@ NS_SWIFT_NAME(WriteBatch)
 // clang-format on
 
 /**
+ * Writes to the document referred to by `document` and only replace the fields
+ * specified under `mergeFields`. Any field that is not specified in `mergeFields`
+ * is ignored and remains untouched. If the document doesn't yet exist,
+ * this method creates it and then sets the data.
+ *
+ * It is an error to include a field in `mergeFields` that does not have a corresponding
+ * value in the `data` dictionary.
+ *
+ * @param data An `NSDictionary` that contains the fields and data to write to the document.
+ * @param document A reference to the document whose data should be overwritten.
+ * @param mergeFields An `NSArray` that contains a list of `NSString` or `FIRFieldPath` elements
+ *     specifying which fields to merge. Fields can contain dots to reference nested fields within
+ *     the document.
+ * @return This `FIRWriteBatch` instance. Used for chaining method calls.
+ */
+// clang-format off
+- (FIRWriteBatch *)setData:(NSDictionary<NSString *, id> *)data
+               forDocument:(FIRDocumentReference *)document
+               mergeFields:(NSArray<id> *)mergeFields
+   NS_SWIFT_NAME(setData(_:forDocument:mergeFields:));
+// clang-format on
+
+/**
  * Updates fields in the document referred to by `document`.
  * If document does not exist, the write batch will fail.
  *