Repository: lambdaclass/cairo-vm (language: rust)

Issue:
Mod Builtin Constraints
Right now, modulo builtin constraints are not always fulfilled. 

- Given a sub mod gate (for example: `x3 = x2 - x1`), then it must always be true that `x3 + x1 - x2 = 0 or p`. 
 
Take, for example, `x2 = p + 1` and `x1 = 1`. Our current implementation will calculate `x3 = 0` as a result of the operation `(x2 - x1) % p`. The expected output is `p`, as it's the minimum value that meets the constraint (`p + 1 - (p + 1) = 0`). (the other possible value is `2p`.
 
See https://github.com/starkware-libs/cairo-lang/pull/196 and cairo-lang implementation as reference.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 6dde21ab95..600b0e52ca 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -600,6 +600,36 @@ jobs:
         key: codecov-cache-test-no_std-extensive_hints-${{ github.sha }}
         fail-on-cache-miss: true
 
+    - name: Fetch results for tests with stdlib (w/mod_builtin; part. 1)
+      uses: actions/cache/restore@v3
+      with:
+        path: lcov-test#1-mod_builtin.info
+        key: codecov-cache-test#1-mod_builtin-${{ github.sha }}
+        fail-on-cache-miss: true
+    - name: Fetch results for tests with stdlib (w/mod_builtin; part. 2)
+      uses: actions/cache/restore@v3
+      with:
+        path: lcov-test#2-mod_builtin.info
+        key: codecov-cache-test#2-mod_builtin-${{ github.sha }}
+        fail-on-cache-miss: true
+    - name: Fetch results for tests with stdlib (w/mod_builtin; part. 3)
+      uses: actions/cache/restore@v3
+      with:
+        path: lcov-test#3-mod_builtin.info
+        key: codecov-cache-test#3-mod_builtin-${{ github.sha }}
+        fail-on-cache-miss: true
+    - name: Fetch results for tests with stdlib (w/mod_builtin; part. 4)
+      uses: actions/cache/restore@v3
+      with:
+        path: lcov-test#4-mod_builtin.info
+        key: codecov-cache-test#4-mod_builtin-${{ github.sha }}
+        fail-on-cache-miss: true
+    - name: Fetch results for tests without stdlib (w/mod_builtin)
+      uses: actions/cache/restore@v3
+      with:
+        path: lcov-no_std-mod_builtin.info
+        key: codecov-cache-test-no_std-mod_builtin-${{ github.sha }}
+        fail-on-cache-miss: true
 
     - name: Upload coverage to codecov.io
       uses: codecov/codecov-action@v3
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3317000d50..1f4001b5a7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 #### Upcoming Changes
 
