1 files changed, 106 insertions, 26 deletions

diff --git a/Firestore/Source/API/FIRDocumentSnapshot.m b/Firestore/Source/API/FIRDocumentSnapshot.m
index b78472e..358ddac 100644
--- a/Firestore/Source/API/FIRDocumentSnapshot.m
+++ b/Firestore/Source/API/FIRDocumentSnapshot.m
@@ -20,11 +20,13 @@
 #import "Firestore/Source/API/FIRFieldPath+Internal.h"
 #import "Firestore/Source/API/FIRFirestore+Internal.h"
 #import "Firestore/Source/API/FIRSnapshotMetadata+Internal.h"
+#import "Firestore/Source/API/FIRSnapshotOptions+Internal.h"
 #import "Firestore/Source/Model/FSTDatabaseID.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #import "Firestore/Source/Model/FSTPath.h"
+#import "Firestore/Source/Util/FSTAssert.h"
 #import "Firestore/Source/Util/FSTUsageValidation.h"
 
 NS_ASSUME_NONNULL_BEGIN
@@ -76,6 +78,37 @@ NS_ASSUME_NONNULL_BEGIN
   return self;
 }
 
+// NSObject Methods
+- (BOOL)isEqual:(nullable id)other {
+  if (other == self) return YES;
+  // self class could be FIRDocumentSnapshot or subtype. So we compare with base type explicitly.
+  if (![other isKindOfClass:[FIRDocumentSnapshot class]]) return NO;
+
+  return [self isEqualToSnapshot:other];
+}
+
+- (BOOL)isEqualToSnapshot:(nullable FIRDocumentSnapshot *)snapshot {
+  if (self == snapshot) return YES;
+  if (snapshot == nil) return NO;
+  if (self.firestore != snapshot.firestore && ![self.firestore isEqual:snapshot.firestore])
+    return NO;
+  if (self.internalKey != snapshot.internalKey && ![self.internalKey isEqual:snapshot.internalKey])
+    return NO;
+  if (self.internalDocument != snapshot.internalDocument &&
+      ![self.internalDocument isEqual:snapshot.internalDocument])
+    return NO;
+  if (self.fromCache != snapshot.fromCache) return NO;
+  return YES;
+}
+
+- (NSUInteger)hash {
+  NSUInteger hash = [self.firestore hash];
+  hash = hash * 31u + [self.internalKey hash];
+  hash = hash * 31u + [self.internalDocument hash];
+  hash = hash * 31u + (self.fromCache ? 1 : 0);
+  return hash;
+}
+
 @dynamic exists;
 
 - (BOOL)exists {
@@ -99,40 +132,48 @@ NS_ASSUME_NONNULL_BEGIN
   return _cachedMetadata;
 }
 
