Allowing field deletes with merge-enabled sets (#404)

7 files changed, 93 insertions, 33 deletions

diff --git a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
index 47ba95b..68692cc 100644
--- a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
@@ -176,6 +176,40 @@
   XCTAssertTrue([document[@"time"] isKindOfClass:[NSDate class]]);
 }
 
+- (void)testCanDeleteFieldUsingMerge {
+  FIRDocumentReference *doc = [[self.db collectionWithPath:@"rooms"] documentWithAutoID];
+
+  NSDictionary<NSString *, id> *initialData = @{
+    @"untouched" : @YES,
+    @"foo" : @"bar",
+    @"nested" : @{@"untouched" : @YES, @"foo" : @"bar"}
+  };
+  NSDictionary<NSString *, id> *mergeData = @{
+    @"foo" : [FIRFieldValue fieldValueForDelete],
+    @"nested" : @{@"foo" : [FIRFieldValue fieldValueForDelete]}
+  };
+
+  [self writeDocumentRef:doc data:initialData];
+
+  XCTestExpectation *completed =
+      [self expectationWithDescription:@"testCanMergeDataWithAnExistingDocumentUsingSet"];
+
+  [doc setData:mergeData
+         options:[FIRSetOptions merge]
+      completion:^(NSError *error) {
+        XCTAssertNil(error);
+        [completed fulfill];
+      }];
+
+  [self awaitExpectations];
+
+  FIRDocumentSnapshot *document = [self readDocumentForRef:doc];
+  XCTAssertEqual(document[@"untouched"], @YES);
+  XCTAssertNil(document[@"foo"]);
+  XCTAssertEqual(document[@"nested.untouched"], @YES);
+  XCTAssertNil(document[@"nested.foo"]);
+}
+
 - (void)testMergeReplacesArrays {
   FIRDocumentReference *doc = [[self.db collectionWithPath:@"rooms"] documentWithAutoID];
 
diff --git a/Firestore/Example/Tests/Integration/API/FIRValidationTests.m b/Firestore/Example/Tests/Integration/API/FIRValidationTests.m
index e4e7cd1..7fa4c3e 100644
--- a/Firestore/Example/Tests/Integration/API/FIRValidationTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRValidationTests.m
@@ -280,7 +280,9 @@
 
 - (void)testSetsWithFieldValueDeleteFail {
   [self expectSet:@{@"foo" : [FIRFieldValue fieldValueForDelete]}
-      toFailWithReason:@"FieldValue.delete() can only be used with updateData()."];
+      toFailWithReason:
+          @"FieldValue.delete() can only be used with updateData() and setData() with "
+          @"SetOptions.merge()."];
 }
 
 - (void)testUpdatesWithNestedFieldValueDeleteFail {
diff --git a/Firestore/Source/API/FIRDocumentReference.m b/Firestore/Source/API/FIRDocumentReference.m
index b750887..5fa6cf7 100644
--- a/Firestore/Source/API/FIRDocumentReference.m
+++ b/Firestore/Source/API/FIRDocumentReference.m
@@ -174,8 +174,9 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)setData:(NSDictionary<NSString *, id> *)documentData
         options:(FIRSetOptions *)options
      completion:(nullable void (^)(NSError *_Nullable error))completion {
-  FSTParsedSetData *parsed =
-      [self.firestore.dataConverter parsedSetData:documentData options:options];
+  FSTParsedSetData *parsed = options.isMerge
+                                 ? [self.firestore.dataConverter parsedMergeData:documentData]
+                                 : [self.firestore.dataConverter parsedSetData:documentData];
   return [self.firestore.client
       writeMutations:[parsed mutationsWithKey:self.key precondition:[FSTPrecondition none]]
           completion:completion];
diff --git a/Firestore/Source/API/FIRTransaction.m b/Firestore/Source/API/FIRTransaction.m
index 02006bb..69e7b12 100644
--- a/Firestore/Source/API/FIRTransaction.m
+++ b/Firestore/Source/API/FIRTransaction.m
@@ -69,7 +69,8 @@ NS_ASSUME_NONNULL_BEGIN
                 forDocument:(FIRDocumentReference *)document
                     options:(FIRSetOptions *)options {
   [self validateReference:document];
-  FSTParsedSetData *parsed = [self.firestore.dataConverter parsedSetData:data options:options];
+  FSTParsedSetData *parsed = options.isMerge ? [self.firestore.dataConverter parsedMergeData:data]
+                                             : [self.firestore.dataConverter parsedSetData:data];
   [self.internalTransaction setData:parsed forDocument:document.key];
   return self;
 }
diff --git a/Firestore/Source/API/FIRWriteBatch.m b/Firestore/Source/API/FIRWriteBatch.m
index 32b6ce8..20559b8 100644
--- a/Firestore/Source/API/FIRWriteBatch.m
+++ b/Firestore/Source/API/FIRWriteBatch.m
@@ -67,7 +67,8 @@ NS_ASSUME_NONNULL_BEGIN
                    options:(FIRSetOptions *)options {
   [self verifyNotCommitted];
   [self validateReference:document];
-  FSTParsedSetData *parsed = [self.firestore.dataConverter parsedSetData:data options:options];
+  FSTParsedSetData *parsed = options.isMerge ? [self.firestore.dataConverter parsedMergeData:data]
+                                             : [self.firestore.dataConverter parsedSetData:data];
   [self.mutations addObjectsFromArray:[parsed mutationsWithKey:document.key
                                                   precondition:[FSTPrecondition none]]];
   return self;
diff --git a/Firestore/Source/API/FSTUserDataConverter.h b/Firestore/Source/API/FSTUserDataConverter.h
index 69d1fa9..2c52ad8 100644
--- a/Firestore/Source/API/FSTUserDataConverter.h
+++ b/Firestore/Source/API/FSTUserDataConverter.h
@@ -110,10 +110,13 @@ typedef id _Nullable (^FSTPreConverterBlock)(id _Nullable);
 - (instancetype)initWithDatabaseID:(FSTDatabaseID *)databaseID
                       preConverter:(FSTPreConverterBlock)preConverter NS_DESIGNATED_INITIALIZER;
 
-/** Parse document data (e.g. from a setData call). */
-- (FSTParsedSetData *)parsedSetData:(id)input options:(FIRSetOptions *)options;
+/** Parse document data from a non-merge setData call.*/
+- (FSTParsedSetData *)parsedSetData:(id)input;
 
-/** Parse "update" data (i.e. from an updateData call). */
+/** Parse document data from a setData call with '[FIRSetOptions merge]'. */
+- (FSTParsedSetData *)parsedMergeData:(id)input;
+
+/** Parse update data from an updateData call. */
 - (FSTParsedUpdateData *)parsedUpdateData:(id)input;
 
 /** Parse a "query value" (e.g. value in a where filter or a value in a cursor bound). */
diff --git a/Firestore/Source/API/FSTUserDataConverter.m b/Firestore/Source/API/FSTUserDataConverter.m
index eb4dbda..aa79a51 100644
--- a/Firestore/Source/API/FSTUserDataConverter.m
+++ b/Firestore/Source/API/FSTUserDataConverter.m
@@ -109,6 +109,7 @@ static NSString *const RESERVED_FIELD_DESIGNATOR = @"__";
  */
 typedef NS_ENUM(NSInteger, FSTUserDataSource) {
   FSTUserDataSourceSet,
+  FSTUserDataSourceMergeSet,
   FSTUserDataSourceUpdate,
   FSTUserDataSourceQueryValue,  // from a where clause or cursor bound.
 };
@@ -158,6 +159,9 @@ typedef NS_ENUM(NSInteger, FSTUserDataSource) {
 - (instancetype)contextForField:(NSString *)fieldName;
 - (instancetype)contextForFieldPath:(FSTFieldPath *)fieldPath;
 - (instancetype)contextForArrayIndex:(NSUInteger)index;
+
+/** Returns true for the non-query parse contexts (Set, MergeSet and Update) */
+- (BOOL)isWrite;
 @end
 
 @implementation FSTParseContext
@@ -231,7 +235,16 @@ typedef NS_ENUM(NSInteger, FSTUserDataSource) {
 }
 
 - (BOOL)isWrite {
-  return _dataSource == FSTUserDataSourceSet || _dataSource == FSTUserDataSourceUpdate;
+  switch (self.dataSource) {
+    case FSTUserDataSourceSet:       // Falls through.
+    case FSTUserDataSourceMergeSet:  // Falls through.
+    case FSTUserDataSourceUpdate:
+      return YES;
+    case FSTUserDataSourceQueryValue:
+      return NO;
+    default:
+      FSTThrowInvalidArgument(@"Unexpected case for FSTUserDataSource: %d", self.dataSource);
+  }
 }
 
 - (void)validatePath {
@@ -288,7 +301,7 @@ typedef NS_ENUM(NSInteger, FSTUserDataSource) {
   return self;
 }
 
-- (FSTParsedSetData *)parsedSetData:(id)input options:(FIRSetOptions *)options {
+- (FSTParsedSetData *)parsedMergeData:(id)input {
   // 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]]) {
@@ -296,27 +309,29 @@ typedef NS_ENUM(NSInteger, FSTUserDataSource) {
   }
 
   FSTParseContext *context =
-      [FSTParseContext contextWithSource:FSTUserDataSourceSet path:[FSTFieldPath emptyPath]];
-
-  __block FSTObjectValue *updateData = [FSTObjectValue objectValue];
+      [FSTParseContext contextWithSource:FSTUserDataSourceMergeSet path:[FSTFieldPath emptyPath]];
+  FSTObjectValue *updateData = (FSTObjectValue *)[self parseData:input context:context];
 
-  [input enumerateKeysAndObjectsUsingBlock:^(NSString *key, id value, BOOL *stop) {
-    // Treat key as a complete field name (don't split on dots, etc.)
-    FSTFieldPath *path = [[FIRFieldPath alloc] initWithFields:@[ key ]].internalValue;
+  return
+      [[FSTParsedSetData alloc] initWithData:updateData
+                                   fieldMask:[[FSTFieldMask alloc] initWithFields:context.fieldMask]
+                             fieldTransforms:context.fieldTransforms];
+}
 
-    value = self.preConverter(value);
+- (FSTParsedSetData *)parsedSetData:(id)input {
+  // 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]]) {
+    FSTThrowInvalidArgument(@"Data to be written must be an NSDictionary.");
+  }
 
-    FSTFieldValue *_Nullable parsedValue =
-        [self parseData:value context:[context contextForFieldPath:path]];
-    if (parsedValue) {
-      updateData = [updateData objectBySettingValue:parsedValue forPath:path];
-    }
-  }];
+  FSTParseContext *context =
+      [FSTParseContext contextWithSource:FSTUserDataSourceSet path:[FSTFieldPath emptyPath]];
+  FSTObjectValue *updateData = (FSTObjectValue *)[self parseData:input context:context];
 
-  return [[FSTParsedSetData alloc]
-         initWithData:updateData
-            fieldMask:options.merge ? [[FSTFieldMask alloc] initWithFields:context.fieldMask] : nil
-      fieldTransforms:context.fieldTransforms];
+  return [[FSTParsedSetData alloc] initWithData:updateData
+                                      fieldMask:nil
+                                fieldTransforms:context.fieldTransforms];
 }
 
 - (FSTParsedUpdateData *)parsedUpdateData:(id)input {
@@ -536,20 +551,23 @@ typedef NS_ENUM(NSInteger, FSTUserDataSource) {
 
   } else if ([input isKindOfClass:[FIRFieldValue class]]) {
     if ([input isKindOfClass:[FSTDeleteFieldValue class]]) {
-      // We shouldn't encounter delete sentinels here. Provide a good error.
-      if (context.dataSource != FSTUserDataSourceUpdate) {
-        FSTThrowInvalidArgument(@"FieldValue.delete() can only be used with updateData().");
-      } else {
+      if (context.dataSource == FSTUserDataSourceMergeSet) {
+        return nil;
+      } else if (context.dataSource == FSTUserDataSourceUpdate) {
         FSTAssert(context.path.length > 0,
                   @"FieldValue.delete() at the top level should have already been handled.");
         FSTThrowInvalidArgument(
             @"FieldValue.delete() can only appear at the top level of your "
              "update data%@",
             [context fieldDescription]);
+      } else {
+        // We shouldn't encounter delete sentinels for queries or non-merge setData calls.
+        FSTThrowInvalidArgument(
+            @"FieldValue.delete() can only be used with updateData() and setData() with "
+            @"SetOptions.merge().");
       }
     } else if ([input isKindOfClass:[FSTServerTimestampFieldValue class]]) {
-      if (context.dataSource != FSTUserDataSourceSet &&
-          context.dataSource != FSTUserDataSourceUpdate) {
+      if (![context isWrite]) {
         FSTThrowInvalidArgument(
             @"FieldValue.serverTimestamp() can only be used with setData() and updateData().");
       }