diff --git a/substrate/primitives/state-machine/src/basic.rs b/substrate/primitives/state-machine/src/basic.rs index a5d573f514..c1745c03da 100644 --- a/substrate/primitives/state-machine/src/basic.rs +++ b/substrate/primitives/state-machine/src/basic.rs @@ -262,9 +262,8 @@ impl Externalities for BasicExternalities { key: Vec, value: Vec, ) { - let previous = self.storage(&key).unwrap_or_default(); - let new = crate::ext::append_to_storage(previous, value).expect("Failed to append to storage"); - self.place_storage(key.clone(), Some(new)); + let previous = self.inner.top.entry(key).or_default(); + crate::ext::append_to_storage(previous, value).expect("Failed to append to storage"); } fn chain_id(&self) -> u64 { 42 } diff --git a/substrate/primitives/state-machine/src/ext.rs b/substrate/primitives/state-machine/src/ext.rs index a92a10bf0e..902763e6e9 100644 --- a/substrate/primitives/state-machine/src/ext.rs +++ b/substrate/primitives/state-machine/src/ext.rs @@ -423,16 +423,12 @@ where let _guard = sp_panic_handler::AbortGuard::force_abort(); self.mark_dirty(); - let current_value = self.overlay - .value_mut(&key) - .map(|v| v.take()) - .unwrap_or_else(|| self.backend.storage(&key).expect(EXT_NOT_ALLOWED_TO_FAIL)) - .unwrap_or_default(); - - self.overlay.set_storage( - key, - Some(append_to_storage(current_value, value).expect(EXT_NOT_ALLOWED_TO_FAIL)), + let backend = &mut self.backend; + let current_value = self.overlay.value_mut_or_insert_with( + &key, + || backend.storage(&key).expect(EXT_NOT_ALLOWED_TO_FAIL).unwrap_or_default() ); + append_to_storage(current_value, value).expect(EXT_NOT_ALLOWED_TO_FAIL); } fn chain_id(&self) -> u64 { @@ -589,14 +585,14 @@ fn extract_length_data(data: &[u8]) -> Result<(u32, usize, usize), &'static str> } pub fn append_to_storage( - mut self_encoded: Vec, + self_encoded: &mut Vec, value: Vec, -) -> Result, &'static str> { +) -> Result<(), &'static str> { // No data present, just encode the given input data. if self_encoded.is_empty() { - Compact::from(1u32).encode_to(&mut self_encoded); + Compact::from(1u32).encode_to(self_encoded); self_encoded.extend(value); - return Ok(self_encoded); + return Ok(()); } let (new_len, encoded_len, encoded_new_len) = extract_length_data(&self_encoded)?; @@ -612,10 +608,10 @@ pub fn append_to_storage( // If old and new encoded len is equal, we don't need to copy the // already encoded data. if encoded_len == encoded_new_len { - replace_len(&mut self_encoded); - append_new_elems(&mut self_encoded); + replace_len(self_encoded); + append_new_elems(self_encoded); - Ok(self_encoded) + Ok(()) } else { let size = encoded_new_len + self_encoded.len() - encoded_len; @@ -627,8 +623,9 @@ pub fn append_to_storage( replace_len(&mut res); res[encoded_new_len..size].copy_from_slice(&self_encoded[encoded_len..]); append_new_elems(&mut res); + *self_encoded = res; - Ok(res) + Ok(()) } } diff --git a/substrate/primitives/state-machine/src/lib.rs b/substrate/primitives/state-machine/src/lib.rs index aa3a3ad573..0577afea00 100644 --- a/substrate/primitives/state-machine/src/lib.rs +++ b/substrate/primitives/state-machine/src/lib.rs @@ -1060,6 +1060,98 @@ mod tests { ); } + #[test] + fn append_storage_works() { + let reference_data = vec![ + b"data1".to_vec(), + b"2".to_vec(), + b"D3".to_vec(), + b"d4".to_vec(), + ]; + let key = b"key".to_vec(); + let mut state = new_in_mem::(); + let backend = state.as_trie_backend().unwrap(); + let mut overlay = OverlayedChanges::default(); + let mut offchain_overlay = OffchainOverlayedChanges::default(); + let mut cache = StorageTransactionCache::default(); + { + let mut ext = Ext::new( + &mut overlay, + &mut offchain_overlay, + &mut cache, + backend, + changes_trie::disabled_state::<_, u64>(), + None, + ); + + ext.storage_append(key.clone(), reference_data[0].encode()); + assert_eq!( + ext.storage(key.as_slice()), + Some(vec![reference_data[0].clone()].encode()), + ); + } + overlay.commit_prospective(); + { + let mut ext = Ext::new( + &mut overlay, + &mut offchain_overlay, + &mut cache, + backend, + changes_trie::disabled_state::<_, u64>(), + None, + ); + + for i in reference_data.iter().skip(1) { + ext.storage_append(key.clone(), i.encode()); + } + assert_eq!( + ext.storage(key.as_slice()), + Some(reference_data.encode()), + ); + } + overlay.discard_prospective(); + { + let ext = Ext::new( + &mut overlay, + &mut offchain_overlay, + &mut cache, + backend, + changes_trie::disabled_state::<_, u64>(), + None, + ); + assert_eq!( + ext.storage(key.as_slice()), + Some(vec![reference_data[0].clone()].encode()), + ); + } + } + + #[test] + fn remove_then_append() { + let key = b"key".to_vec(); + let mut state = new_in_mem::(); + let backend = state.as_trie_backend().unwrap(); + let mut overlay = OverlayedChanges::default(); + let mut offchain_overlay = OffchainOverlayedChanges::default(); + let mut cache = StorageTransactionCache::default(); + + let mut ext = Ext::new( + &mut overlay, + &mut offchain_overlay, + &mut cache, + backend, + changes_trie::disabled_state::<_, u64>(), + None, + ); + + ext.clear_storage(key.as_slice()); + ext.storage_append(key.clone(), b"Item".to_vec().encode()); + assert_eq!( + ext.storage(key.as_slice()), + Some(vec![b"Item".to_vec()].encode()), + ); + } + #[test] fn prove_read_and_proof_check_works() { let child_info = ChildInfo::new_default(b"sub1"); diff --git a/substrate/primitives/state-machine/src/overlayed_changes.rs b/substrate/primitives/state-machine/src/overlayed_changes.rs index f5b3c88ac8..c98da11c27 100644 --- a/substrate/primitives/state-machine/src/overlayed_changes.rs +++ b/substrate/primitives/state-machine/src/overlayed_changes.rs @@ -223,11 +223,37 @@ impl OverlayedChanges { }) } - /// Returns mutable reference to either commited or propsective value. - pub fn value_mut(&mut self, key: &[u8]) -> Option<&mut Option> { - // not using map because of double borrow inside closure - if let Some(entry) = self.prospective.top.get_mut(key) { return Some(&mut entry.value) } - return self.committed.top.get_mut(key).map(|entry| &mut entry.value); + /// Returns mutable reference to current changed value (prospective). + /// If there is no value in the overlay, the default callback is used to initiate + /// the value. + /// Warning this function register a change, so the mutable reference MUST be modified. + #[must_use = "A change was registered, so this value MUST be modified."] + pub fn value_mut_or_insert_with( + &mut self, + key: &[u8], + init: impl Fn() -> StorageValue, + ) -> &mut StorageValue { + let extrinsic_index = self.extrinsic_index(); + let committed = &self.committed.top; + + let mut entry = self.prospective.top.entry(key.to_vec()) + .or_insert_with(|| { + if let Some(overlay_state) = committed.get(key).cloned() { + overlay_state + } else { + OverlayedValue { value: Some(init()), ..Default::default() } + } + }); + + //if was deleted initialise back with empty vec + if entry.value.is_none() { + entry.value = Some(Default::default()); + } + if let Some(extrinsic) = extrinsic_index { + entry.extrinsics.get_or_insert_with(Default::default) + .insert(extrinsic); + } + entry.value.as_mut().expect("Initialized above; qed") } /// Returns a double-Option: None if the key is unknown (i.e. and the query should be referred