Implement a performant bucketing MSM (#68)

* add baseline test * implement bucketing MSM for grumpkin * forge fmt
argumentcomputer · Mar 17, 2024 · b5ebce6 · b5ebce6
1 parent fcb357d
commit b5ebce6
Show file tree

Hide file tree

Showing 4 changed files with 371 additions and 0 deletions.
diff --git a/src/Utilities.sol b/src/Utilities.sol
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 pragma solidity ^0.8.16;
 
+import "@std/Test.sol";
 import "src/blocks/EqPolynomial.sol";
 
 /**

diff --git a/src/blocks/grumpkin/Grumpkin.sol b/src/blocks/grumpkin/Grumpkin.sol
@@ -44,6 +44,10 @@ library Grumpkin {
         uint256 self_z = 1;
         uint256 rhs_z = 1;
 
+        if (is_identity(p1) && is_identity(p2)) {
+            return Identity();
+        }
+
         if (is_identity(p1)) {
             self_x = 0;
             self_y = 1;
@@ -306,4 +310,59 @@ library Grumpkin {
 
         return GrumpkinAffinePoint(x, y);
     }
+
+    function getAt(uint256 segment, uint256 c, uint256 scalar) public returns (uint256) {
+        uint256 skipBits = segment * c;
+        if (skipBits >= 256) {
+            return 0;
+        }
+
+        uint256 res = (scalar >> skipBits) % (1 << c);
+        return res;
+    }
+
+    function multiScalarMulSerial(GrumpkinAffinePoint[] memory bases, uint256[] memory scalars)
+        public
+        returns (GrumpkinAffinePoint memory r)
+    {
+        require(scalars.length == bases.length, "MSM error: length does not match");
+
+        uint256 c;
+        if (bases.length < 4) {
+            c = 1;
+        } else if (bases.length < 32) {
+            c = 3;
+        } else {
+            c = CommonUtilities.log2(bases.length);
+        }
+
+        GrumpkinAffinePoint[] memory buckets = new GrumpkinAffinePoint[]((1 << c) - 1);
+        GrumpkinAffinePoint memory res = Identity();
+
+        uint256 segments = (256 / c) + 1;
+        for (uint256 segment = segments; segment > 0; segment--) {
+            for (uint256 i = 0; i < c; i++) {
+                res = double(res);
+            }
+
+            for (uint256 i = 0; i < buckets.length; i++) {
+                buckets[i] = Identity();
+            }
+
+            for (uint256 i = 0; i < bases.length; i++) {
+                uint256 limb = getAt(segment - 1, c, scalars[i]);
+                if (limb != 0) {
+                    buckets[limb - 1] = add(buckets[limb - 1], bases[i]);
+                }
+            }
+
+            GrumpkinAffinePoint memory runningSum = Identity();
+            for (uint256 i = buckets.length; i > 0; i--) {
+                runningSum = add(buckets[i - 1], runningSum);
+                res = add(res, runningSum);
+            }
+        }
+
+        r = res;
+    }
 }
diff --git a/test/grumpkin-curves-tests.t.sol b/test/grumpkin-curves-tests.t.sol
@@ -249,6 +249,16 @@ contract GrumpkinCurvesContractTests is Test {
         assertEq(a_add_b.y, 0x001a2f39bd6cddd97502bc15b89c9d3376c3bf227a1ef667ba593fc9ed1361cd);
     }
 
+    function testGrumpkinPointsAddition11() public {
+        Grumpkin.GrumpkinAffinePoint memory a = Grumpkin.Identity();
+        Grumpkin.GrumpkinAffinePoint memory b = Grumpkin.Identity();
+
+        Grumpkin.GrumpkinAffinePoint memory a_add_b = Grumpkin.add(a, b);
+
+        assertEq(a_add_b.x, 0);
+        assertEq(a_add_b.y, 0);
+    }
+
     function testGrumpkinScalarMultiplication() public {
         uint256 scalar = 0x29bd9a803cd11224817183fc6bceb32d59926fd9aa37d3cfb1c7845cbf7fae0d;