1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

#![cfg_attr(not(feature = "std"), no_std)]

use frame_support::storage::{with_transaction, TransactionOutcome};
use sp_runtime::DispatchError;
use sp_std::result::Result;

pub mod offchain_worker;
pub mod ordered_set;

pub use offchain_worker::OffchainErr;
pub use ordered_set::OrderedSet;

/// Execute the supplied function in a new storage transaction.
///
/// All changes to storage performed by the supplied function are discarded if
/// the returned outcome is `Result::Err`.
///
/// Transactions can be nested to any depth. Commits happen to the parent
/// transaction.
pub fn with_transaction_result<R>(f: impl FnOnce() -> Result<R, DispatchError>) -> Result<R, DispatchError> {
	with_transaction(|| {
		let res = f();
		if res.is_ok() {
			TransactionOutcome::Commit(res)
		} else {
			TransactionOutcome::Rollback(res)
		}
	})
}

/// Simulate execution of the supplied function in a new storage transaction.
/// Changes to storage performed by the supplied function are always discarded.
pub fn simulate_execution<R>(f: impl FnOnce() -> Result<R, DispatchError>) -> Result<R, DispatchError> {
	with_transaction(|| {
		let res = f();
		TransactionOutcome::Rollback(res)
	})
}

#[cfg(test)]
mod tests {
	use super::*;
	use frame_support::{assert_noop, assert_ok, construct_runtime, derive_impl, pallet_prelude::*};
	use sp_io::TestExternalities;
	use sp_runtime::traits::IdentityLookup;
	use sp_runtime::{DispatchError, DispatchResult};
	use sp_std::result::Result;

	#[allow(dead_code)]
	#[frame_support::pallet]
	pub mod module {
		use super::*;

		#[pallet::config]
		pub trait Config: frame_system::Config {}

		#[pallet::pallet]
		pub struct Pallet<T>(_);

		#[pallet::storage]
		pub type Value<T: Config> = StorageValue<_, u32, ValueQuery>;

		#[pallet::storage]
		pub type Map<T: Config> = StorageMap<_, Twox64Concat, [u8; 4], u32, ValueQuery>;
	}

	use module::*;

	#[derive_impl(frame_system::config_preludes::TestDefaultConfig as frame_system::DefaultConfig)]
	impl frame_system::Config for Runtime {
		type AccountId = u128;
		type Lookup = IdentityLookup<Self::AccountId>;
		type Block = Block;
	}

	impl module::Config for Runtime {}

	type Block = frame_system::mocking::MockBlock<Runtime>;

	construct_runtime!(
		pub enum Runtime {
			System: frame_system,
			TestModule: module,
		}
	);

	#[test]
	fn storage_transaction_basic_commit() {
		TestExternalities::default().execute_with(|| {
			assert_eq!(Value::<Runtime>::get(), 0);
			assert!(!Map::<Runtime>::contains_key(b"val0"));

			assert_ok!(with_transaction_result(|| -> DispatchResult {
				Value::<Runtime>::set(99);
				Map::<Runtime>::insert(b"val0", 99);
				assert_eq!(Value::<Runtime>::get(), 99);
				assert_eq!(Map::<Runtime>::get(b"val0"), 99);
				Ok(())
			}));

			assert_eq!(Value::<Runtime>::get(), 99);
			assert_eq!(Map::<Runtime>::get(b"val0"), 99);
		});
	}

	#[test]
	fn storage_transaction_basic_rollback() {
		TestExternalities::default().execute_with(|| {
			assert_eq!(Value::<Runtime>::get(), 0);
			assert_eq!(Map::<Runtime>::get(b"val0"), 0);

			assert_noop!(
				with_transaction_result(|| -> DispatchResult {
					Value::<Runtime>::set(99);
					Map::<Runtime>::insert(b"val0", 99);
					assert_eq!(Value::<Runtime>::get(), 99);
					assert_eq!(Map::<Runtime>::get(b"val0"), 99);
					Err("test".into())
				}),
				DispatchError::Other("test")
			);

			assert_eq!(Value::<Runtime>::get(), 0);
			assert_eq!(Map::<Runtime>::get(b"val0"), 0);
		});
	}

	#[test]
	fn simulate_execution_works() {
		TestExternalities::default().execute_with(|| {
			assert_eq!(Value::<Runtime>::get(), 0);
			assert_eq!(Map::<Runtime>::get(b"val0"), 0);

			// Roll back on `Err`.
			assert_noop!(
				simulate_execution(|| -> DispatchResult {
					Value::<Runtime>::set(99);
					Map::<Runtime>::insert(b"val0", 99);
					Err(DispatchError::Other("test"))
				}),
				DispatchError::Other("test")
			);
			assert_eq!(Value::<Runtime>::get(), 0);
			assert_eq!(Map::<Runtime>::get(b"val0"), 0);

			// Roll back on `Ok`, but returns `Ok` result.
			assert_ok!(
				simulate_execution(|| -> Result<u32, DispatchError> {
					Value::<Runtime>::set(99);
					Map::<Runtime>::insert(b"val0", 99);
					Ok(99)
				}),
				99
			);
			assert_eq!(Value::<Runtime>::get(), 0);
			assert_eq!(Map::<Runtime>::get(b"val0"), 0);
		});
	}
}