4 files changed, 6 insertions, 71 deletions
diff --git a/Firestore/Source/Model/FSTDocumentKeySet.h b/Firestore/Source/Model/FSTDocumentKeySet.h
deleted file mode 100644
index 80f6624..0000000
--- a/Firestore/Source/Model/FSTDocumentKeySet.h
+++ /dev/null
@@ -1,35 +0,0 @@
-/*
- * 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 <Foundation/Foundation.h>
-
-#import "Firestore/third_party/Immutable/FSTImmutableSortedSet.h"
-
-@class FSTDocumentKey;
-
-NS_ASSUME_NONNULL_BEGIN
-
-/** Convenience type for a set of keys, since they are so common. */
-typedef FSTImmutableSortedSet<FSTDocumentKey *> FSTDocumentKeySet;
-
-@interface FSTImmutableSortedSet (FSTDocumentKey)
-
-/** Returns a new set using the DocumentKeyComparator. */
-+ (FSTDocumentKeySet *)keySet;
-
-@end
-
-NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Model/FSTDocumentKeySet.mm b/Firestore/Source/Model/FSTDocumentKeySet.mm
deleted file mode 100644
index f07b785..0000000
--- a/Firestore/Source/Model/FSTDocumentKeySet.mm
+++ /dev/null
@@ -1,31 +0,0 @@
-/*
- * 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/FSTDocumentKeySet.h"
-
-#import "Firestore/Source/Model/FSTDocumentKey.h"
-
-NS_ASSUME_NONNULL_BEGIN
-
-@implementation FSTImmutableSortedSet (FSTDocumentKey)
-
-+ (instancetype)keySet {
-  return [FSTDocumentKeySet setWithComparator:FSTDocumentKeyComparator];
-}
-
-@end
-
-NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Model/FSTMutationBatch.h b/Firestore/Source/Model/FSTMutationBatch.h
index d3d68a9..f95b145 100644
--- a/Firestore/Source/Model/FSTMutationBatch.h
+++ b/Firestore/Source/Model/FSTMutationBatch.h
@@ -17,10 +17,10 @@
 #import <Foundation/Foundation.h>
 
 #import "Firestore/Source/Core/FSTTypes.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
 #import "Firestore/Source/Model/FSTDocumentVersionDictionary.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 @class FSTMutation;
@@ -85,7 +85,7 @@ extern const FSTBatchID kFSTBatchIDUnknown;
 - (FSTMutationBatch *)toTombstone;
 
 /** Returns the set of unique keys referenced by all mutations in the batch. */
-- (FSTDocumentKeySet *)keys;
+- (firebase::firestore::model::DocumentKeySet)keys;
 
 @property(nonatomic, assign, readonly) FSTBatchID batchID;
 @property(nonatomic, strong, readonly) FIRTimestamp *localWriteTime;
diff --git a/Firestore/Source/Model/FSTMutationBatch.mm b/Firestore/Source/Model/FSTMutationBatch.mm
index 3c03bdc..9926bd4 100644
--- a/Firestore/Source/Model/FSTMutationBatch.mm
+++ b/Firestore/Source/Model/FSTMutationBatch.mm
@@ -28,6 +28,7 @@
 
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -115,10 +116,10 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
 }
 
 // TODO(klimt): This could use NSMutableDictionary instead.
-- (FSTDocumentKeySet *)keys {
-  FSTDocumentKeySet *set = [FSTDocumentKeySet keySet];
+- (DocumentKeySet)keys {
+  DocumentKeySet set;
   for (FSTMutation *mutation in self.mutations) {
-    set = [set setByAddingObject:mutation.key];
+    set = set.insert(mutation.key);
   }
   return set;
 }