Avoid changing overlay committed layer. (#5839)

* Avoid changing overlay committed layer.

* basic test

* Add some tx in the test.

* only update from backend value on missing entry in both layer.
deleted entry is replace by empty vec.

* test and review changes

* additional test and review change

* remove test on changing existing value, it does not always panic
depending on existing content

* Update primitives/state-machine/src/overlayed_changes.rs

* Update primitives/state-machine/src/overlayed_changes.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/primitives/state-machine/src/basic.rs

@@ -262,9 +262,8 @@ impl Externalities for BasicExternalities {
 		key: Vec<u8>,
 		value: Vec<u8>,
 	) {
-		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 }

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<u8>,
+	self_encoded: &mut Vec<u8>,
 	value: Vec<u8>,
-) -> Result<Vec<u8>, &'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(())
 	}
 }
 

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::<BlakeTwo256>();
+		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::<BlakeTwo256>();
+		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");

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<StorageValue>> {
-		// 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