Revert when overflow using pow for u8, u16, u32

sway-lib-std/src/math.sw

@@ -2,7 +2,12 @@
 library;
 
 use ::assert::*;
-use ::flags::{disable_panic_on_overflow, F_UNSAFEMATH_DISABLE_MASK, set_flags};
+use ::flags::{
+    disable_panic_on_overflow,
+    F_UNSAFEMATH_DISABLE_MASK,
+    F_WRAPPING_DISABLE_MASK,
+    set_flags,
+};
 use ::registers::{flags, overflow};
 
 /// Calculates the square root.
@@ -113,27 +118,48 @@ impl Power for u64 {
 
 impl Power for u32 {
     fn pow(self, exponent: u32) -> Self {
-        asm(r1: self, r2: exponent, r3) {
+        let res = asm(r1: self, r2: exponent, r3) {
             exp r3 r1 r2;
-            r3: Self
+            r3: u64
+        };
+        // If panic on wrapping math is enabled, only then revert
+        if flags() & F_WRAPPING_DISABLE_MASK == 0 {
+            assert(res <= Self::max().as_u64());
+        }
+        asm(r1: res) {
+            r1: Self
         }
     }
 }
 
 impl Power for u16 {
     fn pow(self, exponent: u32) -> Self {
-        asm(r1: self, r2: exponent, r3) {
+        let res = asm(r1: self, r2: exponent, r3) {
             exp r3 r1 r2;
-            r3: Self
+            r3: u64
+        };
+        // If panic on wrapping math is enabled, only then revert
+        if flags() & F_WRAPPING_DISABLE_MASK == 0 {
+            assert(res <= Self::max().as_u64());
+        }
+        asm(r1: res) {
+            r1: Self
         }
     }
 }
 
 impl Power for u8 {
     fn pow(self, exponent: u32) -> Self {
-        asm(r1: self, r2: exponent, r3) {
+        let res = asm(r1: self, r2: exponent, r3) {
             exp r3 r1 r2;
-            r3: Self
+            r3: u64
+        };
+        // If panic on wrapping math is enabled, only then revert
+        if flags() & F_WRAPPING_DISABLE_MASK == 0 {
+            assert(res <= Self::max().as_u64());
+        }
+        asm(r1: res) {
+            r1: Self
         }
     }
 }