1
- use crate :: { node:: InMemoryNodeInner , utils:: bytecode_to_factory_dep} ;
2
- use anyhow:: { anyhow, Result } ;
1
+ use crate :: {
2
+ node:: { BlockContext , InMemoryNodeInner } ,
3
+ utils:: bytecode_to_factory_dep,
4
+ } ;
3
5
use ethers:: { abi:: AbiDecode , prelude:: abigen} ;
4
6
use multivm:: {
7
+ interface:: L1BatchEnv ,
5
8
vm_1_3_2:: zk_evm_1_3_3:: {
6
- tracing:: { AfterExecutionData , BeforeExecutionData , VmLocalStateData } ,
9
+ tracing:: { BeforeExecutionData , VmLocalStateData } ,
7
10
zkevm_opcode_defs:: all:: Opcode ,
8
11
zkevm_opcode_defs:: { FatPointer , CALL_IMPLICIT_CALLDATA_FAT_PTR_REGISTER } ,
9
12
} ,
10
- vm_virtual_blocks:: {
11
- DynTracer , ExecutionEndTracer , ExecutionProcessing , HistoryMode , SimpleMemory , VmTracer ,
12
- } ,
13
+ vm_latest:: { DynTracer , HistoryMode , SimpleMemory , VmTracer } ,
14
+ } ;
15
+ use std:: {
16
+ fmt:: Debug ,
17
+ sync:: { Arc , Mutex , RwLock } ,
13
18
} ;
14
- use std:: sync:: { Arc , RwLock } ;
15
- use zksync_basic_types:: { H160 , H256 } ;
19
+ use zksync_basic_types:: { AccountTreeId , H160 , H256 } ;
16
20
use zksync_state:: { StoragePtr , WriteStorage } ;
17
21
use zksync_types:: {
22
+ block:: { pack_block_info, unpack_block_info} ,
18
23
get_code_key, get_nonce_key,
19
24
utils:: { decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance} ,
25
+ StorageKey ,
20
26
} ;
21
27
use zksync_utils:: { h256_to_u256, u256_to_h256} ;
22
28
@@ -27,11 +33,12 @@ const CHEATCODE_ADDRESS: H160 = H160([
27
33
28
34
#[ derive( Clone , Debug , Default ) ]
29
35
pub struct CheatcodeTracer < F > {
30
- factory_deps : F ,
36
+ node_ctx : F ,
31
37
}
32
38
33
- pub trait FactoryDeps {
34
- fn store_factory_dep ( & mut self , hash : H256 , bytecode : Vec < u8 > ) -> Result < ( ) > ;
39
+ pub trait NodeCtx {
40
+ fn set_time ( & mut self , time : u64 ) ;
41
+ fn store_factory_dep ( & mut self , hash : H256 , bytecode : Vec < u8 > ) ;
35
42
}
36
43
37
44
abigen ! (
@@ -40,10 +47,11 @@ abigen!(
40
47
function deal(address who, uint256 newBalance)
41
48
function etch(address who, bytes calldata code)
42
49
function setNonce(address account, uint64 nonce)
50
+ function warp(uint256 timestamp)
43
51
]"#
44
52
) ;
45
53
46
- impl < F : FactoryDeps , S : WriteStorage , H : HistoryMode > DynTracer < S , H > for CheatcodeTracer < F > {
54
+ impl < F : NodeCtx , S : WriteStorage , H : HistoryMode > DynTracer < S , H > for CheatcodeTracer < F > {
47
55
fn before_execution (
48
56
& mut self ,
49
57
state : VmLocalStateData < ' _ > ,
@@ -86,16 +94,11 @@ impl<F: FactoryDeps, S: WriteStorage, H: HistoryMode> DynTracer<S, H> for Cheatc
86
94
}
87
95
}
88
96
89
- impl < F : FactoryDeps + Send , S : WriteStorage , H : HistoryMode > VmTracer < S , H > for CheatcodeTracer < F > { }
90
- impl < F : FactoryDeps , H : HistoryMode > ExecutionEndTracer < H > for CheatcodeTracer < F > { }
91
- impl < F : FactoryDeps , S : WriteStorage , H : HistoryMode > ExecutionProcessing < S , H >
92
- for CheatcodeTracer < F >
93
- {
94
- }
97
+ impl < F : NodeCtx + Send , S : WriteStorage , H : HistoryMode > VmTracer < S , H > for CheatcodeTracer < F > { }
95
98
96
- impl < F : FactoryDeps > CheatcodeTracer < F > {
97
- pub fn new ( factory_deps : F ) -> Self {
98
- Self { factory_deps }
99
+ impl < F : NodeCtx > CheatcodeTracer < F > {
100
+ pub fn new ( node_ctx : F ) -> Self {
101
+ Self { node_ctx }
99
102
}
100
103
101
104
fn dispatch_cheatcode < S : WriteStorage , H : HistoryMode > (
@@ -119,7 +122,7 @@ impl<F: FactoryDeps> CheatcodeTracer<F> {
119
122
let code_key = get_code_key ( & who) ;
120
123
let ( hash, code) = bytecode_to_factory_dep ( code. 0 . into ( ) ) ;
121
124
let hash = u256_to_h256 ( hash) ;
122
- if let Err ( err ) = self . factory_deps . store_factory_dep (
125
+ self . node_ctx . store_factory_dep (
123
126
hash,
124
127
code. iter ( )
125
128
. flat_map ( |entry| {
@@ -128,13 +131,7 @@ impl<F: FactoryDeps> CheatcodeTracer<F> {
128
131
bytes. to_vec ( )
129
132
} )
130
133
. collect ( ) ,
131
- ) {
132
- tracing:: error!(
133
- "Etch cheatcode failed, failed to store factory dep: {:?}" ,
134
- err
135
- ) ;
136
- return ;
137
- }
134
+ ) ;
138
135
storage. borrow_mut ( ) . set_value ( code_key, hash) ;
139
136
}
140
137
SetNonce ( SetNonceCall { account, nonce } ) => {
@@ -170,17 +167,58 @@ impl<F: FactoryDeps> CheatcodeTracer<F> {
170
167
) ;
171
168
storage. set_value ( nonce_key, u256_to_h256 ( enforced_full_nonce) ) ;
172
169
}
170
+ Warp ( WarpCall { timestamp } ) => {
171
+ tracing:: info!( "Setting block timestamp {}" , timestamp) ;
172
+ self . node_ctx . set_time ( timestamp. as_u64 ( ) ) ;
173
+
174
+ let key = StorageKey :: new (
175
+ AccountTreeId :: new ( zksync_types:: SYSTEM_CONTEXT_ADDRESS ) ,
176
+ zksync_types:: CURRENT_VIRTUAL_BLOCK_INFO_POSITION ,
177
+ ) ;
178
+ let mut storage = storage. borrow_mut ( ) ;
179
+ let ( block_number, _) = unpack_block_info ( h256_to_u256 ( storage. read_value ( & key) ) ) ;
180
+ storage. set_value (
181
+ key,
182
+ u256_to_h256 ( pack_block_info ( block_number, timestamp. as_u64 ( ) ) ) ,
183
+ ) ;
184
+ }
173
185
} ;
174
186
}
175
187
}
176
188
177
- impl < T > FactoryDeps for Arc < RwLock < InMemoryNodeInner < T > > > {
178
- fn store_factory_dep ( & mut self , hash : H256 , bytecode : Vec < u8 > ) -> Result < ( ) > {
179
- self . try_write ( )
180
- . map_err ( |e| anyhow ! ( format!( "Failed to grab write lock: {e}" ) ) ) ?
189
+ pub struct CheatcodeNodeContext < T > {
190
+ pub in_memory_node_inner : Arc < RwLock < InMemoryNodeInner < T > > > ,
191
+ pub batch_env : Arc < Mutex < L1BatchEnv > > ,
192
+ pub block_ctx : Arc < Mutex < BlockContext > > ,
193
+ }
194
+
195
+ impl < T > CheatcodeNodeContext < T > {
196
+ pub fn new (
197
+ in_memory_node_inner : Arc < RwLock < InMemoryNodeInner < T > > > ,
198
+ batch_env : Arc < Mutex < L1BatchEnv > > ,
199
+ block_ctx : Arc < Mutex < BlockContext > > ,
200
+ ) -> Self {
201
+ Self {
202
+ in_memory_node_inner,
203
+ batch_env,
204
+ block_ctx,
205
+ }
206
+ }
207
+ }
208
+
209
+ impl < T > NodeCtx for CheatcodeNodeContext < T > {
210
+ fn set_time ( & mut self , time : u64 ) {
211
+ self . in_memory_node_inner . write ( ) . unwrap ( ) . current_timestamp = time;
212
+ self . batch_env . lock ( ) . unwrap ( ) . timestamp = time + 1 ;
213
+ self . block_ctx . lock ( ) . unwrap ( ) . timestamp = time + 1 ;
214
+ }
215
+
216
+ fn store_factory_dep ( & mut self , hash : H256 , bytecode : Vec < u8 > ) {
217
+ self . in_memory_node_inner
218
+ . write ( )
219
+ . unwrap ( )
181
220
. fork_storage
182
- . store_factory_dep ( hash, bytecode) ;
183
- Ok ( ( ) )
221
+ . store_factory_dep ( hash, bytecode)
184
222
}
185
223
}
186
224
0 commit comments