u512 arithmetics

src/math/src/lib.cairo

@@ -15,6 +15,7 @@ pub mod karatsuba;
 pub mod keccak256;
 pub mod lcm_of_n_numbers;
 pub mod mod_arithmetics;
+pub mod u512_arithmetics;
 pub mod perfect_number;
 pub mod sha256;
 pub mod sha512;

src/math/src/tests.cairo

@@ -20,5 +20,6 @@ mod sha256_test;
 mod sha512_test;
 mod test_keccak256;
 mod trigonometry_test;
+mod u512_arithmetics_test;
 mod wad_ray_math_test;
 mod zellers_congruence_test;

src/math/src/tests/u512_arithmetics_test.cairo

@@ -0,0 +1,50 @@
+use core::integer::u512;
+use alexandria_math::u512_arithmetics::{u512_add, u512_sub};
+
+const MAX_128: u128 = 0xffffffffffffffffffffffffffffffff;
+const A: u256 = 9099547013904003590785796930435194473319680151794113978918064868415326638035;
+const B: u256 = 8021715850804026033197027745655159931503181100513576347155970296011118125764;
+
+#[inline(always)]
+fn mu512(limb0: u128, limb1: u128, limb2: u128, limb3: u128) -> u512 {
+    u512 { limb0, limb1, limb2, limb3 }
+}
+
+#[test]
+fn test_u512_add() {
+    assert(
+        u512_add(mu512(1, 2, 3, 4), mu512(5, 6, 7, 8)) == mu512(6, 8, 10, 12), 'incorrect u512 add'
+    );
+    assert(
+        u512_add(mu512(MAX_128, 1, 2, 3), mu512(4, 5, 6, 7)) == mu512(3, 7, 8, 10),
+        'incorrect u512 add'
+    );
+}
+
+#[test]
+fn test_u512_sub() {
+    let sub0 = u512_sub(mu512(5, 6, 7, 8), mu512(1, 2, 3, 4));
+    assert(sub0 == mu512(4, 4, 4, 4), 'incorrect u512 sub');
+
+    let sub1 = u512_sub(mu512(3, 2, 1, MAX_128,), mu512(7, 6, 5, 4));
+    assert(
+        sub1 == mu512(
+            0xfffffffffffffffffffffffffffffffc,
+            0xfffffffffffffffffffffffffffffffb,
+            0xfffffffffffffffffffffffffffffffb,
+            0xfffffffffffffffffffffffffffffffa
+        ),
+        'incorrect u512 sub1'
+    );
+
+    let sub2 = u512_sub(mu512(3, 2, 1, 1), mu512(7, 6, 5, 0));
+    assert(
+        sub2 == mu512(
+            0xfffffffffffffffffffffffffffffffc,
+            0xfffffffffffffffffffffffffffffffb,
+            0xfffffffffffffffffffffffffffffffb,
+            0
+        ),
+        'incorrect u512 sub2'
+    );
+}

src/math/src/u512_arithmetics.cairo

@@ -0,0 +1,87 @@
+use core::result::ResultTrait;
+use core::integer::u512;
+use core::integer::{u128_overflowing_add, u128_overflowing_sub};
+
+pub fn u256_overflow_add(lhs: u256, rhs: u256) -> Result<u256, u256> implicits(RangeCheck) nopanic {
+    let (sum, overflow) = core::integer::u256_overflowing_add(lhs, rhs);
+    if overflow {
+        Result::Err(sum)
+    } else {
+        Result::Ok(sum)
+    }
+}
+
+pub fn u256_overflow_sub(lhs: u256, rhs: u256) -> Result<u256, u256> implicits(RangeCheck) nopanic {
+    let (sum, overflow) = core::integer::u256_overflow_sub(lhs, rhs);
+    if overflow {
+        Result::Err(sum)
+    } else {
+        Result::Ok(sum)
+    }
+}
+
+#[derive(Copy, Drop, Hash, PartialEq, Serde)]
+pub struct u256X2 {
+    low: u256,
+    high: u256,
+}
+
+pub impl U512Intou256X2 of Into<u512, u256X2> {
+    #[inline(always)]
+    fn into(self: u512) -> u256X2 {
+        let u512 { limb0: low, limb1: high, limb2, limb3 } = self;
+        u256X2 { low: u256 { low, high }, high: u256 { low: limb2, high: limb3 } }
+    }
+}
+
+#[inline(always)]
+pub fn u512_add(lhs: u512, rhs: u512) -> u512 {
+    let lhs: u256X2 = lhs.into();
+    let rhs: u256X2 = rhs.into();
+
+    // No overflow allowed
+    let u256 { low: limb2, high: limb3 } = u256_overflow_add(lhs.high, rhs.high)
+        .expect('u512 add overflow');
+
+    match u256_overflow_add(lhs.low, rhs.low) {
+        Result::Ok(u256 { low: limb0, high: limb1 }) => { u512 { limb0, limb1, limb2, limb3 } },
+        Result::Err(u256 { low: limb0,
+        high: limb1 }) => {
+            // Try to move overflow to limb2
+            return match u128_overflowing_add(limb2, 1_u128) {
+                Result::Ok(limb2) => u512 { limb0, limb1, limb2, limb3 },
+                Result::Err(limb2) => {
+                    // Try to move overflow to limb3
+                    let limb3 = u128_overflowing_add(limb3, 1_u128).expect('u512 add overflow');
+                    u512 { limb0, limb1, limb2, limb3 }
+                },
+            };
+        },
+    }
+}
+
+#[inline(always)]
+pub fn u512_sub(lhs: u512, rhs: u512) -> u512 {
+    let lhs: u256X2 = lhs.into();
+    let rhs: u256X2 = rhs.into();
+
+    // No overflow allowed
+    let u256 { low: limb2, high: limb3 } = u256_overflow_sub(lhs.high, rhs.high)
+        .expect('u512 sub overflow');
+
+    match u256_overflow_sub(lhs.low, rhs.low) {
+        Result::Ok(u256 { low: limb0, high: limb1 }) => { u512 { limb0, limb1, limb2, limb3 } },
+        Result::Err(u256 { low: limb0,
+        high: limb1 }) => {
+            // Try to move overflow to limb2
+            return match u128_overflowing_sub(limb2, 1_u128) {
+                Result::Ok(limb2) => u512 { limb0, limb1, limb2, limb3 },
+                Result::Err(limb2) => {
+                    // Try to move overflow to limb3
+                    let limb3 = u128_overflowing_sub(limb3, 1_u128).expect('u512 sub overflow');
+                    u512 { limb0, limb1, limb2, limb3 }
+                },
+            };
+        },
+    }
+}