Make it possible to calculate the storage root as often as you want

frame/authorship/src/lib.rs

@@ -582,7 +582,6 @@ mod tests {
 				&number,
 				&hash,
 				&Default::default(),
-				&Default::default(),
 				Default::default()
 			);
 
@@ -681,7 +680,6 @@ mod tests {
 			System::initialize(
 				&1,
 				&Default::default(),
-				&Default::default(),
 				header.digest(),
 				Default::default(),
 			);

frame/babe/src/mock.rs

@@ -259,7 +259,7 @@ pub fn go_to_block(n: u64, s: u64) {
 
 	let pre_digest = make_secondary_plain_pre_digest(0, s);
 
-	System::initialize(&n, &parent_hash, &Default::default(), &pre_digest, InitKind::Full);
+	System::initialize(&n, &parent_hash, &pre_digest, InitKind::Full);
 	System::set_block_number(n);
 	Timestamp::set_timestamp(n);
 
@@ -447,7 +447,7 @@ pub fn generate_equivocation_proof(
 	let make_header = || {
 		let parent_hash = System::parent_hash();
 		let pre_digest = make_secondary_plain_pre_digest(offender_authority_index, slot_number);
-		System::initialize(&current_block, &parent_hash, &Default::default(), &pre_digest, InitKind::Full);
+		System::initialize(&current_block, &parent_hash, &pre_digest, InitKind::Full);
 		System::set_block_number(current_block);
 		Timestamp::set_timestamp(current_block);
 		System::finalize()

frame/babe/src/tests.rs

@@ -77,7 +77,6 @@ fn first_block_epoch_zero_start() {
 		System::initialize(
 			&1,
 			&Default::default(),
-			&Default::default(),
 			&pre_digest,
 			Default::default(),
 		);
@@ -128,7 +127,6 @@ fn author_vrf_output_for_primary() {
 		System::initialize(
 			&1,
 			&Default::default(),
-			&Default::default(),
 			&primary_pre_digest,
 			Default::default(),
 		);
@@ -155,7 +153,6 @@ fn author_vrf_output_for_secondary_vrf() {
 		System::initialize(
 			&1,
 			&Default::default(),
-			&Default::default(),
 			&secondary_vrf_pre_digest,
 			Default::default(),
 		);
@@ -179,7 +176,6 @@ fn no_author_vrf_output_for_secondary_plain() {
 		System::initialize(
 			&1,
 			&Default::default(),
-			&Default::default(),
 			&secondary_plain_pre_digest,
 			Default::default(),
 		);

frame/contracts/src/tests.rs

@@ -698,7 +698,6 @@ fn initialize_block(number: u64) {
 	System::initialize(
 		&number,
 		&[0u8; 32].into(),
-		&[0u8; 32].into(),
 		&Default::default(),
 		Default::default(),
 	);

frame/executive/src/lib.rs

@@ -130,7 +130,7 @@ use sp_runtime::{
 	transaction_validity::{TransactionValidity, TransactionSource},
 };
 use codec::{Codec, Encode};
-use frame_system::{extrinsics_root, DigestOf};
+use frame_system::DigestOf;
 
 /// Trait that can be used to execute a block.
 pub trait ExecuteBlock<Block: BlockT> {
@@ -213,7 +213,6 @@ where
 		Self::initialize_block_impl(
 			header.number(),
 			header.parent_hash(),
-			header.extrinsics_root(),
 			&digests
 		);
 	}
@@ -231,7 +230,6 @@ where
 	fn initialize_block_impl(
 		block_number: &System::BlockNumber,
 		parent_hash: &System::Hash,
-		extrinsics_root: &System::Hash,
 		digest: &Digest<System::Hash>,
 	) {
 		let mut weight = 0;
@@ -244,7 +242,6 @@ where
 		<frame_system::Module<System>>::initialize(
 			block_number,
 			parent_hash,
-			extrinsics_root,
 			digest,
 			frame_system::InitKind::Full,
 		);
@@ -286,13 +283,8 @@ where
 		assert!(
 			n > System::BlockNumber::zero()
 			&& <frame_system::Module<System>>::block_hash(n - System::BlockNumber::one()) == *header.parent_hash(),
-			"Parent hash should be valid."
+			"Parent hash should be valid.",
 		);
-
-		// Check that transaction trie root represents the transactions.
-		let xts_root = extrinsics_root::<System::Hashing, _>(&block.extrinsics());
-		header.extrinsics_root().check_equal(&xts_root);
-		assert!(header.extrinsics_root() == &xts_root, "Transaction trie root must be valid.");
 	}
 
 	/// Actually execute all transitions for `block`.
@@ -322,8 +314,14 @@ where
 	}
 
 	/// Execute given extrinsics and take care of post-extrinsics book-keeping.
-	fn execute_extrinsics_with_book_keeping(extrinsics: Vec<Block::Extrinsic>, block_number: NumberFor<Block>) {
-		extrinsics.into_iter().for_each(Self::apply_extrinsic_no_note);
+	fn execute_extrinsics_with_book_keeping(
+		extrinsics: Vec<Block::Extrinsic>,
+		block_number: NumberFor<Block>,
+	) {
+		extrinsics.into_iter().for_each(|e| if let Err(e) = Self::apply_extrinsic(e) {
+			let err: &'static str = e.into();
+			panic!(err)
+		});
 
 		// post-extrinsics book-keeping
 		<frame_system::Module<System>>::note_finished_extrinsics();
@@ -341,8 +339,6 @@ where
 		<frame_system::Module<System> as OnFinalize<System::BlockNumber>>::on_finalize(block_number);
 		<AllModules as OnFinalize<System::BlockNumber>>::on_finalize(block_number);
 
-		// set up extrinsics
-		<frame_system::Module<System>>::derive_extrinsics();
 		<frame_system::Module<System>>::finalize()
 	}
 
@@ -354,23 +350,14 @@ where
 		sp_io::init_tracing();
 		let encoded = uxt.encode();
 		let encoded_len = encoded.len();
-		Self::apply_extrinsic_with_len(uxt, encoded_len, Some(encoded))
-	}
-
-	/// Apply an extrinsic inside the block execution function.
-	fn apply_extrinsic_no_note(uxt: Block::Extrinsic) {
-		let l = uxt.encode().len();
-		match Self::apply_extrinsic_with_len(uxt, l, None) {
-			Ok(_) => (),
-			Err(e) => { let err: &'static str = e.into(); panic!(err) },
-		}
+		Self::apply_extrinsic_with_len(uxt, encoded_len, encoded)
 	}
 
 	/// Actually apply an extrinsic given its `encoded_len`; this doesn't note its hash.
 	fn apply_extrinsic_with_len(
 		uxt: Block::Extrinsic,
 		encoded_len: usize,
-		to_note: Option<Vec<u8>>,
+		to_note: Vec<u8>,
 	) -> ApplyExtrinsicResult {
 		sp_tracing::enter_span!(
 			sp_tracing::info_span!("apply_extrinsic",
@@ -382,9 +369,7 @@ where
 		// We don't need to make sure to `note_extrinsic` only after we know it's going to be
 		// executed to prevent it from leaking in storage since at this point, it will either
 		// execute or panic (and revert storage changes).
-		if let Some(encoded) = to_note {
-			<frame_system::Module<System>>::note_extrinsic(encoded);
-		}
+		<frame_system::Module<System>>::note_extrinsic(to_note);
 
 		// AUDIT: Under no circumstances may this function panic from here onwards.
 
@@ -418,6 +403,11 @@ where
 		let storage_root = new_header.state_root();
 		header.state_root().check_equal(&storage_root);
 		assert!(header.state_root() == storage_root, "Storage root must match that calculated.");
+
+		assert!(
+			header.extrinsics_root() == new_header.extrinsics_root(),
+			"Transaction trie root must be valid.",
+		);
 	}
 
 	/// Check a given signed transaction for validity. This doesn't execute any
@@ -462,7 +452,6 @@ where
 		<frame_system::Module<System>>::initialize(
 			header.number(),
 			header.parent_hash(),
-			header.extrinsics_root(),
 			&digests,
 			frame_system::InitKind::Inspection,
 		);
@@ -558,6 +547,12 @@ mod tests {
 				fn offchain_worker(n: T::BlockNumber) {
 					assert_eq!(T::BlockNumber::from(1u32), n);
 				}
+
+				#[weight = 0]
+				fn calculate_storage_root(origin) {
+					let root = sp_io::storage::root();
+					sp_io::storage::set("storage_root".as_bytes(), &root);
+				}
 			}
 		}
 
@@ -1153,4 +1148,29 @@ mod tests {
 			assert_eq!(header.hash(), System::block_hash(1));
 		});
 	}
+
+	#[test]
+	fn calculating_storage_root_twice_works() {
+		let call = Call::Custom(custom::Call::calculate_storage_root());
+		let xt = TestXt::new(call, sign_extra(1, 0, 0));
+
+		let header = new_test_ext(1).execute_with(|| {
+			// Let's build some fake block.
+			Executive::initialize_block(&Header::new(
+				1,
+				H256::default(),
+				H256::default(),
+				[69u8; 32].into(),
+				Digest::default(),
+			));
+
+			Executive::apply_extrinsic(xt.clone()).unwrap().unwrap();
+
+			Executive::finalize_block()
+		});
+
+		new_test_ext(1).execute_with(|| {
+			Executive::execute_block(Block::new(header, vec![xt]));
+		});
+	}
 }

frame/grandpa/src/mock.rs

@@ -359,7 +359,6 @@ pub fn start_session(session_index: SessionIndex) {
 			&(i as u64 + 1),
 			&parent_hash,
 			&Default::default(),
-			&Default::default(),
 			Default::default(),
 		);
 		System::set_block_number((i + 1).into());
@@ -384,7 +383,6 @@ pub fn initialize_block(number: u64, parent_hash: H256) {
 		&number,
 		&parent_hash,
 		&Default::default(),
-		&Default::default(),
 		Default::default(),
 	);
 }

frame/merkle-mountain-range/src/tests.rs

@@ -46,7 +46,6 @@ fn new_block() -> u64 {
 		&number,
 		&hash,
 		&Default::default(),
-		&Default::default(),
 		frame_system::InitKind::Full,
 	);
 	MMR::on_initialize(number)

frame/randomness-collective-flip/src/lib.rs

@@ -205,7 +205,6 @@ mod tests {
 				&i,
 				&parent_hash,
 				&Default::default(),
-				&Default::default(),
 				frame_system::InitKind::Full,
 			);
 			CollectiveFlip::on_initialize(i);

frame/system/src/lib.rs

@@ -107,7 +107,7 @@ use sp_runtime::{
 		self, CheckEqual, AtLeast32Bit, Zero, Lookup, LookupError,
 		SimpleBitOps, Hash, Member, MaybeDisplay, BadOrigin,
 		MaybeSerialize, MaybeSerializeDeserialize, MaybeMallocSizeOf, StaticLookup, One, Bounded,
-		Dispatchable, AtLeast32BitUnsigned
+		Dispatchable, AtLeast32BitUnsigned, Saturating,
 	},
 	offchain::storage_lock::BlockNumberProvider,
 };
@@ -405,9 +405,6 @@ decl_storage! {
 		/// Hash of the previous block.
 		ParentHash get(fn parent_hash) build(|_| hash69()): T::Hash;
 
-		/// Extrinsics root of the current block, also part of the block header.
-		ExtrinsicsRoot get(fn extrinsics_root): T::Hash;
-
 		/// Digest of the current block, also part of the block header.
 		Digest get(fn digest): DigestOf<T>;
 
@@ -989,7 +986,6 @@ impl<T: Config> Module<T> {
 	pub fn initialize(
 		number: &T::BlockNumber,
 		parent_hash: &T::Hash,
-		txs_root: &T::Hash,
 		digest: &DigestOf<T>,
 		kind: InitKind,
 	) {
@@ -1000,7 +996,6 @@ impl<T: Config> Module<T> {
 		<Digest<T>>::put(digest);
 		<ParentHash<T>>::put(parent_hash);
 		<BlockHash<T>>::insert(*number - One::one(), parent_hash);
-		<ExtrinsicsRoot<T>>::put(txs_root);
 
 		// Remove previous block data from storage
 		BlockWeight::kill();
@@ -1017,7 +1012,6 @@ impl<T: Config> Module<T> {
 	/// resulting header for this block.
 	pub fn finalize() -> T::Header {
 		ExecutionPhase::kill();
-		ExtrinsicCount::kill();
 		AllExtrinsicsLen::kill();
 
 		// The following fields
@@ -1034,17 +1028,18 @@ impl<T: Config> Module<T> {
 		let parent_hash = <ParentHash<T>>::get();
 		let mut digest = <Digest<T>>::get();
 
-		let extrinsics_root = <ExtrinsicsRoot<T>>::take();
+		let extrinsics = (0..ExtrinsicCount::take().unwrap_or_default())
+			.map(ExtrinsicData::take)
+			.collect();
+		let extrinsics_root = extrinsics_data_root::<T::Hashing>(extrinsics);
 
 		// move block hash pruning window by one block
-		let block_hash_count = <T::BlockHashCount>::get();
-		if number > block_hash_count {
-			let to_remove = number - block_hash_count - One::one();
+		let block_hash_count = T::BlockHashCount::get();
+		let to_remove = number.saturating_sub(block_hash_count).saturating_sub(One::one());
 
-			// keep genesis hash
-			if to_remove != Zero::zero() {
-				<BlockHash<T>>::remove(to_remove);
-			}
+		// keep genesis hash
+		if !to_remove.is_zero() {
+			<BlockHash<T>>::remove(to_remove);
 		}
 
 		let storage_root = T::Hash::decode(&mut &sp_io::storage::root()[..])
@@ -1138,12 +1133,10 @@ impl<T: Config> Module<T> {
 		Account::<T>::mutate(who, |a| a.nonce += T::Index::one());
 	}
 
-	/// Note what the extrinsic data of the current extrinsic index is. If this
-	/// is called, then ensure `derive_extrinsics` is also called before
-	/// block-building is completed.
+	/// Note what the extrinsic data of the current extrinsic index is.
 	///
-	/// NOTE: This function is called only when the block is being constructed locally.
-	/// `execute_block` doesn't note any extrinsics.
+	/// This is required to be called before applying an extrinsic. The data will used
+	/// in [`Self::finalize`] to calculate the correct extrinsics root.
 	pub fn note_extrinsic(encoded_xt: Vec<u8>) {
 		ExtrinsicData::insert(Self::extrinsic_index().unwrap_or_default(), encoded_xt);
 	}
@@ -1182,14 +1175,6 @@ impl<T: Config> Module<T> {
 		ExecutionPhase::put(Phase::ApplyExtrinsic(0))
 	}
 
-	/// Remove all extrinsic data and save the extrinsics trie root.
-	pub fn derive_extrinsics() {
-		let extrinsics = (0..ExtrinsicCount::get().unwrap_or_default())
-			.map(ExtrinsicData::take).collect();
-		let xts_root = extrinsics_data_root::<T::Hashing>(extrinsics);
-		<ExtrinsicsRoot<T>>::put(xts_root);
-	}
-
 	/// An account is being created.
 	pub fn on_created_account(who: T::AccountId) {
 		T::OnNewAccount::on_new_account(&who);