7 files changed, 23 insertions, 31 deletions
diff --git a/Firestore/Source/Local/FSTLevelDB.h b/Firestore/Source/Local/FSTLevelDB.h
index 95b80a6..a56c133 100644
--- a/Firestore/Source/Local/FSTLevelDB.h
+++ b/Firestore/Source/Local/FSTLevelDB.h
@@ -95,7 +95,7 @@ NS_ASSUME_NONNULL_BEGIN
 + (NSString *)descriptionOfStatus:(leveldb::Status)status;
 
 /** The native db pointer, allocated during start. */
-@property(nonatomic, assign, readonly) std::shared_ptr<leveldb::DB> ptr;
+@property(nonatomic, assign, readonly) leveldb::DB *ptr;
 
 @property(nonatomic, readonly) firebase::firestore::local::LevelDbTransaction *currentTransaction;
 
diff --git a/Firestore/Source/Local/FSTLevelDB.mm b/Firestore/Source/Local/FSTLevelDB.mm
index 9f75a3e..9dc50a2 100644
--- a/Firestore/Source/Local/FSTLevelDB.mm
+++ b/Firestore/Source/Local/FSTLevelDB.mm
@@ -17,6 +17,7 @@
 #import "Firestore/Source/Local/FSTLevelDB.h"
 
 #include <memory>
+#include <utility>
 
 #import "FIRFirestoreErrors.h"
 #import "Firestore/Source/Local/FSTLevelDBMigrations.h"
@@ -60,6 +61,7 @@ using leveldb::WriteOptions;
 
 @implementation FSTLevelDB {
   std::unique_ptr<LevelDbTransaction> _transaction;
+  std::unique_ptr<leveldb::DB> _ptr;
   FSTTransactionRunner _transactionRunner;
 }
 
@@ -82,6 +84,10 @@ using leveldb::WriteOptions;
   return self;
 }
 
+- (leveldb::DB *)ptr {
+  return _ptr.get();
+}
+
 - (const FSTTransactionRunner &)run {
   return _transactionRunner;
 }
diff --git a/Firestore/Source/Local/FSTLevelDBMutationQueue.h b/Firestore/Source/Local/FSTLevelDBMutationQueue.h
index 034738f..911fa37 100644
--- a/Firestore/Source/Local/FSTLevelDBMutationQueue.h
+++ b/Firestore/Source/Local/FSTLevelDBMutationQueue.h
@@ -51,7 +51,7 @@ NS_ASSUME_NONNULL_BEGIN
  * Returns one larger than the largest batch ID that has been stored. If there are no mutations
  * returns 0. Note that batch IDs are global.
  */
-+ (FSTBatchID)loadNextBatchIDFromDB:(std::shared_ptr<leveldb::DB>)db;
++ (FSTBatchID)loadNextBatchIDFromDB:(leveldb::DB *)db;
 
 @end
 
diff --git a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
index 3b4687c..94ab8a5 100644
--- a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
+++ b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
@@ -132,7 +132,7 @@ using leveldb::WriteOptions;
   self.metadata = metadata;
 }
 
