Backport array contains / transform improvements from Web. (#1242)

8 files changed, 189 insertions, 181 deletions

diff --git a/Firestore/Example/Tests/Integration/API/FIRArrayTransformTests.mm b/Firestore/Example/Tests/Integration/API/FIRArrayTransformTests.mm
index 1c82461..0a5b413 100644
--- a/Firestore/Example/Tests/Integration/API/FIRArrayTransformTests.mm
+++ b/Firestore/Example/Tests/Integration/API/FIRArrayTransformTests.mm
@@ -36,7 +36,7 @@
   FIRDocumentReference *_docRef;
 
   // Accumulator used to capture events during the test.
-  FSTEventAccumulator *_accumulator;
+  FSTEventAccumulator<FIRDocumentSnapshot *> *_accumulator;
 
   // Listener registration for a listener maintained during the course of the test.
   id<FIRListenerRegistration> _listenerRegistration;
@@ -64,29 +64,11 @@
 
 #pragma mark - Test Helpers
 
-/** Waits for a snapshot with local writes. */
-- (FIRDocumentSnapshot *)waitForLocalEvent {
-  FIRDocumentSnapshot *snapshot;
-  do {
-    snapshot = [_accumulator awaitEventWithName:@"Local event."];
-  } while (!snapshot.metadata.hasPendingWrites);
-  return snapshot;
-}
-
-/** Waits for a snapshot that has no pending writes */
-- (FIRDocumentSnapshot *)waitForRemoteEvent {
-  FIRDocumentSnapshot *snapshot;
-  do {
-    snapshot = [_accumulator awaitEventWithName:@"Remote event."];
-  } while (snapshot.metadata.hasPendingWrites);
-  return snapshot;
-}
-
 /** Writes some initial data and consumes the events generated. */
 - (void)writeInitialData:(NSDictionary<NSString *, id> *)data {
   [self writeDocumentRef:_docRef data:data];
-  XCTAssertEqualObjects([self waitForLocalEvent].data, data);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, data);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, data);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, data);
 }
 
 #pragma mark - Test Cases
