impl Power for U256

sway-lib-std/src/u128.sw

@@ -321,7 +321,7 @@ impl core::ops::Subtract for U128 {
     }
 }
 impl core::ops::Multiply for U128 {
-    /// Multiply a `U128` with a `U128`. Panics of overflow.
+    /// Multiply a `U128` with a `U128`. Panics if overflow.
     fn multiply(self, other: Self) -> Self {
         // in case both of the `U128` upper parts are bigger than zero,
         // it automatically means overflow, as any `U128` value

sway-lib-std/src/u256.sw

@@ -1,10 +1,11 @@
 //! A 256-bit unsigned integer type.
 library;
 
-use ::assert::assert;
+use ::assert::{assert_eq, assert};
 use ::convert::From;
 use ::result::Result::{self, *};
 use ::u128::U128;
+use ::math::Power;
 
 /// Left shift a `u64` and preserve the overflow amount if any.
 fn lsh_with_carry(word: u64, shift_amount: u64) -> (u64, u64) {
@@ -81,6 +82,41 @@ impl U256 {
         }
     }
 
+    /// Initializes a new `U256` with a value of 1.
+    ///
+    /// ### Examples
+    ///
+    /// ```sway
+    /// use std::u256::U256;
+    ///
+    /// let init_one = U256::one();
+    /// let one_u256 = U256 { a: 0, b: 0, c: 0, d: 1 };
+    ///
+    /// assert(init_one == one_u256);
+    /// ```
+    pub fn one() -> Self {
+        Self {
+            a: 0,
+            b: 0,
+            c: 0,
+            d: 1,
+        }
+    }
+
+    /// Returns true if value is zero.
+    ///
+    /// ### Examples
+    ///
+    /// ```sway
+    /// use std::u256::U256
+    ///
+    /// let zero_u256 = U256::new();
+    /// assert(zero_u256.is_zero());
+    /// ```
+    pub fn is_zero(self) -> bool {
+        self.a == 0 && self.b == 0 && self.c == 0 && self.d == 0
+    }
+
     /// Safely downcast to `u64` without loss of precision.
     /// Returns `Err` if the `number > u64::max()`.
     ///
@@ -570,3 +606,63 @@ impl core::ops::Divide for U256 {
         quotient
     }
 }
+
+impl Power for U256 {
+    /// Fast exponentiation by squaring
+	/// https://en.wikipedia.org/wiki/Exponentiation_by_squaring
+	///
+	/// # Panics
+	///
+	/// Panics if the result overflows the type.
+    fn pow(self, exponent: Self) -> Self {
+        let mut value = self;
+        let mut exp = exponent;
+        let one = Self::one();
+
+        if exp.is_zero() {
+            return one;
+        }
+
+        if exp == one {
+            // Manually clone `self`. Otherwise, we may have a `MemoryOverflow`
+            // issue with code that looks like: `x = x.pow(other)`
+            return Self::from((self.a, self.b, self.c, self.d));
+        }
+
+        while (exp & one).is_zero() {
+            value = value * value;
+            exp >>= 1;
+        }
+
+        if exp == one {
+            return value;
+        }
+
+        let mut acc = value;
+        while exp > one {
+            exp >>= 1;
+            value = value * value;
+            if exp & one == one {
+                acc = acc * value;
+            }
+        }
+        acc
+    }
+}
+
+#[test]
+fn test_pow_u256() {
+    let five = U256::from((0, 0, 0, 5));
+    let two = U256::from((0, 0, 0, 2));
+    let three = U256::from((0, 0, 0, 3));
+    let twenty_eight = U256::from((0, 0, 0, 28));
+    assert_eq(five.pow(two), U256::from((0, 0, 0, 25)));
+    assert_eq(five.pow(three), U256::from((0, 0, 0, 125)));
+    assert_eq(five.pow(twenty_eight), U256::from((0, 359414837200037395, 18446744073709551615, 18446744073709551615)));
+}
+
+#[test]
+fn test_is_zero() {
+    let zero_u256 = U256::new();
+    assert(zero_u256.is_zero());
+}