Implement local and server application of arrayUnion and arrayRemove transforms. (#1101)

3 files changed, 273 insertions, 22 deletions

diff --git a/Firestore/Example/Tests/Model/FSTMutationTests.mm b/Firestore/Example/Tests/Model/FSTMutationTests.mm
index 56bf1c2..0bb7518 100644
--- a/Firestore/Example/Tests/Model/FSTMutationTests.mm
+++ b/Firestore/Example/Tests/Model/FSTMutationTests.mm
@@ -185,7 +185,126 @@ using firebase::firestore::model::TransformOperation;
   }
 }
 
-- (void)testAppliesServerAckedTransformsToDocuments {
+- (void)testAppliesLocalArrayUnionTransformToMissingField {
+  auto baseDoc = @{};
+  auto transform = @{ @"missing" : [FIRFieldValue fieldValueForArrayUnion:@[ @1, @2 ]] };
+  auto expected = @{ @"missing" : @[ @1, @2 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayUnionTransformToNonArrayField {
+  auto baseDoc = @{ @"non-array" : @42 };
+  auto transform = @{ @"non-array" : [FIRFieldValue fieldValueForArrayUnion:@[ @1, @2 ]] };
+  auto expected = @{ @"non-array" : @[ @1, @2 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayUnionTransformWithNonExistingElements {
+  auto baseDoc = @{ @"array" : @[ @1, @3 ] };
+  auto transform = @{ @"array" : [FIRFieldValue fieldValueForArrayUnion:@[ @2, @4 ]] };
+  auto expected = @{ @"array" : @[ @1, @3, @2, @4 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayUnionTransformWithExistingElements {
+  auto baseDoc = @{ @"array" : @[ @1, @3 ] };
+  auto transform = @{ @"array" : [FIRFieldValue fieldValueForArrayUnion:@[ @1, @3 ]] };
+  auto expected = @{ @"array" : @[ @1, @3 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayUnionTransformWithDuplicateExistingElements {
+  // Duplicate entries in your existing array should be preserved.
+  auto baseDoc = @{ @"array" : @[ @1, @2, @2, @3 ] };
+  auto transform = @{ @"array" : [FIRFieldValue fieldValueForArrayUnion:@[ @2 ]] };
+  auto expected = @{ @"array" : @[ @1, @2, @2, @3 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayUnionTransformWithDuplicateUnionElements {
+  // Duplicate entries in your union array should only be added once.
+  auto baseDoc = @{ @"array" : @[ @1, @3 ] };
+  auto transform = @{ @"array" : [FIRFieldValue fieldValueForArrayUnion:@[ @2, @2 ]] };
+  auto expected = @{ @"array" : @[ @1, @3, @2 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayUnionTransformWithNonPrimitiveElements {
+  // Union nested object values (one existing, one not).
+  auto baseDoc = @{ @"array" : @[ @1, @{@"a" : @"b"} ] };
+  auto transform =
+      @{ @"array" : [FIRFieldValue fieldValueForArrayUnion:@[ @{@"a" : @"b"}, @{@"c" : @"d"} ]] };
+  auto expected = @{ @"array" : @[ @1, @{@"a" : @"b"}, @{@"c" : @"d"} ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayUnionTransformWithPartiallyOverlappingElements {
+  // Union objects that partially overlap an existing object.
+  auto baseDoc = @{ @"array" : @[ @1, @{@"a" : @"b", @"c" : @"d"} ] };
+  auto transform =
+      @{ @"array" : [FIRFieldValue fieldValueForArrayUnion:@[ @{@"a" : @"b"}, @{@"c" : @"d"} ]] };
+  auto expected =
+      @{ @"array" : @[ @1, @{@"a" : @"b", @"c" : @"d"}, @{@"a" : @"b"}, @{@"c" : @"d"} ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayRemoveTransformToMissingField {
+  auto baseDoc = @{};
+  auto transform = @{ @"missing" : [FIRFieldValue fieldValueForArrayRemove:@[ @1, @2 ]] };
+  auto expected = @{ @"missing" : @[] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayRemoveTransformToNonArrayField {
+  auto baseDoc = @{ @"non-array" : @42 };
+  auto transform = @{ @"non-array" : [FIRFieldValue fieldValueForArrayRemove:@[ @1, @2 ]] };
+  auto expected = @{ @"non-array" : @[] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayRemoveTransformWithNonExistingElements {
+  auto baseDoc = @{ @"array" : @[ @1, @3 ] };
+  auto transform = @{ @"array" : [FIRFieldValue fieldValueForArrayRemove:@[ @2, @4 ]] };
+  auto expected = @{ @"array" : @[ @1, @3 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayRemoveTransformWithExistingElements {
+  auto baseDoc = @{ @"array" : @[ @1, @2, @3, @4 ] };
+  auto transform = @{ @"array" : [FIRFieldValue fieldValueForArrayRemove:@[ @1, @3 ]] };
+  auto expected = @{ @"array" : @[ @2, @4 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+- (void)testAppliesLocalArrayRemoveTransformWithNonPrimitiveElements {
+  // Remove nested object values (one existing, one not).
+  auto baseDoc = @{ @"array" : @[ @1, @{@"a" : @"b"} ] };
+  auto transform =
+      @{ @"array" : [FIRFieldValue fieldValueForArrayRemove:@[ @{@"a" : @"b"}, @{@"c" : @"d"} ]] };
+  auto expected = @{ @"array" : @[ @1 ] };
+  [self transformBaseDoc:baseDoc with:transform expecting:expected];
+}
+
+// Helper to test a particular transform scenario.
+- (void)transformBaseDoc:(NSDictionary<NSString *, id> *)baseData
+                    with:(NSDictionary<NSString *, id> *)transformData
+               expecting:(NSDictionary<NSString *, id> *)expectedData {
+  FSTDocument *baseDoc = FSTTestDoc("collection/key", 0, baseData, NO);
+
+  FSTMutation *transform = FSTTestTransformMutation(@"collection/key", transformData);
+
+  FSTMaybeDocument *transformedDoc =
+      [transform applyTo:baseDoc baseDocument:baseDoc localWriteTime:_timestamp];
+
+  FSTDocument *expectedDoc = [FSTDocument documentWithData:FSTTestObjectValue(expectedData)
+                                                       key:FSTTestDocKey(@"collection/key")
+                                                   version:FSTTestVersion(0)
+                                         hasLocalMutations:YES];
+
+  XCTAssertEqualObjects(transformedDoc, expectedDoc);
+}
+
+- (void)testAppliesServerAckedServerTimestampTransformToDocuments {
   NSDictionary *docData = @{ @"foo" : @{@"bar" : @"bar-value"}, @"baz" : @"baz-value" };
   FSTDocument *baseDoc = FSTTestDoc("collection/key", 0, docData, NO);
 
@@ -207,6 +326,29 @@ using firebase::firestore::model::TransformOperation;
   XCTAssertEqualObjects(transformedDoc, FSTTestDoc("collection/key", 0, expectedData, NO));
 }
 
+- (void)testAppliesServerAckedArrayTransformsToDocuments {
+  NSDictionary *docData = @{ @"array_1" : @[ @1, @2 ], @"array_2" : @[ @"a", @"b" ] };
+  FSTDocument *baseDoc = FSTTestDoc("collection/key", 0, docData, NO);
+
+  FSTMutation *transform = FSTTestTransformMutation(@"collection/key", @{
+    @"array_1" : [FIRFieldValue fieldValueForArrayUnion:@[ @2, @3 ]],
+    @"array_2" : [FIRFieldValue fieldValueForArrayRemove:@[ @"a", @"c" ]]
+  });
+
+  // Server just sends null transform results for array operations.
+  FSTMutationResult *mutationResult = [[FSTMutationResult alloc]
+       initWithVersion:FSTTestVersion(1)
+      transformResults:@[ [FSTNullValue nullValue], [FSTNullValue nullValue] ]];
+
+  FSTMaybeDocument *transformedDoc = [transform applyTo:baseDoc
+                                           baseDocument:baseDoc
+                                         localWriteTime:_timestamp
+                                         mutationResult:mutationResult];
+
+  NSDictionary *expectedData = @{ @"array_1" : @[ @1, @2, @3 ], @"array_2" : @[ @"b" ] };
+  XCTAssertEqualObjects(transformedDoc, FSTTestDoc("collection/key", 0, expectedData, NO));
+}
+
 - (void)testDeleteDeletes {
   NSDictionary *docData = @{@"foo" : @"bar"};
   FSTDocument *baseDoc = FSTTestDoc("collection/key", 0, docData, NO);
diff --git a/Firestore/Source/Model/FSTMutation.mm b/Firestore/Source/Model/FSTMutation.mm
index 99d2e51..47e34da 100644
--- a/Firestore/Source/Model/FSTMutation.mm
+++ b/Firestore/Source/Model/FSTMutation.mm
@@ -36,6 +36,7 @@
 #include "Firestore/core/src/firebase/firestore/model/precondition.h"
 #include "Firestore/core/src/firebase/firestore/model/transform_operations.h"
 
+using firebase::firestore::model::ArrayTransform;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::FieldMask;
 using firebase::firestore::model::FieldPath;
@@ -50,8 +51,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 @implementation FSTMutationResult
 
-- (instancetype)initWithVersion:(FSTSnapshotVersion *_Nullable)version
-               transformResults:(NSArray<FSTFieldValue *> *_Nullable)transformResults {
+- (instancetype)initWithVersion:(nullable FSTSnapshotVersion *)version
+               transformResults:(nullable NSArray<FSTFieldValue *> *)transformResults {
   if (self = [super init]) {
     _version = version;
     _transformResults = transformResults;
@@ -345,13 +346,18 @@ NS_ASSUME_NONNULL_BEGIN
             [maybeDoc class]);
   FSTDocument *doc = (FSTDocument *)maybeDoc;
 
-  FSTAssert([doc.key isEqual:self.key], @"Can only patch a document with the same key");
+  FSTAssert([doc.key isEqual:self.key], @"Can only transform a document with the same key");
 
   BOOL hasLocalMutations = (mutationResult == nil);
-  NSArray<FSTFieldValue *> *transformResults =
-      mutationResult
-          ? mutationResult.transformResults
-          : [self localTransformResultsWithBaseDocument:baseDoc writeTime:localWriteTime];
+  NSArray<FSTFieldValue *> *transformResults;
+  if (mutationResult) {
+    transformResults =
+        [self serverTransformResultsWithBaseDocument:baseDoc
+                              serverTransformResults:mutationResult.transformResults];
+  } else {
+    transformResults =
+        [self localTransformResultsWithBaseDocument:baseDoc writeTime:localWriteTime];
+  }
   FSTObjectValue *newData = [self transformObject:doc.data transformResults:transformResults];
   return [FSTDocument documentWithData:newData
                                    key:doc.key
@@ -361,6 +367,53 @@ NS_ASSUME_NONNULL_BEGIN
 
 /**
  * Creates an array of "transform results" (a transform result is a field value representing the
+ * result of applying a transform) for use after a FSTTransformMutation has been acknowledged by
+ * the server.
+ *
+ * @param baseDocument The document prior to applying this mutation batch.
+ * @param serverTransformResults The transform results received by the server.
+ * @return The transform results array.
+ */
+- (NSArray<FSTFieldValue *> *)
+serverTransformResultsWithBaseDocument:(nullable FSTMaybeDocument *)baseDocument
+                serverTransformResults:(NSArray<FSTFieldValue *> *)serverTransformResults {
+  NSMutableArray<FSTFieldValue *> *transformResults = [NSMutableArray array];
+  FSTAssert(self.fieldTransforms.size() == serverTransformResults.count,
+            @"server transform result count (%ld) should match field transforms count (%ld)",
+            serverTransformResults.count, self.fieldTransforms.size());
+
+  for (NSUInteger i = 0; i < serverTransformResults.count; i++) {
+    const FieldTransform &fieldTransform = self.fieldTransforms[i];
+    FSTFieldValue *previousValue = nil;
+    if ([baseDocument isMemberOfClass:[FSTDocument class]]) {
+      previousValue = [((FSTDocument *)baseDocument) fieldForPath:fieldTransform.path()];
+    }
+
+    FSTFieldValue *transformResult;
+    // The server just sends null as the transform result for array union / remove operations, so
+    // we have to calculate a result the same as we do for local applications.
+    if (fieldTransform.transformation().type() == TransformOperation::Type::ArrayUnion) {
+      transformResult = [self
+          arrayUnionResultWithElements:ArrayTransform::Elements(fieldTransform.transformation())
+                         previousValue:previousValue];
+
+    } else if (fieldTransform.transformation().type() == TransformOperation::Type::ArrayRemove) {
+      transformResult = [self
+          arrayRemoveResultWithElements:ArrayTransform::Elements(fieldTransform.transformation())
+                          previousValue:previousValue];
+
+    } else {
+      // Just use the server-supplied result.
+      transformResult = serverTransformResults[i];
+    }
+
+    [transformResults addObject:transformResult];
+  }
+  return transformResults;
+}
+
+/**
+ * Creates an array of "transform results" (a transform result is a field value representing the
  * result of applying a transform) for use when applying an FSTTransformMutation locally.
  *
  * @param baseDocument The document prior to applying this mutation batch.
@@ -369,27 +422,81 @@ NS_ASSUME_NONNULL_BEGIN
  * @return The transform results array.
  */
 - (NSArray<FSTFieldValue *> *)localTransformResultsWithBaseDocument:
-                                  (FSTMaybeDocument *_Nullable)baseDocument
+                                  (nullable FSTMaybeDocument *)baseDocument
                                                           writeTime:(FIRTimestamp *)localWriteTime {
   NSMutableArray<FSTFieldValue *> *transformResults = [NSMutableArray array];
   for (const FieldTransform &fieldTransform : self.fieldTransforms) {
+    FSTFieldValue *previousValue = nil;
+    if ([baseDocument isMemberOfClass:[FSTDocument class]]) {
+      previousValue = [((FSTDocument *)baseDocument) fieldForPath:fieldTransform.path()];
+    }
+
+    FSTFieldValue *transformResult;
     if (fieldTransform.transformation().type() == TransformOperation::Type::ServerTimestamp) {
-      FSTFieldValue *previousValue = nil;
+      transformResult =
+          [FSTServerTimestampValue serverTimestampValueWithLocalWriteTime:localWriteTime
+                                                            previousValue:previousValue];
+
+    } else if (fieldTransform.transformation().type() == TransformOperation::Type::ArrayUnion) {
+      transformResult = [self
+          arrayUnionResultWithElements:ArrayTransform::Elements(fieldTransform.transformation())
+                         previousValue:previousValue];
 
-      if ([baseDocument isMemberOfClass:[FSTDocument class]]) {
-        previousValue = [((FSTDocument *)baseDocument) fieldForPath:fieldTransform.path()];
-      }
+    } else if (fieldTransform.transformation().type() == TransformOperation::Type::ArrayRemove) {
+      transformResult = [self
+          arrayRemoveResultWithElements:ArrayTransform::Elements(fieldTransform.transformation())
+                          previousValue:previousValue];
 
-      [transformResults
-          addObject:[FSTServerTimestampValue serverTimestampValueWithLocalWriteTime:localWriteTime
-                                                                      previousValue:previousValue]];
     } else {
       FSTFail(@"Encountered unknown transform: %d type", fieldTransform.transformation().type());
     }
+
+    [transformResults addObject:transformResult];
   }
   return transformResults;
 }
 
+/**
+ * Transforms the provided `previousValue` via the provided `elements`. Used both for local
+ * application and after server acknowledgement.
+ */
+- (FSTFieldValue *)arrayUnionResultWithElements:(const std::vector<FSTFieldValue *> &)elements
+                                  previousValue:(FSTFieldValue *)previousValue {
+  NSMutableArray<FSTFieldValue *> *result = [self coercedFieldValuesArray:previousValue];
+  for (FSTFieldValue *element : elements) {
+    if (![result containsObject:element]) {
+      [result addObject:element];
+    }
+  }
+  return [[FSTArrayValue alloc] initWithValueNoCopy:result];
+}
+
+/**
+ * Transforms the provided `previousValue` via the provided `elements`. Used both for local
+ * application and after server acknowledgement.
+ */
+- (FSTFieldValue *)arrayRemoveResultWithElements:(const std::vector<FSTFieldValue *> &)elements
+                                   previousValue:(FSTFieldValue *)previousValue {
+  NSMutableArray<FSTFieldValue *> *result = [self coercedFieldValuesArray:previousValue];
+  for (FSTFieldValue *element : elements) {
+    [result removeObject:element];
+  }
+  return [[FSTArrayValue alloc] initWithValueNoCopy:result];
+}
+
+/**
+ * Inspects the provided value, returning a mutable copy of the internal array if it's an
+ * FSTArrayValue and an empty mutable array if it's nil or any other type of FSTFieldValue.
+ */
+- (NSMutableArray<FSTFieldValue *> *)coercedFieldValuesArray:(nullable FSTFieldValue *)value {
+  if ([value isMemberOfClass:[FSTArrayValue class]]) {
+    return [NSMutableArray arrayWithArray:((FSTArrayValue *)value).internalValue];
+  } else {
+    // coerce to empty array.
+    return [NSMutableArray array];
+  }
+}
+
 - (FSTObjectValue *)transformObject:(FSTObjectValue *)objectValue
                    transformResults:(NSArray<FSTFieldValue *> *)transformResults {
   FSTAssert(transformResults.count == self.fieldTransforms.size(),
@@ -397,13 +504,8 @@ NS_ASSUME_NONNULL_BEGIN
 
   for (size_t i = 0; i < self.fieldTransforms.size(); i++) {
     const FieldTransform &fieldTransform = self.fieldTransforms[i];
-    const TransformOperation &transform = fieldTransform.transformation();
     const FieldPath &fieldPath = fieldTransform.path();
-    if (transform.type() == TransformOperation::Type::ServerTimestamp) {
-      objectValue = [objectValue objectBySettingValue:transformResults[i] forPath:fieldPath];
-    } else {
-      FSTFail(@"Encountered unknown transform: %d type", transform.type());
-    }
+    objectValue = [objectValue objectBySettingValue:transformResults[i] forPath:fieldPath];
   }
   return objectValue;
 }
diff --git a/Firestore/core/src/firebase/firestore/model/transform_operations.h b/Firestore/core/src/firebase/firestore/model/transform_operations.h
index aad5a9b..2943ea0 100644
--- a/Firestore/core/src/firebase/firestore/model/transform_operations.h
+++ b/Firestore/core/src/firebase/firestore/model/transform_operations.h
@@ -151,6 +151,13 @@ class ArrayTransform : public TransformOperation {
   }
 #endif  // defined(__OBJC__)
 
+  static const std::vector<FSTFieldValue*>& Elements(
+      const TransformOperation& op) {
+    FIREBASE_ASSERT(op.type() == Type::ArrayUnion ||
+                    op.type() == Type::ArrayRemove);
+    return static_cast<const ArrayTransform&>(op).elements();
+  }
+
  private:
   Type type_;
   std::vector<FSTFieldValue*> elements_;