Correctly check for overflow in add, mul and pow

sway-lib-core/src/ops.sw

@@ -59,9 +59,15 @@ impl Add for u32 {
         // any non-64-bit value is compiled to a u64 value under-the-hood
         // constants (like Self::max() below) are also automatically promoted to u64
         let res = __add(self, other);
+        // integer overflow
         if __gt(res, Self::max()) {
-            // integer overflow
-            __revert(0)
+            if panic_on_overflow_is_enabled() {
+                __revert(0)
+            } else {
+                // overflow enabled
+                // res % (Self::max() + 1)
+                __mod(res, __add(Self::max(), 1))
+            }
         } else {
             // no overflow
             res
@@ -73,7 +79,13 @@ impl Add for u16 {
     fn add(self, other: Self) -> Self {
         let res = __add(self, other);
         if __gt(res, Self::max()) {
-            __revert(0)
+            if panic_on_overflow_is_enabled() {
+                __revert(0)
+            } else {
+                // overflow enabled
+                // res % (Self::max() + 1)
+                __mod(res, __add(Self::max(), 1))
+            }
         } else {
             res
         }
@@ -93,7 +105,16 @@ impl Add for u8 {
             input: u64
         };
         if __gt(res_u64, max_u8_u64) {
-            __revert(0)
+            if panic_on_overflow_is_enabled() {
+                __revert(0)
+            } else {
+                // overflow enabled
+                // res % (Self::max() + 1)
+                let res_u64 = __mod(res_u64, __add(max_u8_u64, 1));
+                asm(input: res_u64) {
+                    input: u8
+                }
+            }
         } else {
             asm(input: res_u64) {
                 input: u8
@@ -229,8 +250,14 @@ impl Multiply for u32 {
         // constants (like Self::max() below) are also automatically promoted to u64
         let res = __mul(self, other);
         if __gt(res, Self::max()) {
-            // integer overflow
-            __revert(0)
+            if panic_on_overflow_is_enabled() {
+                // integer overflow
+                __revert(0)
+            } else {
+                // overflow enabled
+                // res % (Self::max() + 1)
+                __mod(res, __add(Self::max(), 1))
+            }
         } else {
             // no overflow
             res
@@ -242,7 +269,11 @@ impl Multiply for u16 {
     fn multiply(self, other: Self) -> Self {
         let res = __mul(self, other);
         if __gt(res, Self::max()) {
-            __revert(0)
+            if panic_on_overflow_is_enabled() {
+                __revert(0)
+            } else {
+                __mod(res, __add(Self::max(), 1))
+            }
         } else {
             res
         }
@@ -262,7 +293,16 @@ impl Multiply for u8 {
             input: u64
         };
         if __gt(res_u64, max_u8_u64) {
-            __revert(0)
+            if panic_on_overflow_is_enabled() {
+                __revert(0)
+            } else {
+                // overflow enabled
+                // res % (Self::max() + 1)
+                let res_u64 = __mod(res_u64, __add(max_u8_u64, 1));
+                asm(input: res_u64) {
+                    input: u8
+                }
+            }
         } else {
             asm(input: res_u64) {
                 input: u8
@@ -1258,3 +1298,19 @@ pub fn ok_str_eq() {
     assert("" != "a");
     assert("a" != "b");
 }
+
+fn flags() -> u64 {
+    asm() {
+        flag
+    }
+}
+
+fn panic_on_overflow_is_enabled() -> bool {
+    __eq(
+        __and(
+            flags(),
+            0b00000000_00000000_00000000_00000000_00000000_00000000_00000000_00000010,
+        ),
+        0,
+    )
+}

sway-lib-std/src/math.sw

@@ -2,6 +2,8 @@
 library;
 
 use ::assert::*;
+use ::revert::revert;
+use ::option::Option::{self, None, Some};
 use ::flags::{
     disable_panic_on_overflow,
     panic_on_overflow_enabled,
@@ -78,12 +80,38 @@ pub trait Power {
     fn pow(self, exponent: u32) -> Self;
 }
 
+fn u256_checked_mul(a: u256, b: u256) -> Option<u256> {
+    let res = u256::zero();
+
+    // The six-bit immediate value is used to select operating mode, as follows:
+
+    // Bits	Short name	Description
+    // ..XXXX	reserved	Reserved and must be zero
+    // .X....	indirect0	Is lhs operand ($rB) indirect or not
+    // X.....	indirect1	Is rhs operand ($rC) indirect or not
+    // As both operands are indirect, 110000 is used, which is 48 in decimal.
+    let of = asm(res: res, a: a, b: b) {
+        wqml res a b i48;
+        of: u64
+    };
+
+    if of != 0 {
+        return None;
+    }
+
+    Some(res)
+}
+
 impl Power for u256 {
     /// Raises self to the power of `exponent`, using exponentiation by squaring.
     ///
-    /// # Panics
+    /// # Additional Information
+    ///
+    /// * If panic on overflow is disabled, and the result overflows, the return value will be 0.
     ///
-    /// Panics if the result overflows the type.
+    /// # Reverts
+    ///
+    /// * Reverts if the result overflows the type, if panic on overflow is enabled.
     fn pow(self, exponent: u32) -> Self {
         let one = 0x0000000000000000000000000000000000000000000000000000000000000001u256;
 
@@ -97,13 +125,28 @@ impl Power for u256 {
 
         while exp > 1 {
             if (exp & 1) == 1 {
-                acc = acc * base;
+                // acc = acc * base;
+                let res = u256_checked_mul(acc, base);
+                acc = match res {
+                    Some(val) => val,
+                    None => return u256::zero(),
+                }
             }
             exp = exp >> 1;
-            base = base * base;
+            // base = base * base;
+            let res = u256_checked_mul(base, base);
+            base = match res {
+                Some(val) => val,
+                None => return u256::zero(),
+            }
         }
 
-        acc * base
+        // acc * base
+        let res = u256_checked_mul(acc, base);
+        match res {
+            Some(val) => val,
+            None => u256::zero(),
+        }
     }
 }
 
@@ -118,14 +161,21 @@ impl Power for u64 {
 
 impl Power for u32 {
     fn pow(self, exponent: u32) -> Self {
-        let res = asm(r1: self, r2: exponent, r3) {
+        let mut res = asm(r1: self, r2: exponent, r3) {
             exp r3 r1 r2;
             r3: u64
         };
-        // If panic on wrapping math is enabled, only then revert
-        if panic_on_overflow_enabled() {
-            assert(res <= Self::max().as_u64());
+
+        if res > Self::max().as_u64() {
+            // If panic on wrapping math is enabled, only then revert
+            if panic_on_overflow_enabled() {
+                revert(0);
+            } else {
+                // Follow spec of returning 0 for overflow
+                res = 0;
+            }
         }
+
         asm(r1: res) {
             r1: Self
         }
@@ -134,14 +184,21 @@ impl Power for u32 {
 
 impl Power for u16 {
     fn pow(self, exponent: u32) -> Self {
-        let res = asm(r1: self, r2: exponent, r3) {
+        let mut res = asm(r1: self, r2: exponent, r3) {
             exp r3 r1 r2;
             r3: u64
         };
-        // If panic on wrapping math is enabled, only then revert
-        if panic_on_overflow_enabled() {
-            assert(res <= Self::max().as_u64());
+
+        if res > Self::max().as_u64() {
+            // If panic on wrapping math is enabled, only then revert
+            if panic_on_overflow_enabled() {
+                revert(0);
+            } else {
+                // Follow spec of returning 0 for overflow
+                res = 0;
+            }
         }
+
         asm(r1: res) {
             r1: Self
         }
@@ -150,14 +207,21 @@ impl Power for u16 {
 
 impl Power for u8 {
     fn pow(self, exponent: u32) -> Self {
-        let res = asm(r1: self, r2: exponent, r3) {
+        let mut res = asm(r1: self, r2: exponent, r3) {
             exp r3 r1 r2;
             r3: u64
         };
-        // If panic on wrapping math is enabled, only then revert
-        if panic_on_overflow_enabled() {
-            assert(res <= Self::max().as_u64());
+
+        if res > Self::max().as_u64() {
+            // If panic on wrapping math is enabled, only then revert
+            if panic_on_overflow_enabled() {
+                revert(0);
+            } else {
+                // Follow spec of returning 0 for overflow
+                res = 0;
+            }
         }
+
         asm(r1: res) {
             r1: Self
         }