Added addap opcodes.

runnair/crates/runner/src/lib.rs

@@ -1 +1,2 @@
 pub mod memory;
+pub mod vm;

runnair/crates/runner/src/memory/mod.rs

@@ -5,12 +5,25 @@ use std::ops::Index;
 use stwo_prover::core::fields::m31::M31;
 use stwo_prover::core::fields::qm31::QM31;
 
-use self::relocatable::{MaybeRelocatable, RelocationTable};
+use self::relocatable::{assert_and_project, MaybeRelocatable, RelocationTable};
 
 pub mod relocatable;
 
-type MaybeRelocatableAddr = MaybeRelocatable<M31>;
-type MaybeRelocatableValue = MaybeRelocatable<QM31>;
+pub type MaybeRelocatableAddr = MaybeRelocatable<M31>;
+pub type MaybeRelocatableValue = MaybeRelocatable<QM31>;
+
+impl From<MaybeRelocatableValue> for MaybeRelocatableAddr {
+    fn from(value: MaybeRelocatableValue) -> Self {
+        match value {
+            MaybeRelocatableValue::Relocatable(relocatable) => {
+                MaybeRelocatableAddr::Relocatable(relocatable)
+            }
+            MaybeRelocatableValue::Absolute(value) => {
+                MaybeRelocatable::Absolute(assert_and_project(value))
+            }
+        }
+    }
+}
 
 #[derive(Debug, Clone, Default)]
 pub struct Memory {
@@ -85,7 +98,7 @@ mod test {
         memory.insert(Relocatable::from((0, 0)), QM31::zero());
         memory.insert(Relocatable::from((1, 1)), Relocatable::from((1, 12)));
 
-        let table = [(0, 1), (1, 1234)].iter().cloned().collect();
+        let table = [(0, M31(1)), (1, M31(1234))].iter().cloned().collect();
 
         memory.relocate(&table);
 

runnair/crates/runner/src/memory/relocatable.rs

@@ -1,26 +1,31 @@
 use std::collections::HashMap;
+use std::ops::{Add, Mul};
+
+use num_traits::Zero;
+use stwo_prover::core::fields::m31::M31;
+use stwo_prover::core::fields::qm31::QM31;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct Relocatable {
     segment: isize,
-    offset: u32,
+    offset: M31,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub enum MaybeRelocatable<T: From<u32>> {
+pub enum MaybeRelocatable<T: From<M31>> {
     Relocatable(Relocatable),
     Absolute(T),
 }
 
-pub type RelocationTable = HashMap<isize, u32>;
+pub type RelocationTable = HashMap<isize, M31>;
 
 impl Relocatable {
-    pub fn relocate(&self, table: &RelocationTable) -> u32 {
+    pub fn relocate(&self, table: &RelocationTable) -> M31 {
         table[&self.segment] + self.offset
     }
 }
 
-impl<T: From<u32> + Copy> MaybeRelocatable<T> {
+impl<T: From<M31> + Copy> MaybeRelocatable<T> {
     pub fn relocate(&self, table: &RelocationTable) -> T {
         match self {
             MaybeRelocatable::Relocatable(relocatable) => relocatable.relocate(table).into(),
@@ -29,20 +34,166 @@ impl<T: From<u32> + Copy> MaybeRelocatable<T> {
     }
 }
 
+impl From<(isize, M31)> for Relocatable {
+    fn from((segment, offset): (isize, M31)) -> Self {
+        Relocatable { segment, offset }
+    }
+}
+
 impl From<(isize, u32)> for Relocatable {
     fn from((segment, offset): (isize, u32)) -> Self {
-        Relocatable { segment, offset }
+        Relocatable {
+            segment,
+            offset: M31(offset),
+        }
     }
 }
 
-impl<T: From<u32>> From<Relocatable> for MaybeRelocatable<T> {
+impl<T: From<M31>> From<Relocatable> for MaybeRelocatable<T> {
     fn from(relocatable: Relocatable) -> Self {
         MaybeRelocatable::Relocatable(relocatable)
     }
 }
 
-impl<T: From<u32>> From<T> for MaybeRelocatable<T> {
+impl<T: From<M31>> From<T> for MaybeRelocatable<T> {
     fn from(value: T) -> Self {
         MaybeRelocatable::Absolute(value)
     }
 }
+
+impl Add<M31> for Relocatable {
+    type Output = Self;
+    fn add(self, rhs: M31) -> Self {
+        Self {
+            segment: self.segment,
+            offset: self.offset + rhs,
+        }
+    }
+}
+
+// TODO(alont): Can this be generalized?
+impl Add<MaybeRelocatable<M31>> for MaybeRelocatable<M31> {
+    type Output = Self;
+    fn add(self, rhs: Self) -> Self {
+        match (self, rhs) {
+            (MaybeRelocatable::Relocatable(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Relocatable(lhs + rhs)
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Relocatable(rhs)) => {
+                MaybeRelocatable::Relocatable(rhs + lhs)
+            }
+            (MaybeRelocatable::Relocatable(_), MaybeRelocatable::Relocatable(_)) => {
+                panic!("Cannot add two relocatables.")
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Absolute(lhs + rhs)
+            }
+        }
+    }
+}
+
+/// Assert that the input is in the base field and return the projection to the base field.
+pub fn assert_and_project(x: QM31) -> M31 {
+    assert!(x.1.is_zero());
+    assert!(x.0 .1.is_zero());
+    x.0 .0
+}
+
+impl Add<MaybeRelocatable<QM31>> for MaybeRelocatable<M31> {
+    type Output = MaybeRelocatable<QM31>;
+    fn add(self, rhs: MaybeRelocatable<QM31>) -> MaybeRelocatable<QM31> {
+        match (self, rhs) {
+            (MaybeRelocatable::Relocatable(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Relocatable(lhs + assert_and_project(rhs))
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Relocatable(rhs)) => {
+                MaybeRelocatable::Relocatable(rhs + lhs)
+            }
+            (MaybeRelocatable::Relocatable(_), MaybeRelocatable::Relocatable(_)) => {
+                panic!("Cannot add two relocatables.")
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Absolute(lhs + rhs)
+            }
+        }
+    }
+}
+
+impl Add<MaybeRelocatable<M31>> for MaybeRelocatable<QM31> {
+    type Output = Self;
+    fn add(self, rhs: MaybeRelocatable<M31>) -> Self {
+        match (self, rhs) {
+            (MaybeRelocatable::Relocatable(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Relocatable(lhs + rhs)
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Relocatable(rhs)) => {
+                MaybeRelocatable::Relocatable(rhs + assert_and_project(lhs))
+            }
+            (MaybeRelocatable::Relocatable(_), MaybeRelocatable::Relocatable(_)) => {
+                panic!("Cannot add two relocatables.")
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Absolute(lhs + rhs)
+            }
+        }
+    }
+}
+
+impl Add<MaybeRelocatable<QM31>> for MaybeRelocatable<QM31> {
+    type Output = Self;
+    fn add(self, rhs: Self) -> Self {
+        match (self, rhs) {
+            (MaybeRelocatable::Relocatable(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Relocatable(lhs + assert_and_project(rhs))
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Relocatable(rhs)) => {
+                MaybeRelocatable::Relocatable(rhs + assert_and_project(lhs))
+            }
+            (MaybeRelocatable::Relocatable(_), MaybeRelocatable::Relocatable(_)) => {
+                panic!("Cannot add two relocatables.")
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Absolute(lhs + rhs)
+            }
+        }
+    }
+}
+
+impl<T: Add<M31, Output = T> + From<M31>> Add<M31> for MaybeRelocatable<T> {
+    type Output = Self;
+    fn add(self, rhs: M31) -> Self {
+        match self {
+            MaybeRelocatable::Relocatable(lhs) => MaybeRelocatable::Relocatable(lhs + rhs),
+            MaybeRelocatable::Absolute(lhs) => MaybeRelocatable::Absolute(lhs + rhs),
+        }
+    }
+}
+
+impl<T: From<M31> + Mul<S, Output = S>, S: From<M31>> Mul<MaybeRelocatable<S>>
+    for MaybeRelocatable<T>
+{
+    type Output = MaybeRelocatable<S>;
+    fn mul(self, rhs: MaybeRelocatable<S>) -> Self::Output {
+        match (self, rhs) {
+            (MaybeRelocatable::Relocatable(_), _) => {
+                panic!("Cannot multiply a relocatable.")
+            }
+            (_, MaybeRelocatable::Relocatable(_)) => {
+                panic!("Cannot multiply a relocatable.")
+            }
+            (MaybeRelocatable::Absolute(lhs), MaybeRelocatable::Absolute(rhs)) => {
+                MaybeRelocatable::Absolute(lhs * rhs)
+            }
+        }
+    }
+}
+
+impl<T: From<M31> + Mul<M31, Output = T>> Mul<M31> for MaybeRelocatable<T> {
+    type Output = Self;
+    fn mul(self, rhs: M31) -> Self {
+        match self {
+            MaybeRelocatable::Relocatable(_) => panic!("Cannot multiply a relocatable."),
+            MaybeRelocatable::Absolute(lhs) => MaybeRelocatable::Absolute(lhs * rhs),
+        }
+    }
+}

runnair/crates/runner/src/vm/add_ap.rs

@@ -0,0 +1,45 @@
+use stwo_prover::core::fields::m31::M31;
+
+use super::{Instruction, State};
+use crate::memory::relocatable::assert_and_project;
+use crate::memory::{MaybeRelocatableValue, Memory};
+
+pub fn addap(state: State, summand: M31) -> State {
+    State {
+        ap: state.ap + summand,
+        fp: state.fp,
+        pc: state.pc,
+    }
+}
+
+macro_rules! addap_with_operand {
+    ($instruction:ident, $operand:ident) => {
+        pub fn $instruction(memory: &mut Memory, state: State, instruction: Instruction) -> State {
+            if let MaybeRelocatableValue::Absolute(summand) =
+                crate::vm::operand::$operand(memory, state, &instruction.args)
+            {
+                addap(state, assert_and_project(summand))
+            } else {
+                panic!("Can't addap by a relocatable.")
+            }
+        }
+    };
+}
+
+addap_with_operand!(addap_add_ap_ap, add_ap_ap);
+addap_with_operand!(addap_add_ap_fp, add_ap_fp);
+addap_with_operand!(addap_add_fp_ap, add_fp_ap);
+addap_with_operand!(addap_add_fp_fp, add_fp_fp);
+addap_with_operand!(addap_add_imm_ap, add_imm_ap);
+addap_with_operand!(addap_add_imm_fp, add_imm_fp);
+addap_with_operand!(addap_deref_ap, deref_ap);
+addap_with_operand!(addap_deref_fp, deref_fp);
+addap_with_operand!(addap_double_deref_ap, double_deref_ap);
+addap_with_operand!(addap_double_deref_fp, double_deref_fp);
+addap_with_operand!(addap_imm, imm);
+addap_with_operand!(addap_mul_ap_ap, mul_ap_ap);
+addap_with_operand!(addap_mul_ap_fp, mul_ap_fp);
+addap_with_operand!(addap_mul_fp_ap, mul_fp_ap);
+addap_with_operand!(addap_mul_fp_fp, mul_fp_fp);
+addap_with_operand!(addap_mul_imm_ap, mul_imm_ap);
+addap_with_operand!(addap_mul_imm_fp, mul_imm_fp);