@@ -62,21 +62,32 @@ - (instancetype)initWithRemoteDocumentCache:(id<FSTRemoteDocumentCache>)remoteDo
62
62
}
63
63
64
64
- (nullable FSTMaybeDocument *)documentForKey : (const DocumentKey &)key {
65
+ NSArray <FSTMutationBatch *> *batches =
66
+ [self .mutationQueue allMutationBatchesAffectingDocumentKey: key];
67
+ return [self documentForKey: key inBatches: batches];
68
+ }
69
+
70
+ // Internal version of documentForKey: which allows reusing `batches`.
71
+ - (nullable FSTMaybeDocument *)documentForKey : (const DocumentKey &)key
72
+ inBatches : (NSArray <FSTMutationBatch *> *)batches {
65
73
FSTMaybeDocument *_Nullable remoteDoc = [self .remoteDocumentCache entryForKey: key];
66
- return [self localDocument: remoteDoc key: key];
74
+ return [self localDocument: remoteDoc key: key inBatches: batches ];
67
75
}
68
76
69
77
- (FSTMaybeDocumentDictionary *)documentsForKeys : (const DocumentKeySet &)keys {
70
78
FSTMaybeDocumentDictionary *results = [FSTMaybeDocumentDictionary maybeDocumentDictionary ];
79
+ NSArray <FSTMutationBatch *> *batches =
80
+ [self .mutationQueue allMutationBatchesAffectingDocumentKeys: keys];
71
81
for (const DocumentKey &key : keys) {
72
82
// TODO(mikelehen): PERF: Consider fetching all remote documents at once rather than one-by-one.
73
- FSTMaybeDocument *maybeDoc = [self documentForKey: key];
83
+ FSTMaybeDocument *maybeDoc = [self documentForKey: key inBatches: batches ];
74
84
// TODO(http://b/32275378): Don't conflate missing / deleted.
75
85
if (!maybeDoc) {
76
86
maybeDoc = [FSTDeletedDocument documentWithKey: key version:SnapshotVersion: :None ()];
77
87
}
78
88
results = [results dictionaryBySettingObject: maybeDoc forKey: key];
79
89
}
90
+
80
91
return results;
81
92
}
82
93
@@ -122,12 +133,13 @@ - (FSTDocumentDictionary *)documentsMatchingCollectionQuery:(FSTQuery *)query {
122
133
}
123
134
124
135
// Now add in results for the matchingKeys.
125
- for (const DocumentKey &key : matchingKeys) {
126
- FSTMaybeDocument *doc = [self documentForKey: key];
127
- if ([doc isKindOfClass: [FSTDocument class ]]) {
128
- results = [results dictionaryBySettingObject: (FSTDocument *)doc forKey: key];
129
- }
130
- }
136
+ FSTMaybeDocumentDictionary *matchingKeysDocs = [self documentsForKeys: matchingKeys];
137
+ [matchingKeysDocs
138
+ enumerateKeysAndObjectsUsingBlock: ^(FSTDocumentKey *key, FSTMaybeDocument *doc, BOOL *stop) {
139
+ if ([doc isKindOfClass: [FSTDocument class ]]) {
140
+ results = [results dictionaryBySettingObject: (FSTDocument *)doc forKey: key];
141
+ }
142
+ }];
131
143
132
144
// Finally, filter out any documents that don't actually match the query. Note that the extra
133
145
// reference here prevents ARC from deallocating the initial unfiltered results while we're
@@ -151,9 +163,8 @@ - (FSTDocumentDictionary *)documentsMatchingCollectionQuery:(FSTQuery *)query {
151
163
* @param documentKey The key of the document (necessary when remoteDocument is nil).
152
164
*/
153
165
- (nullable FSTMaybeDocument *)localDocument : (nullable FSTMaybeDocument *)document
154
- key : (const DocumentKey &)documentKey {
155
- NSArray <FSTMutationBatch *> *batches =
156
- [self .mutationQueue allMutationBatchesAffectingDocumentKey: documentKey];
166
+ key : (const DocumentKey &)documentKey
167
+ inBatches : (NSArray <FSTMutationBatch *> *)batches {
157
168
for (FSTMutationBatch *batch in batches) {
158
169
document = [batch applyTo: document documentKey: documentKey];
159
170
}
@@ -169,10 +180,18 @@ - (nullable FSTMaybeDocument *)localDocument:(nullable FSTMaybeDocument *)docume
169
180
* @return The local view of the documents.
170
181
*/
171
182
- (FSTDocumentDictionary *)localDocuments : (FSTDocumentDictionary *)documents {
183
+ __block DocumentKeySet keySet;
184
+ [documents
185
+ enumerateKeysAndObjectsUsingBlock: ^(FSTDocumentKey *key, FSTDocument *doc, BOOL *stop) {
186
+ keySet = keySet.insert (doc.key );
187
+ }];
188
+ NSArray <FSTMutationBatch *> *batches =
189
+ [self .mutationQueue allMutationBatchesAffectingDocumentKeys: keySet];
190
+
172
191
__block FSTDocumentDictionary *result = documents;
173
192
[documents enumerateKeysAndObjectsUsingBlock: ^(FSTDocumentKey *key, FSTDocument *remoteDocument,
174
193
BOOL *stop) {
175
- FSTMaybeDocument *mutatedDoc = [self localDocument: remoteDocument key: key];
194
+ FSTMaybeDocument *mutatedDoc = [self localDocument: remoteDocument key: key inBatches: batches ];
176
195
if ([mutatedDoc isKindOfClass: [FSTDeletedDocument class ]]) {
177
196
result = [result dictionaryByRemovingObjectForKey: key];
178
197
} else if ([mutatedDoc isKindOfClass: [FSTDocument class ]]) {
0 commit comments