1use alloc::{sync::Arc, vec::Vec};
2
3use miden_air::trace::{chiplets::hasher::STATE_WIDTH, rows::RowIndex};
4use miden_core::{
5 EMPTY_WORD, Felt, ONE, Word, ZERO,
6 crypto::merkle::MerklePath,
7 mast::{
8 BasicBlockNode, JoinNode, LoopNode, MastForest, MastNode, MastNodeExt, MastNodeId,
9 SplitNode,
10 },
11 precompile::PrecompileTranscript,
12 stack::MIN_STACK_DEPTH,
13};
14
15use crate::{
16 chiplets::CircuitEvaluation,
17 continuation_stack::ContinuationStack,
18 decoder::block_stack::{BlockInfo, BlockStack, BlockType, ExecutionContextInfo},
19 fast::{
20 FastProcessor,
21 trace_state::{
22 AceReplay, AdviceReplay, BitwiseReplay, BlockStackReplay, CoreTraceFragmentContext,
23 CoreTraceState, DecoderState, ExecutionContextSystemInfo, ExecutionReplay,
24 HasherRequestReplay, HasherResponseReplay, KernelReplay, MastForestResolutionReplay,
25 MemoryReadsReplay, MemoryWritesReplay, NodeExecutionState, NodeFlags,
26 RangeCheckerReplay, StackOverflowReplay, StackState, SystemState,
27 },
28 tracer::Tracer,
29 },
30 stack::OverflowTable,
31 system::ContextId,
32 utils::{HASH_CYCLE_LEN_FELT, split_u32_into_u16},
33};
34
35#[derive(Debug)]
37struct StateSnapshot {
38 state: CoreTraceState,
39 continuation_stack: ContinuationStack,
40 execution_state: NodeExecutionState,
41 initial_mast_forest: Arc<MastForest>,
42}
43
44pub struct TraceGenerationContext {
45 pub core_trace_contexts: Vec<CoreTraceFragmentContext>,
47
48 pub range_checker_replay: RangeCheckerReplay,
51 pub memory_writes: MemoryWritesReplay,
52 pub bitwise_replay: BitwiseReplay,
53 pub hasher_for_chiplet: HasherRequestReplay,
54 pub kernel_replay: KernelReplay,
55 pub ace_replay: AceReplay,
56
57 pub final_pc_transcript: PrecompileTranscript,
59
60 pub fragment_size: usize,
63}
64
65#[derive(Debug)]
78pub struct ExecutionTracer {
79 state_snapshot: Option<StateSnapshot>,
86
87 pub overflow_table: OverflowTable,
89 pub overflow_replay: StackOverflowReplay,
90
91 pub block_stack: BlockStack,
92 pub block_stack_replay: BlockStackReplay,
93
94 pub hasher_chiplet_shim: HasherChipletShim,
95 pub memory_reads: MemoryReadsReplay,
96 pub advice: AdviceReplay,
97 pub external: MastForestResolutionReplay,
98
99 pub range_checker: RangeCheckerReplay,
102 pub memory_writes: MemoryWritesReplay,
103 pub bitwise: BitwiseReplay,
104 pub kernel: KernelReplay,
105 pub hasher_for_chiplet: HasherRequestReplay,
106 pub ace: AceReplay,
107
108 fragment_contexts: Vec<CoreTraceFragmentContext>,
110
111 fragment_size: usize,
113}
114
115impl ExecutionTracer {
116 pub fn new(fragment_size: usize) -> Self {
118 Self {
119 state_snapshot: None,
120 overflow_table: OverflowTable::default(),
121 overflow_replay: StackOverflowReplay::default(),
122 block_stack: BlockStack::default(),
123 block_stack_replay: BlockStackReplay::default(),
124 hasher_chiplet_shim: HasherChipletShim::default(),
125 memory_reads: MemoryReadsReplay::default(),
126 range_checker: RangeCheckerReplay::default(),
127 memory_writes: MemoryWritesReplay::default(),
128 advice: AdviceReplay::default(),
129 bitwise: BitwiseReplay::default(),
130 kernel: KernelReplay::default(),
131 hasher_for_chiplet: HasherRequestReplay::default(),
132 ace: AceReplay::default(),
133 external: MastForestResolutionReplay::default(),
134 fragment_contexts: Vec::new(),
135 fragment_size,
136 }
137 }
138
139 pub fn into_trace_generation_context(
145 mut self,
146 final_pc_transcript: PrecompileTranscript,
147 ) -> TraceGenerationContext {
148 self.finish_current_fragment_context();
150
151 TraceGenerationContext {
152 core_trace_contexts: self.fragment_contexts,
153 range_checker_replay: self.range_checker,
154 memory_writes: self.memory_writes,
155 bitwise_replay: self.bitwise,
156 kernel_replay: self.kernel,
157 hasher_for_chiplet: self.hasher_for_chiplet,
158 ace_replay: self.ace,
159 final_pc_transcript,
160 fragment_size: self.fragment_size,
161 }
162 }
163
164 fn start_new_fragment_context(
175 &mut self,
176 system_state: SystemState,
177 stack_top: [Felt; MIN_STACK_DEPTH],
178 continuation_stack: ContinuationStack,
179 execution_state: NodeExecutionState,
180 current_forest: Arc<MastForest>,
181 ) {
182 self.finish_current_fragment_context();
184
185 self.state_snapshot = {
187 let decoder_state = {
188 if self.block_stack.is_empty() {
189 DecoderState { current_addr: ZERO, parent_addr: ZERO }
190 } else {
191 let block_info = self.block_stack.peek();
192
193 DecoderState {
194 current_addr: block_info.addr,
195 parent_addr: block_info.parent_addr,
196 }
197 }
198 };
199 let stack = {
200 let stack_depth =
201 MIN_STACK_DEPTH + self.overflow_table.num_elements_in_current_ctx();
202 let last_overflow_addr = self.overflow_table.last_update_clk_in_current_ctx();
203 StackState::new(stack_top, stack_depth, last_overflow_addr)
204 };
205
206 Some(StateSnapshot {
207 state: CoreTraceState {
208 system: system_state,
209 decoder: decoder_state,
210 stack,
211 },
212 continuation_stack,
213 execution_state,
214 initial_mast_forest: current_forest,
215 })
216 };
217 }
218
219 fn record_control_node_start(
220 &mut self,
221 node: &MastNode,
222 processor: &FastProcessor,
223 current_forest: &MastForest,
224 ) {
225 let (ctx_info, block_type) = match node {
226 MastNode::Join(node) => {
227 let child1_hash = current_forest
228 .get_node_by_id(node.first())
229 .expect("join node's first child expected to be in the forest")
230 .digest();
231 let child2_hash = current_forest
232 .get_node_by_id(node.second())
233 .expect("join node's second child expected to be in the forest")
234 .digest();
235 self.hasher_for_chiplet.record_hash_control_block(
236 child1_hash,
237 child2_hash,
238 JoinNode::DOMAIN,
239 node.digest(),
240 );
241
242 (None, BlockType::Join(false))
243 },
244 MastNode::Split(node) => {
245 let child1_hash = current_forest
246 .get_node_by_id(node.on_true())
247 .expect("split node's true child expected to be in the forest")
248 .digest();
249 let child2_hash = current_forest
250 .get_node_by_id(node.on_false())
251 .expect("split node's false child expected to be in the forest")
252 .digest();
253 self.hasher_for_chiplet.record_hash_control_block(
254 child1_hash,
255 child2_hash,
256 SplitNode::DOMAIN,
257 node.digest(),
258 );
259
260 (None, BlockType::Split)
261 },
262 MastNode::Loop(node) => {
263 let body_hash = current_forest
264 .get_node_by_id(node.body())
265 .expect("loop node's body expected to be in the forest")
266 .digest();
267
268 self.hasher_for_chiplet.record_hash_control_block(
269 body_hash,
270 EMPTY_WORD,
271 LoopNode::DOMAIN,
272 node.digest(),
273 );
274
275 let loop_entered = {
276 let condition = processor.stack_get(0);
277 condition == ONE
278 };
279
280 (None, BlockType::Loop(loop_entered))
281 },
282 MastNode::Call(node) => {
283 let callee_hash = current_forest
284 .get_node_by_id(node.callee())
285 .expect("call node's callee expected to be in the forest")
286 .digest();
287
288 self.hasher_for_chiplet.record_hash_control_block(
289 callee_hash,
290 EMPTY_WORD,
291 node.domain(),
292 node.digest(),
293 );
294
295 let exec_ctx = {
296 let overflow_addr = self.overflow_table.last_update_clk_in_current_ctx();
297 ExecutionContextInfo::new(
298 processor.ctx,
299 processor.caller_hash,
300 processor.stack_depth(),
301 overflow_addr,
302 )
303 };
304 let block_type = if node.is_syscall() {
305 BlockType::SysCall
306 } else {
307 BlockType::Call
308 };
309
310 (Some(exec_ctx), block_type)
311 },
312 MastNode::Dyn(dyn_node) => {
313 self.hasher_for_chiplet.record_hash_control_block(
314 EMPTY_WORD,
315 EMPTY_WORD,
316 dyn_node.domain(),
317 dyn_node.digest(),
318 );
319
320 if dyn_node.is_dyncall() {
321 let exec_ctx = {
322 let overflow_addr = self.overflow_table.last_update_clk_in_current_ctx();
323 let stack_depth_after_drop = processor.stack_depth() - 1;
332 ExecutionContextInfo::new(
333 processor.ctx,
334 processor.caller_hash,
335 stack_depth_after_drop,
336 overflow_addr,
337 )
338 };
339 (Some(exec_ctx), BlockType::Dyncall)
340 } else {
341 (None, BlockType::Dyn)
342 }
343 },
344 MastNode::Block(_) => panic!(
345 "`ExecutionTracer::record_basic_block_start()` must be called instead for basic blocks"
346 ),
347 MastNode::External(_) => panic!(
348 "External nodes are guaranteed to be resolved before record_control_node_start is called"
349 ),
350 };
351
352 let block_addr = self.hasher_chiplet_shim.record_hash_control_block();
353 let parent_addr = self.block_stack.push(block_addr, block_type, ctx_info);
354 self.block_stack_replay.record_node_start_parent_addr(parent_addr);
355 }
356
357 fn record_node_end(&mut self, block_info: &BlockInfo) {
359 let flags = NodeFlags::new(
360 block_info.is_loop_body() == ONE,
361 block_info.is_entered_loop() == ONE,
362 block_info.is_call() == ONE,
363 block_info.is_syscall() == ONE,
364 );
365 let (prev_addr, prev_parent_addr) = if self.block_stack.is_empty() {
366 (ZERO, ZERO)
367 } else {
368 let prev_block = self.block_stack.peek();
369 (prev_block.addr, prev_block.parent_addr)
370 };
371 self.block_stack_replay.record_node_end(
372 block_info.addr,
373 flags,
374 prev_addr,
375 prev_parent_addr,
376 );
377 }
378
379 fn record_execution_context(&mut self, ctx_info: ExecutionContextSystemInfo) {
381 self.block_stack_replay.record_execution_context(ctx_info);
382 }
383
384 fn finish_current_fragment_context(&mut self) {
393 if let Some(snapshot) = self.state_snapshot.take() {
394 let hasher_replay = self.hasher_chiplet_shim.extract_replay();
396 let memory_reads_replay = core::mem::take(&mut self.memory_reads);
397 let advice_replay = core::mem::take(&mut self.advice);
398 let external_replay = core::mem::take(&mut self.external);
399 let stack_overflow_replay = core::mem::take(&mut self.overflow_replay);
400 let block_stack_replay = core::mem::take(&mut self.block_stack_replay);
401
402 let trace_state = CoreTraceFragmentContext {
403 state: snapshot.state,
404 replay: ExecutionReplay {
405 hasher: hasher_replay,
406 memory_reads: memory_reads_replay,
407 advice: advice_replay,
408 mast_forest_resolution: external_replay,
409 stack_overflow: stack_overflow_replay,
410 block_stack: block_stack_replay,
411 },
412 continuation: snapshot.continuation_stack,
413 initial_execution_state: snapshot.execution_state,
414 initial_mast_forest: snapshot.initial_mast_forest,
415 };
416
417 self.fragment_contexts.push(trace_state);
418 }
419 }
420}
421
422impl Tracer for ExecutionTracer {
423 fn start_clock_cycle(
426 &mut self,
427 processor: &FastProcessor,
428 execution_state: NodeExecutionState,
429 continuation_stack: &mut ContinuationStack,
430 current_forest: &Arc<MastForest>,
431 ) {
432 if processor.clk.as_usize().is_multiple_of(self.fragment_size) {
434 self.start_new_fragment_context(
435 SystemState::from_processor(processor),
436 processor
437 .stack_top()
438 .try_into()
439 .expect("stack_top expected to be MIN_STACK_DEPTH elements"),
440 continuation_stack.clone(),
441 execution_state.clone(),
442 current_forest.clone(),
443 );
444 }
445
446 match &execution_state {
448 NodeExecutionState::BasicBlock { .. } => {
449 },
451 NodeExecutionState::Start(mast_node_id) => match ¤t_forest[*mast_node_id] {
452 MastNode::Join(_)
453 | MastNode::Split(_)
454 | MastNode::Loop(_)
455 | MastNode::Dyn(_)
456 | MastNode::Call(_) => {
457 self.record_control_node_start(
458 ¤t_forest[*mast_node_id],
459 processor,
460 current_forest,
461 );
462 },
463 MastNode::Block(basic_block_node) => {
464 self.hasher_for_chiplet.record_hash_basic_block(
465 basic_block_node.op_batches().to_vec(),
466 basic_block_node.digest(),
467 );
468 let block_addr =
469 self.hasher_chiplet_shim.record_hash_basic_block(basic_block_node);
470 let parent_addr =
471 self.block_stack.push(block_addr, BlockType::BasicBlock, None);
472 self.block_stack_replay.record_node_start_parent_addr(parent_addr);
473 },
474 MastNode::External(_) => unreachable!(
475 "start_clock_cycle is guaranteed not to be called on external nodes"
476 ),
477 },
478 NodeExecutionState::Respan { node_id: _, batch_index: _ } => {
479 self.block_stack.peek_mut().addr += HASH_CYCLE_LEN_FELT;
480 },
481 NodeExecutionState::LoopRepeat(_) => {
482 },
484 NodeExecutionState::End(_) => {
485 let block_info = self.block_stack.pop();
486 self.record_node_end(&block_info);
487
488 if let Some(ctx_info) = block_info.ctx_info {
489 self.record_execution_context(ExecutionContextSystemInfo {
490 parent_ctx: ctx_info.parent_ctx,
491 parent_fn_hash: ctx_info.parent_fn_hash,
492 });
493 }
494 },
495 }
496 }
497
498 fn record_mast_forest_resolution(&mut self, node_id: MastNodeId, forest: &Arc<MastForest>) {
499 self.external.record_resolution(node_id, forest.clone());
500 }
501
502 fn record_hasher_permute(
503 &mut self,
504 input_state: [Felt; STATE_WIDTH],
505 output_state: [Felt; STATE_WIDTH],
506 ) {
507 self.hasher_for_chiplet.record_permute_input(input_state);
508 self.hasher_chiplet_shim.record_permute_output(output_state);
509 }
510
511 fn record_hasher_build_merkle_root(
512 &mut self,
513 node: Word,
514 path: Option<&MerklePath>,
515 index: Felt,
516 output_root: Word,
517 ) {
518 let path = path.expect("execution tracer expects a valid Merkle path");
519 self.hasher_chiplet_shim.record_build_merkle_root(path, output_root);
520 self.hasher_for_chiplet.record_build_merkle_root(node, path.clone(), index);
521 }
522
523 fn record_hasher_update_merkle_root(
524 &mut self,
525 old_value: Word,
526 new_value: Word,
527 path: Option<&MerklePath>,
528 index: Felt,
529 old_root: Word,
530 new_root: Word,
531 ) {
532 let path = path.expect("execution tracer expects a valid Merkle path");
533 self.hasher_chiplet_shim.record_update_merkle_root(path, old_root, new_root);
534 self.hasher_for_chiplet.record_update_merkle_root(
535 old_value,
536 new_value,
537 path.clone(),
538 index,
539 );
540 }
541
542 fn record_memory_read_element(
543 &mut self,
544 element: Felt,
545 addr: Felt,
546 ctx: ContextId,
547 clk: RowIndex,
548 ) {
549 self.memory_reads.record_read_element(element, addr, ctx, clk);
550 }
551
552 fn record_memory_read_word(&mut self, word: Word, addr: Felt, ctx: ContextId, clk: RowIndex) {
553 self.memory_reads.record_read_word(word, addr, ctx, clk);
554 }
555
556 fn record_memory_write_element(
557 &mut self,
558 element: Felt,
559 addr: Felt,
560 ctx: ContextId,
561 clk: RowIndex,
562 ) {
563 self.memory_writes.record_write_element(element, addr, ctx, clk);
564 }
565
566 fn record_memory_write_word(&mut self, word: Word, addr: Felt, ctx: ContextId, clk: RowIndex) {
567 self.memory_writes.record_write_word(word, addr, ctx, clk);
568 }
569
570 fn record_advice_pop_stack(&mut self, value: Felt) {
571 self.advice.record_pop_stack(value);
572 }
573
574 fn record_advice_pop_stack_word(&mut self, word: Word) {
575 self.advice.record_pop_stack_word(word);
576 }
577
578 fn record_advice_pop_stack_dword(&mut self, words: [Word; 2]) {
579 self.advice.record_pop_stack_dword(words);
580 }
581
582 fn record_u32and(&mut self, a: Felt, b: Felt) {
583 self.bitwise.record_u32and(a, b);
584 }
585
586 fn record_u32xor(&mut self, a: Felt, b: Felt) {
587 self.bitwise.record_u32xor(a, b);
588 }
589
590 fn record_u32_range_checks(&mut self, clk: RowIndex, u32_lo: Felt, u32_hi: Felt) {
591 let (t1, t0) = split_u32_into_u16(u32_lo.as_int());
592 let (t3, t2) = split_u32_into_u16(u32_hi.as_int());
593
594 self.range_checker.record_range_check_u32(clk, [t0, t1, t2, t3]);
595 }
596
597 fn record_kernel_proc_access(&mut self, proc_hash: Word) {
598 self.kernel.record_kernel_proc_access(proc_hash);
599 }
600
601 fn record_circuit_evaluation(&mut self, clk: RowIndex, circuit_eval: CircuitEvaluation) {
602 self.ace.record_circuit_evaluation(clk, circuit_eval);
603 }
604
605 fn increment_clk(&mut self) {
606 self.overflow_table.advance_clock();
607 }
608
609 fn increment_stack_size(&mut self, processor: &FastProcessor) {
610 let new_overflow_value = processor.stack_get(15);
611 self.overflow_table.push(new_overflow_value);
612 }
613
614 fn decrement_stack_size(&mut self) {
615 if let Some(popped_value) = self.overflow_table.pop() {
617 let new_overflow_addr = self.overflow_table.last_update_clk_in_current_ctx();
618 self.overflow_replay.record_pop_overflow(popped_value, new_overflow_addr);
619 }
620 }
621
622 fn start_context(&mut self) {
623 self.overflow_table.start_context();
624 }
625
626 fn restore_context(&mut self) {
627 self.overflow_table.restore_context();
628 self.overflow_replay.record_restore_context_overflow_addr(
629 MIN_STACK_DEPTH + self.overflow_table.num_elements_in_current_ctx(),
630 self.overflow_table.last_update_clk_in_current_ctx(),
631 );
632 }
633}
634
635const NUM_HASHER_ROWS_PER_PERMUTATION: u32 = 8;
641
642#[derive(Debug)]
649pub struct HasherChipletShim {
650 addr: u32,
654 replay: HasherResponseReplay,
656}
657
658impl HasherChipletShim {
659 pub fn new() -> Self {
661 Self {
662 addr: 1,
663 replay: HasherResponseReplay::default(),
664 }
665 }
666
667 pub fn record_hash_control_block(&mut self) -> Felt {
669 let block_addr = self.addr.into();
670
671 self.replay.record_block_address(block_addr);
672 self.addr += NUM_HASHER_ROWS_PER_PERMUTATION;
673
674 block_addr
675 }
676
677 pub fn record_hash_basic_block(&mut self, basic_block_node: &BasicBlockNode) -> Felt {
679 let block_addr = self.addr.into();
680
681 self.replay.record_block_address(block_addr);
682 self.addr += NUM_HASHER_ROWS_PER_PERMUTATION * basic_block_node.num_op_batches() as u32;
683
684 block_addr
685 }
686 pub fn record_permute_output(&mut self, hashed_state: [Felt; 12]) {
688 self.replay.record_permute(self.addr.into(), hashed_state);
689 self.addr += NUM_HASHER_ROWS_PER_PERMUTATION;
690 }
691
692 pub fn record_build_merkle_root(&mut self, path: &MerklePath, computed_root: Word) {
694 self.replay.record_build_merkle_root(self.addr.into(), computed_root);
695 self.addr += NUM_HASHER_ROWS_PER_PERMUTATION * path.depth() as u32;
696 }
697
698 pub fn record_update_merkle_root(&mut self, path: &MerklePath, old_root: Word, new_root: Word) {
700 self.replay.record_update_merkle_root(self.addr.into(), old_root, new_root);
701
702 self.addr += 2 * NUM_HASHER_ROWS_PER_PERMUTATION * path.depth() as u32;
704 }
705
706 pub fn extract_replay(&mut self) -> HasherResponseReplay {
707 core::mem::take(&mut self.replay)
708 }
709}
710
711impl Default for HasherChipletShim {
712 fn default() -> Self {
713 Self::new()
714 }
715}