@@ -97,8 +79,8 @@
                       @"array" : [FIRFieldValue fieldValueForArrayUnion:@[ @1, @2 ]]
                     }];
   id expected = @{ @"array" : @[ @1, @2 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testAppendToArrayViaUpdate {
@@ -110,8 +92,8 @@
                      }];
 
   id expected = @{ @"array" : @[ @1, @3, @2, @4 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testAppendToArrayViaMergeSet {
@@ -123,8 +105,8 @@
                     }];
 
   id expected = @{ @"array" : @[ @1, @3, @2, @4 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testAppendObjectToArrayViaUpdate {
@@ -137,8 +119,8 @@
                      }];
 
   id expected = @{ @"array" : @[ @{@"a" : @"hi"}, @{@"a" : @"bye"} ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testRemoveFromArrayViaUpdate {
@@ -150,8 +132,8 @@
                      }];
 
   id expected = @{ @"array" : @[ @3, @3 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testRemoveFromArrayViaMergeSet {
@@ -163,8 +145,8 @@
                     }];
 
   id expected = @{ @"array" : @[ @3, @3 ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 - (void)testRemoveObjectFromArrayViaUpdate {
@@ -176,8 +158,8 @@
                      }];
 
   id expected = @{ @"array" : @[ @{@"a" : @"bye"} ] };
-  XCTAssertEqualObjects([self waitForLocalEvent].data, expected);
-  XCTAssertEqualObjects([self waitForRemoteEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitLocalEvent].data, expected);
+  XCTAssertEqualObjects([_accumulator awaitRemoteEvent].data, expected);
 }
 
 @end
diff --git a/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.mm b/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.mm
index e1ad3d2..6a755c6 100644
--- a/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.mm
+++ b/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.mm
@@ -90,24 +90,6 @@
   XCTAssertEqualObjects(initialDataSnap.data, _initialData);
 }
 
-/** Waits for a snapshot with local writes. */
-- (FIRDocumentSnapshot *)waitForLocalEvent {
-  FIRDocumentSnapshot *snapshot;
-  do {
-    snapshot = [_accumulator awaitEventWithName:@"Local event."];
-  } while (!snapshot.metadata.hasPendingWrites);
-  return snapshot;
-}
-
-/** Waits for a snapshot that has no pending writes */
-- (FIRDocumentSnapshot *)waitForRemoteEvent {
-  FIRDocumentSnapshot *snapshot;
-  do {
-    snapshot = [_accumulator awaitEventWithName:@"Remote event."];
-  } while (snapshot.metadata.hasPendingWrites);
-  return snapshot;
-}
-
 /** Verifies a snapshot containing _setData but with NSNull for the timestamps. */
 - (void)verifyTimestampsAreNullInSnapshot:(FIRDocumentSnapshot *)snapshot {
   XCTAssertEqualObjects(snapshot.data, [self expectedDataWithTimestamp:[NSNull null]]);
@@ -170,41 +152,42 @@
 
 - (void)testServerTimestampsWorkViaSet {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsAreNullInSnapshot:[self waitForLocalEvent]];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsAreNullInSnapshot:[_accumulator awaitLocalEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWorkViaUpdate {
   [self writeInitialData];
   [self updateDocumentRef:_docRef data:_updateData];
-  [self verifyTimestampsAreNullInSnapshot:[self waitForLocalEvent]];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsAreNullInSnapshot:[_accumulator awaitLocalEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWithEstimatedValue {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsAreEstimatedInSnapshot:[self waitForLocalEvent]];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsAreEstimatedInSnapshot:[_accumulator awaitLocalEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWithPreviousValue {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:nil];
-  FIRDocumentSnapshot *remoteSnapshot = [self waitForRemoteEvent];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent] fromPreviousSnapshot:nil];
+  FIRDocumentSnapshot *remoteSnapshot = [_accumulator awaitRemoteEvent];
 
   [_docRef updateData:_updateData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:remoteSnapshot];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent]
+              fromPreviousSnapshot:remoteSnapshot];
 
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWithPreviousValueOfDifferentType {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:nil];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent] fromPreviousSnapshot:nil];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  FIRDocumentSnapshot *localSnapshot = [self waitForLocalEvent];
+  FIRDocumentSnapshot *localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects([localSnapshot valueForField:@"a"], [NSNull null]);
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
@@ -213,7 +196,7 @@
       [[localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorEstimate]
           isKindOfClass:[FIRTimestamp class]]);
 
-  FIRDocumentSnapshot *remoteSnapshot = [self waitForRemoteEvent];
+  FIRDocumentSnapshot *remoteSnapshot = [_accumulator awaitRemoteEvent];
   XCTAssertTrue([[remoteSnapshot valueForField:@"a"] isKindOfClass:[FIRTimestamp class]]);
   XCTAssertTrue([
       [remoteSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious]
@@ -225,55 +208,55 @@
 
 - (void)testServerTimestampsWithConsecutiveUpdates {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:nil];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent] fromPreviousSnapshot:nil];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 
   [self disableNetwork];
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  FIRDocumentSnapshot *localSnapshot = [self waitForLocalEvent];
+  FIRDocumentSnapshot *localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
       @42);
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  localSnapshot = [self waitForLocalEvent];
+  localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
       @42);
 
   [self enableNetwork];
 
-  FIRDocumentSnapshot *remoteSnapshot = [self waitForRemoteEvent];
+  FIRDocumentSnapshot *remoteSnapshot = [_accumulator awaitRemoteEvent];
   XCTAssertTrue([[remoteSnapshot valueForField:@"a"] isKindOfClass:[FIRTimestamp class]]);
 }
 
 - (void)testServerTimestampsPreviousValueFromLocalMutation {
   [self writeDocumentRef:_docRef data:_setData];
-  [self verifyTimestampsInSnapshot:[self waitForLocalEvent] fromPreviousSnapshot:nil];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifyTimestampsInSnapshot:[_accumulator awaitLocalEvent] fromPreviousSnapshot:nil];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 
   [self disableNetwork];
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  FIRDocumentSnapshot *localSnapshot = [self waitForLocalEvent];
+  FIRDocumentSnapshot *localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
       @42);
 
   [_docRef updateData:@{ @"a" : @1337 }];
-  localSnapshot = [self waitForLocalEvent];
+  localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects([localSnapshot valueForField:@"a"], @1337);
 
   [_docRef updateData:@{@"a" : [FIRFieldValue fieldValueForServerTimestamp]}];
-  localSnapshot = [self waitForLocalEvent];
+  localSnapshot = [_accumulator awaitLocalEvent];
   XCTAssertEqualObjects(
       [localSnapshot valueForField:@"a" serverTimestampBehavior:FIRServerTimestampBehaviorPrevious],
       @1337);
 
   [self enableNetwork];
 
