add functions that compute packed reduced qm31 arithmetics to math_utils

ohad-nir-starkware · ohad-nir-starkware · commit ea9f93fa4ca2 · 2025-02-12T10:05:12.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* feat: add functions that compute packed reduced qm31 arithmetics to `math_utils` [#1944](https://github.com/lambdaclass/cairo-vm/pull/1944)
+
 * feat: set `encoded_instruction` to be u128 for opcode_extensions to come [#1940](https://github.com/lambdaclass/cairo-vm/pull/1940)
 
 * feat: add `get_u32_range` to `impl VirtualMachine` add `get_u32` and `get_u32_range` to `impl Memory` [#1936](https://github.com/lambdaclass/cairo-vm/pull/1936)
diff --git a/vm/src/math_utils/mod.rs b/vm/src/math_utils/mod.rs
@@ -24,6 +24,8 @@ lazy_static! {
             .collect::<Vec<_>>();
 }
 
+const STWO_PRIME: u128 = (1 << 31) - 1;
+
 /// Returns the `n`th (up to the `251`th power) power of 2 as a [`Felt252`]
 /// in constant time.
 /// It silently returns `1` if the input is out of bounds.
@@ -66,6 +68,181 @@ pub fn signed_felt(felt: Felt252) -> BigInt {
     }
 }
 
+/// Reads four u128 coordinates from a single Felt252.
+/// Returns an error if the input has over 144 bits or any coordinate is unreduced.
+fn qm31_packed_reduced_read_coordinates(felt: Felt252) -> Result<[u128; 4], MathError> {
+    let limbs = felt.to_le_digits();
+    if limbs[3] != 0 || limbs[2] >= 1 << 16 {
+        return Err(MathError::QM31UpackingError(Box::new(felt)));
+    }
+    let coordinates = [
+        (limbs[0] & ((1 << 36) - 1)) as u128,
+        ((limbs[0] >> 36) + ((limbs[1] & ((1 << 8) - 1)) << 28)) as u128,
+        ((limbs[1] >> 8) & ((1 << 36) - 1)) as u128,
+        ((limbs[1] >> 44) + (limbs[2] << 20)) as u128,
+    ];
+    for x in coordinates.iter() {
+        if *x >= STWO_PRIME {
+            return Err(MathError::QM31UnreducedError(Box::new(felt)));
+        }
+    }
+    Ok(coordinates)
+}
+
+/// Reduces four u128 coordinates and packs them into a single Felt252.
+fn qm31_coordinates_to_packed_reduced(coordinates: [u128; 4]) -> Felt252 {
+    let bytes_part1 =
+        ((coordinates[0] % STWO_PRIME) + ((coordinates[1] % STWO_PRIME) << 36)).to_le_bytes();
+    let bytes_part2 =
+        ((coordinates[2] % STWO_PRIME) + ((coordinates[3] % STWO_PRIME) << 36)).to_le_bytes();
+    let mut result_bytes = [0u8; 32];
+    result_bytes[0..9].copy_from_slice(&bytes_part1[0..9]);
+    result_bytes[9..18].copy_from_slice(&bytes_part2[0..9]);
+    Felt252::from_bytes_le(&result_bytes)
+}
+
+/// Computes the addition of two QM31 elements in reduced form.
+/// Returns an error if either operand is not reduced.
+pub fn qm31_packed_reduced_add(felt1: Felt252, felt2: Felt252) -> Result<Felt252, MathError> {
+    let coordinates1 = qm31_packed_reduced_read_coordinates(felt1)?;
+    let coordinates2 = qm31_packed_reduced_read_coordinates(felt2)?;
+    let result_unreduced_coordinates = [
+        coordinates1[0] + coordinates2[0],
+        coordinates1[1] + coordinates2[1],
+        coordinates1[2] + coordinates2[2],
+        coordinates1[3] + coordinates2[3],
+    ];
+    Ok(qm31_coordinates_to_packed_reduced(
+        result_unreduced_coordinates,
+    ))
+}
+
+/// Computes the negative of a QM31 element in reduced form.
+/// Returns an error if the input is not reduced.
+pub fn qm31_packed_reduced_neg(felt: Felt252) -> Result<Felt252, MathError> {
+    let coordinates = qm31_packed_reduced_read_coordinates(felt)?;
+    Ok(qm31_coordinates_to_packed_reduced([
+        STWO_PRIME - coordinates[0],
+        STWO_PRIME - coordinates[1],
+        STWO_PRIME - coordinates[2],
+        STWO_PRIME - coordinates[3],
+    ]))
+}
+
+/// Computes the subtraction of two QM31 elements in reduced form.
+/// Returns an error if either operand is not reduced.
+pub fn qm31_packed_reduced_sub(felt1: Felt252, felt2: Felt252) -> Result<Felt252, MathError> {
+    let coordinates1 = qm31_packed_reduced_read_coordinates(felt1)?;
+    let coordinates2 = qm31_packed_reduced_read_coordinates(felt2)?;
+    let result_unreduced_coordinates = [
+        STWO_PRIME + coordinates1[0] - coordinates2[0],
+        STWO_PRIME + coordinates1[1] - coordinates2[1],
+        STWO_PRIME + coordinates1[2] - coordinates2[2],
+        STWO_PRIME + coordinates1[3] - coordinates2[3],
+    ];
+    Ok(qm31_coordinates_to_packed_reduced(
+        result_unreduced_coordinates,
+    ))
+}
+
+/// Computes the multiplication of two QM31 elements in reduced form.
+/// Returns an error if either operand is not reduced.
+pub fn qm31_packed_reduced_mul(felt1: Felt252, felt2: Felt252) -> Result<Felt252, MathError> {
+    let coordinates1 = qm31_packed_reduced_read_coordinates(felt1)?;
+    let coordinates2 = qm31_packed_reduced_read_coordinates(felt2)?;
+    let result_unreduced_coordinates = [
+        coordinates1[0] * coordinates2[0] + STWO_PRIME * STWO_PRIME
+            - coordinates1[1] * coordinates2[1]
+            + 2 * (coordinates1[2] * coordinates2[2] + STWO_PRIME * STWO_PRIME
+                - coordinates1[3] * coordinates2[3])
+            - coordinates1[2] * coordinates2[3]
+            - coordinates1[3] * coordinates2[2],
+        coordinates1[0] * coordinates2[1]
+            + coordinates2[0] * coordinates1[1]
+            + 2 * (coordinates1[2] * coordinates2[3] + coordinates1[3] * coordinates2[2])
+            + coordinates1[2] * coordinates2[2]
+            - coordinates1[3] * coordinates2[3],
+        coordinates1[0] * coordinates2[2] + STWO_PRIME * STWO_PRIME
+            - coordinates1[1] * coordinates2[3]
+            + coordinates1[2] * coordinates2[0]
+            - coordinates1[3] * coordinates2[1],
+        coordinates1[0] * coordinates2[3]
+            + coordinates1[1] * coordinates2[2]
+            + coordinates1[2] * coordinates2[1]
+            + coordinates1[3] * coordinates2[0],
+    ];
+    Ok(qm31_coordinates_to_packed_reduced(
+        result_unreduced_coordinates,
+    ))
+}
+
+/// Computes `v^(STWO_PRIME-2) modulo STWO_PRIME`.
+pub fn pow2147483645(v: u128) -> u128 {
+    let t0 = (sqn(v, 2) * v) % STWO_PRIME;
+    let t1 = (sqn(t0, 1) * t0) % STWO_PRIME;
+    let t2 = (sqn(t1, 3) * t0) % STWO_PRIME;
+    let t3 = (sqn(t2, 1) * t0) % STWO_PRIME;
+    let t4 = (sqn(t3, 8) * t3) % STWO_PRIME;
+    let t5 = (sqn(t4, 8) * t3) % STWO_PRIME;
+    (sqn(t5, 7) * t2) % STWO_PRIME
+}
+
+/// Computes `v^(2*n) modulo STWO_PRIME`.
+fn sqn(v: u128, n: usize) -> u128 {
+    let mut u = v;
+    for _ in 0..n {
+        u = (u * u) % STWO_PRIME;
+    }
+    u
+}
+
+/// Computes the inverse of a QM31 element in reduced form.
+/// Returns an error if the denominator is zero or either operand is not reduced.
+pub fn qm31_packed_reduced_inv(felt: Felt252) -> Result<Felt252, MathError> {
+    if felt.is_zero() {
+        return Err(MathError::DividedByZero);
+    }
+    let coordinates = qm31_packed_reduced_read_coordinates(felt)?;
+
+    let b2_r = (coordinates[2] * coordinates[2] + STWO_PRIME * STWO_PRIME
+        - coordinates[3] * coordinates[3])
+        % STWO_PRIME;
+    let b2_i = (2 * coordinates[2] * coordinates[3]) % STWO_PRIME;
+
+    let denom_r = (coordinates[0] * coordinates[0] + STWO_PRIME * STWO_PRIME
+        - coordinates[1] * coordinates[1]
+        + 2 * STWO_PRIME
+        - 2 * b2_r
+        + b2_i)
+        % STWO_PRIME;
+    let denom_i =
+        (2 * coordinates[0] * coordinates[1] + 3 * STWO_PRIME - 2 * b2_i - b2_r) % STWO_PRIME;
+
+    let denom_norm_squared = (denom_r * denom_r + denom_i * denom_i) % STWO_PRIME;
+    let denom_norm_inverse_squared = pow2147483645(denom_norm_squared);
+
+    let denom_inverse_r = (denom_r * denom_norm_inverse_squared) % STWO_PRIME;
+    let denom_inverse_i = ((STWO_PRIME - denom_i) * denom_norm_inverse_squared) % STWO_PRIME;
+
+    Ok(qm31_coordinates_to_packed_reduced([
+        coordinates[0] * denom_inverse_r + STWO_PRIME * STWO_PRIME
+            - coordinates[1] * denom_inverse_i,
+        coordinates[0] * denom_inverse_i + coordinates[1] * denom_inverse_r,
+        coordinates[3] * denom_inverse_i + STWO_PRIME * STWO_PRIME
+            - coordinates[2] * denom_inverse_r,
+        2 * STWO_PRIME * STWO_PRIME
+            - coordinates[2] * denom_inverse_i
+            - coordinates[3] * denom_inverse_r,
+    ]))
+}
+
+/// Computes the division of two QM31 elements in reduced form.
+/// Returns an error if the input is zero.
+pub fn qm31_packed_reduced_div(felt1: Felt252, felt2: Felt252) -> Result<Felt252, MathError> {
+    let felt2_inv = qm31_packed_reduced_inv(felt2)?;
+    qm31_packed_reduced_mul(felt1, felt2_inv)
+}
+
 ///Returns the integer square root of the nonnegative integer n.
 ///This is the floor of the exact square root of n.
 ///Unlike math.sqrt(), this function doesn't have rounding error issues.
@@ -946,6 +1123,142 @@ mod tests {
         )
     }
 
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn qm31_packed_reduced_read_coordinates_over_144_bits() {
+        let mut felt_bytes = [0u8; 32];
+        felt_bytes[18] = 1;
+        let felt = Felt252::from_bytes_le(&felt_bytes);
+        assert_matches!(
+            qm31_packed_reduced_read_coordinates(felt),
+            Err(MathError::QM31UpackingError(bx)) if *bx == felt
+        );
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn qm31_packed_reduced_read_coordinates_unreduced() {
+        let mut felt_bytes = [0u8; 32];
+        felt_bytes[0] = 0xff;
+        felt_bytes[1] = 0xff;
+        felt_bytes[2] = 0xff;
+        felt_bytes[3] = (1 << 7) - 1;
+        let felt = Felt252::from_bytes_le(&felt_bytes);
+        assert_matches!(
+            qm31_packed_reduced_read_coordinates(felt),
+            Err(MathError::QM31UnreducedError(bx)) if *bx == felt
+        );
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn qm31_packed_reduced_add_test() {
+        let x_coordinates = [1414213562, 1732050807, 1618033988, 1234567890];
+        let y_coordinates = [1234567890, 1414213562, 1732050807, 1618033988];
+        let x = qm31_coordinates_to_packed_reduced(x_coordinates);
+        let y = qm31_coordinates_to_packed_reduced(y_coordinates);
+        let res = qm31_packed_reduced_add(x, y).unwrap();
+        let res_coordinates = qm31_packed_reduced_read_coordinates(res);
+        assert_eq!(
+            res_coordinates,
+            Ok([
+                (1414213562 + 1234567890) % STWO_PRIME,
+                (1732050807 + 1414213562) % STWO_PRIME,
+                (1618033988 + 1732050807) % STWO_PRIME,
+                (1234567890 + 1618033988) % STWO_PRIME,
+            ])
+        );
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn qm31_packed_reduced_neg_test() {
+        let x_coordinates = [1749652895, 834624081, 1930174752, 2063872165];
+        let x = qm31_coordinates_to_packed_reduced(x_coordinates);
+        let res = qm31_packed_reduced_neg(x).unwrap();
+        let res_coordinates = qm31_packed_reduced_read_coordinates(res);
+        assert_eq!(
+            res_coordinates,
+            Ok([
+                STWO_PRIME - x_coordinates[0],
+                STWO_PRIME - x_coordinates[1],
+                STWO_PRIME - x_coordinates[2],
+                STWO_PRIME - x_coordinates[3]
+            ])
+        );
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn qm31_packed_reduced_sub_test() {
+        let x_coordinates = [
+            (1414213562 + 1234567890) % STWO_PRIME,
+            (1732050807 + 1414213562) % STWO_PRIME,
+            (1618033988 + 1732050807) % STWO_PRIME,
+            (1234567890 + 1618033988) % STWO_PRIME,
+        ];
+        let y_coordinates = [1414213562, 1732050807, 1618033988, 1234567890];
+        let x = qm31_coordinates_to_packed_reduced(x_coordinates);
+        let y = qm31_coordinates_to_packed_reduced(y_coordinates);
+        let res = qm31_packed_reduced_sub(x, y).unwrap();
+        let res_coordinates = qm31_packed_reduced_read_coordinates(res);
+        assert_eq!(
+            res_coordinates,
+            Ok([1234567890, 1414213562, 1732050807, 1618033988])
+        );
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn qm31_packed_reduced_mul_test() {
+        let x_coordinates = [1414213562, 1732050807, 1618033988, 1234567890];
+        let y_coordinates = [1259921049, 1442249570, 1847759065, 2094551481];
+        let x = qm31_coordinates_to_packed_reduced(x_coordinates);
+        let y = qm31_coordinates_to_packed_reduced(y_coordinates);
+        let res = qm31_packed_reduced_mul(x, y).unwrap();
+        let res_coordinates = qm31_packed_reduced_read_coordinates(res);
+        assert_eq!(
+            res_coordinates,
+            Ok([947980980, 1510986506, 623360030, 1260310989])
+        );
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn qm31_packed_reduced_inv_test() {
+        let x_coordinates = [1259921049, 1442249570, 1847759065, 2094551481];
+        let x = qm31_coordinates_to_packed_reduced(x_coordinates);
+        let res = qm31_packed_reduced_inv(x).unwrap();
+        assert_eq!(qm31_packed_reduced_mul(x, res), Ok(Felt252::from(1)));
+
+        let x_coordinates = [1, 2, 3, 4];
+        let x = qm31_coordinates_to_packed_reduced(x_coordinates);
+        let res = qm31_packed_reduced_inv(x).unwrap();
+        assert_eq!(qm31_packed_reduced_mul(x, res), Ok(Felt252::from(1)));
+
+        let x_coordinates = [1749652895, 834624081, 1930174752, 2063872165];
+        let x = qm31_coordinates_to_packed_reduced(x_coordinates);
+        let res = qm31_packed_reduced_inv(x).unwrap();
+        assert_eq!(qm31_packed_reduced_mul(x, res), Ok(Felt252::from(1)));
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn qm31_packed_reduced_div_test() {
+        let x_coordinates = [1259921049, 1442249570, 1847759065, 2094551481];
+        let y_coordinates = [1414213562, 1732050807, 1618033988, 1234567890];
+        let xy_coordinates = [947980980, 1510986506, 623360030, 1260310989];
+        let x = qm31_coordinates_to_packed_reduced(x_coordinates);
+        let y = qm31_coordinates_to_packed_reduced(y_coordinates);
+        let xy = qm31_coordinates_to_packed_reduced(xy_coordinates);
+
+        let res = qm31_packed_reduced_div(xy, y).unwrap();
+        assert_eq!(res, x);
+
+        let res = qm31_packed_reduced_div(xy, x).unwrap();
+        assert_eq!(res, y);
+    }
+
     #[cfg(feature = "std")]
     proptest! {
         #[test]
diff --git a/vm/src/types/errors/math_errors.rs b/vm/src/types/errors/math_errors.rs
@@ -65,6 +65,10 @@ pub enum MathError {
         "Operation failed: divmod({}, {}, {}), igcdex({}, {}) != 1 ", (*.0).0, (*.0).1, (*.0).2, (*.0).1, (*.0).2
     )]
     DivModIgcdexNotZero(Box<(BigInt, BigInt, BigInt)>),
+    #[error("Numbers over 144 bit cannot be packed QM31 elements: {0})")]
+    QM31UpackingError(Box<Felt252>),
+    #[error("Number is not a packing of a QM31 in reduced form: {0})")]
+    QM31UnreducedError(Box<Felt252>),
 }
 
 #[cfg(test)]
