Source code for zisk/state-machines/main/src/main_planner.rs

 1//! The `MainPlanner` module defines a planner for the Main State Machine.
 2//!
 3//! It generates execution plans for segments of the main trace, mapping each segment
 4//! to a specific `Plan` instance.
 5
 6use std::any::Any;
 7
 8use asm_runner::MinimalTraces;
 9use fields::PrimeField64;
10use zisk_common::{BusDeviceMetrics, CheckPoint, ChunkId, InstanceType, Metrics, Plan, SegmentId};
11use zisk_pil::{MainTrace, MAIN_AIR_IDS, ZISK_AIRGROUP_ID};
12
13use crate::MainCounter;
14
15/// The `MainPlanner` struct generates execution plans for the Main State Machine.
16///
17/// It organizes the execution flow by creating a `Plan` instance for each segment

18/// of the main trace, associating it with the corresponding segment ID.

19pub struct MainPlanner {}

20
21impl MainPlanner {
22    /// Generates execution plans for the Main State Machine.
23    ///
24    /// This method creates a `Plan` for each segment of the provided traces, associating
25    /// the segment ID with the corresponding execution plan.
26    ///
27    /// # Arguments
28    /// * `min_traces` - A slice of `EmuTrace` instances representing the segments to be planned.
29    /// * `main_counters` - A vector of main counters, each associated with a specific chunk ID.
30    /// * `min_traces_size` - The size of the minimal traces.
31    ///
32    /// # Returns

33    /// A vector of `Plan` instances, each corresponding to a segment of the main trace.
34    pub fn plan<F: PrimeField64>(
35        min_traces: &MinimalTraces,
36        main_counters: Vec<(ChunkId, Box<dyn BusDeviceMetrics>)>,
37        min_traces_size: u64,
38    ) -> (Vec<Plan>, Vec<(u64, u32)>) {
39        let min_traces = match min_traces {
40            MinimalTraces::AsmEmuTrace(asm_min_traces) => &asm_min_traces.vec_chunks,
41            MinimalTraces::EmuTrace(vec_chunks) => vec_chunks,
42            MinimalTraces::None => {
43                panic!("Minimal traces are required for planning the main state machine.");
44            }
45        };
46
47        let num_rows = MainTrace::<F>::NUM_ROWS as u64;
48
49        let mut publics = Vec::new();
50
51        main_counters.iter().for_each(|(_, counter)| {
52            let reg_counter = Metrics::as_any(&**counter).downcast_ref::<MainCounter>().unwrap();
53            publics.extend_from_slice(&reg_counter.publics);
54        });
55
56        assert!(num_rows.is_power_of_two());
57        assert!(min_traces_size.is_power_of_two());
58        assert!(num_rows >= min_traces_size);
59
60        // This is the number of minimal traces wrapped in a main trace
61        let num_within = num_rows / min_traces_size;
62        let num_instances = (min_traces.len() as f64 / num_within as f64).ceil() as usize;
63
64        let plans: Vec<Plan> = (0..num_instances)
65            .map(|segment_id| {
66                Plan::new(
67                    ZISK_AIRGROUP_ID,
68                    MAIN_AIR_IDS[0],
69                    Some(SegmentId(segment_id)),
70                    InstanceType::Instance,
71                    CheckPoint::Single(ChunkId(segment_id)),
72                    Some(Box::new(segment_id == num_instances - 1) as Box<dyn Any>),
73                )
74            })
75            .collect();
76
77        (plans, publics)

78    }

79}