-  FIRDocumentSnapshot *remoteSnapshot = [self waitForRemoteEvent];
+  FIRDocumentSnapshot *remoteSnapshot = [_accumulator awaitRemoteEvent];
   XCTAssertTrue([[remoteSnapshot valueForField:@"a"] isKindOfClass:[FIRTimestamp class]]);
 }
 
@@ -282,7 +265,7 @@
     [transaction setData:_setData forDocument:_docRef];
   }];
 
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsWorkViaTransactionUpdate {
@@ -290,7 +273,7 @@
   [self runTransactionBlock:^(FIRTransaction *transaction) {
     [transaction updateData:_updateData forDocument:_docRef];
   }];
-  [self verifySnapshotWithResolvedTimestamps:[self waitForRemoteEvent]];
+  [self verifySnapshotWithResolvedTimestamps:[_accumulator awaitRemoteEvent]];
 }
 
 - (void)testServerTimestampsFailViaUpdateOnNonexistentDocument {
diff --git a/Firestore/Example/Tests/Integration/API/FIRValidationTests.mm b/Firestore/Example/Tests/Integration/API/FIRValidationTests.mm
index 6d10aba..8af8d15 100644
--- a/Firestore/Example/Tests/Integration/API/FIRValidationTests.mm
+++ b/Firestore/Example/Tests/Integration/API/FIRValidationTests.mm
@@ -507,8 +507,8 @@
   FSTAssertThrows([collection queryWhereFieldPath:[FIRFieldPath documentID] isEqualTo:@1], reason);
 
   reason =
-      @"Invalid query. You can't do arrayContains queries on document ID since document IDs are "
-      @"not arrays.";
+      @"Invalid query. You can't perform arrayContains queries on document ID since document IDs "
+       "are not arrays.";
   FSTAssertThrows([collection queryWhereFieldPath:[FIRFieldPath documentID] arrayContains:@1],
                   reason);
 }
diff --git a/Firestore/Example/Tests/Util/FSTEventAccumulator.h b/Firestore/Example/Tests/Util/FSTEventAccumulator.h
index baa501b..58b802b 100644
--- a/Firestore/Example/Tests/Util/FSTEventAccumulator.h
+++ b/Firestore/Example/Tests/Util/FSTEventAccumulator.h
@@ -25,15 +25,21 @@ NS_ASSUME_NONNULL_BEGIN
 
 typedef void (^FSTValueEventHandler)(id _Nullable, NSError *_Nullable error);
 
-@interface FSTEventAccumulator : NSObject
+@interface FSTEventAccumulator <EventType> : NSObject
 
 + (instancetype)accumulatorForTest:(XCTestCase *)testCase;
 
 - (instancetype)init NS_UNAVAILABLE;
 
-- (id)awaitEventWithName:(NSString *)name;
+- (EventType)awaitEventWithName:(NSString *)name;
 
-- (NSArray<id> *)awaitEvents:(NSUInteger)events name:(NSString *)name;
+- (NSArray<EventType> *)awaitEvents:(NSUInteger)events name:(NSString *)name;
+
+/** Waits for a latency compensated local snapshot. */
+- (EventType)awaitLocalEvent;
+
+/** Waits for a snapshot that has no pending writes */
+- (EventType)awaitRemoteEvent;
 
 @property(nonatomic, strong, readonly) FSTValueEventHandler valueEventHandler;
 
diff --git a/Firestore/Example/Tests/Util/FSTEventAccumulator.mm b/Firestore/Example/Tests/Util/FSTEventAccumulator.mm
index 623ba2d..3ab6035 100644
--- a/Firestore/Example/Tests/Util/FSTEventAccumulator.mm
+++ b/Firestore/Example/Tests/Util/FSTEventAccumulator.mm
@@ -18,6 +18,9 @@
 
 #import <XCTest/XCTest.h>
 
+#import "Firestore/Source/Public/FIRDocumentSnapshot.h"
+#import "Firestore/Source/Public/FIRQuerySnapshot.h"
+#import "Firestore/Source/Public/FIRSnapshotMetadata.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
 #import "Firestore/Example/Tests/Util/XCTestCase+Await.h"
@@ -68,6 +71,31 @@ NS_ASSUME_NONNULL_BEGIN
   return events[0];
 }
 
