1use crate::error::StateRootTaskError;
11use alloy_evm::block::OnStateHook;
12use alloy_primitives::{keccak256, map::B256Map, B256};
13use reth_trie::{
14 prefix_set::TriePrefixSetsMut, updates::TrieUpdates, HashedPostState, HashedStorage,
15 MultiProofTargetsV2, ProofV2Target,
16};
17use revm::state::EvmState;
18use std::{fmt, sync::Arc};
19use tracing::trace;
20
21#[derive(Debug)]
23pub enum StateRootMessage {
24 PrefetchProofs(MultiProofTargetsV2),
26 StateUpdate(EvmState),
28 HashedStateUpdate(HashedPostState),
30 FinishedStateUpdates,
35}
36
37#[derive(Debug, Clone)]
40pub struct StateRootComputeOutcome {
41 pub state_root: B256,
43 pub trie_updates: Arc<TrieUpdates>,
45 pub changed_paths: Option<Arc<TriePrefixSetsMut>>,
47 #[cfg(feature = "trie-debug")]
50 pub debug_recorders: Vec<(Option<B256>, reth_trie_sparse::debug_recorder::TrieDebugRecorder)>,
51}
52
53#[derive(Debug)]
61pub struct StateRootHandle {
62 cached_trie_state_root: B256,
64 hint: Option<StateRootHintStream>,
66 authoritative: Option<StateRootUpdateStream>,
73 cancel_guard: StateRootTaskCancelGuard,
75 state_root_rx:
77 Option<std::sync::mpsc::Receiver<Result<StateRootComputeOutcome, StateRootTaskError>>>,
78 hashed_state_rx: Option<std::sync::mpsc::Receiver<HashedPostState>>,
80}
81
82impl StateRootHandle {
83 pub fn new(
85 cached_trie_state_root: B256,
86 updates_tx: crossbeam_channel::Sender<StateRootMessage>,
87 cancel_guard: StateRootTaskCancelGuard,
88 state_root_rx: std::sync::mpsc::Receiver<
89 Result<StateRootComputeOutcome, StateRootTaskError>,
90 >,
91 hashed_state_rx: std::sync::mpsc::Receiver<HashedPostState>,
92 ) -> Self {
93 let sink: Arc<dyn StateRootSink> = Arc::new(SparseTrieStateRootSink::new(updates_tx));
94 Self {
95 cached_trie_state_root,
96 hint: Some(StateRootHintStream::new(Arc::clone(&sink))),
97 authoritative: Some(StateRootUpdateStream::new(sink)),
98 cancel_guard,
99 state_root_rx: Some(state_root_rx),
100 hashed_state_rx: Some(hashed_state_rx),
101 }
102 }
103
104 pub const fn cached_trie_state_root(&self) -> B256 {
106 self.cached_trie_state_root
107 }
108
109 pub const fn take_hint_stream(&mut self) -> StateRootHintStream {
115 self.hint.take().expect("hint stream already taken")
116 }
117
118 pub fn take_execution_hook(&mut self) -> StateRootUpdateHook {
127 self.take_hashed_update_stream().into_state_hook()
128 }
129
130 pub const fn take_hashed_update_stream(&mut self) -> StateRootUpdateStream {
140 self.authoritative.take().expect("authoritative update capability already taken")
141 }
142
143 pub fn state_root(&mut self) -> Result<StateRootComputeOutcome, StateRootTaskError> {
149 self.state_root_rx
150 .take()
151 .expect("state_root already taken")
152 .recv()
153 .map_err(|_| StateRootTaskError::Other("sparse trie task dropped".to_string()))?
154 }
155
156 pub const fn take_state_root_rx(
162 &mut self,
163 ) -> std::sync::mpsc::Receiver<Result<StateRootComputeOutcome, StateRootTaskError>> {
164 self.state_root_rx.take().expect("state_root already taken")
165 }
166
167 pub const fn take_hashed_state_rx(&mut self) -> std::sync::mpsc::Receiver<HashedPostState> {
173 self.hashed_state_rx.take().expect("hashed_state already taken")
174 }
175
176 pub fn into_payload_state_root_handle(mut self) -> PayloadStateRootHandle {
181 let hook = self.take_execution_hook();
182 PayloadStateRootHandle {
183 name: "sparse-trie",
184 hook: Some(hook),
185 cancel_guard: Some(self.cancel_guard),
186 state_root_rx: self.state_root_rx.take(),
187 hashed_state_rx: self.hashed_state_rx.take(),
188 }
189 }
190}
191
192#[derive(Debug)]
198pub struct StateRootTaskCancelGuard(#[allow(dead_code)] crossbeam_channel::Sender<()>);
199
200impl StateRootTaskCancelGuard {
201 pub fn channel() -> (Self, crossbeam_channel::Receiver<()>) {
203 let (tx, rx) = crossbeam_channel::bounded(0);
204 (Self(tx), rx)
205 }
206}
207
208pub struct PayloadStateRootHandle {
210 name: &'static str,
211 hook: Option<StateRootUpdateHook>,
213 cancel_guard: Option<StateRootTaskCancelGuard>,
215 state_root_rx:
216 Option<std::sync::mpsc::Receiver<Result<StateRootComputeOutcome, StateRootTaskError>>>,
217 hashed_state_rx: Option<std::sync::mpsc::Receiver<HashedPostState>>,
218}
219
220impl fmt::Debug for PayloadStateRootHandle {
221 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
222 f.debug_struct("PayloadStateRootHandle")
223 .field("name", &self.name)
224 .field("has_hook", &self.hook.is_some())
225 .field("has_cancel_guard", &self.cancel_guard.is_some())
226 .field("has_state_root_rx", &self.state_root_rx.is_some())
227 .field("has_hashed_state_rx", &self.hashed_state_rx.is_some())
228 .finish()
229 }
230}
231
232impl PayloadStateRootHandle {
233 pub const fn new(
238 name: &'static str,
239 hook: Option<StateRootUpdateHook>,
240 state_root_rx: std::sync::mpsc::Receiver<
241 Result<StateRootComputeOutcome, StateRootTaskError>,
242 >,
243 hashed_state_rx: Option<std::sync::mpsc::Receiver<HashedPostState>>,
244 ) -> Self {
245 Self { name, hook, cancel_guard: None, state_root_rx: Some(state_root_rx), hashed_state_rx }
246 }
247
248 pub const fn name(&self) -> &'static str {
250 self.name
251 }
252
253 pub const fn take_state_hook(&mut self) -> StateRootUpdateHook {
259 self.hook.take().expect("payload state root task missing execution hook")
260 }
261
262 pub fn state_root(&mut self) -> Result<StateRootComputeOutcome, StateRootTaskError> {
268 self.state_root_rx
269 .take()
270 .expect("state_root already taken")
271 .recv()
272 .map_err(|_| StateRootTaskError::Other("state root task dropped".to_string()))?
273 }
274
275 pub const fn try_take_hashed_state_rx(
278 &mut self,
279 ) -> Option<std::sync::mpsc::Receiver<HashedPostState>> {
280 self.hashed_state_rx.take()
281 }
282}
283
284#[derive(Debug, Clone, Default)]
290pub struct StateAccessHint {
291 pub accounts: Vec<B256>,
293 pub storages: B256Map<Vec<B256>>,
295}
296
297impl From<MultiProofTargetsV2> for StateAccessHint {
298 fn from(targets: MultiProofTargetsV2) -> Self {
299 Self {
300 accounts: targets.account_targets.into_iter().map(|target| target.key()).collect(),
301 storages: targets
302 .storage_targets
303 .into_iter()
304 .map(|(account, slots)| {
305 (account, slots.into_iter().map(|target| target.key()).collect())
306 })
307 .collect(),
308 }
309 }
310}
311
312impl From<StateAccessHint> for MultiProofTargetsV2 {
313 fn from(hint: StateAccessHint) -> Self {
314 Self {
315 account_targets: hint.accounts.into_iter().map(ProofV2Target::from).collect(),
316 storage_targets: hint
317 .storages
318 .into_iter()
319 .map(|(account, slots)| {
320 (account, slots.into_iter().map(ProofV2Target::from).collect())
321 })
322 .collect(),
323 }
324 }
325}
326
327pub trait StateRootSink: Send + Sync + 'static {
329 fn on_access_hint(&self, _hint: StateAccessHint) {}
331
332 fn on_state_update(&self, state: EvmState);
334
335 fn on_hashed_state_update(&self, state: HashedPostState);
337
338 fn on_updates_finished(&self);
340}
341
342#[derive(Clone)]
344pub struct StateRootHintStream {
345 inner: Arc<dyn StateRootSink>,
346}
347
348impl fmt::Debug for StateRootHintStream {
349 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350 f.debug_struct("StateRootHintStream").finish_non_exhaustive()
351 }
352}
353
354impl StateRootHintStream {
355 pub fn new(inner: Arc<dyn StateRootSink>) -> Self {
357 Self { inner }
358 }
359
360 pub fn on_access_hint(&self, hint: StateAccessHint) {
362 self.inner.on_access_hint(hint);
363 }
364}
365
366pub struct StateRootUpdateStream {
378 inner: Arc<dyn StateRootSink>,
379}
380
381impl fmt::Debug for StateRootUpdateStream {
382 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
383 f.debug_struct("StateRootUpdateStream").finish_non_exhaustive()
384 }
385}
386
387impl StateRootUpdateStream {
388 pub fn new(inner: Arc<dyn StateRootSink>) -> Self {
390 Self { inner }
391 }
392
393 pub fn on_hashed_state_update(&self, state: HashedPostState) {
395 self.inner.on_hashed_state_update(state);
396 }
397
398 pub fn finish(self) {
400 self.inner.on_updates_finished();
401 }
402
403 pub fn into_state_hook(self) -> StateRootUpdateHook {
408 StateRootUpdateHook { inner: self.inner }
409 }
410}
411
412pub struct StateRootUpdateHook {
425 inner: Arc<dyn StateRootSink>,
426}
427
428impl fmt::Debug for StateRootUpdateHook {
429 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
430 f.debug_struct("StateRootUpdateHook").finish_non_exhaustive()
431 }
432}
433
434impl OnStateHook for StateRootUpdateHook {
435 fn on_state(&mut self, state: EvmState) {
436 self.inner.on_state_update(state);
437 }
438}
439
440impl Drop for StateRootUpdateHook {
441 fn drop(&mut self) {
442 if std::thread::panicking() {
445 return;
446 }
447 self.inner.on_updates_finished();
448 }
449}
450
451#[derive(Debug, Clone)]
452struct SparseTrieStateRootSink {
453 sender: crossbeam_channel::Sender<StateRootMessage>,
454}
455
456impl SparseTrieStateRootSink {
457 const fn new(sender: crossbeam_channel::Sender<StateRootMessage>) -> Self {
458 Self { sender }
459 }
460}
461
462impl StateRootSink for SparseTrieStateRootSink {
463 fn on_access_hint(&self, hint: StateAccessHint) {
464 let _ = self.sender.send(StateRootMessage::PrefetchProofs(hint.into()));
465 }
466
467 fn on_state_update(&self, state: EvmState) {
468 let _ = self.sender.send(StateRootMessage::StateUpdate(state));
469 }
470
471 fn on_hashed_state_update(&self, state: HashedPostState) {
472 let _ = self.sender.send(StateRootMessage::HashedStateUpdate(state));
473 }
474
475 fn on_updates_finished(&self) {
476 let _ = self.sender.send(StateRootMessage::FinishedStateUpdates);
477 }
478}
479
480pub fn evm_state_to_hashed_post_state(update: EvmState) -> HashedPostState {
482 let mut hashed_state = HashedPostState::with_capacity(update.len());
483
484 for (address, account) in update {
485 if account.is_touched() {
486 let hashed_address = keccak256(address);
487 trace!(target: "trie::parallel::sparse", ?address, ?hashed_address, "Adding account to state update");
488
489 let destroyed = account.is_selfdestructed();
490 if account.info != account.original_info() {
491 let info = if destroyed { None } else { Some(account.info.into()) };
492 hashed_state.accounts.insert(hashed_address, info);
493 }
494
495 let mut changed_storage_iter = account
496 .storage
497 .into_iter()
498 .filter(|(_slot, value)| value.is_changed())
499 .map(|(slot, value)| (keccak256(B256::from(slot)), value.present_value))
500 .peekable();
501
502 if destroyed {
503 hashed_state.storages.insert(hashed_address, HashedStorage::new(true));
504 } else if changed_storage_iter.peek().is_some() {
505 hashed_state
506 .storages
507 .insert(hashed_address, HashedStorage::from_iter(false, changed_storage_iter));
508 }
509 }
510 }
511
512 hashed_state
513}
514
515#[cfg(test)]
516mod tests {
517 use super::*;
518 use std::sync::atomic::{AtomicUsize, Ordering};
519
520 #[derive(Default)]
521 struct CountingSink {
522 access_hints: AtomicUsize,
523 state_updates: AtomicUsize,
524 hashed_state_updates: AtomicUsize,
525 finished_updates: AtomicUsize,
526 }
527
528 impl StateRootSink for CountingSink {
529 fn on_access_hint(&self, hint: StateAccessHint) {
530 assert_eq!(hint.accounts, vec![B256::repeat_byte(0x01)]);
531 assert_eq!(
532 hint.storages.get(&B256::repeat_byte(0x02)),
533 Some(&vec![B256::repeat_byte(0x03)])
534 );
535 self.access_hints.fetch_add(1, Ordering::Relaxed);
536 }
537
538 fn on_state_update(&self, state: EvmState) {
539 assert!(state.is_empty());
540 self.state_updates.fetch_add(1, Ordering::Relaxed);
541 }
542
543 fn on_hashed_state_update(&self, state: HashedPostState) {
544 assert!(state.accounts.is_empty());
545 assert!(state.storages.is_empty());
546 self.hashed_state_updates.fetch_add(1, Ordering::Relaxed);
547 }
548
549 fn on_updates_finished(&self) {
550 self.finished_updates.fetch_add(1, Ordering::Relaxed);
551 }
552 }
553
554 #[test]
555 fn state_access_hint_converts_to_sparse_targets() {
556 let account = B256::repeat_byte(0x01);
557 let storage_account = B256::repeat_byte(0x02);
558 let storage_slot = B256::repeat_byte(0x03);
559
560 let mut storages = B256Map::default();
561 storages.insert(storage_account, vec![storage_slot]);
562 let hint = StateAccessHint { accounts: vec![account], storages };
563
564 let targets = MultiProofTargetsV2::from(hint);
565 assert_eq!(targets.account_targets.len(), 1);
566 assert_eq!(targets.account_targets[0].key(), account);
567 assert_eq!(targets.storage_targets.len(), 1);
568 assert_eq!(targets.storage_targets[&storage_account].len(), 1);
569 assert_eq!(targets.storage_targets[&storage_account][0].key(), storage_slot);
570
571 let hint = StateAccessHint::from(targets);
572 assert_eq!(hint.accounts, vec![account]);
573 assert_eq!(hint.storages.len(), 1);
574 assert_eq!(hint.storages[&storage_account], vec![storage_slot]);
575 }
576
577 #[test]
578 fn state_root_capabilities_forward_to_sink() {
579 let sink = Arc::new(CountingSink::default());
580
581 let hint_stream = StateRootHintStream::new(sink.clone());
582 let mut storages = B256Map::default();
583 storages.insert(B256::repeat_byte(0x02), vec![B256::repeat_byte(0x03)]);
584 hint_stream
585 .on_access_hint(StateAccessHint { accounts: vec![B256::repeat_byte(0x01)], storages });
586
587 let updates = StateRootUpdateStream::new(sink.clone());
588 updates.on_hashed_state_update(HashedPostState::default());
589 updates.finish();
590
591 {
592 let mut hook = StateRootUpdateStream::new(sink.clone()).into_state_hook();
593 hook.on_state(EvmState::default());
594 }
595
596 assert_eq!(sink.access_hints.load(Ordering::Relaxed), 1);
597 assert_eq!(sink.state_updates.load(Ordering::Relaxed), 1);
598 assert_eq!(sink.hashed_state_updates.load(Ordering::Relaxed), 1);
599 assert_eq!(sink.finished_updates.load(Ordering::Relaxed), 2);
600 }
601
602 #[test]
605 fn hook_dropped_during_panic_does_not_finish_stream() {
606 let sink = Arc::new(CountingSink::default());
607 let hook = StateRootUpdateStream::new(sink.clone()).into_state_hook();
608
609 let result = std::thread::spawn(move || {
610 let _hook = hook;
611 panic!("execution died mid-block");
612 })
613 .join();
614
615 assert!(result.is_err());
616 assert_eq!(sink.finished_updates.load(Ordering::Relaxed), 0);
617 }
618
619 #[test]
622 #[should_panic(expected = "authoritative update capability already taken")]
623 fn authoritative_capability_can_only_be_taken_once() {
624 let (updates_tx, _updates_rx) = crossbeam_channel::unbounded();
625 let (cancel_guard, _cancel_rx) = StateRootTaskCancelGuard::channel();
626 let (_state_root_tx, state_root_rx) = std::sync::mpsc::channel();
627 let (_hashed_state_tx, hashed_state_rx) = std::sync::mpsc::channel();
628 let mut handle = StateRootHandle::new(
629 B256::ZERO,
630 updates_tx,
631 cancel_guard,
632 state_root_rx,
633 hashed_state_rx,
634 );
635
636 let _hook = handle.take_execution_hook();
637 let _ = handle.take_hashed_update_stream();
638 }
639
640 #[test]
645 fn payload_state_root_handle_lifecycle() {
646 let sink = Arc::new(CountingSink::default());
647 let hook = StateRootUpdateStream::new(sink.clone()).into_state_hook();
648
649 let (state_root_tx, state_root_rx) = std::sync::mpsc::channel();
650 let (hashed_state_tx, hashed_state_rx) = std::sync::mpsc::channel();
651 let mut handle =
652 PayloadStateRootHandle::new("test", Some(hook), state_root_rx, Some(hashed_state_rx));
653
654 assert_eq!(handle.name(), "test");
655
656 {
657 let mut hook = handle.take_state_hook();
658 hook.on_state(EvmState::default());
659 }
660 assert_eq!(sink.state_updates.load(Ordering::Relaxed), 1);
661 assert_eq!(sink.finished_updates.load(Ordering::Relaxed), 1);
662
663 hashed_state_tx.send(HashedPostState::default()).unwrap();
664 let rx = handle.try_take_hashed_state_rx().expect("first take returns the receiver");
665 assert!(rx.recv().is_ok());
666 assert!(handle.try_take_hashed_state_rx().is_none(), "second take returns None");
667
668 state_root_tx
669 .send(Ok(StateRootComputeOutcome {
670 state_root: B256::repeat_byte(0x42),
671 trie_updates: Arc::new(TrieUpdates::default()),
672 changed_paths: None,
673 #[cfg(feature = "trie-debug")]
674 debug_recorders: Vec::new(),
675 }))
676 .unwrap();
677 let outcome = handle.state_root().expect("outcome is delivered");
678 assert_eq!(outcome.state_root, B256::repeat_byte(0x42));
679 }
680}