diff --git a/framework/base/src/storage/mappers/vec_mapper.rs b/framework/base/src/storage/mappers/vec_mapper.rs index 80f8c6018a..22b8763ece 100644 --- a/framework/base/src/storage/mappers/vec_mapper.rs +++ b/framework/base/src/storage/mappers/vec_mapper.rs @@ -209,6 +209,16 @@ where self.set_unchecked(index, item); } + /// Syntactic sugar, to more compactly express a get, update and set in one line. + /// Takes whatever lies in storage, apples the given closure and saves the final value back to storage. + /// Propagates the return value of the given function. + pub fn update R>(&mut self, index: usize, f: F) -> R { + let mut value = self.get(index); + let result = f(&mut value); + self.set(index, &value); + result + } + /// Keeping `set_unchecked` private on purpose, so developers don't write out of index limits by accident. fn set_unchecked(&self, index: usize, item: &T) { storage_set(self.item_key(index).as_ref(), item); diff --git a/framework/scenario/tests/test_vec_mapper.rs b/framework/scenario/tests/test_vec_mapper.rs index f7bc6587e4..92002fcc80 100644 --- a/framework/scenario/tests/test_vec_mapper.rs +++ b/framework/scenario/tests/test_vec_mapper.rs @@ -79,6 +79,20 @@ fn test_vec_swap_remove() { check_vec(&vect, vec![42, 45]); } +#[test] +fn test_vec_update() { + let mut vect = create_vec(); + + vect.extend_from_slice(&[42, 43, 44, 45, 46]); + + assert_eq!(vect.len(), 5); + assert_eq!(vect.get(5), 46); + vect.update(5, |item| *item = 42); + assert_eq!(vect.len(), 5); + assert_eq!(vect.get(5), 42); + check_vec(&vect, vec![42, 43, 44, 45, 42]); +} + #[test] fn test_vec_iter_processing() { let mut vect = create_vec();