1 files changed, 197 insertions, 0 deletions
diff --git a/Firestore/Source/Model/FSTDocumentSet.mm b/Firestore/Source/Model/FSTDocumentSet.mm
new file mode 100644
index 0000000..c4c0f49
--- /dev/null
+++ b/Firestore/Source/Model/FSTDocumentSet.mm
@@ -0,0 +1,197 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "Firestore/Source/Model/FSTDocumentSet.h"
+
+#import "Firestore/Source/Model/FSTDocument.h"
+#import "Firestore/Source/Model/FSTDocumentKey.h"
+#import "Firestore/third_party/Immutable/FSTImmutableSortedSet.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+/**
+ * The type of the index of the documents in an FSTDocumentSet.
+ * @see FSTDocumentSet#index
+ */
+typedef FSTImmutableSortedDictionary<FSTDocumentKey *, FSTDocument *> IndexType;
+
+/**
+ * The type of the main collection of documents in an FSTDocumentSet.
+ * @see FSTDocumentSet#sortedSet
+ */
+typedef FSTImmutableSortedSet<FSTDocument *> SetType;
+
+@interface FSTDocumentSet ()
+
+- (instancetype)initWithIndex:(IndexType *)index set:(SetType *)sortedSet NS_DESIGNATED_INITIALIZER;
+
+/**
+ * An index of the documents in the FSTDocumentSet, indexed by document key. The index
+ * exists to guarantee the uniqueness of document keys in the set and to allow lookup and removal
+ * of documents by key.
+ */
+@property(nonatomic, strong, readonly) IndexType *index;
+
+/**
+ * The main collection of documents in the FSTDocumentSet. The documents are ordered by a
+ * comparator supplied from a query. The SetType collection exists in addition to the index to
+ * allow ordered traversal of the FSTDocumentSet.
+ */
+@property(nonatomic, strong, readonly) SetType *sortedSet;
+@end
+
+@implementation FSTDocumentSet
+
++ (instancetype)documentSetWithComparator:(NSComparator)comparator {
+  IndexType *index =
+      [FSTImmutableSortedDictionary dictionaryWithComparator:FSTDocumentKeyComparator];
+  SetType *set = [FSTImmutableSortedSet setWithComparator:comparator];
+  return [[FSTDocumentSet alloc] initWithIndex:index set:set];
+}
+
+- (instancetype)initWithIndex:(IndexType *)index set:(SetType *)sortedSet {
+  self = [super init];
+  if (self) {
+    _index = index;
+    _sortedSet = sortedSet;
+  }
+  return self;
+}
+
+- (BOOL)isEqual:(id)other {
+  if (other == self) {
+    return YES;
+  }
+  if (![other isMemberOfClass:[FSTDocumentSet class]]) {
+    return NO;
+  }
+
+  FSTDocumentSet *otherSet = (FSTDocumentSet *)other;
+  if ([self count] != [otherSet count]) {
+    return NO;
+  }
+
+  NSEnumerator<FSTDocument *> *selfIter = [self.sortedSet objectEnumerator];
+  NSEnumerator<FSTDocument *> *otherIter = [otherSet.sortedSet objectEnumerator];
+
+  FSTDocument *selfDoc = [selfIter nextObject];
+  FSTDocument *otherDoc = [otherIter nextObject];
+  while (selfDoc) {
+    if (![selfDoc isEqual:otherDoc]) {
+      return NO;
+    }
+    selfDoc = [selfIter nextObject];
+    otherDoc = [otherIter nextObject];
+  }
+  return YES;
+}
+
+- (NSUInteger)hash {
+  NSUInteger hash = 0;
+  for (FSTDocument *doc in self.sortedSet.objectEnumerator) {
+    hash = 31 * hash + [doc hash];
+  }
+  return hash;
+}
+
+- (NSString *)description {
+  return [self.sortedSet description];
+}
+
+- (NSUInteger)count {
+  return [self.index count];
+}
+
+- (BOOL)isEmpty {
+  return [self.index isEmpty];
+}
+
+- (BOOL)containsKey:(FSTDocumentKey *)key {
+  return [self.index objectForKey:key] != nil;
+}
+
+- (FSTDocument *_Nullable)documentForKey:(FSTDocumentKey *)key {
+  return [self.index objectForKey:key];
+}
+
+- (FSTDocument *_Nullable)firstDocument {
+  return [self.sortedSet firstObject];
+}
+
+- (FSTDocument *_Nullable)lastDocument {
+  return [self.sortedSet lastObject];
+}
+
+- (FSTDocument *_Nullable)predecessorDocumentForKey:(FSTDocumentKey *)key {
+  FSTDocument *doc = [self.index objectForKey:key];
+  if (!doc) {
+    @throw [NSException exceptionWithName:NSInvalidArgumentException
+                                   reason:[NSString stringWithFormat:@"Key %@ does not exist", key]
+                                 userInfo:nil];
+  }
+  return [self.sortedSet predecessorObject:doc];
+}
+
+- (NSUInteger)indexOfKey:(FSTDocumentKey *)key {
+  FSTDocument *doc = [self.index objectForKey:key];
+  return doc ? [self.sortedSet indexOfObject:doc] : NSNotFound;
+}
+
+- (NSEnumerator<FSTDocument *> *)documentEnumerator {
+  return [self.sortedSet objectEnumerator];
+}
+
+- (NSArray *)arrayValue {
+  NSMutableArray<FSTDocument *> *result = [NSMutableArray arrayWithCapacity:self.count];
+  for (FSTDocument *doc in self.documentEnumerator) {
+    [result addObject:doc];
+  }
+  return result;
+}
+
+- (FSTMaybeDocumentDictionary *)dictionaryValue {
+  return self.index;
+}
+
+- (instancetype)documentSetByAddingDocument:(FSTDocument *_Nullable)document {
+  // TODO(mcg): look into making document nonnull.
+  if (!document) {
+    return self;
+  }
+
+  // Remove any prior mapping of the document's key before adding, preventing sortedSet from
+  // accumulating values that aren't in the index.
+  FSTDocumentSet *removed = [self documentSetByRemovingKey:document.key];
+
+  IndexType *index = [removed.index dictionaryBySettingObject:document forKey:document.key];
+  SetType *set = [removed.sortedSet setByAddingObject:document];
+  return [[FSTDocumentSet alloc] initWithIndex:index set:set];
+}
+
+- (instancetype)documentSetByRemovingKey:(FSTDocumentKey *)key {
+  FSTDocument *doc = [self.index objectForKey:key];
+  if (!doc) {
+    return self;
+  }
+
+  IndexType *index = [self.index dictionaryByRemovingObjectForKey:key];
+  SetType *set = [self.sortedSet setByRemovingObject:doc];
+  return [[FSTDocumentSet alloc] initWithIndex:index set:set];
+}
+
+@end
+
+NS_ASSUME_NONNULL_END