feat: MemoryAccessAdapter

compiler/src/asm/builder.rs

@@ -16,7 +16,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmBu
     }
 
     pub fn compile_isa_with_options(self, options: CompilerOptions) -> Program<F> {
-        let mut compiler = AsmCompiler::new(1);
+        let mut compiler = AsmCompiler::new(options.word_size);
         compiler.build(self.operations);
         let asm_code = compiler.code();
         convert_program(asm_code, options)

compiler/src/asm/compiler.rs

@@ -12,18 +12,18 @@ use crate::{
 };
 
 /// The memory location for the top of memory
-pub const MEMORY_TOP: i32 = (1 << 29) - 4;
+pub const MEMORY_TOP: u32 = (1 << 29) - 4;
+
+// The memory location for the start of the heap.
+pub(crate) const HEAP_START_ADDRESS: i32 = 1 << 24;
 
 /// The heap pointer address.
-pub(crate) const HEAP_PTR: i32 = MEMORY_TOP - 4;
+pub(crate) const HEAP_PTR: i32 = HEAP_START_ADDRESS - 4;
 /// Utility register.
-pub(crate) const A0: i32 = MEMORY_TOP - 8;
+pub(crate) const A0: i32 = HEAP_START_ADDRESS - 8;
 
 /// The memory location for the top of the stack.
-pub(crate) const STACK_TOP: i32 = MEMORY_TOP - 100;
-
-// The memory location for the start of the heap.
-pub(crate) const HEAP_START_ADDRESS: usize = 4;
+pub(crate) const STACK_TOP: i32 = HEAP_START_ADDRESS - 64;
 
 /// The assembly compiler.
 // #[derive(Debug, Clone, Default)]
@@ -41,21 +41,24 @@ pub struct AsmCompiler<F, EF> {
 impl<F> Var<F> {
     /// Gets the frame pointer for a var.
     pub const fn fp(&self) -> i32 {
-        STACK_TOP - self.0 as i32
+        // Vars are stored in stack positions 1, 2, 9, 10, 17, 18, ...
+        STACK_TOP - (8 * (self.0 / 2) + 1 + (self.0 % 2)) as i32
     }
 }
 
 impl<F> Felt<F> {
     /// Gets the frame pointer for a felt.
     pub const fn fp(&self) -> i32 {
-        STACK_TOP - self.0 as i32
+        // Felts are stored in stack positions 3, 4, 11, 12, 19, 20, ...
+        STACK_TOP - (((self.0 >> 1) << 3) + 3 + (self.0 & 1)) as i32
     }
 }
 
 impl<F, EF> Ext<F, EF> {
     /// Gets the frame pointer for an extension element
     pub const fn fp(&self) -> i32 {
-        STACK_TOP - self.0 as i32
+        // Exts are stored in stack positions 5-8, 13-16, 21-24, ...
+        STACK_TOP - 8 * self.0 as i32
     }
 }
 
@@ -94,7 +97,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
     pub fn build(&mut self, operations: TracedVec<DslIr<AsmConfig<F, EF>>>) {
         if self.block_label().is_zero() {
             // Initialize the heap pointer value.
-            let heap_start = F::from_canonical_usize(HEAP_START_ADDRESS);
+            let heap_start = F::from_canonical_u32(HEAP_START_ADDRESS as u32);
             self.push(AsmInstruction::ImmF(HEAP_PTR, heap_start), None);
             // Jump over the TRAP instruction we are about to add.
             self.push(AsmInstruction::j(self.trap_label + F::one()), None);
@@ -668,25 +671,25 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
         size: usize,
         debug_info: Option<DebugInfo>,
     ) {
+        let word_size = self.word_size;
+        let align = |x: usize| (x + word_size - 1) / word_size * word_size;
         // Load the current heap ptr address to the stack value and advance the heap ptr.
-        let size = F::from_canonical_usize(size);
         let len = len.into();
         match len {
             RVar::Const(len) => {
                 self.push(
                     AsmInstruction::CopyF(ptr.fp(), HEAP_PTR),
                     debug_info.clone(),
                 );
-                self.push(
-                    AsmInstruction::AddFI(HEAP_PTR, HEAP_PTR, len * size),
-                    debug_info,
-                );
+                let inc = F::from_canonical_usize(align((len.as_canonical_u32() as usize) * size));
+                self.push(AsmInstruction::AddFI(HEAP_PTR, HEAP_PTR, inc), debug_info);
             }
             RVar::Val(len) => {
                 self.push(
                     AsmInstruction::CopyF(ptr.fp(), HEAP_PTR),
                     debug_info.clone(),
                 );
+                let size = F::from_canonical_usize(align(size));
                 self.push(
                     AsmInstruction::MulFI(A0, len.fp(), size),
                     debug_info.clone(),
@@ -1042,8 +1045,10 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
     fn assign_exti(&mut self, dst: i32, imm: EF, debug_info: Option<DebugInfo>) {
         let imm = imm.as_base_slice();
         for i in 0..EF::D {
-            let j = (i * self.word_size) as i32;
-            self.push(AsmInstruction::ImmF(dst + j, imm[i]), debug_info.clone());
+            self.push(
+                AsmInstruction::ImmF(dst + i as i32, imm[i]),
+                debug_info.clone(),
+            );
         }
     }
 
@@ -1059,11 +1064,11 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
                 for i in 0..EF::D {
                     self.push(
                         AsmInstruction::LoadFI(
-                            val.fp() + (i * self.word_size) as i32,
+                            val.fp() + i as i32,
                             addr,
                             index,
                             size,
-                            offset + F::from_canonical_usize(i * self.word_size),
+                            offset + F::from_canonical_usize(i),
                         ),
                         debug_info.clone(),
                     )
@@ -1073,11 +1078,11 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
                 for i in 0..EF::D {
                     self.push(
                         AsmInstruction::LoadF(
-                            val.fp() + (i * self.word_size) as i32,
+                            val.fp() + i as i32,
                             addr,
                             index,
                             size,
-                            offset + F::from_canonical_usize(i * self.word_size),
+                            offset + F::from_canonical_usize(i),
                         ),
                         debug_info.clone(),
                     )
@@ -1098,11 +1103,11 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
                 for i in 0..EF::D {
                     self.push(
                         AsmInstruction::StoreFI(
-                            val.fp() + (i * self.word_size) as i32,
+                            val.fp() + i as i32,
                             addr,
                             index,
                             size,
-                            offset + F::from_canonical_usize(i * self.word_size),
+                            offset + F::from_canonical_usize(i),
                         ),
                         debug_info.clone(),
                     )
@@ -1112,11 +1117,11 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
                 for i in 0..EF::D {
                     self.push(
                         AsmInstruction::StoreF(
-                            val.fp() + (i * self.word_size) as i32,
+                            val.fp() + i as i32,
                             addr,
                             index,
                             size,
-                            offset + F::from_canonical_usize(i * self.word_size),
+                            offset + F::from_canonical_usize(i),
                         ),
                         debug_info.clone(),
                     )
@@ -1134,7 +1139,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
     ) {
         let rhs = rhs.as_base_slice();
         for i in 0..EF::D {
-            let j = (i * self.word_size) as i32;
+            let j = i as i32;
             self.push(
                 AsmInstruction::AddFI(dst.fp() + j, lhs.fp() + j, rhs[i]),
                 debug_info.clone(),
@@ -1151,7 +1156,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
     ) {
         let lhs = lhs.as_base_slice();
         for i in 0..EF::D {
-            let j = (i * self.word_size) as i32;
+            let j = i as i32;
             self.push(
                 AsmInstruction::SubFIN(dst.fp() + j, lhs[i], rhs.fp() + j),
                 debug_info.clone(),
@@ -1171,7 +1176,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
             debug_info.clone(),
         );
         for i in 1..EF::D {
-            let j = (i * self.word_size) as i32;
+            let j = i as i32;
             self.push(
                 AsmInstruction::CopyF(dst.fp() + j, lhs.fp() + j),
                 debug_info.clone(),
@@ -1191,7 +1196,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
             debug_info.clone(),
         );
         for i in 1..EF::D {
-            let j = (i * self.word_size) as i32;
+            let j = i as i32;
             self.push(
                 AsmInstruction::CopyF(dst.fp() + j, lhs.fp() + j),
                 debug_info.clone(),
@@ -1214,7 +1219,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
         );
 
         for i in 1..EF::D {
-            let j = (i * self.word_size) as i32;
+            let j = i as i32;
             self.push(
                 AsmInstruction::ImmF(dst.fp() + j, rhs[i]),
                 debug_info.clone(),
@@ -1230,7 +1235,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
         debug_info: Option<DebugInfo>,
     ) {
         for i in 0..EF::D {
-            let j = (i * self.word_size) as i32;
+            let j = i as i32;
             self.push(
                 AsmInstruction::MulF(dst.fp() + j, lhs.fp() + j, rhs.fp()),
                 debug_info.clone(),
@@ -1246,7 +1251,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
         debug_info: Option<DebugInfo>,
     ) {
         for i in 0..EF::D {
-            let j = (i * self.word_size) as i32;
+            let j = i as i32;
             self.push(
                 AsmInstruction::MulFI(dst.fp() + j, lhs.fp() + j, rhs),
                 debug_info.clone(),
@@ -1262,7 +1267,7 @@ impl<F: PrimeField32 + TwoAdicField, EF: ExtensionField<F> + TwoAdicField> AsmCo
         debug_info: Option<DebugInfo>,
     ) {
         for i in 0..EF::D {
-            let j = (i * self.word_size) as i32;
+            let j = i as i32;
             self.push(
                 AsmInstruction::DivF(dst.fp() + j, lhs.fp() + j, rhs.fp()),
                 debug_info.clone(),

compiler/src/conversion/mod.rs

@@ -8,6 +8,8 @@ use crate::asm::{AsmInstruction, AssemblyCode};
 
 #[derive(Clone, Copy, Debug)]
 pub struct CompilerOptions {
+    // The compiler will ensure that the heap pointer is aligned to be a multiple of `word_size`.
+    pub word_size: usize,
     pub compile_prints: bool,
     pub enable_cycle_tracker: bool,
     pub field_arithmetic_enabled: bool,
@@ -17,6 +19,7 @@ pub struct CompilerOptions {
 impl Default for CompilerOptions {
     fn default() -> Self {
         CompilerOptions {
+            word_size: 8,
             compile_prints: true,
             enable_cycle_tracker: false,
             field_arithmetic_enabled: true,

compiler/src/ir/builder.rs

@@ -106,7 +106,9 @@ pub struct BuilderFlags {
 /// Can compile to both assembly and a set of constraints.
 #[derive(Debug, Clone, Default)]
 pub struct Builder<C: Config> {
-    pub(crate) stack_ptr: u32,
+    pub(crate) var_count: u32,
+    pub(crate) felt_count: u32,
+    pub(crate) ext_count: u32,
     pub operations: TracedVec<DslIr<C>>,
     pub(crate) nb_public_values: Option<Var<C::N>>,
     pub(crate) witness_var_count: u32,
@@ -120,13 +122,17 @@ pub struct Builder<C: Config> {
 impl<C: Config> Builder<C> {
     /// Creates a new builder with a given number of counts for each type.
     pub fn new_sub_builder(
-        stack_ptr: u32,
+        var_count: u32,
+        felt_count: u32,
+        ext_count: u32,
         nb_public_values: Option<Var<C::N>>,
         flags: BuilderFlags,
         bigint_repr_size: u32,
     ) -> Self {
         Self {
-            stack_ptr,
+            var_count,
+            felt_count,
+            ext_count,
             // Witness counts are only used when the target is a gnark circuit.  And sub-builders are
             // not used when the target is a gnark circuit, so it's fine to set the witness counts to 0.
             witness_var_count: 0,
@@ -346,7 +352,9 @@ impl<C: Config> Builder<C> {
     /// Evaluate a block of operations repeatedly (until a break).
     pub fn do_loop(&mut self, mut f: impl FnMut(&mut Builder<C>) -> Result<(), BreakLoop>) {
         let mut loop_body_builder = Builder::<C>::new_sub_builder(
-            self.stack_ptr,
+            self.var_count,
+            self.felt_count,
+            self.ext_count,
             self.nb_public_values,
             self.flags,
             self.bigint_repr_size,
@@ -637,7 +645,9 @@ impl<'a, C: Config> IfBuilder<'a, C> {
 
         // Execute the `then` block and collect the instructions.
         let mut f_builder = Builder::<C>::new_sub_builder(
-            self.builder.stack_ptr,
+            self.builder.var_count,
+            self.builder.felt_count,
+            self.builder.ext_count,
             self.builder.nb_public_values,
             self.builder.flags,
             self.builder.bigint_repr_size,
@@ -714,7 +724,9 @@ impl<'a, C: Config> IfBuilder<'a, C> {
             "Cannot use dynamic branch in static mode"
         );
         let mut then_builder = Builder::<C>::new_sub_builder(
-            self.builder.stack_ptr,
+            self.builder.var_count,
+            self.builder.felt_count,
+            self.builder.ext_count,
             self.builder.nb_public_values,
             self.builder.flags,
             self.builder.bigint_repr_size,
@@ -725,7 +737,9 @@ impl<'a, C: Config> IfBuilder<'a, C> {
         let then_instructions = then_builder.operations;
 
         let mut else_builder = Builder::<C>::new_sub_builder(
-            self.builder.stack_ptr,
+            self.builder.var_count,
+            self.builder.felt_count,
+            self.builder.ext_count,
             self.builder.nb_public_values,
             self.builder.flags,
             self.builder.bigint_repr_size,
@@ -897,7 +911,9 @@ impl<'a, C: Config> RangeBuilder<'a, C> {
         let step_size = C::N::from_canonical_usize(self.step_size);
         let loop_variable: Var<C::N> = self.builder.uninit();
         let mut loop_body_builder = Builder::<C>::new_sub_builder(
-            self.builder.stack_ptr,
+            self.builder.var_count,
+            self.builder.felt_count,
+            self.builder.ext_count,
             self.builder.nb_public_values,
             self.builder.flags,
             self.builder.bigint_repr_size,

compiler/src/ir/types.rs

@@ -410,8 +410,8 @@ impl<C: Config> Variable<C> for Var<C::N> {
     type Expression = SymbolicVar<C::N>;
 
     fn uninit(builder: &mut Builder<C>) -> Self {
-        builder.stack_ptr += 1;
-        Var(builder.stack_ptr, PhantomData)
+        builder.var_count += 1;
+        Var(builder.var_count, PhantomData)
     }
 
     fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
@@ -780,8 +780,8 @@ impl<C: Config> Variable<C> for Felt<C::F> {
     type Expression = SymbolicFelt<C::F>;
 
     fn uninit(builder: &mut Builder<C>) -> Self {
-        builder.stack_ptr += 1;
-        Felt(builder.stack_ptr, PhantomData)
+        builder.felt_count += 1;
+        Felt(builder.felt_count, PhantomData)
     }
 
     fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {
@@ -1209,8 +1209,8 @@ impl<C: Config> Variable<C> for Ext<C::F, C::EF> {
     type Expression = SymbolicExt<C::F, C::EF>;
 
     fn uninit(builder: &mut Builder<C>) -> Self {
-        builder.stack_ptr += <Self as MemVariable<C>>::size_of() as u32;
-        Ext(builder.stack_ptr, PhantomData)
+        builder.ext_count += 1;
+        Ext(builder.ext_count, PhantomData)
     }
 
     fn assign(&self, src: Self::Expression, builder: &mut Builder<C>) {