Making DocumentSnapshot nullable

7 files changed, 110 insertions, 62 deletions

diff --git a/Firestore/Source/API/FIRDocumentChange.m b/Firestore/Source/API/FIRDocumentChange.m
index 970dc90..d1d9999 100644
--- a/Firestore/Source/API/FIRDocumentChange.m
+++ b/Firestore/Source/API/FIRDocumentChange.m
@@ -57,11 +57,11 @@ NS_ASSUME_NONNULL_BEGIN
     NSUInteger index = 0;
     NSMutableArray<FIRDocumentChange *> *changes = [NSMutableArray array];
     for (FSTDocumentViewChange *change in snapshot.documentChanges) {
-      FIRDocumentSnapshot *document =
-          [FIRDocumentSnapshot snapshotWithFirestore:firestore
-                                         documentKey:change.document.key
-                                            document:change.document
-                                           fromCache:snapshot.isFromCache];
+      FIRQueryDocumentSnapshot *document =
+          [FIRQueryDocumentSnapshot snapshotWithFirestore:firestore
+                                              documentKey:change.document.key
+                                                 document:change.document
+                                                fromCache:snapshot.isFromCache];
       FSTAssert(change.type == FSTDocumentViewChangeTypeAdded,
                 @"Invalid event type for first snapshot");
       FSTAssert(!lastDocument ||
@@ -79,11 +79,11 @@ NS_ASSUME_NONNULL_BEGIN
     FSTDocumentSet *indexTracker = snapshot.oldDocuments;
     NSMutableArray<FIRDocumentChange *> *changes = [NSMutableArray array];
     for (FSTDocumentViewChange *change in snapshot.documentChanges) {
-      FIRDocumentSnapshot *document =
-          [FIRDocumentSnapshot snapshotWithFirestore:firestore
-                                         documentKey:change.document.key
-                                            document:change.document
-                                           fromCache:snapshot.isFromCache];
+      FIRQueryDocumentSnapshot *document =
+          [FIRQueryDocumentSnapshot snapshotWithFirestore:firestore
+                                              documentKey:change.document.key
+                                                 document:change.document
+                                                fromCache:snapshot.isFromCache];
 
       NSUInteger oldIndex = NSNotFound;
       NSUInteger newIndex = NSNotFound;
@@ -112,7 +112,7 @@ NS_ASSUME_NONNULL_BEGIN
 @implementation FIRDocumentChange
 
 - (instancetype)initWithType:(FIRDocumentChangeType)type
-                    document:(FIRDocumentSnapshot *)document
+                    document:(FIRQueryDocumentSnapshot *)document
                     oldIndex:(NSUInteger)oldIndex
                     newIndex:(NSUInteger)newIndex {
   if (self = [super init]) {
diff --git a/Firestore/Source/API/FIRDocumentSnapshot.m b/Firestore/Source/API/FIRDocumentSnapshot.m
index b78472e..0aba639 100644
--- a/Firestore/Source/API/FIRDocumentSnapshot.m
+++ b/Firestore/Source/API/FIRDocumentSnapshot.m
@@ -25,6 +25,7 @@
 #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
@@ -99,18 +100,8 @@ 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);
-  }
-
-  return [self convertedObject:[self.internalDocument data]];
+- (nullable NSDictionary<NSString *, id> *)data {
+  return self.internalDocument == nil ? nil : [self convertedObject:[self.internalDocument data]];
 }
 
 - (nullable id)objectForKeyedSubscript:(id)key {
@@ -125,7 +116,7 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTFieldValue *fieldValue = [[self.internalDocument data] valueForPath:fieldPath.internalValue];
-  return [self convertedValue:fieldValue];
+  return fieldValue == nil ? nil : [self convertedValue:fieldValue];
 }
 
 - (id)convertedValue:(FSTFieldValue *)value {
@@ -172,4 +163,34 @@ NS_ASSUME_NONNULL_BEGIN
 
 @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;
+}
+
+@end
+
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/API/FIRQuery.m b/Firestore/Source/API/FIRQuery.m
index 9ab854a..70bd51e 100644
--- a/Firestore/Source/API/FIRQuery.m
+++ b/Firestore/Source/API/FIRQuery.m
@@ -259,8 +259,9 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
 - (FIRQuery *)queryFilteredUsingComparisonPredicate:(NSPredicate *)predicate {
   NSComparisonPredicate *comparison = (NSComparisonPredicate *)predicate;
   if (comparison.comparisonPredicateModifier != NSDirectPredicateModifier) {
-    FSTThrowInvalidArgument(@"Invalid query. Predicate cannot have an "
-                            "aggregate modifier.");
+    FSTThrowInvalidArgument(
+        @"Invalid query. Predicate cannot have an "
+         "aggregate modifier.");
   }
   NSString *path;
   id value = nil;
@@ -279,8 +280,7 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
         return [self queryWhereField:path isGreaterThan:value];
       case NSGreaterThanOrEqualToPredicateOperatorType:
         return [self queryWhereField:path isGreaterThanOrEqualTo:value];
-      default:
-        ;  // Fallback below to throw assertion.
+      default:;  // Fallback below to throw assertion.
     }
   } else if ([comparison.leftExpression expressionType] == NSConstantValueExpressionType &&
              [comparison.rightExpression expressionType] == NSKeyPathExpressionType) {
@@ -297,18 +297,19 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
         return [self queryWhereField:path isLessThan:value];
       case NSGreaterThanOrEqualToPredicateOperatorType:
         return [self queryWhereField:path isLessThanOrEqualTo:value];
-      default:
-        ;  // Fallback below to throw assertion.
+      default:;  // Fallback below to throw assertion.
     }
   } else {
-    FSTThrowInvalidArgument(@"Invalid query. Predicate comparisons must "
-                            "include a key path and a constant.");
+    FSTThrowInvalidArgument(
+        @"Invalid query. Predicate comparisons must "
+         "include a key path and a constant.");
   }
   // Fallback cases of unsupported comparison operator.
   switch (comparison.predicateOperatorType) {
     case NSCustomSelectorPredicateOperatorType:
-      FSTThrowInvalidArgument(@"Invalid query. Custom predicate filters are "
-                              "not supported.");
+      FSTThrowInvalidArgument(
+          @"Invalid query. Custom predicate filters are "
+           "not supported.");
       break;
     default:
       FSTThrowInvalidArgument(@"Invalid query. Operator type %lu is not supported.",
@@ -318,10 +319,10 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
 
 - (FIRQuery *)queryFilteredUsingCompoundPredicate:(NSPredicate *)predicate {
   NSCompoundPredicate *compound = (NSCompoundPredicate *)predicate;
-  if (compound.compoundPredicateType != NSAndPredicateType ||
-      compound.subpredicates.count == 0) {
-    FSTThrowInvalidArgument(@"Invalid query. Only compound queries using AND "
-                            "are supported.");
+  if (compound.compoundPredicateType != NSAndPredicateType || compound.subpredicates.count == 0) {
+    FSTThrowInvalidArgument(
+        @"Invalid query. Only compound queries using AND "
+         "are supported.");
   }
   FIRQuery *query = self;
   for (NSPredicate *pred in compound.subpredicates) {
@@ -335,16 +336,19 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
     return [self queryFilteredUsingComparisonPredicate:predicate];
   } else if ([predicate isKindOfClass:[NSCompoundPredicate class]]) {
     return [self queryFilteredUsingCompoundPredicate:predicate];
-  } else if ([predicate isKindOfClass:
-              [[NSPredicate predicateWithBlock:
-                ^BOOL(id obj, NSDictionary *bindings) { return true; }] class]]) {
-    FSTThrowInvalidArgument(@"Invalid query. Block-based predicates are not "
-                            "supported. Please use predicateWithFormat to "
-                            "create predicates instead.");
+  } else if ([predicate isKindOfClass:[[NSPredicate
+                                          predicateWithBlock:^BOOL(id obj, NSDictionary *bindings) {
+                                            return true;
+                                          }] class]]) {
+    FSTThrowInvalidArgument(
+        @"Invalid query. Block-based predicates are not "
+         "supported. Please use predicateWithFormat to "
+         "create predicates instead.");
   } else {
-    FSTThrowInvalidArgument(@"Invalid query. Expect comparison or compound of "
-                            "comparison predicate. Please use "
-                            "predicateWithFormat to create predicates.");
+    FSTThrowInvalidArgument(
+        @"Invalid query. Expect comparison or compound of "
+         "comparison predicate. Please use "
+         "predicateWithFormat to create predicates.");
   }
 }
 
diff --git a/Firestore/Source/API/FIRQuerySnapshot.m b/Firestore/Source/API/FIRQuerySnapshot.m
index 6bc6761..5e1af9a 100644
--- a/Firestore/Source/API/FIRQuerySnapshot.m
+++ b/Firestore/Source/API/FIRQuerySnapshot.m
@@ -57,7 +57,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 @implementation FIRQuerySnapshot {
   // Cached value of the documents property.
-  NSArray<FIRDocumentSnapshot *> *_documents;
+  NSArray<FIRQueryDocumentSnapshot *> *_documents;
 
   // Cached value of the documentChanges property.
   NSArray<FIRDocumentChange *> *_documentChanges;
@@ -93,18 +93,18 @@ NS_ASSUME_NONNULL_BEGIN
   return self.snapshot.documents.count;
 }
 
-- (NSArray<FIRDocumentSnapshot *> *)documents {
+- (NSArray<FIRQueryDocumentSnapshot *> *)documents {
   if (!_documents) {
     FSTDocumentSet *documentSet = self.snapshot.documents;
     FIRFirestore *firestore = self.firestore;
     BOOL fromCache = self.metadata.fromCache;
 
-    NSMutableArray<FIRDocumentSnapshot *> *result = [NSMutableArray array];
+    NSMutableArray<FIRQueryDocumentSnapshot *> *result = [NSMutableArray array];
     for (FSTDocument *document in documentSet.documentEnumerator) {
-      [result addObject:[FIRDocumentSnapshot snapshotWithFirestore:firestore
-                                                       documentKey:document.key
-                                                          document:document
-                                                         fromCache:fromCache]];
+      [result addObject:[FIRQueryDocumentSnapshot snapshotWithFirestore:firestore
+                                                            documentKey:document.key
+                                                               document:document
+                                                              fromCache:fromCache]];
     }
 
     _documents = result;
diff --git a/Firestore/Source/Public/FIRDocumentChange.h b/Firestore/Source/Public/FIRDocumentChange.h
index 022c81b..4717067 100644
--- a/Firestore/Source/Public/FIRDocumentChange.h
+++ b/Firestore/Source/Public/FIRDocumentChange.h
@@ -18,7 +18,7 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
-@class FIRDocumentSnapshot;
+@class FIRQueryDocumentSnapshot;
 
 /** An enumeration of document change types. */
 typedef NS_ENUM(NSInteger, FIRDocumentChangeType) {
@@ -47,7 +47,7 @@ NS_SWIFT_NAME(DocumentChange)
 @property(nonatomic, readonly) FIRDocumentChangeType type;
 
 /** The document affected by this change. */
-@property(nonatomic, strong, readonly) FIRDocumentSnapshot *document;
+@property(nonatomic, strong, readonly) FIRQueryDocumentSnapshot *document;
 
 /**
  * The index of the changed document in the result set immediately prior to this FIRDocumentChange
diff --git a/Firestore/Source/Public/FIRDocumentSnapshot.h b/Firestore/Source/Public/FIRDocumentSnapshot.h
index 3e67c25..9706edd 100644
--- a/Firestore/Source/Public/FIRDocumentSnapshot.h
+++ b/Firestore/Source/Public/FIRDocumentSnapshot.h
@@ -46,21 +46,44 @@ NS_SWIFT_NAME(DocumentSnapshot)
 @property(nonatomic, strong, readonly) FIRSnapshotMetadata *metadata;
 
 /**
- * Retrieves all fields in the document as an `NSDictionary`.
+ * Retrieves all fields in the document as an `NSDictionary`. Returns `nil` if the document doesn't
+ * exist.
  *
- * @return An `NSDictionary` containing all fields in the document.
+ * @return An `NSDictionary` containing all fields in the document or `nil` if the document doesn't
+ *      exist.
  */
-- (NSDictionary<NSString *, id> *)data;
+- (nullable NSDictionary<NSString *, id> *)data;
 
 /**
  * Retrieves a specific field from the document.
  *
  * @param key The field to retrieve.
  *
- * @return The value contained in the field or `nil` if the field doesn't exist.
+ * @return The value contained in the field or `nil` if the document or field doesn't exist.
  */
 - (nullable id)objectForKeyedSubscript:(id)key;
 
 @end
 
+/**
+ * A `FIRDocumentSnapshot` contains data read from a document in your Firestore database. The
+ * document is guaranteed to exist and its data can be extracted with the `data` property or by using
+ * subscript syntax to access a specific field.
+ */
+NS_SWIFT_NAME(QueryDocumentSnapshot)
+@interface FIRQueryDocumentSnapshot : FIRDocumentSnapshot
+
+/**   */
+- (instancetype)init
+    __attribute__((unavailable("FIRQueryDocumentSnapshot cannot be created directly.")));
+
+/**
+ * Retrieves all fields in the document as an `NSDictionary`.
+ *
+ * @return An `NSDictionary` containing all fields in the document.
+ */
+- (NSDictionary<NSString *, id> *)data;
+
+@end
+
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Public/FIRQuerySnapshot.h b/Firestore/Source/Public/FIRQuerySnapshot.h
index c49a07a..1266832 100644
--- a/Firestore/Source/Public/FIRQuerySnapshot.h
+++ b/Firestore/Source/Public/FIRQuerySnapshot.h
@@ -19,8 +19,8 @@
 NS_ASSUME_NONNULL_BEGIN
 
 @class FIRDocumentChange;
-@class FIRDocumentSnapshot;
 @class FIRQuery;
+@class FIRQueryDocumentSnapshot;
 @class FIRSnapshotMetadata;
 
 /**
@@ -50,7 +50,7 @@ NS_SWIFT_NAME(QuerySnapshot)
 @property(nonatomic, readonly) NSInteger count;
 
 /** An Array of the `FIRDocumentSnapshots` that make up this document set. */
-@property(nonatomic, strong, readonly) NSArray<FIRDocumentSnapshot *> *documents;
+@property(nonatomic, strong, readonly) NSArray<FIRQueryDocumentSnapshot *> *documents;
 
 /**
  * An array of the documents that changed since the last snapshot. If this is the first snapshot,