@@ -7,7 +7,8 @@ use polkadot_sdk::{
77 } ,
88 pallet_revive:: {
99 self , AccountId32Mapper , AccountInfo , AddressMapper , BytecodeType , ContractInfo ,
10- ExecConfig , Executable , HoldReason , Pallet , ResourceMeter ,
10+ ExecConfig , Executable , HoldReason , Pallet , ResourceMeter , TRANSIENT_STORAGE_LIMIT ,
11+ TransientStorage ,
1112 } ,
1213 sp_core:: { self , H160 , H256 } ,
1314 sp_externalities:: Externalities ,
@@ -17,13 +18,15 @@ use polkadot_sdk::{
1718} ;
1819use revive_env:: { Balances , BlockAuthor , ExtBuilder , NativeToEthRatio , Runtime , System , Timestamp } ;
1920use std:: {
21+ cell:: RefCell ,
2022 fmt:: Debug ,
2123 sync:: { Arc , Mutex } ,
2224} ;
2325
2426pub ( crate ) struct Inner {
2527 pub externalities : TestExternalities ,
2628 pub depth : usize ,
29+ pub transient_storage : TransientStorage < Runtime > ,
2730}
2831
2932#[ derive( Default ) ]
@@ -39,6 +42,7 @@ impl Default for Inner {
3942 ) ] )
4043 . build ( ) ,
4144 depth : 0 ,
45+ transient_storage : TransientStorage :: new ( TRANSIENT_STORAGE_LIMIT ) ,
4246 }
4347 }
4448}
@@ -52,8 +56,10 @@ impl Debug for TestEnv {
5256impl Clone for TestEnv {
5357 fn clone ( & self ) -> Self {
5458 let mut inner: Inner = Default :: default ( ) ;
55- inner. externalities . backend = self . 0 . lock ( ) . unwrap ( ) . externalities . as_backend ( ) ;
56- inner. depth = self . 0 . lock ( ) . unwrap ( ) . depth ;
59+ let mut data = self . 0 . lock ( ) . unwrap ( ) ;
60+ inner. externalities . backend = data. externalities . as_backend ( ) ;
61+ inner. transient_storage = data. transient_storage . clone ( ) ;
62+ inner. depth = data. depth ;
5763 Self ( Arc :: new ( Mutex :: new ( inner) ) )
5864 }
5965}
@@ -67,22 +73,43 @@ impl TestEnv {
6773 let mut state = self . 0 . lock ( ) . unwrap ( ) ;
6874 state. depth += 1 ;
6975 state. externalities . ext ( ) . storage_start_transaction ( ) ;
76+ state. transient_storage . start_transaction ( ) ;
7077 }
7178
7279 pub fn revert ( & mut self , depth : usize ) {
7380 let mut state = self . 0 . lock ( ) . unwrap ( ) ;
7481 while state. depth > depth + 1 {
7582 state. externalities . ext ( ) . storage_rollback_transaction ( ) . unwrap ( ) ;
83+ state. transient_storage . rollback_transaction ( ) ;
7684 state. depth -= 1 ;
7785 }
7886 state. externalities . ext ( ) . storage_rollback_transaction ( ) . unwrap ( ) ;
87+ state. transient_storage . rollback_transaction ( ) ;
7988 state. externalities . ext ( ) . storage_start_transaction ( ) ;
89+ state. transient_storage . start_transaction ( ) ;
8090 }
8191
8292 pub fn execute_with < R , F : FnOnce ( ) -> R > ( & mut self , f : F ) -> R {
8393 self . 0 . lock ( ) . unwrap ( ) . externalities . execute_with ( f)
8494 }
8595
96+ pub fn execute_with_transient_storage <
97+ R ,
98+ F : FnOnce ( RefCell < TransientStorage < Runtime > > ) -> ( R , RefCell < TransientStorage < Runtime > > ) ,
99+ > (
100+ & mut self ,
101+ f : F ,
102+ ) -> R {
103+ let mut data = self . 0 . lock ( ) . unwrap ( ) ;
104+ let mut temp = TransientStorage :: new ( 0 ) ;
105+ std:: mem:: swap ( & mut data. transient_storage , & mut temp) ;
106+ let transient_storage = RefCell :: new ( temp) ;
107+ let to_use = transient_storage. clone ( ) ;
108+ let ( result, transient_storage) = data. externalities . execute_with ( || f ( to_use) ) ;
109+ data. transient_storage = transient_storage. into_inner ( ) ;
110+ result
111+ }
112+
86113 pub fn get_nonce ( & mut self , account : Address ) -> u32 {
87114 self . 0 . lock ( ) . unwrap ( ) . externalities . execute_with ( || {
88115 System :: account_nonce ( AccountId32Mapper :: < Runtime > :: to_fallback_account_id (
0 commit comments