-- (NSDictionary<NSString *, id> *)data {
-  FSTDocument *document = self.internalDocument;
-
-  if (!document) {
-    FSTThrowInvalidUsage(
-        @"NonExistentDocumentException",
-        @"Document '%@' doesn't exist. "
-        @"Check document.exists to make sure the document exists before calling document.data.",
-        self.internalKey);
-  }
+- (nullable NSDictionary<NSString *, id> *)data {
+  return [self dataWithOptions:[FIRSnapshotOptions defaultOptions]];
+}
 
-  return [self convertedObject:[self.internalDocument data]];
+- (nullable NSDictionary<NSString *, id> *)dataWithOptions:(FIRSnapshotOptions *)options {
+  return self.internalDocument == nil
+             ? nil
+             : [self convertedObject:[self.internalDocument data]
+                             options:[FSTFieldValueOptions optionsForSnapshotOptions:options]];
 }
 
-- (nullable id)objectForKeyedSubscript:(id)key {
+- (nullable id)valueForField:(id)field {
+  return [self valueForField:field options:[FIRSnapshotOptions defaultOptions]];
+}
+
+- (nullable id)valueForField:(id)field options:(FIRSnapshotOptions *)options {
   FIRFieldPath *fieldPath;
 
-  if ([key isKindOfClass:[NSString class]]) {
-    fieldPath = [FIRFieldPath pathWithDotSeparatedString:key];
-  } else if ([key isKindOfClass:[FIRFieldPath class]]) {
-    fieldPath = key;
+  if ([field isKindOfClass:[NSString class]]) {
+    fieldPath = [FIRFieldPath pathWithDotSeparatedString:field];
+  } else if ([field isKindOfClass:[FIRFieldPath class]]) {
+    fieldPath = field;
   } else {
     FSTThrowInvalidArgument(@"Subscript key must be an NSString or FIRFieldPath.");
   }
 
   FSTFieldValue *fieldValue = [[self.internalDocument data] valueForPath:fieldPath.internalValue];
-  return [self convertedValue:fieldValue];
+  return fieldValue == nil
+             ? nil
+             : [self convertedValue:fieldValue
+                            options:[FSTFieldValueOptions optionsForSnapshotOptions:options]];
 }
 
-- (id)convertedValue:(FSTFieldValue *)value {
+- (nullable id)objectForKeyedSubscript:(id)key {
+  return [self valueForField:key];
+}
+
+- (id)convertedValue:(FSTFieldValue *)value options:(FSTFieldValueOptions *)options {
   if ([value isKindOfClass:[FSTObjectValue class]]) {
-    return [self convertedObject:(FSTObjectValue *)value];
+    return [self convertedObject:(FSTObjectValue *)value options:options];
   } else if ([value isKindOfClass:[FSTArrayValue class]]) {
-    return [self convertedArray:(FSTArrayValue *)value];
+    return [self convertedArray:(FSTArrayValue *)value options:options];
   } else if ([value isKindOfClass:[FSTReferenceValue class]]) {
     FSTReferenceValue *ref = (FSTReferenceValue *)value;
     FSTDatabaseID *refDatabase = ref.databaseID;
@@ -146,30 +187,69 @@ NS_ASSUME_NONNULL_BEGIN
           self.reference.path, refDatabase.projectID, refDatabase.databaseID, database.projectID,
           database.databaseID);
     }
-    return [FIRDocumentReference referenceWithKey:ref.value firestore:self.firestore];
+    return [FIRDocumentReference referenceWithKey:[ref valueWithOptions:options]
+                                        firestore:self.firestore];
   } else {
-    return value.value;
+    return [value valueWithOptions:options];
   }
 }
 
-- (NSDictionary<NSString *, id> *)convertedObject:(FSTObjectValue *)objectValue {
+- (NSDictionary<NSString *, id> *)convertedObject:(FSTObjectValue *)objectValue
+                                          options:(FSTFieldValueOptions *)options {
   NSMutableDictionary *result = [NSMutableDictionary dictionary];
   [objectValue.internalValue
       enumerateKeysAndObjectsUsingBlock:^(NSString *key, FSTFieldValue *value, BOOL *stop) {
-        result[key] = [self convertedValue:value];
+        result[key] = [self convertedValue:value options:options];
       }];
   return result;
 }
 
-- (NSArray<id> *)convertedArray:(FSTArrayValue *)arrayValue {
+- (NSArray<id> *)convertedArray:(FSTArrayValue *)arrayValue
+                        options:(FSTFieldValueOptions *)options {
   NSArray<FSTFieldValue *> *internalValue = arrayValue.internalValue;
   NSMutableArray *result = [NSMutableArray arrayWithCapacity:internalValue.count];
   [internalValue enumerateObjectsUsingBlock:^(id value, NSUInteger idx, BOOL *stop) {
-    [result addObject:[self convertedValue:value]];
+    [result addObject:[self convertedValue:value options:options]];
   }];
   return result;
 }
 
 @end
 
+@interface FIRQueryDocumentSnapshot ()
+
+- (instancetype)initWithFirestore:(FIRFirestore *)firestore
+                      documentKey:(FSTDocumentKey *)documentKey
+                         document:(FSTDocument *)document
+                        fromCache:(BOOL)fromCache NS_DESIGNATED_INITIALIZER;
+
+@end
+
+@implementation FIRQueryDocumentSnapshot
+
+- (instancetype)initWithFirestore:(FIRFirestore *)firestore
+                      documentKey:(FSTDocumentKey *)documentKey
+                         document:(FSTDocument *)document
+                        fromCache:(BOOL)fromCache {
+  self = [super initWithFirestore:firestore
+                      documentKey:documentKey
+                         document:document
+                        fromCache:fromCache];
+  return self;
+}
+
+- (NSDictionary<NSString *, id> *)data {
+  NSDictionary<NSString *, id> *data = [super data];
+  FSTAssert(data, @"Document in a QueryDocumentSnapshot should exist");
+  return data;
+}
+
+- (NSDictionary<NSString *, id> *)dataWithOptions:(FIRSnapshotOptions *)options {
+  NSDictionary<NSString *, id> *data = [super dataWithOptions:options];
+  FSTAssert(data, @"Document in a QueryDocumentSnapshot should exist");
+  return data;
+}
+
+@end
+
 NS_ASSUME_NONNULL_END