1 files changed, 7 insertions, 81 deletions

diff --git a/Firestore/Source/Model/FSTMutation.mm b/Firestore/Source/Model/FSTMutation.mm
index fdf6014..ee6ef9c 100644
--- a/Firestore/Source/Model/FSTMutation.mm
+++ b/Firestore/Source/Model/FSTMutation.mm
@@ -392,30 +392,15 @@ serverTransformResultsWithBaseDocument:(nullable FSTMaybeDocument *)baseDocument
 
   for (NSUInteger i = 0; i < serverTransformResults.count; i++) {
     const FieldTransform &fieldTransform = self.fieldTransforms[i];
+    const TransformOperation &transform = fieldTransform.transformation();
+
     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];
+    [transformResults
+        addObject:transform.applyToRemoteDocument(previousValue, serverTransformResults[i])];
   }
   return transformResults;
 }
@@ -434,77 +419,18 @@ serverTransformResultsWithBaseDocument:(nullable FSTMaybeDocument *)baseDocument
                                                           writeTime:(FIRTimestamp *)localWriteTime {
   NSMutableArray<FSTFieldValue *> *transformResults = [NSMutableArray array];
   for (const FieldTransform &fieldTransform : self.fieldTransforms) {
+    const TransformOperation &transform = fieldTransform.transformation();
+
     FSTFieldValue *previousValue = nil;
     if ([baseDocument isMemberOfClass:[FSTDocument class]]) {
       previousValue = [((FSTDocument *)baseDocument) fieldForPath:fieldTransform.path()];
     }
 
-    FSTFieldValue *transformResult;
-    if (fieldTransform.transformation().type() == TransformOperation::Type::ServerTimestamp) {
-      transformResult =
-          [FSTServerTimestampValue serverTimestampValueWithLocalWriteTime:localWriteTime
-                                                            previousValue:previousValue];
-
-    } else 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 {
-      FSTFail(@"Encountered unknown transform: %d type", fieldTransform.transformation().type());
-    }
-
-    [transformResults addObject:transformResult];
+    [transformResults addObject:transform.applyToLocalView(previousValue, localWriteTime)];
   }
   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(),