referrerpolicy=no-referrer-when-downgrade

sp_state_machine/overlayed_changes/
changeset.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License
17
18//! Houses the code that implements the transactional overlay storage.
19
20use super::{Extrinsics, StorageKey, StorageValue};
21
22#[cfg(not(feature = "std"))]
23use alloc::collections::btree_set::BTreeSet as Set;
24use codec::{Compact, CompactLen};
25#[cfg(feature = "std")]
26use std::collections::HashSet as Set;
27
28use crate::{ext::StorageAppend, warn};
29use alloc::{
30	collections::{btree_map::BTreeMap, btree_set::BTreeSet},
31	vec::Vec,
32};
33use core::hash::Hash;
34use smallvec::SmallVec;
35
36const PROOF_OVERLAY_NON_EMPTY: &str = "\
37	An OverlayValue is always created with at least one transaction and dropped as soon
38	as the last transaction is removed; qed";
39
40type DirtyKeysSets<K> = SmallVec<[Set<K>; 5]>;
41type Transactions<V> = SmallVec<[InnerValue<V>; 5]>;
42
43/// Error returned when trying to commit or rollback while no transaction is open or
44/// when the runtime is trying to close a transaction started by the client.
45#[derive(Debug)]
46#[cfg_attr(test, derive(PartialEq))]
47pub struct NoOpenTransaction;
48
49/// Error when calling `enter_runtime` when already being in runtime execution mode.
50#[derive(Debug)]
51#[cfg_attr(test, derive(PartialEq))]
52pub struct AlreadyInRuntime;
53
54/// Error when calling `exit_runtime` when not being in runtime execution mode.
55#[derive(Debug)]
56#[cfg_attr(test, derive(PartialEq))]
57pub struct NotInRuntime;
58
59/// Describes in which mode the node is currently executing.
60#[derive(Debug, Clone, Copy)]
61pub enum ExecutionMode {
62	/// Executing in client mode: Removal of all transactions possible.
63	Client,
64	/// Executing in runtime mode: Transactions started by the client are protected.
65	Runtime,
66}
67
68#[derive(Debug, Default, Clone)]
69#[cfg_attr(test, derive(PartialEq))]
70struct InnerValue<V> {
71	/// Current value. None if value has been deleted.
72	value: V,
73	/// The set of extrinsic indices where the values has been changed.
74	extrinsics: Extrinsics,
75}
76
77/// An overlay that contains all versions of a value for a specific key.
78#[derive(Debug, Clone)]
79#[cfg_attr(test, derive(PartialEq))]
80pub struct OverlayedEntry<V> {
81	/// The individual versions of that value.
82	/// One entry per transactions during that the value was actually written.
83	transactions: Transactions<V>,
84}
85
86impl<V> Default for OverlayedEntry<V> {
87	fn default() -> Self {
88		Self { transactions: SmallVec::new() }
89	}
90}
91
92/// History of value, with removal support.
93pub type OverlayedValue = OverlayedEntry<StorageEntry>;
94
95/// Content in an overlay for a given transactional depth.
96#[derive(Debug, Clone, Default)]
97#[cfg_attr(test, derive(PartialEq))]
98pub enum StorageEntry {
99	/// The storage entry should be set to the stored value.
100	Set(StorageValue),
101	/// The storage entry should be removed.
102	#[default]
103	Remove,
104	/// The storage entry was appended to.
105	///
106	/// This assumes that the storage entry is encoded as a SCALE list. This means that it is
107	/// prefixed with a `Compact<u32>` that reprensents the length, followed by all the encoded
108	/// elements.
109	Append {
110		/// The value of the storage entry.
111		///
112		/// This may or may not be prefixed by the length, depending on the materialized length.
113		data: StorageValue,
114		/// Current number of elements stored in data.
115		current_length: u32,
116		/// The number of elements as stored in the prefixed length in `data`.
117		///
118		/// If `None`, than `data` is not yet prefixed with the length.
119		materialized_length: Option<u32>,
120		/// The size of `data` in the parent transactional layer.
121		///
122		/// Only set when the parent layer is in  `Append` state.
123		parent_size: Option<usize>,
124	},
125}
126
127impl StorageEntry {
128	/// Convert to an [`Option<StorageValue>`].
129	pub(super) fn to_option(mut self) -> Option<StorageValue> {
130		self.materialize_in_place();
131		match self {
132			StorageEntry::Append { data, .. } | StorageEntry::Set(data) => Some(data),
133			StorageEntry::Remove => None,
134		}
135	}
136
137	/// Return as an [`Option<StorageValue>`].
138	fn as_option(&mut self) -> Option<&StorageValue> {
139		self.materialize_in_place();
140		match self {
141			StorageEntry::Append { data, .. } | StorageEntry::Set(data) => Some(data),
142			StorageEntry::Remove => None,
143		}
144	}
145
146	/// Materialize the internal state and cache the resulting materialized value.
147	fn materialize_in_place(&mut self) {
148		if let StorageEntry::Append { data, materialized_length, current_length, .. } = self {
149			let current_length = *current_length;
150			if materialized_length.map_or(false, |m| m == current_length) {
151				return
152			}
153			StorageAppend::new(data).replace_length(*materialized_length, current_length);
154			*materialized_length = Some(current_length);
155		}
156	}
157
158	/// Materialize the internal state.
159	#[cfg(test)]
160	pub(crate) fn materialize(&self) -> Option<alloc::borrow::Cow<[u8]>> {
161		use alloc::borrow::Cow;
162
163		match self {
164			StorageEntry::Append { data, materialized_length, current_length, .. } => {
165				let current_length = *current_length;
166				if materialized_length.map_or(false, |m| m == current_length) {
167					Some(Cow::Borrowed(data.as_ref()))
168				} else {
169					let mut data = data.clone();
170					StorageAppend::new(&mut data)
171						.replace_length(*materialized_length, current_length);
172
173					Some(data.into())
174				}
175			},
176			StorageEntry::Remove => None,
177			StorageEntry::Set(e) => Some(Cow::Borrowed(e.as_ref())),
178		}
179	}
180}
181
182/// Change set for basic key value with extrinsics index recording and removal support.
183pub type OverlayedChangeSet = OverlayedMap<StorageKey, StorageEntry>;
184
185/// Holds a set of changes with the ability modify them using nested transactions.
186#[derive(Debug, Clone)]
187pub struct OverlayedMap<K, V> {
188	/// Stores the changes that this overlay constitutes.
189	changes: BTreeMap<K, OverlayedEntry<V>>,
190	/// Stores which keys are dirty per transaction. Needed in order to determine which
191	/// values to merge into the parent transaction on commit. The length of this vector
192	/// therefore determines how many nested transactions are currently open (depth).
193	dirty_keys: DirtyKeysSets<K>,
194	/// The number of how many transactions beginning from the first transactions are started
195	/// by the client. Those transactions are protected against close (commit, rollback)
196	/// when in runtime mode.
197	num_client_transactions: usize,
198	/// Determines whether the node is using the overlay from the client or the runtime.
199	execution_mode: ExecutionMode,
200}
201
202impl<K, V> Default for OverlayedMap<K, V> {
203	fn default() -> Self {
204		Self {
205			changes: BTreeMap::new(),
206			dirty_keys: SmallVec::new(),
207			num_client_transactions: Default::default(),
208			execution_mode: Default::default(),
209		}
210	}
211}
212
213#[cfg(not(substrate_runtime))]
214impl From<sp_core::storage::StorageMap> for OverlayedMap<StorageKey, StorageEntry> {
215	fn from(storage: sp_core::storage::StorageMap) -> Self {
216		Self {
217			changes: storage
218				.into_iter()
219				.map(|(k, v)| {
220					(
221						k,
222						OverlayedEntry {
223							transactions: SmallVec::from_iter([InnerValue {
224								value: StorageEntry::Set(v),
225								extrinsics: Default::default(),
226							}]),
227						},
228					)
229				})
230				.collect(),
231			dirty_keys: Default::default(),
232			num_client_transactions: 0,
233			execution_mode: ExecutionMode::Client,
234		}
235	}
236}
237
238impl Default for ExecutionMode {
239	fn default() -> Self {
240		Self::Client
241	}
242}
243
244impl<V> OverlayedEntry<V> {
245	/// The value as seen by the current transaction.
246	pub fn value_ref(&self) -> &V {
247		&self.transactions.last().expect(PROOF_OVERLAY_NON_EMPTY).value
248	}
249
250	/// The value as seen by the current transaction.
251	pub fn into_value(mut self) -> V {
252		self.transactions.pop().expect(PROOF_OVERLAY_NON_EMPTY).value
253	}
254
255	/// Unique list of extrinsic indices which modified the value.
256	pub fn extrinsics(&self) -> BTreeSet<u32> {
257		let mut set = BTreeSet::new();
258		self.transactions
259			.iter()
260			.for_each(|t| t.extrinsics.copy_extrinsics_into(&mut set));
261		set
262	}
263
264	/// Mutable reference to the most recent version.
265	fn value_mut(&mut self) -> &mut V {
266		&mut self.transactions.last_mut().expect(PROOF_OVERLAY_NON_EMPTY).value
267	}
268
269	/// Remove the last version and return it.
270	fn pop_transaction(&mut self) -> InnerValue<V> {
271		self.transactions.pop().expect(PROOF_OVERLAY_NON_EMPTY)
272	}
273
274	/// Mutable reference to the set which holds the indices for the **current transaction only**.
275	fn transaction_extrinsics_mut(&mut self) -> &mut Extrinsics {
276		&mut self.transactions.last_mut().expect(PROOF_OVERLAY_NON_EMPTY).extrinsics
277	}
278
279	/// Writes a new version of a value.
280	///
281	/// This makes sure that the old version is not overwritten and can be properly
282	/// rolled back when required.
283	fn set_offchain(&mut self, value: V, first_write_in_tx: bool, at_extrinsic: Option<u32>) {
284		if first_write_in_tx || self.transactions.is_empty() {
285			self.transactions.push(InnerValue { value, extrinsics: Default::default() });
286		} else {
287			*self.value_mut() = value;
288		}
289
290		if let Some(extrinsic) = at_extrinsic {
291			self.transaction_extrinsics_mut().insert(extrinsic);
292		}
293	}
294}
295
296/// Restore the `current_data` from an [`StorageEntry::Append`] back to the parent.
297///
298/// When creating a new transaction layer from an appended entry, the `data` will be moved to
299/// prevent extra allocations. So, we need to move back the `data` to the parent layer when there is
300/// a roll back or the entry is set to some different value. This functions puts back the data to
301/// the `parent` and truncates any extra elements that got added in the current layer.
302///
303/// The current and the `parent` layer need to be [`StorageEntry::Append`] or otherwise the function
304/// is a no-op.
305fn restore_append_to_parent(
306	parent: &mut StorageEntry,
307	mut current_data: Vec<u8>,
308	current_materialized: Option<u32>,
309	mut target_parent_size: usize,
310) {
311	match parent {
312		StorageEntry::Append {
313			data: parent_data,
314			materialized_length: parent_materialized,
315			..
316		} => {
317			// Forward the materialized length to the parent with the data. Next time when
318			// materializing the value, the length will be corrected. This prevents doing a
319			// potential allocation here.
320
321			let prev = parent_materialized.map(|l| Compact::<u32>::compact_len(&l)).unwrap_or(0);
322			let new = current_materialized.map(|l| Compact::<u32>::compact_len(&l)).unwrap_or(0);
323			let delta = new.abs_diff(prev);
324			if prev >= new {
325				target_parent_size -= delta;
326			} else {
327				target_parent_size += delta;
328			}
329			*parent_materialized = current_materialized;
330
331			// Truncate the data to remove any extra elements
332			current_data.truncate(target_parent_size);
333			*parent_data = current_data;
334		},
335		_ => {
336			// No value or a simple value, no need to restore
337		},
338	}
339}
340
341impl OverlayedEntry<StorageEntry> {
342	/// Writes a new version of a value.
343	///
344	/// This makes sure that the old version is not overwritten and can be properly
345	/// rolled back when required.
346	fn set(
347		&mut self,
348		value: Option<StorageValue>,
349		first_write_in_tx: bool,
350		at_extrinsic: Option<u32>,
351	) {
352		let value = value.map_or_else(|| StorageEntry::Remove, StorageEntry::Set);
353
354		if first_write_in_tx || self.transactions.is_empty() {
355			self.transactions.push(InnerValue { value, extrinsics: Default::default() });
356		} else {
357			let mut old_value = self.value_mut();
358
359			let set_prev = if let StorageEntry::Append {
360				data,
361				current_length: _,
362				materialized_length,
363				parent_size,
364			} = &mut old_value
365			{
366				parent_size
367					.map(|parent_size| (core::mem::take(data), *materialized_length, parent_size))
368			} else {
369				None
370			};
371
372			*old_value = value;
373
374			if let Some((data, current_materialized, parent_size)) = set_prev {
375				let transactions = self.transactions.len();
376
377				debug_assert!(transactions >= 2);
378				let parent = self
379					.transactions
380					.get_mut(transactions - 2)
381					.expect("`set_prev` is only `Some(_)`, if the value came from parent; qed");
382				restore_append_to_parent(
383					&mut parent.value,
384					data,
385					current_materialized,
386					parent_size,
387				);
388			}
389		}
390
391		if let Some(extrinsic) = at_extrinsic {
392			self.transaction_extrinsics_mut().insert(extrinsic);
393		}
394	}
395
396	/// Append content to a value, updating a prefixed compact encoded length.
397	///
398	/// This makes sure that the old version is not overwritten and can be properly
399	/// rolled back when required.
400	/// This avoid copying value from previous transaction.
401	fn append(
402		&mut self,
403		element: StorageValue,
404		first_write_in_tx: bool,
405		init: impl Fn() -> StorageValue,
406		at_extrinsic: Option<u32>,
407	) {
408		if self.transactions.is_empty() {
409			let mut init_value = init();
410
411			let mut append = StorageAppend::new(&mut init_value);
412
413			// Either the init value is a SCALE list like value to that the `element` gets appended
414			// or the value is reset to `[element]`.
415			let (data, current_length, materialized_length) =
416				if let Some(len) = append.extract_length() {
417					append.append_raw(element);
418
419					(init_value, len + 1, Some(len))
420				} else {
421					(element, 1, None)
422				};
423
424			self.transactions.push(InnerValue {
425				value: StorageEntry::Append {
426					data,
427					current_length,
428					materialized_length,
429					parent_size: None,
430				},
431				extrinsics: Default::default(),
432			});
433		} else if first_write_in_tx {
434			let parent = self.value_mut();
435			let (data, current_length, materialized_length, parent_size) = match parent {
436				StorageEntry::Remove => (element, 1, None, None),
437				StorageEntry::Append { data, current_length, materialized_length, .. } => {
438					let parent_len = data.len();
439					let mut data_buf = core::mem::take(data);
440					StorageAppend::new(&mut data_buf).append_raw(element);
441					(data_buf, *current_length + 1, *materialized_length, Some(parent_len))
442				},
443				StorageEntry::Set(prev) => {
444					// For compatibility: append if there is a encoded length, overwrite
445					// with value otherwhise.
446					if let Some(current_length) = StorageAppend::new(prev).extract_length() {
447						// The `prev` is cloned here, but it could be optimized to not do the clone
448						// here as it is done for `Append` above.
449						let mut data = prev.clone();
450						StorageAppend::new(&mut data).append_raw(element);
451						(data, current_length + 1, Some(current_length), None)
452					} else {
453						// overwrite, same as empty case.
454						(element, 1, None, None)
455					}
456				},
457			};
458
459			self.transactions.push(InnerValue {
460				value: StorageEntry::Append {
461					data,
462					current_length,
463					materialized_length,
464					parent_size,
465				},
466				extrinsics: Default::default(),
467			});
468		} else {
469			// not first transaction write
470			let old_value = self.value_mut();
471			let replace = match old_value {
472				StorageEntry::Remove => Some((element, 1, None)),
473				StorageEntry::Set(data) => {
474					// Note that when the data here is not initialized with append,
475					// and still starts with a valid compact u32 we can have totally broken
476					// encoding.
477					let mut append = StorageAppend::new(data);
478
479					// For compatibility: append if there is a encoded length, overwrite
480					// with value otherwhise.
481					if let Some(current_length) = append.extract_length() {
482						append.append_raw(element);
483						Some((core::mem::take(data), current_length + 1, Some(current_length)))
484					} else {
485						Some((element, 1, None))
486					}
487				},
488				StorageEntry::Append { data, current_length, .. } => {
489					StorageAppend::new(data).append_raw(element);
490					*current_length += 1;
491					None
492				},
493			};
494
495			if let Some((data, current_length, materialized_length)) = replace {
496				*old_value = StorageEntry::Append {
497					data,
498					current_length,
499					materialized_length,
500					parent_size: None,
501				};
502			}
503		}
504
505		if let Some(extrinsic) = at_extrinsic {
506			self.transaction_extrinsics_mut().insert(extrinsic);
507		}
508	}
509
510	/// The value as seen by the current transaction.
511	pub fn value(&mut self) -> Option<&StorageValue> {
512		self.value_mut().as_option()
513	}
514}
515
516/// Inserts a key into the dirty set.
517///
518/// Returns true iff we are currently have at least one open transaction and if this
519/// is the first write to the given key that transaction.
520fn insert_dirty<K: Ord + Hash>(set: &mut DirtyKeysSets<K>, key: K) -> bool {
521	set.last_mut().map(|dk| dk.insert(key)).unwrap_or_default()
522}
523
524impl<K: Ord + Hash + Clone, V> OverlayedMap<K, V> {
525	/// Create a new changeset at the same transaction state but without any contents.
526	///
527	/// This changeset might be created when there are already open transactions.
528	/// We need to catch up here so that the child is at the same transaction depth.
529	pub fn spawn_child(&self) -> Self {
530		use core::iter::repeat;
531		Self {
532			changes: Default::default(),
533			dirty_keys: repeat(Set::new()).take(self.transaction_depth()).collect(),
534			num_client_transactions: self.num_client_transactions,
535			execution_mode: self.execution_mode,
536		}
537	}
538
539	/// True if no changes at all are contained in the change set.
540	pub fn is_empty(&self) -> bool {
541		self.changes.is_empty()
542	}
543
544	/// Get an optional reference to the value stored for the specified key.
545	pub fn get<Q>(&mut self, key: &Q) -> Option<&mut OverlayedEntry<V>>
546	where
547		K: core::borrow::Borrow<Q>,
548		Q: Ord + ?Sized,
549	{
550		self.changes.get_mut(key)
551	}
552
553	/// Set a new value for the specified key.
554	///
555	/// Can be rolled back or committed when called inside a transaction.
556	pub fn set_offchain(&mut self, key: K, value: V, at_extrinsic: Option<u32>) {
557		let overlayed = self.changes.entry(key.clone()).or_default();
558		overlayed.set_offchain(value, insert_dirty(&mut self.dirty_keys, key), at_extrinsic);
559	}
560
561	/// Get a list of all changes as seen by current transaction.
562	pub fn changes(&self) -> impl Iterator<Item = (&K, &OverlayedEntry<V>)> {
563		self.changes.iter()
564	}
565
566	/// Get a list of all changes as seen by current transaction.
567	pub fn changes_mut(&mut self) -> impl Iterator<Item = (&K, &mut OverlayedEntry<V>)> {
568		self.changes.iter_mut()
569	}
570
571	/// Get a list of all changes as seen by current transaction, consumes
572	/// the overlay.
573	pub fn into_changes(self) -> impl Iterator<Item = (K, OverlayedEntry<V>)> {
574		self.changes.into_iter()
575	}
576
577	/// Consume this changeset and return all committed changes.
578	///
579	/// Panics:
580	/// Panics if there are open transactions: `transaction_depth() > 0`
581	pub fn drain_committed(self) -> impl Iterator<Item = (K, V)> {
582		assert!(self.transaction_depth() == 0, "Drain is not allowed with open transactions.");
583		self.changes.into_iter().map(|(k, mut v)| (k, v.pop_transaction().value))
584	}
585
586	/// Returns the current nesting depth of the transaction stack.
587	///
588	/// A value of zero means that no transaction is open and changes are committed on write.
589	pub fn transaction_depth(&self) -> usize {
590		self.dirty_keys.len()
591	}
592
593	/// Call this before transferring control to the runtime.
594	///
595	/// This protects all existing transactions from being removed by the runtime.
596	/// Calling this while already inside the runtime will return an error.
597	pub fn enter_runtime(&mut self) -> Result<(), AlreadyInRuntime> {
598		if let ExecutionMode::Runtime = self.execution_mode {
599			return Err(AlreadyInRuntime)
600		}
601		self.execution_mode = ExecutionMode::Runtime;
602		self.num_client_transactions = self.transaction_depth();
603		Ok(())
604	}
605
606	/// Call this when control returns from the runtime.
607	///
608	/// This rollbacks all dangling transaction left open by the runtime.
609	/// Calling this while already outside the runtime will return an error.
610	pub fn exit_runtime_offchain(&mut self) -> Result<(), NotInRuntime> {
611		if let ExecutionMode::Client = self.execution_mode {
612			return Err(NotInRuntime)
613		}
614		self.execution_mode = ExecutionMode::Client;
615		if self.has_open_runtime_transactions() {
616			warn!(
617				"{} storage transactions are left open by the runtime. Those will be rolled back.",
618				self.transaction_depth() - self.num_client_transactions,
619			);
620		}
621		while self.has_open_runtime_transactions() {
622			self.rollback_transaction_offchain()
623				.expect("The loop condition checks that the transaction depth is > 0; qed");
624		}
625		Ok(())
626	}
627
628	/// Start a new nested transaction.
629	///
630	/// This allows to either commit or roll back all changes that were made while this
631	/// transaction was open. Any transaction must be closed by either `commit_transaction`
632	/// or `rollback_transaction` before this overlay can be converted into storage changes.
633	///
634	/// Changes made without any open transaction are committed immediately.
635	pub fn start_transaction(&mut self) {
636		self.dirty_keys.push(Default::default());
637	}
638
639	/// Rollback the last transaction started by `start_transaction`.
640	///
641	/// Any changes made during that transaction are discarded. Returns an error if
642	/// there is no open transaction that can be rolled back.
643	pub fn rollback_transaction_offchain(&mut self) -> Result<(), NoOpenTransaction> {
644		self.close_transaction_offchain(true)
645	}
646
647	/// Commit the last transaction started by `start_transaction`.
648	///
649	/// Any changes made during that transaction are committed. Returns an error if
650	/// there is no open transaction that can be committed.
651	pub fn commit_transaction_offchain(&mut self) -> Result<(), NoOpenTransaction> {
652		self.close_transaction_offchain(false)
653	}
654
655	fn close_transaction_offchain(&mut self, rollback: bool) -> Result<(), NoOpenTransaction> {
656		// runtime is not allowed to close transactions started by the client
657		if matches!(self.execution_mode, ExecutionMode::Runtime) &&
658			!self.has_open_runtime_transactions()
659		{
660			return Err(NoOpenTransaction)
661		}
662
663		for key in self.dirty_keys.pop().ok_or(NoOpenTransaction)? {
664			let overlayed = self.changes.get_mut(&key).expect(
665				"\
666				A write to an OverlayedValue is recorded in the dirty key set. Before an
667				OverlayedValue is removed, its containing dirty set is removed. This
668				function is only called for keys that are in the dirty set. qed\
669			",
670			);
671
672			if rollback {
673				overlayed.pop_transaction();
674
675				// We need to remove the key as an `OverlayValue` with no transactions
676				// violates its invariant of always having at least one transaction.
677				if overlayed.transactions.is_empty() {
678					self.changes.remove(&key);
679				}
680			} else {
681				let has_predecessor = if let Some(dirty_keys) = self.dirty_keys.last_mut() {
682					// Not the last tx: Did the previous tx write to this key?
683					!dirty_keys.insert(key)
684				} else {
685					// Last tx: Is there already a value in the committed set?
686					// Check against one rather than empty because the current tx is still
687					// in the list as it is popped later in this function.
688					overlayed.transactions.len() > 1
689				};
690
691				// We only need to merge if there is an pre-existing value. It may be a value from
692				// the previous transaction or a value committed without any open transaction.
693				if has_predecessor {
694					let dropped_tx = overlayed.pop_transaction();
695					*overlayed.value_mut() = dropped_tx.value;
696					overlayed.transaction_extrinsics_mut().extend(dropped_tx.extrinsics);
697				}
698			}
699		}
700
701		Ok(())
702	}
703
704	fn has_open_runtime_transactions(&self) -> bool {
705		self.transaction_depth() > self.num_client_transactions
706	}
707}
708
709impl OverlayedChangeSet {
710	/// Rollback the last transaction started by `start_transaction`.
711	///
712	/// Any changes made during that transaction are discarded. Returns an error if
713	/// there is no open transaction that can be rolled back.
714	pub fn rollback_transaction(&mut self) -> Result<(), NoOpenTransaction> {
715		self.close_transaction(true)
716	}
717
718	/// Commit the last transaction started by `start_transaction`.
719	///
720	/// Any changes made during that transaction are committed. Returns an error if
721	/// there is no open transaction that can be committed.
722	pub fn commit_transaction(&mut self) -> Result<(), NoOpenTransaction> {
723		self.close_transaction(false)
724	}
725
726	fn close_transaction(&mut self, rollback: bool) -> Result<(), NoOpenTransaction> {
727		// runtime is not allowed to close transactions started by the client
728		if matches!(self.execution_mode, ExecutionMode::Runtime) &&
729			!self.has_open_runtime_transactions()
730		{
731			return Err(NoOpenTransaction)
732		}
733
734		for key in self.dirty_keys.pop().ok_or(NoOpenTransaction)? {
735			let overlayed = self.changes.get_mut(&key).expect(
736				"\
737				A write to an OverlayedValue is recorded in the dirty key set. Before an
738				OverlayedValue is removed, its containing dirty set is removed. This
739				function is only called for keys that are in the dirty set. qed\
740			",
741			);
742
743			if rollback {
744				match overlayed.pop_transaction().value {
745					StorageEntry::Append {
746						data,
747						materialized_length,
748						parent_size: Some(parent_size),
749						..
750					} => {
751						debug_assert!(!overlayed.transactions.is_empty());
752						restore_append_to_parent(
753							overlayed.value_mut(),
754							data,
755							materialized_length,
756							parent_size,
757						);
758					},
759					_ => (),
760				}
761
762				// We need to remove the key as an `OverlayValue` with no transactions
763				// violates its invariant of always having at least one transaction.
764				if overlayed.transactions.is_empty() {
765					self.changes.remove(&key);
766				}
767			} else {
768				let has_predecessor = if let Some(dirty_keys) = self.dirty_keys.last_mut() {
769					// Not the last tx: Did the previous tx write to this key?
770					!dirty_keys.insert(key)
771				} else {
772					// Last tx: Is there already a value in the committed set?
773					// Check against one rather than empty because the current tx is still
774					// in the list as it is popped later in this function.
775					overlayed.transactions.len() > 1
776				};
777
778				// We only need to merge if there is an pre-existing value. It may be a value from
779				// the previous transaction or a value committed without any open transaction.
780				if has_predecessor {
781					let mut committed_tx = overlayed.pop_transaction();
782					let mut merge_appends = false;
783
784					// consecutive appends need to keep past `parent_size` value.
785					if let StorageEntry::Append { parent_size, .. } = &mut committed_tx.value {
786						if parent_size.is_some() {
787							let parent = overlayed.value_mut();
788							if let StorageEntry::Append { parent_size: keep_me, .. } = parent {
789								merge_appends = true;
790								*parent_size = *keep_me;
791							}
792						}
793					}
794
795					if merge_appends {
796						*overlayed.value_mut() = committed_tx.value;
797					} else {
798						let removed = core::mem::replace(overlayed.value_mut(), committed_tx.value);
799						// The transaction being commited is not an append operation. However, the
800						// value being overwritten in the previous transaction might be an append
801						// that needs to be merged with its parent. We only need to handle `Append`
802						// here because `Set` and `Remove` can directly overwrite previous
803						// operations.
804						if let StorageEntry::Append {
805							parent_size, data, materialized_length, ..
806						} = removed
807						{
808							if let Some(parent_size) = parent_size {
809								let transactions = overlayed.transactions.len();
810
811								// info from replaced head so len is at least one
812								// and parent_size implies a parent transaction
813								// so length is at least two.
814								debug_assert!(transactions >= 2);
815								if let Some(parent) =
816									overlayed.transactions.get_mut(transactions - 2)
817								{
818									restore_append_to_parent(
819										&mut parent.value,
820										data,
821										materialized_length,
822										parent_size,
823									)
824								}
825							}
826						}
827					}
828
829					overlayed.transaction_extrinsics_mut().extend(committed_tx.extrinsics);
830				}
831			}
832		}
833
834		Ok(())
835	}
836
837	/// Call this when control returns from the runtime.
838	///
839	/// This commits all dangling transaction left open by the runtime.
840	/// Calling this while already outside the runtime will return an error.
841	pub fn exit_runtime(&mut self) -> Result<(), NotInRuntime> {
842		if matches!(self.execution_mode, ExecutionMode::Client) {
843			return Err(NotInRuntime)
844		}
845
846		self.execution_mode = ExecutionMode::Client;
847		if self.has_open_runtime_transactions() {
848			warn!(
849				"{} storage transactions are left open by the runtime. Those will be rolled back.",
850				self.transaction_depth() - self.num_client_transactions,
851			);
852		}
853		while self.has_open_runtime_transactions() {
854			self.rollback_transaction()
855				.expect("The loop condition checks that the transaction depth is > 0; qed");
856		}
857
858		Ok(())
859	}
860
861	/// Set a new value for the specified key.
862	///
863	/// Can be rolled back or committed when called inside a transaction.
864	pub fn set(&mut self, key: StorageKey, value: Option<StorageValue>, at_extrinsic: Option<u32>) {
865		let overlayed = self.changes.entry(key.clone()).or_default();
866		overlayed.set(value, insert_dirty(&mut self.dirty_keys, key), at_extrinsic);
867	}
868
869	/// Append bytes to an existing content.
870	pub fn append_storage(
871		&mut self,
872		key: StorageKey,
873		value: StorageValue,
874		init: impl Fn() -> StorageValue,
875		at_extrinsic: Option<u32>,
876	) {
877		let overlayed = self.changes.entry(key.clone()).or_default();
878		let first_write_in_tx = insert_dirty(&mut self.dirty_keys, key);
879		overlayed.append(value, first_write_in_tx, init, at_extrinsic);
880	}
881
882	/// Set all values to deleted which are matched by the predicate.
883	///
884	/// Can be rolled back or committed when called inside a transaction.
885	pub fn clear_where(
886		&mut self,
887		predicate: impl Fn(&[u8], &OverlayedValue) -> bool,
888		at_extrinsic: Option<u32>,
889	) -> u32 {
890		let mut count = 0;
891		for (key, val) in self.changes.iter_mut().filter(|(k, v)| predicate(k, v)) {
892			if matches!(val.value_ref(), StorageEntry::Set(..) | StorageEntry::Append { .. }) {
893				count += 1;
894			}
895			val.set(None, insert_dirty(&mut self.dirty_keys, key.clone()), at_extrinsic);
896		}
897		count
898	}
899
900	/// Get the iterator over all changes that follow the supplied `key`.
901	pub fn changes_after(
902		&mut self,
903		key: &[u8],
904	) -> impl Iterator<Item = (&[u8], &mut OverlayedValue)> {
905		use core::ops::Bound;
906		let range = (Bound::Excluded(key), Bound::Unbounded);
907		self.changes.range_mut::<[u8], _>(range).map(|(k, v)| (k.as_slice(), v))
908	}
909}
910
911#[cfg(test)]
912mod test {
913	use super::*;
914	use pretty_assertions::assert_eq;
915
916	type Changes<'a> = Vec<(&'a [u8], (Option<&'a [u8]>, Vec<u32>))>;
917	type Drained<'a> = Vec<(&'a [u8], Option<&'a [u8]>)>;
918
919	fn assert_changes(is: &mut OverlayedChangeSet, expected: &Changes) {
920		let is: Changes = is
921			.changes_mut()
922			.map(|(k, v)| {
923				let extrinsics = v.extrinsics().into_iter().collect();
924				(k.as_ref(), (v.value().map(AsRef::as_ref), extrinsics))
925			})
926			.collect();
927		assert_eq!(&is, expected);
928	}
929
930	fn assert_drained_changes(is: OverlayedChangeSet, expected: Changes) {
931		let is = is.drain_committed().map(|(k, v)| (k, v.to_option())).collect::<Vec<_>>();
932		let expected = expected
933			.iter()
934			.map(|(k, v)| (k.to_vec(), v.0.map(From::from)))
935			.collect::<Vec<_>>();
936		assert_eq!(is, expected);
937	}
938
939	fn assert_drained(is: OverlayedChangeSet, expected: Drained) {
940		let is = is.drain_committed().map(|(k, v)| (k, v.to_option())).collect::<Vec<_>>();
941		let expected = expected
942			.iter()
943			.map(|(k, v)| (k.to_vec(), v.map(From::from)))
944			.collect::<Vec<_>>();
945		assert_eq!(is, expected);
946	}
947
948	#[test]
949	fn no_transaction_works() {
950		let mut changeset = OverlayedChangeSet::default();
951		assert_eq!(changeset.transaction_depth(), 0);
952
953		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
954		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
955		changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(9));
956
957		assert_drained(changeset, vec![(b"key0", Some(b"val0-1")), (b"key1", Some(b"val1"))]);
958	}
959
960	#[test]
961	fn transaction_works() {
962		let mut changeset = OverlayedChangeSet::default();
963		assert_eq!(changeset.transaction_depth(), 0);
964
965		// no transaction: committed on set
966		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
967		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
968		changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(10));
969
970		changeset.start_transaction();
971		assert_eq!(changeset.transaction_depth(), 1);
972
973		// we will commit that later
974		changeset.set(b"key42".to_vec(), Some(b"val42".to_vec()), Some(42));
975		changeset.set(b"key99".to_vec(), Some(b"val99".to_vec()), Some(99));
976
977		changeset.start_transaction();
978		assert_eq!(changeset.transaction_depth(), 2);
979
980		// we will roll that back
981		changeset.set(b"key42".to_vec(), Some(b"val42-rolled".to_vec()), Some(421));
982		changeset.set(b"key7".to_vec(), Some(b"val7-rolled".to_vec()), Some(77));
983		changeset.set(b"key0".to_vec(), Some(b"val0-rolled".to_vec()), Some(1000));
984		changeset.set(b"key5".to_vec(), Some(b"val5-rolled".to_vec()), None);
985
986		// changes contain all changes not only the committed ones.
987		let all_changes: Changes = vec![
988			(b"key0", (Some(b"val0-rolled"), vec![1, 10, 1000])),
989			(b"key1", (Some(b"val1"), vec![1])),
990			(b"key42", (Some(b"val42-rolled"), vec![42, 421])),
991			(b"key5", (Some(b"val5-rolled"), vec![])),
992			(b"key7", (Some(b"val7-rolled"), vec![77])),
993			(b"key99", (Some(b"val99"), vec![99])),
994		];
995		assert_changes(&mut changeset, &all_changes);
996
997		// this should be no-op
998		changeset.start_transaction();
999		assert_eq!(changeset.transaction_depth(), 3);
1000		changeset.start_transaction();
1001		assert_eq!(changeset.transaction_depth(), 4);
1002		changeset.rollback_transaction().unwrap();
1003		assert_eq!(changeset.transaction_depth(), 3);
1004		changeset.commit_transaction().unwrap();
1005		assert_eq!(changeset.transaction_depth(), 2);
1006		assert_changes(&mut changeset, &all_changes);
1007
1008		// roll back our first transactions that actually contains something
1009		changeset.rollback_transaction().unwrap();
1010		assert_eq!(changeset.transaction_depth(), 1);
1011
1012		let rolled_back: Changes = vec![
1013			(b"key0", (Some(b"val0-1"), vec![1, 10])),
1014			(b"key1", (Some(b"val1"), vec![1])),
1015			(b"key42", (Some(b"val42"), vec![42])),
1016			(b"key99", (Some(b"val99"), vec![99])),
1017		];
1018		assert_changes(&mut changeset, &rolled_back);
1019
1020		changeset.commit_transaction().unwrap();
1021		assert_eq!(changeset.transaction_depth(), 0);
1022		assert_changes(&mut changeset, &rolled_back);
1023
1024		assert_drained_changes(changeset, rolled_back);
1025	}
1026
1027	#[test]
1028	fn transaction_commit_then_rollback_works() {
1029		let mut changeset = OverlayedChangeSet::default();
1030		assert_eq!(changeset.transaction_depth(), 0);
1031
1032		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
1033		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
1034		changeset.set(b"key0".to_vec(), Some(b"val0-1".to_vec()), Some(10));
1035
1036		changeset.start_transaction();
1037		assert_eq!(changeset.transaction_depth(), 1);
1038
1039		changeset.set(b"key42".to_vec(), Some(b"val42".to_vec()), Some(42));
1040		changeset.set(b"key99".to_vec(), Some(b"val99".to_vec()), Some(99));
1041
1042		changeset.start_transaction();
1043		assert_eq!(changeset.transaction_depth(), 2);
1044
1045		changeset.set(b"key42".to_vec(), Some(b"val42-rolled".to_vec()), Some(421));
1046		changeset.set(b"key7".to_vec(), Some(b"val7-rolled".to_vec()), Some(77));
1047		changeset.set(b"key0".to_vec(), Some(b"val0-rolled".to_vec()), Some(1000));
1048		changeset.set(b"key5".to_vec(), Some(b"val5-rolled".to_vec()), None);
1049
1050		let all_changes: Changes = vec![
1051			(b"key0", (Some(b"val0-rolled"), vec![1, 10, 1000])),
1052			(b"key1", (Some(b"val1"), vec![1])),
1053			(b"key42", (Some(b"val42-rolled"), vec![42, 421])),
1054			(b"key5", (Some(b"val5-rolled"), vec![])),
1055			(b"key7", (Some(b"val7-rolled"), vec![77])),
1056			(b"key99", (Some(b"val99"), vec![99])),
1057		];
1058		assert_changes(&mut changeset, &all_changes);
1059
1060		// this should be no-op
1061		changeset.start_transaction();
1062		assert_eq!(changeset.transaction_depth(), 3);
1063		changeset.start_transaction();
1064		assert_eq!(changeset.transaction_depth(), 4);
1065		changeset.rollback_transaction().unwrap();
1066		assert_eq!(changeset.transaction_depth(), 3);
1067		changeset.commit_transaction().unwrap();
1068		assert_eq!(changeset.transaction_depth(), 2);
1069		assert_changes(&mut changeset, &all_changes);
1070
1071		changeset.commit_transaction().unwrap();
1072		assert_eq!(changeset.transaction_depth(), 1);
1073
1074		assert_changes(&mut changeset, &all_changes);
1075
1076		changeset.rollback_transaction().unwrap();
1077		assert_eq!(changeset.transaction_depth(), 0);
1078
1079		let rolled_back: Changes =
1080			vec![(b"key0", (Some(b"val0-1"), vec![1, 10])), (b"key1", (Some(b"val1"), vec![1]))];
1081		assert_changes(&mut changeset, &rolled_back);
1082
1083		assert_drained_changes(changeset, rolled_back);
1084	}
1085
1086	#[test]
1087	fn append_works() {
1088		use codec::Encode;
1089		let mut changeset = OverlayedChangeSet::default();
1090		assert_eq!(changeset.transaction_depth(), 0);
1091		let init = || vec![b"valinit".to_vec()].encode();
1092
1093		// committed set
1094		let val0 = vec![b"val0".to_vec()].encode();
1095		changeset.set(b"key0".to_vec(), Some(val0.clone()), Some(0));
1096		changeset.set(b"key1".to_vec(), None, Some(1));
1097		let all_changes: Changes =
1098			vec![(b"key0", (Some(val0.as_slice()), vec![0])), (b"key1", (None, vec![1]))];
1099
1100		assert_changes(&mut changeset, &all_changes);
1101		changeset.append_storage(b"key3".to_vec(), b"-modified".to_vec().encode(), init, Some(3));
1102		let val3 = vec![b"valinit".to_vec(), b"-modified".to_vec()].encode();
1103		let all_changes: Changes = vec![
1104			(b"key0", (Some(val0.as_slice()), vec![0])),
1105			(b"key1", (None, vec![1])),
1106			(b"key3", (Some(val3.as_slice()), vec![3])),
1107		];
1108		assert_changes(&mut changeset, &all_changes);
1109
1110		changeset.start_transaction();
1111		assert_eq!(changeset.transaction_depth(), 1);
1112		changeset.start_transaction();
1113		assert_eq!(changeset.transaction_depth(), 2);
1114
1115		// non existing value -> init value should be returned
1116		changeset.append_storage(b"key3".to_vec(), b"-twice".to_vec().encode(), init, Some(15));
1117
1118		// non existing value -> init value should be returned
1119		changeset.append_storage(b"key2".to_vec(), b"-modified".to_vec().encode(), init, Some(2));
1120		// existing value should be reuse on append
1121		changeset.append_storage(b"key0".to_vec(), b"-modified".to_vec().encode(), init, Some(10));
1122
1123		// should work for deleted keys
1124		changeset.append_storage(
1125			b"key1".to_vec(),
1126			b"deleted-modified".to_vec().encode(),
1127			init,
1128			Some(20),
1129		);
1130		let val0_2 = vec![b"val0".to_vec(), b"-modified".to_vec()].encode();
1131		let val3_2 = vec![b"valinit".to_vec(), b"-modified".to_vec(), b"-twice".to_vec()].encode();
1132		let val1 = vec![b"deleted-modified".to_vec()].encode();
1133		let all_changes: Changes = vec![
1134			(b"key0", (Some(val0_2.as_slice()), vec![0, 10])),
1135			(b"key1", (Some(val1.as_slice()), vec![1, 20])),
1136			(b"key2", (Some(val3.as_slice()), vec![2])),
1137			(b"key3", (Some(val3_2.as_slice()), vec![3, 15])),
1138		];
1139		assert_changes(&mut changeset, &all_changes);
1140
1141		changeset.start_transaction();
1142		let val3_3 =
1143			vec![b"valinit".to_vec(), b"-modified".to_vec(), b"-twice".to_vec(), b"-2".to_vec()]
1144				.encode();
1145		changeset.append_storage(b"key3".to_vec(), b"-2".to_vec().encode(), init, Some(21));
1146		let all_changes2: Changes = vec![
1147			(b"key0", (Some(val0_2.as_slice()), vec![0, 10])),
1148			(b"key1", (Some(val1.as_slice()), vec![1, 20])),
1149			(b"key2", (Some(val3.as_slice()), vec![2])),
1150			(b"key3", (Some(val3_3.as_slice()), vec![3, 15, 21])),
1151		];
1152		assert_changes(&mut changeset, &all_changes2);
1153		changeset.rollback_transaction().unwrap();
1154
1155		assert_changes(&mut changeset, &all_changes);
1156		changeset.start_transaction();
1157		let val3_4 = vec![
1158			b"valinit".to_vec(),
1159			b"-modified".to_vec(),
1160			b"-twice".to_vec(),
1161			b"-thrice".to_vec(),
1162		]
1163		.encode();
1164		changeset.append_storage(b"key3".to_vec(), b"-thrice".to_vec().encode(), init, Some(25));
1165		let all_changes: Changes = vec![
1166			(b"key0", (Some(val0_2.as_slice()), vec![0, 10])),
1167			(b"key1", (Some(val1.as_slice()), vec![1, 20])),
1168			(b"key2", (Some(val3.as_slice()), vec![2])),
1169			(b"key3", (Some(val3_4.as_slice()), vec![3, 15, 25])),
1170		];
1171		assert_changes(&mut changeset, &all_changes);
1172		changeset.commit_transaction().unwrap();
1173		changeset.commit_transaction().unwrap();
1174		assert_eq!(changeset.transaction_depth(), 1);
1175		assert_changes(&mut changeset, &all_changes);
1176
1177		changeset.rollback_transaction().unwrap();
1178		assert_eq!(changeset.transaction_depth(), 0);
1179		let rolled_back: Changes = vec![
1180			(b"key0", (Some(val0.as_slice()), vec![0])),
1181			(b"key1", (None, vec![1])),
1182			(b"key3", (Some(val3.as_slice()), vec![3])),
1183		];
1184		assert_changes(&mut changeset, &rolled_back);
1185		assert_drained_changes(changeset, rolled_back);
1186	}
1187
1188	#[test]
1189	fn clear_works() {
1190		let mut changeset = OverlayedChangeSet::default();
1191
1192		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
1193		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
1194		changeset.set(b"del1".to_vec(), Some(b"delval1".to_vec()), Some(3));
1195		changeset.set(b"del2".to_vec(), Some(b"delval2".to_vec()), Some(4));
1196
1197		changeset.start_transaction();
1198
1199		changeset.clear_where(|k, _| k.starts_with(b"del"), Some(5));
1200
1201		assert_changes(
1202			&mut changeset,
1203			&vec![
1204				(b"del1", (None, vec![3, 5])),
1205				(b"del2", (None, vec![4, 5])),
1206				(b"key0", (Some(b"val0"), vec![1])),
1207				(b"key1", (Some(b"val1"), vec![2])),
1208			],
1209		);
1210
1211		changeset.rollback_transaction().unwrap();
1212
1213		assert_changes(
1214			&mut changeset,
1215			&vec![
1216				(b"del1", (Some(b"delval1"), vec![3])),
1217				(b"del2", (Some(b"delval2"), vec![4])),
1218				(b"key0", (Some(b"val0"), vec![1])),
1219				(b"key1", (Some(b"val1"), vec![2])),
1220			],
1221		);
1222	}
1223
1224	#[test]
1225	fn next_change_works() {
1226		let mut changeset = OverlayedChangeSet::default();
1227
1228		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(0));
1229		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(1));
1230		changeset.set(b"key2".to_vec(), Some(b"val2".to_vec()), Some(2));
1231
1232		changeset.start_transaction();
1233
1234		changeset.set(b"key3".to_vec(), Some(b"val3".to_vec()), Some(3));
1235		changeset.set(b"key4".to_vec(), Some(b"val4".to_vec()), Some(4));
1236		changeset.set(b"key11".to_vec(), Some(b"val11".to_vec()), Some(11));
1237
1238		assert_eq!(changeset.changes_after(b"key0").next().unwrap().0, b"key1");
1239		assert_eq!(
1240			changeset.changes_after(b"key0").next().unwrap().1.value(),
1241			Some(&b"val1".to_vec())
1242		);
1243		assert_eq!(changeset.changes_after(b"key1").next().unwrap().0, b"key11");
1244		assert_eq!(
1245			changeset.changes_after(b"key1").next().unwrap().1.value(),
1246			Some(&b"val11".to_vec())
1247		);
1248		assert_eq!(changeset.changes_after(b"key11").next().unwrap().0, b"key2");
1249		assert_eq!(
1250			changeset.changes_after(b"key11").next().unwrap().1.value(),
1251			Some(&b"val2".to_vec())
1252		);
1253		assert_eq!(changeset.changes_after(b"key2").next().unwrap().0, b"key3");
1254		assert_eq!(
1255			changeset.changes_after(b"key2").next().unwrap().1.value(),
1256			Some(&b"val3".to_vec())
1257		);
1258		assert_eq!(changeset.changes_after(b"key3").next().unwrap().0, b"key4");
1259		assert_eq!(
1260			changeset.changes_after(b"key3").next().unwrap().1.value(),
1261			Some(&b"val4".to_vec())
1262		);
1263		assert_eq!(changeset.changes_after(b"key4").next(), None);
1264
1265		changeset.rollback_transaction().unwrap();
1266
1267		assert_eq!(changeset.changes_after(b"key0").next().unwrap().0, b"key1");
1268		assert_eq!(
1269			changeset.changes_after(b"key0").next().unwrap().1.value(),
1270			Some(&b"val1".to_vec())
1271		);
1272		assert_eq!(changeset.changes_after(b"key1").next().unwrap().0, b"key2");
1273		assert_eq!(
1274			changeset.changes_after(b"key1").next().unwrap().1.value(),
1275			Some(&b"val2".to_vec())
1276		);
1277		assert_eq!(changeset.changes_after(b"key11").next().unwrap().0, b"key2");
1278		assert_eq!(
1279			changeset.changes_after(b"key11").next().unwrap().1.value(),
1280			Some(&b"val2".to_vec())
1281		);
1282		assert_eq!(changeset.changes_after(b"key2").next(), None);
1283		assert_eq!(changeset.changes_after(b"key3").next(), None);
1284		assert_eq!(changeset.changes_after(b"key4").next(), None);
1285	}
1286
1287	#[test]
1288	fn no_open_tx_commit_errors() {
1289		let mut changeset = OverlayedChangeSet::default();
1290		assert_eq!(changeset.transaction_depth(), 0);
1291		assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
1292	}
1293
1294	#[test]
1295	fn no_open_tx_rollback_errors() {
1296		let mut changeset = OverlayedChangeSet::default();
1297		assert_eq!(changeset.transaction_depth(), 0);
1298		assert_eq!(changeset.rollback_transaction(), Err(NoOpenTransaction));
1299	}
1300
1301	#[test]
1302	fn unbalanced_transactions_errors() {
1303		let mut changeset = OverlayedChangeSet::default();
1304		changeset.start_transaction();
1305		changeset.commit_transaction().unwrap();
1306		assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
1307	}
1308
1309	#[test]
1310	#[should_panic]
1311	fn drain_with_open_transaction_panics() {
1312		let mut changeset = OverlayedChangeSet::default();
1313		changeset.start_transaction();
1314		let _ = changeset.drain_committed();
1315	}
1316
1317	#[test]
1318	fn runtime_cannot_close_client_tx() {
1319		let mut changeset = OverlayedChangeSet::default();
1320		changeset.start_transaction();
1321		changeset.enter_runtime().unwrap();
1322		changeset.start_transaction();
1323		changeset.commit_transaction().unwrap();
1324		assert_eq!(changeset.commit_transaction(), Err(NoOpenTransaction));
1325		assert_eq!(changeset.rollback_transaction(), Err(NoOpenTransaction));
1326	}
1327
1328	#[test]
1329	fn exit_runtime_closes_runtime_tx() {
1330		let mut changeset = OverlayedChangeSet::default();
1331
1332		changeset.start_transaction();
1333
1334		changeset.set(b"key0".to_vec(), Some(b"val0".to_vec()), Some(1));
1335
1336		changeset.enter_runtime().unwrap();
1337		changeset.start_transaction();
1338		changeset.set(b"key1".to_vec(), Some(b"val1".to_vec()), Some(2));
1339		changeset.exit_runtime().unwrap();
1340
1341		changeset.commit_transaction().unwrap();
1342		assert_eq!(changeset.transaction_depth(), 0);
1343
1344		assert_drained(changeset, vec![(b"key0", Some(b"val0"))]);
1345	}
1346
1347	#[test]
1348	fn enter_exit_runtime_fails_when_already_in_requested_mode() {
1349		let mut changeset = OverlayedChangeSet::default();
1350
1351		assert_eq!(changeset.exit_runtime(), Err(NotInRuntime));
1352		assert_eq!(changeset.enter_runtime(), Ok(()));
1353		assert_eq!(changeset.enter_runtime(), Err(AlreadyInRuntime));
1354		assert_eq!(changeset.exit_runtime(), Ok(()));
1355		assert_eq!(changeset.exit_runtime(), Err(NotInRuntime));
1356	}
1357
1358	#[test]
1359	fn restore_append_to_parent() {
1360		use codec::{Compact, Encode};
1361		let mut changeset = OverlayedChangeSet::default();
1362		let key: Vec<u8> = b"akey".into();
1363
1364		let from = 50; // 1 byte len
1365		let to = 100; // 2 byte len
1366		for i in 0..from {
1367			changeset.append_storage(key.clone(), vec![i], Default::default, None);
1368		}
1369
1370		// materialized
1371		let encoded = changeset.get(&key).unwrap().value().unwrap();
1372		let encoded_from_len = Compact(from as u32).encode();
1373		assert_eq!(encoded_from_len.len(), 1);
1374		assert!(encoded.starts_with(&encoded_from_len[..]));
1375		let encoded_from = encoded.clone();
1376
1377		changeset.start_transaction();
1378
1379		for i in from..to {
1380			changeset.append_storage(key.clone(), vec![i], Default::default, None);
1381		}
1382
1383		// materialized
1384		let encoded = changeset.get(&key).unwrap().value().unwrap();
1385		let encoded_to_len = Compact(to as u32).encode();
1386		assert_eq!(encoded_to_len.len(), 2);
1387		assert!(encoded.starts_with(&encoded_to_len[..]));
1388
1389		changeset.rollback_transaction().unwrap();
1390
1391		let encoded = changeset.get(&key).unwrap().value().unwrap();
1392		assert_eq!(&encoded_from, encoded);
1393	}
1394
1395	/// First we have some `Set` operation with a valid SCALE list. Then we append data and rollback
1396	/// afterwards.
1397	#[test]
1398	fn restore_initial_set_after_append_to_parent() {
1399		use codec::{Compact, Encode};
1400		let mut changeset = OverlayedChangeSet::default();
1401		let key: Vec<u8> = b"akey".into();
1402
1403		let initial_data = vec![1u8; 50].encode();
1404
1405		changeset.set(key.clone(), Some(initial_data.clone()), None);
1406
1407		changeset.start_transaction();
1408
1409		// Append until we require 2 bytes for the length prefix.
1410		for i in 0..50 {
1411			changeset.append_storage(key.clone(), vec![i], Default::default, None);
1412		}
1413
1414		// Materialize the value.
1415		let encoded = changeset.get(&key).unwrap().value().unwrap();
1416		let encoded_to_len = Compact(100u32).encode();
1417		assert_eq!(encoded_to_len.len(), 2);
1418		assert!(encoded.starts_with(&encoded_to_len[..]));
1419
1420		changeset.rollback_transaction().unwrap();
1421
1422		let encoded = changeset.get(&key).unwrap().value().unwrap();
1423		assert_eq!(&initial_data, encoded);
1424	}
1425}