Skip to content
This repository was archived by the owner on Nov 15, 2023. It is now read-only.

Commit 6d43761

Browse files
bkchrbkchr
authored
Optimize next_storage_key (#8956)
* Optimize `next_storage_key` - Do not rely on recursion - Use an iterator over the overlay to not always call the same method * Fix bug
1 parent fc29e14 commit 6d43761

File tree

3 files changed

+136
-68
lines changed

3 files changed

+136
-68
lines changed

‎primitives/state-machine/src/ext.rs

Lines changed: 92 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ use sp_externalities::{
3232
};
3333
use codec::{Decode, Encode, EncodeAppend};
3434

35-
use sp_std::{fmt, any::{Any, TypeId}, vec::Vec, vec, boxed::Box};
35+
use sp_std::{fmt, any::{Any, TypeId}, vec::Vec, vec, boxed::Box, cmp::Ordering};
3636
use crate::{warn, trace, log_error};
3737
#[cfg(feature = "std")]
3838
use crate::changes_trie::State as ChangesTrieState;
@@ -323,16 +323,37 @@ where
323323
}
324324

325325
fn next_storage_key(&self, key: &[u8]) -> Option<StorageKey> {
326-
let next_backend_key = self.backend.next_storage_key(key).expect(EXT_NOT_ALLOWED_TO_FAIL);
327-
let next_overlay_key_change = self.overlay.next_storage_key_change(key);
328-
329-
match (next_backend_key, next_overlay_key_change) {
330-
(Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key),
331-
(backend_key, None) => backend_key,
332-
(_, Some(overlay_key)) => if overlay_key.1.value().is_some() {
333-
Some(overlay_key.0.to_vec())
334-
} else {
335-
self.next_storage_key(&overlay_key.0[..])
326+
let mut next_backend_key = self.backend.next_storage_key(key).expect(EXT_NOT_ALLOWED_TO_FAIL);
327+
let mut overlay_changes = self.overlay.iter_after(key).peekable();
328+
329+
match (&next_backend_key, overlay_changes.peek()) {
330+
(_, None) => next_backend_key,
331+
(Some(_), Some(_)) => {
332+
while let Some(overlay_key) = overlay_changes.next() {
333+
let cmp = next_backend_key.as_deref().map(|v| v.cmp(&overlay_key.0));
334+
335+
// If `backend_key` is less than the `overlay_key`, we found out next key.
336+
if cmp == Some(Ordering::Less) {
337+
return next_backend_key
338+
} else if overlay_key.1.value().is_some() {
339+
// If there exists a value for the `overlay_key` in the overlay
340+
// (aka the key is still valid), it means we have found our next key.
341+
return Some(overlay_key.0.to_vec())
342+
} else if cmp == Some(Ordering::Equal) {
343+
// If the `backend_key` and `overlay_key` are equal, it means that we need
344+
// to search for the next backend key, because the overlay has overwritten
345+
// this key.
346+
next_backend_key = self.backend.next_storage_key(
347+
&overlay_key.0,
348+
).expect(EXT_NOT_ALLOWED_TO_FAIL);
349+
}
350+
}
351+
352+
next_backend_key
353+
},
354+
(None, Some(_)) => {
355+
// Find the next overlay key that has a value attached.
356+
overlay_changes.find_map(|k| k.1.value().as_ref().map(|_| k.0.to_vec()))
336357
},
337358
}
338359
}
@@ -342,24 +363,43 @@ where
342363
child_info: &ChildInfo,
343364
key: &[u8],
344365
) -> Option<StorageKey> {
345-
let next_backend_key = self.backend
366+
let mut next_backend_key = self.backend
346367
.next_child_storage_key(child_info, key)
347368
.expect(EXT_NOT_ALLOWED_TO_FAIL);
348-
let next_overlay_key_change = self.overlay.next_child_storage_key_change(
369+
let mut overlay_changes = self.overlay.child_iter_after(
349370
child_info.storage_key(),
350371
key
351-
);
372+
).peekable();
373+
374+
match (&next_backend_key, overlay_changes.peek()) {
375+
(_, None) => next_backend_key,
376+
(Some(_), Some(_)) => {
377+
while let Some(overlay_key) = overlay_changes.next() {
378+
let cmp = next_backend_key.as_deref().map(|v| v.cmp(&overlay_key.0));
379+
380+
// If `backend_key` is less than the `overlay_key`, we found out next key.
381+
if cmp == Some(Ordering::Less) {
382+
return next_backend_key
383+
} else if overlay_key.1.value().is_some() {
384+
// If there exists a value for the `overlay_key` in the overlay
385+
// (aka the key is still valid), it means we have found our next key.
386+
return Some(overlay_key.0.to_vec())
387+
} else if cmp == Some(Ordering::Equal) {
388+
// If the `backend_key` and `overlay_key` are equal, it means that we need
389+
// to search for the next backend key, because the overlay has overwritten
390+
// this key.
391+
next_backend_key = self.backend.next_child_storage_key(
392+
child_info,
393+
&overlay_key.0,
394+
).expect(EXT_NOT_ALLOWED_TO_FAIL);
395+
}
396+
}
352397

353-
match (next_backend_key, next_overlay_key_change) {
354-
(Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key),
355-
(backend_key, None) => backend_key,
356-
(_, Some(overlay_key)) => if overlay_key.1.value().is_some() {
357-
Some(overlay_key.0.to_vec())
358-
} else {
359-
self.next_child_storage_key(
360-
child_info,
361-
&overlay_key.0[..],
362-
)
398+
next_backend_key
399+
},
400+
(None, Some(_)) => {
401+
// Find the next overlay key that has a value attached.
402+
overlay_changes.find_map(|k| k.1.value().as_ref().map(|_| k.0.to_vec()))
363403
},
364404
}
365405
}
@@ -971,6 +1011,34 @@ mod tests {
9711011
assert_eq!(ext.next_storage_key(&[40]), Some(vec![50]));
9721012
}
9731013

1014+
#[test]
1015+
fn next_storage_key_works_with_a_lot_empty_values_in_overlay() {
1016+
let mut cache = StorageTransactionCache::default();
1017+
let mut overlay = OverlayedChanges::default();
1018+
overlay.set_storage(vec![20], None);
1019+
overlay.set_storage(vec![21], None);
1020+
overlay.set_storage(vec![22], None);
1021+
overlay.set_storage(vec![23], None);
1022+
overlay.set_storage(vec![24], None);
1023+
overlay.set_storage(vec![25], None);
1024+
overlay.set_storage(vec![26], None);
1025+
overlay.set_storage(vec![27], None);
1026+
overlay.set_storage(vec![28], None);
1027+
overlay.set_storage(vec![29], None);
1028+
let backend = Storage {
1029+
top: map![
1030+
vec![30] => vec![30]
1031+
],
1032+
children_default: map![]
1033+
}.into();
1034+
1035+
let ext = TestExt::new(&mut overlay, &mut cache, &backend, None, None);
1036+
1037+
assert_eq!(ext.next_storage_key(&[5]), Some(vec![30]));
1038+
1039+
drop(ext);
1040+
}
1041+
9741042
#[test]
9751043
fn next_child_storage_key_works() {
9761044
let child_info = ChildInfo::new_default(b"Child1");

‎primitives/state-machine/src/overlayed_changes/changeset.rs

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -426,11 +426,11 @@ impl OverlayedChangeSet {
426426
}
427427
}
428428

429-
/// Get the change that is next to the supplied key.
430-
pub fn next_change(&self, key: &[u8]) -> Option<(&[u8], &OverlayedValue)> {
429+
/// Get the iterator over all changes that follow the supplied `key`.
430+
pub fn changes_after(&self, key: &[u8]) -> impl Iterator<Item = (&[u8], &OverlayedValue)> {
431431
use sp_std::ops::Bound;
432432
let range = (Bound::Excluded(key), Bound::Unbounded);
433-
self.changes.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v))
433+
self.changes.range::<[u8], _>(range).map(|(k, v)| (k.as_slice(), v))
434434
}
435435
}
436436

@@ -707,29 +707,29 @@ mod test {
707707
changeset.set(b"key4".to_vec(), Some(b"val4".to_vec()), Some(4));
708708
changeset.set(b"key11".to_vec(), Some(b"val11".to_vec()), Some(11));
709709

710-
assert_eq!(changeset.next_change(b"key0").unwrap().0, b"key1");
711-
assert_eq!(changeset.next_change(b"key0").unwrap().1.value(), Some(&b"val1".to_vec()));
712-
assert_eq!(changeset.next_change(b"key1").unwrap().0, b"key11");
713-
assert_eq!(changeset.next_change(b"key1").unwrap().1.value(), Some(&b"val11".to_vec()));
714-
assert_eq!(changeset.next_change(b"key11").unwrap().0, b"key2");
715-
assert_eq!(changeset.next_change(b"key11").unwrap().1.value(), Some(&b"val2".to_vec()));
716-
assert_eq!(changeset.next_change(b"key2").unwrap().0, b"key3");
717-
assert_eq!(changeset.next_change(b"key2").unwrap().1.value(), Some(&b"val3".to_vec()));
718-
assert_eq!(changeset.next_change(b"key3").unwrap().0, b"key4");
719-
assert_eq!(changeset.next_change(b"key3").unwrap().1.value(), Some(&b"val4".to_vec()));
720-
assert_eq!(changeset.next_change(b"key4"), None);
710+
assert_eq!(changeset.changes_after(b"key0").next().unwrap().0, b"key1");
711+
assert_eq!(changeset.changes_after(b"key0").next().unwrap().1.value(), Some(&b"val1".to_vec()));
712+
assert_eq!(changeset.changes_after(b"key1").next().unwrap().0, b"key11");
713+
assert_eq!(changeset.changes_after(b"key1").next().unwrap().1.value(), Some(&b"val11".to_vec()));
714+
assert_eq!(changeset.changes_after(b"key11").next().unwrap().0, b"key2");
715+
assert_eq!(changeset.changes_after(b"key11").next().unwrap().1.value(), Some(&b"val2".to_vec()));
716+
assert_eq!(changeset.changes_after(b"key2").next().unwrap().0, b"key3");
717+
assert_eq!(changeset.changes_after(b"key2").next().unwrap().1.value(), Some(&b"val3".to_vec()));
718+
assert_eq!(changeset.changes_after(b"key3").next().unwrap().0, b"key4");
719+
assert_eq!(changeset.changes_after(b"key3").next().unwrap().1.value(), Some(&b"val4".to_vec()));
720+
assert_eq!(changeset.changes_after(b"key4").next(), None);
721721

722722
changeset.rollback_transaction().unwrap();
723723

724-
assert_eq!(changeset.next_change(b"key0").unwrap().0, b"key1");
725-
assert_eq!(changeset.next_change(b"key0").unwrap().1.value(), Some(&b"val1".to_vec()));
726-
assert_eq!(changeset.next_change(b"key1").unwrap().0, b"key2");
727-
assert_eq!(changeset.next_change(b"key1").unwrap().1.value(), Some(&b"val2".to_vec()));
728-
assert_eq!(changeset.next_change(b"key11").unwrap().0, b"key2");
729-
assert_eq!(changeset.next_change(b"key11").unwrap().1.value(), Some(&b"val2".to_vec()));
730-
assert_eq!(changeset.next_change(b"key2"), None);
731-
assert_eq!(changeset.next_change(b"key3"), None);
732-
assert_eq!(changeset.next_change(b"key4"), None);
724+
assert_eq!(changeset.changes_after(b"key0").next().unwrap().0, b"key1");
725+
assert_eq!(changeset.changes_after(b"key0").next().unwrap().1.value(), Some(&b"val1".to_vec()));
726+
assert_eq!(changeset.changes_after(b"key1").next().unwrap().0, b"key2");
727+
assert_eq!(changeset.changes_after(b"key1").next().unwrap().1.value(), Some(&b"val2".to_vec()));
728+
assert_eq!(changeset.changes_after(b"key11").next().unwrap().0, b"key2");
729+
assert_eq!(changeset.changes_after(b"key11").next().unwrap().1.value(), Some(&b"val2".to_vec()));
730+
assert_eq!(changeset.changes_after(b"key2").next(), None);
731+
assert_eq!(changeset.changes_after(b"key3").next(), None);
732+
assert_eq!(changeset.changes_after(b"key4").next(), None);
733733

734734
}
735735

‎primitives/state-machine/src/overlayed_changes/mod.rs

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -669,24 +669,24 @@ impl OverlayedChanges {
669669
})
670670
}
671671

