1 files changed, 11 insertions, 9 deletions
diff --git a/Firestore/Source/Local/FSTReferenceSet.mm b/Firestore/Source/Local/FSTReferenceSet.mm
index 14f5d47..6b34725 100644
--- a/Firestore/Source/Local/FSTReferenceSet.mm
+++ b/Firestore/Source/Local/FSTReferenceSet.mm
@@ -17,10 +17,12 @@
 #import "Firestore/Source/Local/FSTReferenceSet.h"
 
 #import "Firestore/Source/Local/FSTDocumentReference.h"
+#import "Firestore/third_party/Immutable/FSTImmutableSortedSet.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -68,20 +70,20 @@ NS_ASSUME_NONNULL_BEGIN
   self.referencesByID = [self.referencesByID setByAddingObject:reference];
 }
 
-- (void)addReferencesToKeys:(FSTDocumentKeySet *)keys forID:(int)ID {
-  [keys enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
+- (void)addReferencesToKeys:(const DocumentKeySet &)keys forID:(int)ID {
+  for (const DocumentKey &key : keys) {
     [self addReferenceToKey:key forID:ID];
-  }];
+  }
 }
 
 - (void)removeReferenceToKey:(const DocumentKey &)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) {
+- (void)removeReferencesToKeys:(const DocumentKeySet &)keys forID:(int)ID {
+  for (const DocumentKey &key : keys) {
     [self removeReferenceToKey:key forID:ID];
-  }];
+  }
 }
 
 - (void)removeReferencesForID:(int)ID {
@@ -109,17 +111,17 @@ NS_ASSUME_NONNULL_BEGIN
   [self.garbageCollector addPotentialGarbageKey:reference.key];
 }
 
-- (FSTDocumentKeySet *)referencedKeysForID:(int)ID {
+- (DocumentKeySet)referencedKeysForID:(int)ID {
   FSTDocumentReference *start =
       [[FSTDocumentReference alloc] initWithKey:DocumentKey::Empty() ID:ID];
   FSTDocumentReference *end =
       [[FSTDocumentReference alloc] initWithKey:DocumentKey::Empty() ID:(ID + 1)];
 
-  __block FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
+  __block DocumentKeySet keys;
   [self.referencesByID enumerateObjectsFrom:start
                                          to:end
                                  usingBlock:^(FSTDocumentReference *reference, BOOL *stop) {
-                                   keys = [keys setByAddingObject:reference.key];
+                                   keys = keys.insert(reference.key);
                                  }];
   return keys;
 }