-+ (FSTBatchID)loadNextBatchIDFromDB:(std::shared_ptr<DB>)db {
++ (FSTBatchID)loadNextBatchIDFromDB:(DB *)db {
   // TODO(gsoltis): implement Prev() and SeekToLast() on LevelDbTransaction::Iterator, then port
   // this to a transaction.
   std::unique_ptr<Iterator> it(db->NewIterator(LevelDbTransaction::DefaultReadOptions()));
diff --git a/Firestore/Source/Local/FSTLevelDBQueryCache.h b/Firestore/Source/Local/FSTLevelDBQueryCache.h
index 2cd6758..f756d9e 100644
--- a/Firestore/Source/Local/FSTLevelDBQueryCache.h
+++ b/Firestore/Source/Local/FSTLevelDBQueryCache.h
@@ -36,7 +36,7 @@ NS_ASSUME_NONNULL_BEGIN
  * Retrieves the global singleton metadata row from the given database, if it exists.
  * TODO(gsoltis): remove this method once fully ported to transactions.
  */
-+ (nullable FSTPBTargetGlobal *)readTargetMetadataFromDB:(std::shared_ptr<leveldb::DB>)db;
++ (nullable FSTPBTargetGlobal *)readTargetMetadataFromDB:(leveldb::DB *)db;
 
 /**
  * Retrieves the global singleton metadata row using the given transaction, if it exists.
diff --git a/Firestore/Source/Local/FSTLevelDBQueryCache.mm b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
index f28370a..31c2a2e 100644
--- a/Firestore/Source/Local/FSTLevelDBQueryCache.mm
+++ b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
@@ -85,7 +85,7 @@ using firebase::firestore::model::DocumentKeySet;
   return proto;
 }
 
-+ (nullable FSTPBTargetGlobal *)readTargetMetadataFromDB:(std::shared_ptr<DB>)db {
++ (nullable FSTPBTargetGlobal *)readTargetMetadataFromDB:(DB *)db {
   std::string key = [FSTLevelDBTargetGlobalKey key];
   std::string value;
   Status status = db->Get([FSTLevelDB standardReadOptions], key, &value);
diff --git a/Firestore/Source/Local/FSTLocalStore.mm b/Firestore/Source/Local/FSTLocalStore.mm
index 7469c71..6aab78e 100644
--- a/Firestore/Source/Local/FSTLocalStore.mm
+++ b/Firestore/Source/Local/FSTLocalStore.mm
@@ -268,46 +268,32 @@ NS_ASSUME_NONNULL_BEGIN
     FSTListenSequenceNumber sequenceNumber = [self.listenSequence next];
     id<FSTQueryCache> queryCache = self.queryCache;
 
-    [remoteEvent.targetChanges enumerateKeysAndObjectsUsingBlock:^(
-                                   NSNumber *targetIDNumber, FSTTargetChange *change, BOOL *stop) {
-      FSTTargetID targetID = targetIDNumber.intValue;
+    for (const auto &entry : remoteEvent.targetChanges) {
+      FSTTargetID targetID = entry.first;
+      FSTBoxedTargetID *boxedTargetID = @(targetID);
+      FSTTargetChange *change = entry.second;
 
       // Do not ref/unref unassigned targetIDs - it may lead to leaks.
-      FSTQueryData *queryData = self.targetIDs[targetIDNumber];
+      FSTQueryData *queryData = self.targetIDs[boxedTargetID];
       if (!queryData) {
-        return;
+        continue;
       }
 
+      [queryCache removeMatchingKeys:change.removedDocuments forTargetID:targetID];
+      [queryCache addMatchingKeys:change.addedDocuments forTargetID:targetID];
+
       // Update the resume token if the change includes one. Don't clear any preexisting value.
       // Bump the sequence number as well, so that documents being removed now are ordered later
       // than documents that were previously removed from this target.
       NSData *resumeToken = change.resumeToken;
       if (resumeToken.length > 0) {
-        queryData = [queryData queryDataByReplacingSnapshotVersion:change.snapshotVersion
+        queryData = [queryData queryDataByReplacingSnapshotVersion:remoteEvent.snapshotVersion
                                                        resumeToken:resumeToken
                                                     sequenceNumber:sequenceNumber];
-        self.targetIDs[targetIDNumber] = queryData;
+        self.targetIDs[boxedTargetID] = queryData;
         [self.queryCache updateQueryData:queryData];
       }
-
-      FSTTargetMapping *mapping = change.mapping;
-      if (mapping) {
-        // First make sure that all references are deleted.
-        if ([mapping isKindOfClass:[FSTResetMapping class]]) {
-          FSTResetMapping *reset = (FSTResetMapping *)mapping;
-          [queryCache removeMatchingKeysForTargetID:targetID];
-          [queryCache addMatchingKeys:reset.documents forTargetID:targetID];
-
-        } else if ([mapping isKindOfClass:[FSTUpdateMapping class]]) {
-          FSTUpdateMapping *update = (FSTUpdateMapping *)mapping;
-          [queryCache removeMatchingKeys:update.removedDocuments forTargetID:targetID];
-          [queryCache addMatchingKeys:update.addedDocuments forTargetID:targetID];
-
-        } else {
-          HARD_FAIL("Unknown mapping type: %s", mapping);
-        }
-      }
-    }];
+    }
 
     // TODO(klimt): This could probably be an NSMutableDictionary.
     DocumentKeySet changedDocKeys;