+- (id)awaitLocalEvent {
+  id event;
+  do {
+    event = [self awaitEventWithName:@"Local Event"];
+  } while (![self hasPendingWrites:event]);
+  return event;
+}
+
+- (id)awaitRemoteEvent {
+  id event;
+  do {
+    event = [self awaitEventWithName:@"Remote Event"];
+  } while ([self hasPendingWrites:event]);
+  return event;
+}
+
+- (BOOL)hasPendingWrites:(id)event {
+  if ([event isKindOfClass:[FIRDocumentSnapshot class]]) {
+    return ((FIRDocumentSnapshot *)event).metadata.hasPendingWrites;
+  } else {
+    FSTAssert([event isKindOfClass:[FIRQuerySnapshot class]], @"Unexpected event: %@", event);
+    return ((FIRQuerySnapshot *)event).metadata.hasPendingWrites;
+  }
+}
+
 - (void (^)(id _Nullable, NSError *_Nullable))valueEventHandler {
   return ^void(id _Nullable value, NSError *_Nullable error) {
     // We can't store nil in the _events array, but these are still interesting to tests so store
diff --git a/Firestore/Source/API/FIRQuery.mm b/Firestore/Source/API/FIRQuery.mm
index 32d8327..c83af9c 100644
--- a/Firestore/Source/API/FIRQuery.mm
+++ b/Firestore/Source/API/FIRQuery.mm
@@ -456,8 +456,8 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
   if (fieldPath.IsKeyFieldPath()) {
     if (filterOperator == FSTRelationFilterOperatorArrayContains) {
       FSTThrowInvalidArgument(
-          @"Invalid query. You can't do arrayContains queries on document ID since document IDs "
-          @"are not arrays.");
+          @"Invalid query. You can't perform arrayContains queries on document ID since document "
+           "IDs are not arrays.");
     }
     if ([value isKindOfClass:[NSString class]]) {
       NSString *documentKey = (NSString *)value;
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(),
diff --git a/Firestore/core/src/firebase/firestore/model/transform_operations.h b/Firestore/core/src/firebase/firestore/model/transform_operations.h
index a1c2de0..abc8d1b 100644
--- a/Firestore/core/src/firebase/firestore/model/transform_operations.h
+++ b/Firestore/core/src/firebase/firestore/model/transform_operations.h
@@ -20,6 +20,8 @@
 #include <utility>
 #include <vector>
 
+#include "Firestore/core/include/firebase/firestore/timestamp.h"
+
 #if defined(__OBJC__)
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #endif
@@ -49,11 +51,28 @@ class TransformOperation {
   /** Returns the actual type. */
   virtual Type type() const = 0;
 
+// TODO(mikelehen): apply methods require FSTFieldValue which is Obj-C only.
+#if defined(__OBJC__)
+  /**
+   * Computes the local transform result against the provided `previousValue`,
+   * optionally using the provided localWriteTime.
+   */
+  virtual FSTFieldValue *applyToLocalView(
+      FSTFieldValue *previousValue, FIRTimestamp *localWriteTime) const = 0;
+
+  /**
+   * Computes a final transform result after the transform has been acknowledged
+   * by the server, potentially using the server-provided transformResult.
+   */
+  virtual FSTFieldValue *applyToRemoteDocument(
+      FSTFieldValue *previousValue, FSTFieldValue *transformResult) const = 0;
+#endif
+
   /** Returns whether the two are equal. */
-  virtual bool operator==(const TransformOperation& other) const = 0;
+  virtual bool operator==(const TransformOperation &other) const = 0;
 
   /** Returns whether the two are not equal. */
-  bool operator!=(const TransformOperation& other) const {
+  bool operator!=(const TransformOperation &other) const {
     return !operator==(other);
   }
 
@@ -74,12 +93,28 @@ class ServerTimestampTransform : public TransformOperation {
     return Type::ServerTimestamp;
   }
 
-  bool operator==(const TransformOperation& other) const override {
+// TODO(mikelehen): apply methods require FSTFieldValue which is Obj-C only.
+#if defined(__OBJC__)
+  FSTFieldValue *applyToLocalView(FSTFieldValue *previousValue,
+                                  FIRTimestamp *localWriteTime) const override {
+    return [FSTServerTimestampValue
+        serverTimestampValueWithLocalWriteTime:localWriteTime
+                                 previousValue:previousValue];
+  }
+
+  FSTFieldValue *applyToRemoteDocument(
+      FSTFieldValue *previousValue,
+      FSTFieldValue *transformResult) const override {
+    return transformResult;
+  }
+#endif
+
+  bool operator==(const TransformOperation &other) const override {
     // All ServerTimestampTransform objects are equal.
     return other.type() == Type::ServerTimestamp;
   }
 
-  static const ServerTimestampTransform& Get() {
+  static const ServerTimestampTransform &Get() {
     static ServerTimestampTransform shared_instance;
     return shared_instance;
   }
@@ -102,13 +137,14 @@ class ServerTimestampTransform : public TransformOperation {
 // TODO(mikelehen): ArrayTransform can only be used from Obj-C until we switch
 // to using FieldValue instead of FSTFieldValue.
 #if defined(__OBJC__)
+
 /**
  * Transforms an array via a union or remove operation (for convenience, we use
  * this class for both Type::ArrayUnion and Type::ArrayRemove).
  */
 class ArrayTransform : public TransformOperation {
  public:
-  ArrayTransform(Type type, std::vector<FSTFieldValue*> elements)
+  ArrayTransform(Type type, std::vector<FSTFieldValue *> elements)
       : type_(type), elements_(std::move(elements)) {
   }
 
@@ -119,15 +155,29 @@ class ArrayTransform : public TransformOperation {
     return type_;
   }
 
-  const std::vector<FSTFieldValue*>& elements() const {
+  FSTFieldValue *applyToLocalView(FSTFieldValue *previousValue,
+                                  FIRTimestamp *localWriteTime) const override {
+    return apply(previousValue);
+  }
+
+  FSTFieldValue *applyToRemoteDocument(
+      FSTFieldValue *previousValue,
+      FSTFieldValue *transformResult) const override {
+    // The server just sends null as the transform result for array operations,
+    // so we have to calculate a result the same as we do for local
+    // applications.
+    return apply(previousValue);
+  }
+
+  const std::vector<FSTFieldValue *> &elements() const {
     return elements_;
   }
 
-  bool operator==(const TransformOperation& other) const override {
+  bool operator==(const TransformOperation &other) const override {
     if (other.type() != type()) {
       return false;
     }
-    auto array_transform = static_cast<const ArrayTransform&>(other);
+    auto array_transform = static_cast<const ArrayTransform &>(other);
     if (array_transform.elements_.size() != elements_.size()) {
       return false;
     }
@@ -145,23 +195,56 @@ class ArrayTransform : public TransformOperation {
   NSUInteger Hash() const override {
     NSUInteger result = 37;
     result = 31 * result + (type() == Type::ArrayUnion ? 1231 : 1237);
-    for (FSTFieldValue* element : elements_) {
+    for (FSTFieldValue *element : elements_) {
       result = 31 * result + [element hash];
     }
     return result;
   }
 #endif  // defined(__OBJC__)
 
-  static const std::vector<FSTFieldValue*>& Elements(
-      const TransformOperation& op) {
+  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();
+    return static_cast<const ArrayTransform &>(op).elements();
   }
 
  private:
   Type type_;
-  std::vector<FSTFieldValue*> elements_;
+  std::vector<FSTFieldValue *> elements_;
+
+  /**
+   * 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.
+   */
+  static NSMutableArray<FSTFieldValue *> *CoercedFieldValuesArray(
+      FSTFieldValue *value) {
+    if ([value isMemberOfClass:[FSTArrayValue class]]) {
+      return [NSMutableArray
+          arrayWithArray:reinterpret_cast<FSTArrayValue *>(value)
+                             .internalValue];
+    } else {
+      // coerce to empty array.
+      return [NSMutableArray array];
+    }
+  }
+
+  FSTFieldValue *apply(FSTFieldValue *previousValue) const {
+    NSMutableArray<FSTFieldValue *> *result =
+        ArrayTransform::CoercedFieldValuesArray(previousValue);
+    for (FSTFieldValue *element : elements_) {
+      if (type_ == Type::ArrayUnion) {
+        if (![result containsObject:element]) {
+          [result addObject:element];
+        }
+      } else {
+        FIREBASE_ASSERT(type_ == Type::ArrayRemove);
+        [result removeObject:element];
+      }
+    }
+    return [[FSTArrayValue alloc] initWithValueNoCopy:result];
+  }
 };
 #endif