Source code for zisk/state-machines/mem/src/mem_align_write_byte_instance.rs

 1use crate::{mem_align_byte_sm::MemAlignByteSM, MemAlignCollector};
 2use mem_common::MemAlignCheckPoint;
 3
 4use crate::mem_align_byte_sm::{MemAlignWriteByteTraceRowType, MemAlignWriteByteTraceType};
 5use fields::PrimeField64;
 6use proofman_common::{AirInstance, ProofCtx, ProofmanResult, SetupCtx};
 7use std::{collections::HashMap, sync::Arc};
 8use zisk_common::{
 9    BusDevice, CheckPoint, ChunkId, Instance, InstanceCtx, InstanceType, PayloadType,
10};

11
12pub struct MemAlignWriteByteInstance<F: PrimeField64> {
13    /// Instance context
14    ictx: InstanceCtx,
15
16    /// Checkpoint data for this memory align instance.
17    checkpoint: HashMap<ChunkId, MemAlignCheckPoint>,
18
19    mem_align_byte_sm: Arc<MemAlignByteSM<F>>,

20}

21

22impl<F: PrimeField64> MemAlignWriteByteInstance<F> {
23    pub fn new(mem_align_sm: Arc<MemAlignByteSM<F>>, mut ictx: InstanceCtx) -> Self {
24        let meta = ictx.plan.meta.take().expect("Expected metadata in ictx.plan.meta");
25
26        let checkpoint = *meta
27            .downcast::<HashMap<ChunkId, MemAlignCheckPoint>>()
28            .expect("Failed to downcast ictx.plan.meta to expected type");
29
30        Self { ictx, checkpoint, mem_align_byte_sm: mem_align_sm }

31    }

32
33    pub fn build_mem_align_write_byte_collector(&self, chunk_id: ChunkId) -> MemAlignCollector {
34        MemAlignCollector::new(&self.checkpoint[&chunk_id])

35    }

36}

37
38impl<F: PrimeField64> Instance<F> for MemAlignWriteByteInstance<F> {
39    fn compute_witness(
40        &self,
41        _pctx: &ProofCtx<F>,
42        _sctx: &SetupCtx<F>,
43        collectors: Vec<(usize, Box<dyn BusDevice<PayloadType>>)>,
44        trace_buffer: Vec<F>,
45    ) -> ProofmanResult<Option<AirInstance<F>>> {
46        let mut total_rows = 0;
47        let inputs: Vec<_> = collectors
48            .into_iter()
49            .map(|(_, collector)| {
50                let collector = collector.as_any().downcast::<MemAlignCollector>().unwrap();
51
52                total_rows += collector.count();
53
54                collector.inputs
55            })
56            .collect();
57        Ok(Some(
58            self.mem_align_byte_sm
59                .compute_witness::<MemAlignWriteByteTraceType<F>, MemAlignWriteByteTraceRowType<F>>(
60                    &inputs,
61                    total_rows as usize,
62                    trace_buffer,
63                )?,
64        ))
65    }
66
67    fn check_point(&self) -> &CheckPoint {
68        &self.ictx.plan.check_point
69    }
70
71    fn instance_type(&self) -> InstanceType {
72        InstanceType::Instance
73    }
74
75    fn as_any(&self) -> &dyn std::any::Any {
76        self
77    }
78
79    /// Builds an input collector for the instance.
80    ///
81    /// # Arguments
82    /// * `chunk_id` - The chunk ID associated with the input collector.
83    ///
84    /// # Returns
85    /// An `Option` containing the input collector for the instance.
86    fn build_inputs_collector(&self, chunk_id: ChunkId) -> Option<Box<dyn BusDevice<PayloadType>>> {
87        Some(Box::new(MemAlignCollector::new(&self.checkpoint[&chunk_id])))
88    }

89}