1 files changed, 135 insertions, 0 deletions
diff --git a/Firestore/Source/Local/FSTReferenceSet.mm b/Firestore/Source/Local/FSTReferenceSet.mm
new file mode 100644
index 0000000..2acd64b
--- /dev/null
+++ b/Firestore/Source/Local/FSTReferenceSet.mm
@@ -0,0 +1,135 @@
+/*
+ * 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/Local/FSTReferenceSet.h"
+
+#import "Firestore/Source/Local/FSTDocumentReference.h"
+#import "Firestore/Source/Model/FSTDocumentKey.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+#pragma mark - FSTReferenceSet
+
+@interface FSTReferenceSet ()
+
+/** A set of outstanding references to a document sorted by key. */
+@property(nonatomic, strong) FSTImmutableSortedSet<FSTDocumentReference *> *referencesByKey;
+
+/** A set of outstanding references to a document sorted by target ID (or batch ID). */
+@property(nonatomic, strong) FSTImmutableSortedSet<FSTDocumentReference *> *referencesByID;
+
+@end
+
+@implementation FSTReferenceSet
+
+#pragma mark - Initializer
+
+- (instancetype)init {
+  self = [super init];
+  if (self) {
+    _referencesByKey =
+        [FSTImmutableSortedSet setWithComparator:FSTDocumentReferenceComparatorByKey];
+    _referencesByID = [FSTImmutableSortedSet setWithComparator:FSTDocumentReferenceComparatorByID];
+  }
+  return self;
+}
+
+#pragma mark - Testing helper methods
+
+- (BOOL)isEmpty {
+  return [self.referencesByKey isEmpty];
+}
+
+- (NSUInteger)count {
+  return self.referencesByKey.count;
+}
+
+#pragma mark - Public methods
+
+- (void)addReferenceToKey:(FSTDocumentKey *)key forID:(int)ID {
+  FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:ID];
+  self.referencesByKey = [self.referencesByKey setByAddingObject:reference];
+  self.referencesByID = [self.referencesByID setByAddingObject:reference];
+}
+
+- (void)addReferencesToKeys:(FSTDocumentKeySet *)keys forID:(int)ID {
+  [keys enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
+    [self addReferenceToKey:key forID:ID];
+  }];
+}
+
+- (void)removeReferenceToKey:(FSTDocumentKey *)key forID:(int)ID {
+  [self removeReference:[[FSTDocumentReference alloc] initWithKey:key ID:ID]];
+}
+
+- (void)removeReferencesToKeys:(FSTDocumentKeySet *)keys forID:(int)ID {
+  [keys enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
+    [self removeReferenceToKey:key forID:ID];
+  }];
+}
+
+- (void)removeReferencesForID:(int)ID {
+  FSTDocumentKey *emptyKey = [FSTDocumentKey keyWithSegments:@[]];
+  FSTDocumentReference *start = [[FSTDocumentReference alloc] initWithKey:emptyKey ID:ID];
+  FSTDocumentReference *end = [[FSTDocumentReference alloc] initWithKey:emptyKey ID:(ID + 1)];
+
+  [self.referencesByID enumerateObjectsFrom:start
+                                         to:end
+                                 usingBlock:^(FSTDocumentReference *reference, BOOL *stop) {
+                                   [self removeReference:reference];
+                                 }];
+}
+
+- (void)removeAllReferences {
+  for (FSTDocumentReference *reference in self.referencesByKey.objectEnumerator) {
+    [self removeReference:reference];
+  }
+}
+
+- (void)removeReference:(FSTDocumentReference *)reference {
+  self.referencesByKey = [self.referencesByKey setByRemovingObject:reference];
+  self.referencesByID = [self.referencesByID setByRemovingObject:reference];
+  [self.garbageCollector addPotentialGarbageKey:reference.key];
+}
+
+- (FSTDocumentKeySet *)referencedKeysForID:(int)ID {
+  FSTDocumentKey *emptyKey = [FSTDocumentKey keyWithSegments:@[]];
+  FSTDocumentReference *start = [[FSTDocumentReference alloc] initWithKey:emptyKey ID:ID];
+  FSTDocumentReference *end = [[FSTDocumentReference alloc] initWithKey:emptyKey ID:(ID + 1)];
+
+  __block FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
+  [self.referencesByID enumerateObjectsFrom:start
+                                         to:end
+                                 usingBlock:^(FSTDocumentReference *reference, BOOL *stop) {
+                                   keys = [keys setByAddingObject:reference.key];
+                                 }];
+  return keys;
+}
+
+- (BOOL)containsKey:(FSTDocumentKey *)key {
+  // Create a reference with a zero ID as the start position to find any document reference with
+  // this key.
+  FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:0];
+
+  NSEnumerator<FSTDocumentReference *> *enumerator =
+      [self.referencesByKey objectEnumeratorFrom:reference];
+  FSTDocumentKey *_Nullable firstKey = [enumerator nextObject].key;
+  return [firstKey isEqual:reference.key];
+}
+
+@end
+
+NS_ASSUME_NONNULL_END