+* fix: [#1841](https://github.com/lambdaclass/cairo-vm/pull/1841):
+  * Fix modulo builtin to comply with prover constraints
+
 * feat(BREAKING): [#1824](https://github.com/lambdaclass/cairo-vm/pull/1824):
     * Add support for dynamic layout
     * CLI change(BREAKING): The flag `cairo_layout_params_file` must be specified when using dynamic layout.
diff --git a/vm/src/air_private_input.rs b/vm/src/air_private_input.rs
index b76575eaff..421c1b3fa0 100644
--- a/vm/src/air_private_input.rs
+++ b/vm/src/air_private_input.rs
@@ -126,6 +126,8 @@ pub struct ModInputInstance {
     pub values_ptr: usize,
     pub offsets_ptr: usize,
     pub n: usize,
+    #[serde(deserialize_with = "mod_input_instance_batch_serde::deserialize")]
+    #[serde(serialize_with = "mod_input_instance_batch_serde::serialize")]
     pub batch: BTreeMap<usize, ModInputMemoryVars>,
 }
 
@@ -205,6 +207,88 @@ impl AirPrivateInputSerializable {
     }
 }
 
+mod mod_input_instance_batch_serde {
+    use super::*;
+
+    use serde::{Deserializer, Serializer};
+
+    pub(crate) fn serialize<S: Serializer>(
+        value: &BTreeMap<usize, ModInputMemoryVars>,
+        s: S,
+    ) -> Result<S::Ok, S::Error> {
+        let value = value.iter().map(|v| v.1).collect::<Vec<_>>();
+
+        value.serialize(s)
+    }
+
+    pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
+        d: D,
+    ) -> Result<BTreeMap<usize, ModInputMemoryVars>, D::Error> {
+        let value = Vec::<ModInputMemoryVars>::deserialize(d)?;
+
+        Ok(value.into_iter().enumerate().collect())
+    }
+
+    #[cfg(feature = "std")]
+    #[test]
+    fn test_serde() {
+        let input_value = vec![
+            (
+                0,
+                ModInputMemoryVars {
+                    a_offset: 5,
+                    b_offset: 5,
+                    c_offset: 5,
+                    a0: Felt252::from(5u32),
+                    a1: Felt252::from(5u32),
+                    a2: Felt252::from(5u32),
+                    a3: Felt252::from(5u32),
+                    b0: Felt252::from(5u32),
+                    b1: Felt252::from(5u32),
+                    b2: Felt252::from(5u32),
+                    b3: Felt252::from(5u32),
+                    c0: Felt252::from(5u32),
+                    c1: Felt252::from(5u32),
+                    c2: Felt252::from(5u32),
+                    c3: Felt252::from(5u32),
+                },
+            ),
+            (
+                1,
+                ModInputMemoryVars {
+                    a_offset: 7,
+                    b_offset: 7,
+                    c_offset: 7,
+                    a0: Felt252::from(7u32),
+                    a1: Felt252::from(7u32),
+                    a2: Felt252::from(7u32),
+                    a3: Felt252::from(7u32),
+                    b0: Felt252::from(7u32),
+                    b1: Felt252::from(7u32),
+                    b2: Felt252::from(7u32),
+                    b3: Felt252::from(7u32),
+                    c0: Felt252::from(7u32),
+                    c1: Felt252::from(7u32),
+                    c2: Felt252::from(7u32),
+                    c3: Felt252::from(7u32),
+                },
+            ),
+        ]
+        .into_iter()
+        .collect::<BTreeMap<usize, _>>();
+
+        let bytes = Vec::new();
+        let mut serializer = serde_json::Serializer::new(bytes);
+        serialize(&input_value, &mut serializer).unwrap();
+        let bytes = serializer.into_inner();
+
+        let mut deserializer = serde_json::Deserializer::from_slice(&bytes);
+        let output_value = deserialize(&mut deserializer).unwrap();
+
+        assert_eq!(input_value, output_value);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::types::layout_name::LayoutName;
diff --git a/vm/src/vm/runners/builtin_runner/modulo.rs b/vm/src/vm/runners/builtin_runner/modulo.rs
index 572ccf1897..0f5bb7719b 100644
--- a/vm/src/vm/runners/builtin_runner/modulo.rs
+++ b/vm/src/vm/runners/builtin_runner/modulo.rs
@@ -21,7 +21,7 @@ use crate::{
     },
     Felt252,
 };
-use core::{fmt::Display, ops::Shl};
+use core::ops::Shl;
 use num_bigint::BigUint;
 use num_integer::div_ceil;
 use num_integer::Integer;
@@ -47,6 +47,7 @@ pub struct ModBuiltinRunner {
     // Precomputed powers used for reading and writing values that are represented as n_words words of word_bit_len bits each.
     shift: BigUint,
     shift_powers: [BigUint; N_WORDS],
+    k_bound: BigUint,
 }
 
 #[derive(Debug, Clone)]
@@ -55,21 +56,11 @@ pub enum ModBuiltinType {
     Add,
 }
 
-#[derive(Debug)]
-pub enum Operation {
-    Mul,
-    Add,
-    Sub,
-    DivMod(BigUint),
-}
-
-impl Display for Operation {
-    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+impl ModBuiltinType {
+    pub(crate) fn operation_string(&self) -> &'static str {
         match self {
-            Operation::Mul => "*".fmt(f),
-            Operation::Add => "+".fmt(f),
-            Operation::Sub => "-".fmt(f),
-            Operation::DivMod(_) => "/".fmt(f),
+            ModBuiltinType::Mul => "*",
+            ModBuiltinType::Add => "+",
         }
     }
 }
@@ -85,17 +76,28 @@ struct Inputs {
 
 impl ModBuiltinRunner {
     pub(crate) fn new_add_mod(instance_def: &ModInstanceDef, included: bool) -> Self {
-        Self::new(instance_def.clone(), included, ModBuiltinType::Add)
+        Self::new(
+            instance_def.clone(),
+            included,
+            ModBuiltinType::Add,
+            Some(2u32.into()),
+        )
     }
 
     pub(crate) fn new_mul_mod(instance_def: &ModInstanceDef, included: bool) -> Self {
-        Self::new(instance_def.clone(), included, ModBuiltinType::Mul)
+        Self::new(instance_def.clone(), included, ModBuiltinType::Mul, None)
     }
 
-    fn new(instance_def: ModInstanceDef, included: bool, builtin_type: ModBuiltinType) -> Self {
+    fn new(
+        instance_def: ModInstanceDef,
+        included: bool,
+        builtin_type: ModBuiltinType,
+        k_bound: Option<BigUint>,
+    ) -> Self {
         let shift = BigUint::one().shl(instance_def.word_bit_len);
         let shift_powers = core::array::from_fn(|i| shift.pow(i as u32));
         let zero_segment_size = core::cmp::max(N_WORDS, instance_def.batch_size * 3);
+        let int_lim = BigUint::from(2_u32).pow(N_WORDS as u32 * instance_def.word_bit_len);
         Self {
             builtin_type,
             base: 0,
@@ -106,6 +108,7 @@ impl ModBuiltinRunner {
             zero_segment_size,
             shift,
             shift_powers,
+            k_bound: k_bound.unwrap_or(int_lim),
         }
     }
 
@@ -435,13 +438,13 @@ impl ModBuiltinRunner {
 
     // Fills a value in the values table, if exactly one value is missing.
     // Returns true on success or if all values are already known.
+    //
+    // The builtin type (add or mul) determines which operation to perform
     fn fill_value(
         &self,
         memory: &mut Memory,
         inputs: &Inputs,
         index: usize,
-        op: &Operation,
-        inv_op: &Operation,
     ) -> Result<bool, RunnerError> {
         let mut addresses = Vec::new();
         let mut values = Vec::new();
@@ -458,19 +461,19 @@ impl ModBuiltinRunner {
         match (a, b, c) {
             // Deduce c from a and b and write it to memory.
             (Some(a), Some(b), None) => {
-                let value = apply_op(a, b, op)?.mod_floor(&inputs.p);
+                let value = self.apply_operation(a, b, &inputs.p)?;
                 self.write_n_words_value(memory, addresses[2], value)?;
                 Ok(true)
             }
             // Deduce b from a and c and write it to memory.
             (Some(a), None, Some(c)) => {
-                let value = apply_op(c, a, inv_op)?.mod_floor(&inputs.p);
+                let value = self.deduce_operand(a, c, &inputs.p)?;
                 self.write_n_words_value(memory, addresses[1], value)?;
                 Ok(true)
             }
             // Deduce a from b and c and write it to memory.
             (None, Some(b), Some(c)) => {
-                let value = apply_op(c, b, inv_op)?.mod_floor(&inputs.p);
+                let value = self.deduce_operand(b, c, &inputs.p)?;
                 self.write_n_words_value(memory, addresses[0], value)?;
                 Ok(true)
             }
@@ -539,44 +542,33 @@ impl ModBuiltinRunner {
                 Default::default()
             };
 
-        //  Get one of the builtin runners - the rest of this function doesn't depend on batch_size.
-        let mod_runner = if let Some((_, add_mod, _)) = add_mod {
-            add_mod
-        } else {
-            mul_mod.unwrap().1
-        };
         // Fill the values table.
         let mut add_mod_index = 0;
         let mut mul_mod_index = 0;
-        // Create operation here to avoid cloning p in the loop
-        let div_operation = Operation::DivMod(mul_mod_inputs.p.clone());
+
         while add_mod_index < add_mod_n || mul_mod_index < mul_mod_n {
-            if add_mod_index < add_mod_n
-                && mod_runner.fill_value(
-                    memory,
-                    &add_mod_inputs,
-                    add_mod_index,
-                    &Operation::Add,
-                    &Operation::Sub,
-                )?
-            {
-                add_mod_index += 1;
-            } else if mul_mod_index < mul_mod_n
-                && mod_runner.fill_value(
-                    memory,
-                    &mul_mod_inputs,
-                    mul_mod_index,
-                    &Operation::Mul,
-                    &div_operation,
-                )?
-            {
-                mul_mod_index += 1;
-            } else {
-                return Err(RunnerError::FillMemoryCoudNotFillTable(
-                    add_mod_index,
-                    mul_mod_index,
-                ));
+            if add_mod_index < add_mod_n {
+                if let Some((_, add_mod_runner, _)) = add_mod {
+                    if add_mod_runner.fill_value(memory, &add_mod_inputs, add_mod_index)? {
+                        add_mod_index += 1;
+                        continue;
+                    }
+                }
             }
+
+            if mul_mod_index < mul_mod_n {
+                if let Some((_, mul_mod_runner, _)) = mul_mod {
+                    if mul_mod_runner.fill_value(memory, &mul_mod_inputs, mul_mod_index)? {
+                        mul_mod_index += 1;
+                    }
+                    continue;
+                }
+            }
+
+            return Err(RunnerError::FillMemoryCoudNotFillTable(
+                add_mod_index,
+                mul_mod_index,
+            ));
         }
         Ok(())
     }
@@ -625,14 +617,11 @@ impl ModBuiltinRunner {
                     inputs.offsets_ptr,
                     index_in_batch,
                 )?;
-                let op = match self.builtin_type {
-                    ModBuiltinType::Add => Operation::Add,
-                    ModBuiltinType::Mul => Operation::Mul,
-                };
-                let a_op_b = apply_op(&a, &b, &op)?.mod_floor(&inputs.p);
-                if a_op_b != c.mod_floor(&inputs.p) {
+                let a_op_b = self.apply_operation(&a, &b, &inputs.p)?;
+                if a_op_b.mod_floor(&inputs.p) != c.mod_floor(&inputs.p) {
                     // Build error string
                     let p = inputs.p;
+                    let op = self.builtin_type.operation_string();
                     let error_string = format!("Expected a {op} b == c (mod p). Got: instance={instance}, batch={index_in_batch}, p={p}, a={a}, b={b}, c={c}.");
                     return Err(RunnerError::ModBuiltinSecurityCheck(Box::new((
                         self.name(),
@@ -667,23 +656,145 @@ impl ModBuiltinRunner {
         self.shift_powers = core::array::from_fn(|i| self.shift.pow(i as u32));
         self.zero_segment_size = core::cmp::max(N_WORDS, batch_size * 3);
     }
-}
 
-fn apply_op(lhs: &BigUint, rhs: &BigUint, op: &Operation) -> Result<BigUint, MathError> {
-    Ok(match op {
-        Operation::Mul => lhs * rhs,
-        Operation::Add => lhs + rhs,
-        Operation::Sub => lhs - rhs,
-        Operation::DivMod(ref p) => div_mod_unsigned(lhs, rhs, p)?,
-    })
+    // Calculates the result of `lhs OP rhs`
+    //
+    // The builtin type (add or mul) determines the OP
+    pub(crate) fn apply_operation(
+        &self,
+        lhs: &BigUint,
+        rhs: &BigUint,
+        prime: &BigUint,
+    ) -> Result<BigUint, MathError> {
+        let full_value = match self.builtin_type {
+            ModBuiltinType::Mul => lhs * rhs,
+            ModBuiltinType::Add => lhs + rhs,
+        };
+
+        let value = if full_value < &self.k_bound * prime {
+            full_value.mod_floor(prime)
+        } else {
+            full_value - (&self.k_bound - 1u32) * prime
+        };
+
+        Ok(value)
+    }
+
+    // Given `known OP unknown = result (mod p)`, it deduces `unknown`
+    //
+    // The builtin type (add or mul) determines the OP
+    pub(crate) fn deduce_operand(
+        &self,
+        known: &BigUint,
+        result: &BigUint,
+        prime: &BigUint,
+    ) -> Result<BigUint, MathError> {
+        let value = match self.builtin_type {
+            ModBuiltinType::Add => {
+                if known <= result {
+                    result - known
+                } else {
+                    result + prime - known
+                }
+            }
+            ModBuiltinType::Mul => div_mod_unsigned(result, known, prime)?,
+        };
+        Ok(value)
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use super::*;
+
+    #[test]
+    fn apply_operation_add() {
+        let builtin = ModBuiltinRunner::new_add_mod(&ModInstanceDef::new(Some(8), 8, 8), true);
+
+        assert_eq!(
+            builtin
+                .apply_operation(
+                    &BigUint::from(2u32),
+                    &BigUint::from(3u32),
+                    &BigUint::from(7u32)
+                )
+                .unwrap(),
+            BigUint::from(5u32)
+        );
+
+        assert_eq!(
+            builtin
+                .apply_operation(
+                    &BigUint::from(5u32),
+                    &BigUint::from(5u32),
+                    &BigUint::from(5u32)
+                )
+                .unwrap(),
+            BigUint::from(5u32)
+        );
+    }
+
+    #[test]
+    fn apply_operation_mul() {
+        let builtin = ModBuiltinRunner::new_mul_mod(&ModInstanceDef::new(Some(8), 8, 8), true);
+
+        assert_eq!(
+            builtin
+                .apply_operation(
+                    &BigUint::from(2u32),
+                    &BigUint::from(3u32),
+                    &BigUint::from(7u32)
+                )
+                .unwrap(),
+            BigUint::from(6u32)
+        );
+    }
+
+    #[test]
+    fn deduce_operand_add() {
+        let builtin = ModBuiltinRunner::new_add_mod(&ModInstanceDef::new(Some(8), 8, 8), true);
+
+        assert_eq!(
+            builtin
+                .deduce_operand(
+                    &BigUint::from(2u32),
+                    &BigUint::from(5u32),
+                    &BigUint::from(7u32)
+                )
+                .unwrap(),
+            BigUint::from(3u32)
+        );
+        assert_eq!(
+            builtin
+                .deduce_operand(
+                    &BigUint::from(5u32),
+                    &BigUint::from(2u32),
+                    &BigUint::from(7u32)
+                )
+                .unwrap(),
+            BigUint::from(4u32)
+        );
+    }
+
+    #[test]
+    fn deduce_operand_mul() {
+        let builtin = ModBuiltinRunner::new_mul_mod(&ModInstanceDef::new(Some(8), 8, 8), true);
+
+        assert_eq!(
+            builtin
+                .deduce_operand(
+                    &BigUint::from(2u32),
+                    &BigUint::from(1u32),
+                    &BigUint::from(7u32)
+                )
+                .unwrap(),
+            BigUint::from(4u32)
+        );
+    }
+
     #[test]
     #[cfg(feature = "mod_builtin")]
     fn test_air_private_input_all_cairo() {
-        use super::*;
         use crate::{
             air_private_input::{ModInput, ModInputInstance, ModInputMemoryVars, PrivateInput},
             hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor,
@@ -735,8 +846,8 @@ mod tests {
                                 a2: Felt252::ZERO,
                                 a3: Felt252::ZERO,
                                 b_offset: 12,
-                                b0: Felt252::ZERO,
-                                b1: Felt252::ZERO,
+                                b0: Felt252::ONE,
+                                b1: Felt252::ONE,
                                 b2: Felt252::ZERO,
                                 b3: Felt252::ZERO,
                                 c_offset: 4,
@@ -798,8 +909,8 @@ mod tests {
                             0,
                             ModInputMemoryVars {
                                 a_offset: 12,
-                                a0: Felt252::ZERO,
-                                a1: Felt252::ZERO,
+                                a0: Felt252::ONE,
+                                a1: Felt252::ONE,
                                 a2: Felt252::ZERO,
                                 a3: Felt252::ZERO,
                                 b_offset: 8,


Test patch:
diff --git a/vm/src/tests/compare_outputs_dynamic_layouts.sh b/vm/src/tests/compare_outputs_dynamic_layouts.sh
index 7077edaa44..2bc512e3a1 100755
--- a/vm/src/tests/compare_outputs_dynamic_layouts.sh
+++ b/vm/src/tests/compare_outputs_dynamic_layouts.sh
@@ -82,6 +82,9 @@ CASES=(
     "cairo_programs/proof_programs/sha256.json;double_all_cairo"
     "cairo_programs/proof_programs/keccak.json;all_cairo"
     "cairo_programs/proof_programs/keccak.json;double_all_cairo"
+    "cairo_programs/mod_builtin_feature/proof/mod_builtin.json;all_cairo"
+    "cairo_programs/mod_builtin_feature/proof/mod_builtin_failure.json;all_cairo"
+    "cairo_programs/mod_builtin_feature/proof/apply_poly.json;all_cairo"
 )
 
 passed_tests=0