update branch new-arch-mem with new-arch changes

0xPolygonHermez · Jan 7, 2025 · 0a27aa5 · 0a27aa5
1 parent f3111f2
commit 0a27aa5
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 13 deletions.
diff --git a/state-machines/mem/src/mem_align_planner.rs b/state-machines/mem/src/mem_align_planner.rs
@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use crate::{MemCounters, MemPlanCalculator};
-use sm_common::{CheckPoint, ChunkId, InstanceType, Plan};
+use sm_common::{CheckPoint, ChunkId, CollectInfoSkip, InstanceType, Plan};
 use zisk_pil::{MEM_ALIGN_AIR_IDS, ZISK_AIRGROUP_ID};
 
 pub struct MemAlignPlanner<'a> {
@@ -79,16 +79,17 @@ impl<'a> MemAlignPlanner<'a> {
             return;
         }
         // TODO: add multi chunk_id, with skip
+        let chunks = std::mem::take(&mut self.current_chunks);
         let instance = Plan::new(
             ZISK_AIRGROUP_ID,
             MEM_ALIGN_AIR_IDS[0],
             Some(self.instances.len()),
             InstanceType::Instance,
-            Some(CheckPoint::new(self.current_chunks[0], self.current_skip as u64)),
+            CheckPoint::Multiple(chunks),
+            Some(Box::new(CollectInfoSkip::new(self.current_skip as u64))),
             None,
         );
         self.instances.push(instance);
-        self.current_chunks.clear();
     }
     fn open_new_instance(&mut self, next_instance_skip: u32, use_current_chunk_id: bool) {
         self.current_skip = next_instance_skip;

diff --git a/state-machines/mem/src/mem_module_planner.rs b/state-machines/mem/src/mem_module_planner.rs
@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use crate::{MemCounters, MemPlanCalculator, UsesCounter, MEMORY_MAX_DIFF};
-use sm_common::{ChunkId, InstanceType, Plan};
+use sm_common::{CheckPoint, ChunkId, InstanceType, Plan};
 
 pub struct MemModulePlanner<'a> {
     airgroup_id: usize,
@@ -163,16 +163,17 @@ impl<'a> MemModulePlanner<'a> {
         // }
 
         let checkpoint = std::mem::take(&mut self.current_checkpoint);
+        let chunks = std::mem::take(&mut self.current_checkpoint_chunks);
         let instance = Plan::new(
             self.airgroup_id,
             self.air_id,
             Some(self.instances.len()),
             InstanceType::Instance,
+            CheckPoint::Multiple(chunks),
             None,
             Some(Box::new(checkpoint)),
         );
         self.instances.push(instance);
-        self.current_checkpoint_chunks.clear();
     }
     fn set_current_chunk_id(&mut self, chunk_id: ChunkId) {
         if self.current_chunk_id == Some(chunk_id) && !self.current_checkpoint_chunks.is_empty() {

diff --git a/state-machines/mem/src/mem_proxy.rs b/state-machines/mem/src/mem_proxy.rs
@@ -6,9 +6,7 @@ use crate::{
 };
 use p3_field::PrimeField;
 use pil_std_lib::Std;
-use sm_common::{
-    BusDeviceInstance, BusDeviceMetrics, ComponentProvider, InstanceExpanderCtx, Planner,
-};
+use sm_common::{BusDeviceInstance, BusDeviceMetrics, ComponentBuilder, InstanceCtx, Planner};
 use zisk_core::ZiskRequiredMemory;
 
 pub struct MemProxy<F: PrimeField> {
@@ -49,20 +47,20 @@ impl<F: PrimeField> MemProxy<F> {
     }
 }
 
-impl<F: PrimeField> ComponentProvider<F> for MemProxy<F> {
-    fn get_counter(&self) -> Box<dyn BusDeviceMetrics> {
+impl<F: PrimeField> ComponentBuilder<F> for MemProxy<F> {
+    fn build_counter(&self) -> Box<dyn BusDeviceMetrics> {
         Box::new(MemCounters::new())
         // Box::new(MemCounters::new(OPERATION_BUS_ID, vec![zisk_core::ZiskOperationType::Arith]))
     }
 
-    fn get_planner(&self) -> Box<dyn Planner> {
+    fn build_planner(&self) -> Box<dyn Planner> {
         Box::new(MemPlanner::new())
     }
 
-    fn get_instance(&self, _iectx: InstanceExpanderCtx) -> Box<dyn BusDeviceInstance<F>> {
+    fn build_inputs_collector(&self, _iectx: InstanceCtx) -> Box<dyn BusDeviceInstance<F>> {
         unimplemented!("get_instance for MemProxy");
     }
-    fn get_inputs_generator(&self) -> Option<Box<dyn BusDeviceInstance<F>>> {
+    fn build_inputs_generator(&self) -> Option<Box<dyn BusDeviceInstance<F>>> {
         unimplemented!("get_instance for MemProxy");
     }
 }