672-
/// Returns the next (in lexicographic order) storage key in the overlayed alongside its value.
673-
/// If no value is next then `None` is returned.
674-
pub fn next_storage_key_change(&self, key: &[u8]) -> Option<(&[u8], &OverlayedValue)> {
675-
self.top.next_change(key)
672+
/// Returns an iterator over the keys (in lexicographic order) following `key` (excluding `key`)
673+
/// alongside its value.
674+
pub fn iter_after(&self, key: &[u8]) -> impl Iterator<Item = (&[u8], &OverlayedValue)> {
675+
self.top.changes_after(key)
676676
}
677677

678-
/// Returns the next (in lexicographic order) child storage key in the overlayed alongside its
679-
/// value. If no value is next then `None` is returned.
680-
pub fn next_child_storage_key_change(
678+
/// Returns an iterator over the keys (in lexicographic order) following `key` (excluding `key`)
679+
/// alongside its value for the given `storage_key` child.
680+
pub fn child_iter_after(
681681
&self,
682682
storage_key: &[u8],
683683
key: &[u8]
684-
) -> Option<(&[u8], &OverlayedValue)> {
684+
) -> impl Iterator<Item = (&[u8], &OverlayedValue)> {
685685
self.children
686686
.get(storage_key)
687-
.and_then(|(overlay, _)|
688-
overlay.next_change(key)
689-
)
687+
.map(|(overlay, _)| overlay.changes_after(key))
688+
.into_iter()
689+
.flatten()
690690
}
691691

692692
/// Read only access ot offchain overlay.
@@ -988,28 +988,28 @@ mod tests {
988988
overlay.set_storage(vec![30], None);
989989

990990
// next_prospective < next_committed
991-
let next_to_5 = overlay.next_storage_key_change(&[5]).unwrap();
991+
let next_to_5 = overlay.iter_after(&[5]).next().unwrap();
992992
assert_eq!(next_to_5.0.to_vec(), vec![10]);
993993
assert_eq!(next_to_5.1.value(), Some(&vec![10]));
994994

995995
// next_committed < next_prospective
996-
let next_to_10 = overlay.next_storage_key_change(&[10]).unwrap();
996+
let next_to_10 = overlay.iter_after(&[10]).next().unwrap();
997997
assert_eq!(next_to_10.0.to_vec(), vec![20]);
998998
assert_eq!(next_to_10.1.value(), Some(&vec![20]));
999999

10001000
// next_committed == next_prospective
1001-
let next_to_20 = overlay.next_storage_key_change(&[20]).unwrap();
1001+
let next_to_20 = overlay.iter_after(&[20]).next().unwrap();
10021002
assert_eq!(next_to_20.0.to_vec(), vec![30]);
10031003
assert_eq!(next_to_20.1.value(), None);
10041004

10051005
// next_committed, no next_prospective
1006-
let next_to_30 = overlay.next_storage_key_change(&[30]).unwrap();
1006+
let next_to_30 = overlay.iter_after(&[30]).next().unwrap();
10071007
assert_eq!(next_to_30.0.to_vec(), vec![40]);
10081008
assert_eq!(next_to_30.1.value(), Some(&vec![40]));
10091009

10101010
overlay.set_storage(vec![50], Some(vec![50]));
10111011
// next_prospective, no next_committed
1012-
let next_to_40 = overlay.next_storage_key_change(&[40]).unwrap();
1012+
let next_to_40 = overlay.iter_after(&[40]).next().unwrap();
10131013
assert_eq!(next_to_40.0.to_vec(), vec![50]);
10141014
assert_eq!(next_to_40.1.value(), Some(&vec![50]));
10151015
}
@@ -1029,28 +1029,28 @@ mod tests {
10291029
overlay.set_child_storage(child_info, vec![30], None);
10301030

10311031
// next_prospective < next_committed
1032-
let next_to_5 = overlay.next_child_storage_key_change(child, &[5]).unwrap();
1032+
let next_to_5 = overlay.child_iter_after(child, &[5]).next().unwrap();
10331033
assert_eq!(next_to_5.0.to_vec(), vec![10]);
10341034
assert_eq!(next_to_5.1.value(), Some(&vec![10]));
10351035

10361036
// next_committed < next_prospective
1037-
let next_to_10 = overlay.next_child_storage_key_change(child, &[10]).unwrap();
1037+
let next_to_10 = overlay.child_iter_after(child, &[10]).next().unwrap();
10381038
assert_eq!(next_to_10.0.to_vec(), vec![20]);
10391039
assert_eq!(next_to_10.1.value(), Some(&vec![20]));
10401040

10411041
// next_committed == next_prospective
1042-
let next_to_20 = overlay.next_child_storage_key_change(child, &[20]).unwrap();
1042+
let next_to_20 = overlay.child_iter_after(child, &[20]).next().unwrap();
10431043
assert_eq!(next_to_20.0.to_vec(), vec![30]);
10441044
assert_eq!(next_to_20.1.value(), None);
10451045

10461046
// next_committed, no next_prospective
1047-
let next_to_30 = overlay.next_child_storage_key_change(child, &[30]).unwrap();
1047+
let next_to_30 = overlay.child_iter_after(child, &[30]).next().unwrap();
10481048
assert_eq!(next_to_30.0.to_vec(), vec![40]);
10491049
assert_eq!(next_to_30.1.value(), Some(&vec![40]));
10501050

10511051
overlay.set_child_storage(child_info, vec![50], Some(vec![50]));
10521052
// next_prospective, no next_committed
1053-
let next_to_40 = overlay.next_child_storage_key_change(child, &[40]).unwrap();
1053+
let next_to_40 = overlay.child_iter_after(child, &[40]).next().unwrap();
10541054
assert_eq!(next_to_40.0.to_vec(), vec![50]);
10551055
assert_eq!(next_to_40.1.value(), Some(&vec![50]));
10561056
}

0 commit comments

Comments
 (0)