refactor big mul, add big mul up and tests

invariant-labs · Apr 2, 2024 · baefab3 · baefab3
1 parent f8a0d20
commit baefab3
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 25 deletions.
diff --git a/contracts/math/uints.ral b/contracts/math/uints.ral
@@ -92,49 +92,73 @@ Contract Uints () {
         }
 
     pub fn bigAdd512(a: U512, b: U512) -> U512 {
-        let (r0, overflow) = overflowingAdd256(a.lower, b.lower)
-        let r1 = wrappingAdd256(wrappingAdd256(a.higher, b.higher), overflow)
+        let (r0, overflow) = overflowingAdd(a.lower, b.lower)
+        let r1 = wrappingAdd(wrappingAdd(a.higher, b.higher), overflow)
         return U512 { higher: r1, lower: r0 }
     }
 
-    pub fn bigMul256(a: U256, b: U256) -> U512 {
-        let a0 = low256(a)
-        let a1 = high256(a)
-        let b0 = low256(b)
-        let b1 = high256(b)
+    pub fn bigMul(a: U256, b: U256, denominator: U256) -> U512 {
+        let aLower = low128(a)
+        let aHigher = high128(a)
+        let bLower = low128(b)
+        let bHigher = high128(b)
+
+        let rLowerLower = aLower * bLower
+        let rHigherHigher = aHigher * bHigher
+        let rHigherLower = aHigher * bLower
+        let rLowerHigher = aLower * bHigher
+
+        let mut result = bigAdd512(U512 { higher: rHigherHigher, lower: rLowerLower }, U512 { higher: high128(rHigherLower), lower: low128(rHigherLower) << 128 })
+        result = bigAdd512(result, U512 { higher: high128(rLowerHigher), lower: low128(rLowerHigher) << 128 })
+
+        result = bigDiv(result, denominator, 1)
+
+        return result
+    }
+
+    pub fn bigMulUp(a: U256, b: U256, denominator: U256) -> U512 {
+        let a0 = low128(a)
+        let a1 = high128(a)
+        let b0 = low128(b)
+        let b1 = high128(b)
 
         let r0 = a0 * b0
         let r1 = a1 * b1
         let r2 = a1 * b0
         let r3 = a0 * b1
 
-        let mut u512 = bigAdd512(U512 { higher: r1, lower: r0 }, U512 { higher: high256(r2), lower: low256(r2) << 128 })
-        u512 = bigAdd512(u512, U512 { higher: high256(r3), lower: low256(r3) << 128 })
+        let mut u512 = bigAdd512(U512 { higher: r1, lower: r0 }, U512 { higher: high128(r2), lower: low128(r2) << 128 })
+        u512 = bigAdd512(u512, U512 { higher: high128(r3), lower: low128(r3) << 128 })
+
+        u512 = bigAdd512(u512, U512 { higher: 0, lower: 10 ** denominator - 1 })
+        result = bigDiv(result, denominator, 1)
 
         return u512
     }
 
-    pub fn overflowingAdd256(a: U256, b: U256) -> (U256, U256) {
+    pub fn overflowingAdd(a: U256, b: U256) -> (U256, U256) {
         if (u256Max!() - a >= b) {
             return a + b, 0
         } else {
             return b + (u256Max!() - a) - 1, 1
         }
     }
 
-    pub fn wrappingAdd256(a: U256, b: U256) -> U256 {
-        if (u256Max!() - a >= b) {
+    pub fn wrappingAdd(a: U256, b: U256) -> U256 {
+        let u256Max = u256Max!()
+
+        if (u256Max - a >= b) {
             return a + b
         } else {
-            return b + (u256Max!() - a) - 1
+            return b + (u256Max - a) - 1
         }
     }
 
-    pub fn low256(a: U256) -> U256 {
+    pub fn low128(a: U256) -> U256 {
         return a << 128 >> 128
     }
 
-    pub fn high256(a: U256) -> U256 {
+    pub fn high128(a: U256) -> U256 {
         return a >> 128
     }
 }                           

diff --git a/test/uints.test.ts b/test/uints.test.ts
@@ -178,65 +178,89 @@ describe('uints tests', () => {
     }
   })
 
-  test('big mul 256', async () => {
+  test('big mul', async () => {
     const uints = await deployUints(sender)
     {
       const a = 123n
       const b = 2n
-      const result = (await uints.contractInstance.methods.bigMul256({ args: { a, b } })).returns
+      const result = (await uints.contractInstance.methods.bigMul({ args: { a, b, denominator: 1n } })).returns
       expect(result).toStrictEqual({ higher: 0n, lower: 246n })
       // expected: 246
       // real: 246
     }
     {
       const a = 340282366920938463463374607431768211457n
       const b = 340282366920938463463374607431768211457n
-      const result = (await uints.contractInstance.methods.bigMul256({ args: { a, b } })).returns
+      const result = (await uints.contractInstance.methods.bigMul({ args: { a, b, denominator: 1n } })).returns
       expect(result).toStrictEqual({ higher: 1n, lower: 680564733841876926926749214863536422913n })
       // expected: 115792089237316195423570985008687907853950549399482440966384333222776666062849
       //     real: 115792089237316195423570985008687907853950549399482440966384333222776666062849
     }
     {
       const a = 115792089237316195423570985008687907853269984665640564039457584007913129639935n
       const b = 115792089237316195423570985008687907853269984665640564039457584007913129639935n
-      const result = (await uints.contractInstance.methods.bigMul256({ args: { a, b } })).returns
+      const result = (await uints.contractInstance.methods.bigMul({ args: { a, b, denominator: 1n } })).returns
       expect(result).toStrictEqual({
         higher: 115792089237316195423570985008687907853269984665640564039457584007913129639934n,
         lower: 680564733841876926926749214863536422909n
       })
       // expected: 13407807929942597099574024998205846127479365820592393377723561443721764030073315392623399665776056285720014482370779510884422601683867654778417822746804225
       //     real: 13407807929942597099574024998205846127479365820592393377723561443721764030073315392623399665776056285720014482370780191449156443560794581527632686283227133
     }
+    {
+      const a = 500n
+      const b = 0n
+      const result = (await uints.contractInstance.methods.bigMul({ args: { a, b, denominator: 1n } })).returns
+      expect(result).toStrictEqual({ higher: 0n, lower: 0n })
+    }
+    {
+      const a = 100n
+      const b = 100n
+      const result = (await uints.contractInstance.methods.bigMul({ args: { a, b, denominator: 100n } })).returns
+      expect(result).toStrictEqual({ higher: 0n, lower: 100n })
+    }
+    {
+      const a = 30n
+      const b = 1n
+      const result = (await uints.contractInstance.methods.bigMul({ args: { a, b, denominator: 10n } })).returns
+      expect(result).toStrictEqual({ higher: 0n, lower: 3n })
+    }
+    {
+      const a = 500n
+      const b = 4000n
+      const result = (await uints.contractInstance.methods.bigMul({ args: { a, b, denominator: 1000n } })).returns
+      expect(result).toStrictEqual({ higher: 0n, lower: 2000n })
+    }
   })
 
-  test('overflowing add 256', async () => {
+  test('overflowing add', async () => {
     const uints = await deployUints(sender)
     {
       const a = 10n
       const b = 20n
-      const result = (await uints.contractInstance.methods.overflowingAdd256({ args: { a, b } })).returns
+      const result = (await uints.contractInstance.methods.overflowingAdd({ args: { a, b } })).returns
       expect(result).toStrictEqual([30n, 0n])
     }
     {
       const a = 115792089237316195423570985008687907853269984665640564039457584007913129639935n
       const b = 20n
-      const result = (await uints.contractInstance.methods.overflowingAdd256({ args: { a, b } })).returns
+      const result = (await uints.contractInstance.methods.overflowingAdd({ args: { a, b } })).returns
       expect(result).toStrictEqual([19n, 1n])
     }
   })
 
-  test('wrapping add 256', async () => {
+  test('wrapping add', async () => {
     const uints = await deployUints(sender)
     {
       const a = 10n
       const b = 20n
-      const result = (await uints.contractInstance.methods.wrappingAdd256({ args: { a, b } })).returns
+      const result = (await uints.contractInstance.methods.wrappingAdd({ args: { a, b } })).returns
       expect(result).toStrictEqual(30n)
     }
     {
       const a = 115792089237316195423570985008687907853269984665640564039457584007913129639935n
       const b = 20n
-      const result = (await uints.contractInstance.methods.wrappingAdd256({ args: { a, b } })).returns
+      const result = (await uints.contractInstance.methods.wrappingAdd({ args: { a, b } })).returns
       expect(result).toStrictEqual(19n)
     }
   })