Adding mergeFields support (#1141)

5 files changed, 86 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 {