From 6a0f1fdaa78a08ef6b87f54f9eb2603bbb1da1ed Mon Sep 17 00:00:00 2001
From: Alexander Golovanov <Golovanov12345@gmail.com>
Date: Tue, 21 Jan 2025 01:20:21 +0100
Subject: [PATCH] Fix

---
 extensions/native/circuit/src/extension.rs    |  4 ++--
 .../native/circuit/src/loadstore/core.rs      | 19 ++++++++++---------
 .../native/circuit/src/loadstore/tests.rs     |  2 +-
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/extensions/native/circuit/src/extension.rs b/extensions/native/circuit/src/extension.rs
index 37c9c840e0..b500bf510e 100644
--- a/extensions/native/circuit/src/extension.rs
+++ b/extensions/native/circuit/src/extension.rs
@@ -117,7 +117,7 @@ impl<F: PrimeField32> VmExtension<F> for Native {
                 memory_bridge,
                 NativeLoadStoreOpcode::CLASS_OFFSET,
             ),
-            NativeLoadStoreCoreChip::new(),
+            NativeLoadStoreCoreChip::new(NativeLoadStoreOpcode::CLASS_OFFSET),
             offline_memory.clone(),
         );
         load_store_chip.core.set_streams(builder.streams().clone());
@@ -134,7 +134,7 @@ impl<F: PrimeField32> VmExtension<F> for Native {
                 memory_bridge,
                 NativeLoadStore4Opcode::CLASS_OFFSET,
             ),
-            NativeLoadStoreCoreChip::new(),
+            NativeLoadStoreCoreChip::new(NativeLoadStore4Opcode::CLASS_OFFSET),
             offline_memory.clone(),
         );
         block_load_store_chip
diff --git a/extensions/native/circuit/src/loadstore/core.rs b/extensions/native/circuit/src/loadstore/core.rs
index 53e8b73225..3a6af1543c 100644
--- a/extensions/native/circuit/src/loadstore/core.rs
+++ b/extensions/native/circuit/src/loadstore/core.rs
@@ -47,7 +47,9 @@ pub struct NativeLoadStoreCoreRecord<F, const NUM_CELLS: usize> {
 }
 
 #[derive(Clone, Debug)]
-pub struct NativeLoadStoreCoreAir<const NUM_CELLS: usize> {}
+pub struct NativeLoadStoreCoreAir<const NUM_CELLS: usize> {
+    pub offset: usize,
+}
 
 impl<F: Field, const NUM_CELLS: usize> BaseAir<F> for NativeLoadStoreCoreAir<NUM_CELLS> {
     fn width(&self) -> usize {
@@ -107,7 +109,7 @@ where
     }
 
     fn start_offset(&self) -> usize {
-        NativeLoadStoreOpcode::CLASS_OFFSET
+        self.offset
     }
 }
 
@@ -118,9 +120,9 @@ pub struct NativeLoadStoreCoreChip<F: Field, const NUM_CELLS: usize> {
 }
 
 impl<F: Field, const NUM_CELLS: usize> NativeLoadStoreCoreChip<F, NUM_CELLS> {
-    pub fn new() -> Self {
+    pub fn new(offset: usize) -> Self {
         Self {
-            air: NativeLoadStoreCoreAir::<NUM_CELLS> {},
+            air: NativeLoadStoreCoreAir::<NUM_CELLS> { offset },
             streams: OnceLock::new(),
         }
     }
@@ -131,7 +133,7 @@ impl<F: Field, const NUM_CELLS: usize> NativeLoadStoreCoreChip<F, NUM_CELLS> {
 
 impl<F: Field, const NUM_CELLS: usize> Default for NativeLoadStoreCoreChip<F, NUM_CELLS> {
     fn default() -> Self {
-        Self::new()
+        Self::new(NativeLoadStoreOpcode::CLASS_OFFSET)
     }
 }
 
@@ -151,9 +153,8 @@ where
         reads: I::Reads,
     ) -> Result<(AdapterRuntimeContext<F, I>, Self::Record)> {
         let Instruction { opcode, .. } = *instruction;
-        let local_opcode = NativeLoadStoreOpcode::from_usize(
-            opcode.local_opcode_idx(NativeLoadStoreOpcode::CLASS_OFFSET),
-        );
+        let local_opcode =
+            NativeLoadStoreOpcode::from_usize(opcode.local_opcode_idx(self.air.offset));
         let (pointer_read, data_read) = reads.into();
 
         let data_write = if local_opcode == NativeLoadStoreOpcode::HINT_STOREW {
@@ -181,7 +182,7 @@ where
     fn get_opcode_name(&self, opcode: usize) -> String {
         format!(
             "{:?}",
-            NativeLoadStoreOpcode::from_usize(opcode - NativeLoadStoreOpcode::CLASS_OFFSET)
+            NativeLoadStoreOpcode::from_usize(opcode - self.air.offset)
         )
     }
 
diff --git a/extensions/native/circuit/src/loadstore/tests.rs b/extensions/native/circuit/src/loadstore/tests.rs
index e67f080d43..ea6346719d 100644
--- a/extensions/native/circuit/src/loadstore/tests.rs
+++ b/extensions/native/circuit/src/loadstore/tests.rs
@@ -38,7 +38,7 @@ fn setup() -> (StdRng, VmChipTestBuilder<F>, NativeLoadStoreChip<F, 1>) {
         tester.memory_bridge(),
         NativeLoadStoreOpcode::CLASS_OFFSET,
     );
-    let mut inner = NativeLoadStoreCoreChip::new();
+    let mut inner = NativeLoadStoreCoreChip::new(NativeLoadStoreOpcode::CLASS_OFFSET);
     inner.set_streams(Arc::new(Mutex::new(Streams::default())));
     let chip = NativeLoadStoreChip::<F, 1>::new(adapter, inner, tester.offline_memory_mutex_arc());
     (rng, tester, chip)