From 41a30a9b419e63ff6a03e6074d9852879c3459d0 Mon Sep 17 00:00:00 2001
From: Ori Ziv <oriz@starkware.co>
Date: Sun, 7 May 2023 19:38:19 +0300
Subject: [PATCH] Added widemul for u256.

commit-id:7c544b0c
---
 corelib/src/integer.cairo           | 37 +++++++++++++++++++++++++++++
 corelib/src/test/integer_test.cairo | 19 +++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/corelib/src/integer.cairo b/corelib/src/integer.cairo
index f608bd027ea..9fc2506cc67 100644
--- a/corelib/src/integer.cairo
+++ b/corelib/src/integer.cairo
@@ -1080,6 +1080,43 @@ impl U256BitNot of BitNot<u256> {
     }
 }
 
+#[derive(Copy, Drop, PartialEq, Serde)]
+struct u512 {
+    limb0: u128,
+    limb1: u128,
+    limb2: u128,
+    limb3: u128,
+}
+
+// Returns the result of u128 addition, including an overflow word.
+fn u128_add_with_carry(a: u128, b: u128) -> (u128, u128) nopanic {
+    match u128_overflowing_add(a, b) {
+        Result::Ok(v) => (v, 0),
+        Result::Err(v) => (v, 1),
+    }
+}
+
+fn u256_wide_mul(a: u256, b: u256) -> u512 nopanic {
+    let (limb1, limb0) = u128_wide_mul(a.low, b.low);
+    let (limb2, limb1_part) = u128_wide_mul(a.low, b.high);
+    let (limb1, limb1_overflow0) = u128_add_with_carry(limb1, limb1_part);
+    let (limb2_part, limb1_part) = u128_wide_mul(a.high, b.low);
+    let (limb1, limb1_overflow1) = u128_add_with_carry(limb1, limb1_part);
+    let (limb2, limb2_overflow) = u128_add_with_carry(limb2, limb2_part);
+    let (limb3, limb2_part) = u128_wide_mul(a.high, b.high);
+    // No overflow since no limb4.
+    let limb3 = u128_wrapping_add(limb3, limb2_overflow);
+    let (limb2, limb2_overflow) = u128_add_with_carry(limb2, limb2_part);
+    // No overflow since no limb4.
+    let limb3 = u128_wrapping_add(limb3, limb2_overflow);
+    // No overflow possible in this addition since both operands are 0/1.
+    let limb1_overflow = u128_wrapping_add(limb1_overflow0, limb1_overflow1);
+    let (limb2, limb2_overflow) = u128_add_with_carry(limb2, limb1_overflow);
+    // No overflow since no limb4.
+    let limb3 = u128_wrapping_add(limb3, limb2_overflow);
+    u512 { limb0, limb1, limb2, limb3 }
+}
+
 /// Bounded
 trait BoundedInt<T> {
     fn min() -> T nopanic;
diff --git a/corelib/src/test/integer_test.cairo b/corelib/src/test/integer_test.cairo
index f675b9986f8..8ffeddd9861 100644
--- a/corelib/src/test/integer_test.cairo
+++ b/corelib/src/test/integer_test.cairo
@@ -685,6 +685,25 @@ fn test_u256_mul_overflow_2() {
     pow_2_127() * 0x200000000000000000000000000000000;
 }
 
+use integer::{u512, u256_wide_mul};
+
+#[test]
+fn test_u256_wide_mul() {
+    assert(u256_wide_mul(0, 0) == u512 { limb0: 0, limb1: 0, limb2: 0, limb3: 0 }, '0 * 0 != 0');
+    assert(
+        u256_wide_mul(
+            0x1001001001001001001001001001001001001001001001001001,
+            0x1000100010001000100010001000100010001000100010001000100010001
+        ) == u512 {
+            limb0: 0x33233223222222122112111111011001,
+            limb1: 0x54455445544554454444443443343333,
+            limb2: 0x21222222322332333333433443444444,
+            limb3: 0x1001101111112112
+        },
+        'long calculation failed'
+    );
+}
+
 #[test]
 fn test_min() {
     let min_u8: u8 = BoundedInt::min();