🔧 fix: try fns

interest-protocol · Dec 23, 2024 · 4b8daa7 · 4b8daa7
1 parent 6c7f3b6
commit 4b8daa7
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 42 deletions.
diff --git a/sources/fixed18.move b/sources/fixed18.move
@@ -6,6 +6,8 @@ use interest_math::uint_macro as macro;
 
 const FIXED_18_BASE: u256 = 1_000_000_000_000_000_000;
 
+const MAX_U256: u256 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF;
+
 // === Structs ===
 
 public struct Fixed18 has copy, drop, store { value: u256 }
@@ -102,23 +104,33 @@ public fun u256_to_fixed18_up(x: u256, decimals: u8): Fixed18 {
 
 // === Try Functions ===
 
+public fun try_add(x: Fixed18, y: Fixed18): (bool, Fixed18) {
+    let (pred, value) = macro::try_add!(x.value, y.value, MAX_U256);
+    (pred, Fixed18 { value })
+}
+
+public fun try_sub(x: Fixed18, y: Fixed18): (bool, Fixed18) {
+    let (pred, value) = macro::try_sub!(x.value, y.value);
+    (pred, Fixed18 { value })
+}
+
 public fun try_mul_down(x: Fixed18, y: Fixed18): (bool, Fixed18) {
-    let (pred, value) = macro::try_mul_div_down!(x.value, y.value, FIXED_18_BASE);
+    let (pred, value) = macro::try_mul_div_down!(x.value, y.value, FIXED_18_BASE, MAX_U256);
     (pred, Fixed18 { value })
 }
 
 public fun try_mul_up(x: Fixed18, y: Fixed18): (bool, Fixed18) {
-    let (pred, value) = macro::try_mul_div_up!(x.value, y.value, FIXED_18_BASE);
+    let (pred, value) = macro::try_mul_div_up!(x.value, y.value, FIXED_18_BASE, MAX_U256);
     (pred, Fixed18 { value })
 }
 
 public fun try_div_down(x: Fixed18, y: Fixed18): (bool, Fixed18) {
-    let (pred, value) = macro::try_mul_div_down!(x.value, FIXED_18_BASE, y.value);
+    let (pred, value) = macro::try_mul_div_down!(x.value, FIXED_18_BASE, y.value, MAX_U256);
     (pred, Fixed18 { value })
 }
 
 public fun try_div_up(x: Fixed18, y: Fixed18): (bool, Fixed18) {
-    let (pred, value) = macro::try_mul_div_up!(x.value, FIXED_18_BASE, y.value);
+    let (pred, value) = macro::try_mul_div_up!(x.value, FIXED_18_BASE, y.value, MAX_U256);
     (pred, Fixed18 { value })
 }
 

diff --git a/sources/uint/u128.move b/sources/uint/u128.move
@@ -17,8 +17,8 @@ public fun try_sub(x: u128, y: u128): (bool, u128) {
 }
 
 public fun try_mul(x: u128, y: u128): (bool, u128) {
-    let (pred, c) = macro::try_mul!(x, y);
-    if (!pred || c > MAX_U128) (false, 0) else (true, (c as u128))
+    let (pred, r) = macro::try_mul!(x, y, MAX_U128);
+    (pred, r as u128)
 }
 
 public fun try_div_down(x: u128, y: u128): (bool, u128) {
@@ -30,13 +30,11 @@ public fun try_div_up(x: u128, y: u128): (bool, u128) {
 }
 
 public fun try_mul_div_down(x: u128, y: u128, z: u128): (bool, u128) {
-    let (pred, r) = macro::try_mul_div_down!(x, y, z);
-    if (!pred || r > MAX_U128) (false, 0) else (true, (r as u128))
+    macro::try_mul_div_down!(x, y, z, MAX_U128)
 }
 
 public fun try_mul_div_up(x: u128, y: u128, z: u128): (bool, u128) {
-    let (pred, r) = macro::try_mul_div_up!(x, y, z);
-    if (!pred || r > MAX_U128) (false, 0) else (true, (r as u128))
+    macro::try_mul_div_up!(x, y, z, MAX_U128)
 }
 
 public fun try_mod(x: u128, y: u128): (bool, u128) {

diff --git a/sources/uint/u256.move b/sources/uint/u256.move
@@ -2,7 +2,7 @@ module interest_math::u256;
 
 use interest_math::uint_macro as macro;
 
-// === Constants === 
+// === Constants ===
 
 const MAX_U256: u256 = 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff;
 
@@ -22,7 +22,7 @@ public fun try_sub(x: u256, y: u256): (bool, u256) {
 }
 
 public fun try_mul(x: u256, y: u256): (bool, u256) {
-    macro::try_mul!(x, y)
+    macro::try_mul!(x, y, MAX_U256)
 }
 
 public fun try_div_down(x: u256, y: u256): (bool, u256) {
@@ -34,11 +34,11 @@ public fun try_div_up(x: u256, y: u256): (bool, u256) {
 }
 
 public fun try_mul_div_down(x: u256, y: u256, z: u256): (bool, u256) {
-    macro::try_mul_div_down!(x, y, z)
+    macro::try_mul_div_down!(x, y, z, MAX_U256)
 }
 
 public fun try_mul_div_up(x: u256, y: u256, z: u256): (bool, u256) {
-    macro::try_mul_div_up!(x, y, z)
+    macro::try_mul_div_up!(x, y, z, MAX_U256)
 }
 
 public fun try_mod(x: u256, y: u256): (bool, u256) {

diff --git a/sources/uint/u64.move b/sources/uint/u64.move
@@ -1,6 +1,6 @@
 module interest_math::u64;
 
-use interest_math::{uint_macro as macro, i256::{Self, I256}};
+use interest_math::{i256::{Self, I256}, uint_macro as macro};
 
 // === Constants ===
 
@@ -12,15 +12,15 @@ const WRAPPING_MAX: u256 = MAX_U64 + 1;
 // === WrappingFunctions ===
 
 public fun wrapping_add(x: u64, y: u64): u64 {
-    (wrap(i256::from_u64(x).add( i256::from_u64(y)), WRAPPING_MAX) as u64)
+    (wrap(i256::from_u64(x).add(i256::from_u64(y)), WRAPPING_MAX) as u64)
 }
 
 public fun wrapping_sub(x: u64, y: u64): u64 {
-    (wrap(i256::from_u64(x).sub( i256::from_u64(y)), WRAPPING_MAX) as u64)
+    (wrap(i256::from_u64(x).sub(i256::from_u64(y)), WRAPPING_MAX) as u64)
 }
 
 public fun wrapping_mul(x: u64, y: u64): u64 {
-    (wrap(i256::from_u64(x).mul( i256::from_u64(y)), WRAPPING_MAX) as u64)
+    (wrap(i256::from_u64(x).mul(i256::from_u64(y)), WRAPPING_MAX) as u64)
 }
 
 // === Try Functions ===
@@ -34,8 +34,8 @@ public fun try_sub(x: u64, y: u64): (bool, u64) {
 }
 
 public fun try_mul(x: u64, y: u64): (bool, u64) {
-    let (pred, r) = macro::try_mul!(x, y);
-    if (!pred || r > MAX_U64) (false, 0) else (true, (r as u64))
+    let (pred, r) = macro::try_mul!(x, y, MAX_U64);
+    (pred, r as u64)
 }
 
 public fun try_div_down(x: u64, y: u64): (bool, u64) {
@@ -47,13 +47,11 @@ public fun try_div_up(x: u64, y: u64): (bool, u64) {
 }
 
 public fun try_mul_div_down(x: u64, y: u64, z: u64): (bool, u64) {
-    let (pred, r) = macro::try_mul_div_down!(x, y, z);
-    if (!pred || r > MAX_U64) (false, 0) else (true, (r as u64))
+    macro::try_mul_div_down!(x, y, z, MAX_U64)
 }
 
 public fun try_mul_div_up(x: u64, y: u64, z: u64): (bool, u64) {
-    let (pred, r) = macro::try_mul_div_up!(x, y, z);
-    if (!pred || r > MAX_U64) (false, 0) else (true, (r as u64))
+    macro::try_mul_div_up!(x, y, z, MAX_U64)
 }
 
 public fun try_mod(x: u64, y: u64): (bool, u64) {
@@ -168,10 +166,10 @@ public fun max_value(): u64 {
     (MAX_U64 as u64)
 }
 
-// === Private Functions === 
+// === Private Functions ===
 
 fun wrap(self: I256, max: u256): u256 {
     let max = i256::from_u256(max);
 
-    i256::to_u256(if (self.is_negative()) self.add( max) else self.sub( max.mul( self.div( max))))
+    i256::to_u256(if (self.is_negative()) self.add(max) else self.sub(max.mul(self.div(max))))
 }
diff --git a/sources/uint/uint.move b/sources/uint/uint.move
@@ -1,22 +1,31 @@
 module interest_math::uint_macro;
 
 public(package) macro fun try_add<$T>($x: _, $y: _, $max: u256): (bool, $T) {
-    let r = ($x as u256) + ($y as u256);
-    if (r > $max) (false, 0) else (true, (r as $T))
+    let x = $x as u256;
+    let y = $y as u256;
+    let max = $max as u256;
+
+    if (x == max && y != 0) return (false, 0 as $T);
+
+    let rem = max - x;
+    if (y > rem) return (false, 0 as $T);
+
+    (true, (x + y) as $T)
 }
 
 public(package) macro fun try_sub($x: _, $y: _): (bool, _) {
     if ($y > $x) (false, 0) else (true, $x - $y)
 }
 
-public(package) macro fun try_mul($x: _, $y: _): (bool, u256) {
+public(package) macro fun try_mul($x: _, $y: _, $max: u256): (bool, u256) {
     let x = $x as u256;
     let y = $y as u256;
+    let max = $max as u256;
 
     if (y == 0) return (true, 0);
-    if (
-        x > 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff / y
-    ) (false, 0) else (true, x * y)
+    if (x > max / y) return (false, 0);
+
+    (true, (x * y))
 }
 
 public(package) macro fun try_div_down($x: _, $y: _): (bool, _) {
@@ -27,28 +36,46 @@ public(package) macro fun try_div_up($x: _, $y: _): (bool, _) {
     if ($y == 0) (false, 0) else (true, div_up!($x, $y))
 }
 
-public(package) macro fun try_mul_div_down($x: _, $y: _, $z: _): (bool, u256) {
+public(package) macro fun try_mul_div_down<$T>($x: _, $y: _, $z: _, $max: u256): (bool, $T) {
     let x = $x as u256;
     let y = $y as u256;
     let z = $z as u256;
+    let max = $max as u256;
+
+    if (z == 0) return (false, 0 as $T);
+    let (pred, _) = try_mul!(
+        x,
+        y,
+        0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff,
+    );
+    if (!pred) return (false, 0 as $T);
+
+    let r = mul_div_down!<u256>(x, y, z);
 
-    if (z == 0) return (false, 0);
-    let (pred, _) = try_mul!(x, y);
-    if (!pred) return (false, 0);
+    if (r > max) return (false, 0 as $T);
 
-    (true, mul_div_down!<u256>(x, y, z))
+    (true, r as $T)
 }
 
-public(package) macro fun try_mul_div_up($x: _, $y: _, $z: _): (bool, u256) {
+public(package) macro fun try_mul_div_up<$T>($x: _, $y: _, $z: _, $max: u256): (bool, $T) {
     let x = $x as u256;
     let y = $y as u256;
     let z = $z as u256;
+    let max = $max as u256;
+
+    if (z == 0) return (false, 0 as $T);
+    let (pred, _) = try_mul!(
+        x,
+        y,
+        0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff,
+    );
+    if (!pred) return (false, 0 as $T);
+
+    let r = mul_div_up!<u256>(x, y, z);
 
-    if (z == 0) return (false, 0);
-    let (pred, _) = try_mul!(x, y);
-    if (!pred) return (false, 0);
+    if (r > max) return (false, 0 as $T);
 
-    (true, mul_div_up!<u256>(x, y, z))
+    (true, r as $T)
 }
 
 public(package) macro fun try_mod($x: _, $y: _): (bool, _) {

diff --git a/tests/fixed18.move b/tests/fixed18.move
@@ -14,6 +14,8 @@ use interest_math::fixed18::{
     base,
     add,
     sub,
+    try_add,
+    try_sub
 };
 use sui::test_utils::assert_eq;
 
@@ -75,6 +77,28 @@ fun test_convert_functions() {
     );
 }
 
+#[test]
+fun test_try_add() {
+    let (pred, r) = try_add(3u256.from(), 5u256.from());
+    assert_eq(pred, true);
+    assert_eq(r.raw_value(), 8 * FIXED_18_BASE);
+
+    let (pred, r) = try_add(MAX_U256.from_raw(), MAX_U256.from_raw());
+    assert_eq(pred, false);
+    assert_eq(r.raw_value(), 0);
+}
+
+#[test]
+fun test_try_sub() {
+    let (pred, r) = try_sub(5u256.from(), 3u256.from());
+    assert_eq(pred, true);
+    assert_eq(r.raw_value(), 2 * FIXED_18_BASE);
+
+    let (pred, r) = try_sub(3u256.from(), 5u256.from());
+    assert_eq(pred, false);
+    assert_eq(r.raw_value(), 0);
+}
+
 #[test]
 fun test_try_mul_down() {
     let (pred, r) = try_mul_down(3u256.from(), 5u256.from());