feat: add benchmark

.gitignore

@@ -2,4 +2,6 @@
 
 zig-out/
 
-test/bun/node_modules/
+test/bun/node_modules/
+
+test/bun/benchmark_data/

src/compute_indices.zig → src/committee_indices.zig

@@ -5,19 +5,16 @@ const native_endian = builtin.target.cpu.arch.endian();
 const Allocator = std.mem.Allocator;
 
 pub const SEED_SIZE = 32;
-const U32U32HashMap = std.AutoHashMap(u32, u32);
-const U8SliceByU32 = std.AutoHashMap(u32, []const u8);
-// note that AutoHashMap always copy data in put() api
-// so value should be a pointer instead of U8SliceByU32 so that it can be freed
-const U8SliceByU8ByU32 = std.AutoHashMap(u32, *U8SliceByU32);
+// model this as []?[32]u8 does not show better performance
+const U8ArrayByU32 = std.AutoHashMap(u32, [32]u8);
 
 /// a Zig implementation of https://github.com/ChainSafe/swap-or-not-shuffle/pull/5
 pub const ComputeShuffledIndex = struct {
     // this ComputeShuffledIndex is always init() and deinit() inside consumer's function so use arena allocator here
     // to improve performance and implify deinit()
     arena: std.heap.ArenaAllocator,
-    pivot_by_index: U32U32HashMap,
-    source_by_position_by_index: U8SliceByU8ByU32,
+    pivot_by_index: []?u32,
+    source_by_position_by_index: []?*U8ArrayByU32,
     // 32 bytes seed + 1 byte i
     pivot_buffer: [33]u8,
     // 32 bytes seed + 1 byte i + 4 bytes positionDiv
@@ -38,10 +35,12 @@ pub const ComputeShuffledIndex = struct {
             return error.InvalidRounds;
         }
 
-        const arena = std.heap.ArenaAllocator.init(parent_allocator);
+        var arena = std.heap.ArenaAllocator.init(parent_allocator);
 
-        const pivot_by_index = U32U32HashMap.init(parent_allocator);
-        const source_by_position_by_index = U8SliceByU8ByU32.init(parent_allocator);
+        const pivot_by_index = try arena.allocator().alloc(?u32, @intCast(rounds));
+        @memset(pivot_by_index, null);
+        const source_by_position_by_index = try arena.allocator().alloc(?*U8ArrayByU32, @intCast(rounds));
+        @memset(source_by_position_by_index, null);
 
         var pivot_buffer: [33]u8 = [_]u8{0} ** 33;
         var source_buffer: [37]u8 = [_]u8{0} ** 37;
@@ -60,18 +59,16 @@ pub const ComputeShuffledIndex = struct {
     }
 
     pub fn deinit(self: *ComputeShuffledIndex) void {
-        self.pivot_by_index.deinit();
-
-        var it = self.source_by_position_by_index.iterator();
-        while (it.next()) |entry| {
-            var source_by_position = entry.value_ptr.*;
-            // no need to loop through values and free the sources inside source_by_position thanks to arena
-            source_by_position.deinit();
-            // we create() source_by_position in the below get() api
-            // but no need to destroy() it thanks to arena
+        // pivot_by_index is deinit() by arena allocator
+
+        for (0..self.rounds) |i| {
+            const source_by_position = self.source_by_position_by_index[@intCast(i)];
+            if (source_by_position) |item| {
+                item.deinit();
+            }
         }
 
-        self.source_by_position_by_index.deinit();
+        // source_by_position_by_index is deinit() by arena allocator
 
         // this needs to be the last step
         self.arena.deinit();
@@ -82,26 +79,29 @@ pub const ComputeShuffledIndex = struct {
         const allocator = self.arena.allocator();
 
         for (0..self.rounds) |i| {
-            var pivot = self.pivot_by_index.get(@intCast(i));
+            var pivot = self.pivot_by_index[@intCast(i)];
             if (pivot == null) {
                 self.pivot_buffer[SEED_SIZE] = @intCast(i % 256);
                 var digest = [_]u8{0} ** 32;
                 Sha256.hash(self.pivot_buffer[0..], digest[0..], .{});
                 const u64Slice = std.mem.bytesAsSlice(u64, digest[0..8]);
                 const u64_value = u64Slice[0];
                 const le_value = if (native_endian == .big) @byteSwap(u64_value) else u64_value;
-                pivot = @intCast(le_value % self.index_count);
+                const _pivot: u32 = @intCast(le_value % self.index_count);
+                self.pivot_by_index[@intCast(i)] = _pivot;
+                pivot = _pivot;
             }
 
             const flip = (pivot.? + self.index_count - permuted) % self.index_count;
             const position = @max(permuted, flip);
             const position_div: u32 = position / 256;
 
-            var source_by_position = self.source_by_position_by_index.get(@intCast(i));
+            var source_by_position = self.source_by_position_by_index[@intCast(i)];
             if (source_by_position == null) {
-                const _source_by_position = try allocator.create(U8SliceByU32);
-                _source_by_position.* = U8SliceByU32.init(allocator);
-                try self.source_by_position_by_index.put(@intCast(i), _source_by_position);
+                const _source_by_position = try allocator.create(U8ArrayByU32);
+                _source_by_position.* = U8ArrayByU32.init(allocator);
+                try _source_by_position.*.ensureTotalCapacity(256);
+                self.source_by_position_by_index[@intCast(i)] = _source_by_position;
                 source_by_position = _source_by_position;
             }
 
@@ -111,10 +111,8 @@ pub const ComputeShuffledIndex = struct {
                 const u32Slice = std.mem.bytesAsSlice(u32, self.source_buffer[SEED_SIZE + 1 ..]);
                 u32Slice[0] = if (native_endian == .big) @byteSwap(position_div) else position_div;
 
-                const _source = try allocator.alloc(u8, 32);
-                var hash = [_]u8{0} ** 32;
-                Sha256.hash(self.source_buffer[0..], &hash, .{});
-                @memcpy(_source, hash[0..]);
+                var _source: [32]u8 = undefined;
+                Sha256.hash(self.source_buffer[0..], &_source, .{});
                 try source_by_position.?.put(position_div, _source);
                 source = _source;
             }
@@ -180,8 +178,9 @@ fn getCommitteeIndices(allocator: Allocator, seed: []const u8, active_indices: [
 
     var compute_shuffled_index = try ComputeShuffledIndex.init(allocator, seed, @intCast(active_indices.len), rounds);
     defer compute_shuffled_index.deinit();
-    var shuffled_result = U32U32HashMap.init(allocator);
-    defer shuffled_result.deinit();
+    var shuffled_result = try allocator.alloc(?u32, @intCast(active_indices.len));
+    defer allocator.free(shuffled_result);
+    @memset(shuffled_result, null);
 
     var i: u32 = 0;
     var cached_hash_input = [_]u8{0} ** (32 + 8);
@@ -192,10 +191,10 @@ fn getCommitteeIndices(allocator: Allocator, seed: []const u8, active_indices: [
 
     while (next_committee_index < out.len) {
         const index: u32 = @intCast(i % active_indices.len);
-        var shuffled_index = shuffled_result.get(index);
+        var shuffled_index = shuffled_result[index];
         if (shuffled_index == null) {
             const _shuffled_index = try compute_shuffled_index.get(index);
-            try shuffled_result.put(index, _shuffled_index);
+            shuffled_result[index] = _shuffled_index;
             shuffled_index = _shuffled_index;
         }
         const candidate_index = active_indices[@intCast(shuffled_index.?)];

src/root.zig

@@ -1,8 +1,8 @@
 const std = @import("std");
 const testing = std.testing;
 pub const PubkeyIndexMap = @import("./pubkey_index_map.zig").PubkeyIndexMap;
-pub const ComputeIndices = @import("./compute_indices.zig");
-pub const ComputeShuffledIndex = ComputeIndices.ComputeShuffledIndex;
+pub const CommitteeIndices = @import("./committee_indices.zig");
+pub const ComputeShuffledIndex = CommitteeIndices.ComputeShuffledIndex;
 
 export fn add(a: i32, b: i32) i32 {
     return a + b;

src/root_c_abi.zig

@@ -4,7 +4,7 @@ pub const PubkeyIndexMap = @import("pubkey_index_map.zig").PubkeyIndexMap;
 const PUBKEY_INDEX_MAP_KEY_SIZE = @import("pubkey_index_map.zig").PUBKEY_INDEX_MAP_KEY_SIZE;
 const innerShuffleList = @import("shuffle.zig").innerShuffleList;
 const SEED_SIZE = @import("shuffle.zig").SEED_SIZE;
-const compute_indices = @import("compute_indices.zig");
+const committee_indices = @import("committee_indices.zig");
 const ErrorCode = @import("error.zig").ErrorCode;
 const NOT_FOUND_INDEX = @import("error.zig").NOT_FOUND_INDEX;
 const ERROR_INDEX = @import("error.zig").ERROR_INDEX;
@@ -264,26 +264,26 @@ export fn doShuffleList(active_indices: [*c]u32, len: usize, seed: [*c]u8, seed_
 export fn computeProposerIndexElectra(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, max_effective_balance_electra: u64, effective_balance_increment: u32, rounds: u32) u32 {
     const allocator = gpa.allocator();
     // TODO: is it better to define a Result struct with code and value
-    const proposer_index = compute_indices.computeProposerIndexElectra(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], max_effective_balance_electra, effective_balance_increment, rounds) catch return ERROR_INDEX;
+    const proposer_index = committee_indices.computeProposerIndexElectra(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], max_effective_balance_electra, effective_balance_increment, rounds) catch return ERROR_INDEX;
     return proposer_index;
 }
 
-export fn computeProposerIndex(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, rand_byte_count: compute_indices.ByteCount, max_effective_balance: u64, effective_balance_increment: u32, rounds: u32) u32 {
+export fn computeProposerIndex(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, rand_byte_count: committee_indices.ByteCount, max_effective_balance: u64, effective_balance_increment: u32, rounds: u32) u32 {
     const allocator = gpa.allocator();
     // TODO: is it better to define a Result struct with code and value
-    const proposer_index = compute_indices.computeProposerIndex(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], rand_byte_count, max_effective_balance, effective_balance_increment, rounds) catch return ERROR_INDEX;
+    const proposer_index = committee_indices.computeProposerIndex(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], rand_byte_count, max_effective_balance, effective_balance_increment, rounds) catch return ERROR_INDEX;
     return proposer_index;
 }
 
 export fn computeSyncCommitteeIndicesElectra(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, max_effective_balance_electra: u64, effective_balance_increment: u32, rounds: u32, out: [*c]u32, out_len: usize) c_uint {
     const allocator = gpa.allocator();
-    compute_indices.computeSyncCommitteeIndicesElectra(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], max_effective_balance_electra, effective_balance_increment, rounds, out[0..out_len]) catch return ErrorCode.Error;
+    committee_indices.computeSyncCommitteeIndicesElectra(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], max_effective_balance_electra, effective_balance_increment, rounds, out[0..out_len]) catch return ErrorCode.Error;
     return ErrorCode.Success;
 }
 
-export fn computeSyncCommitteeIndices(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, rand_byte_count: compute_indices.ByteCount, max_effective_balance: u64, effective_balance_increment: u32, rounds: u32, out: [*c]u32, out_len: usize) c_uint {
+export fn computeSyncCommitteeIndices(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, rand_byte_count: committee_indices.ByteCount, max_effective_balance: u64, effective_balance_increment: u32, rounds: u32, out: [*c]u32, out_len: usize) c_uint {
     const allocator = gpa.allocator();
-    compute_indices.computeSyncCommitteeIndices(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], rand_byte_count, max_effective_balance, effective_balance_increment, rounds, out[0..out_len]) catch return ErrorCode.Error;
+    committee_indices.computeSyncCommitteeIndices(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], rand_byte_count, max_effective_balance, effective_balance_increment, rounds, out[0..out_len]) catch return ErrorCode.Error;
     return ErrorCode.Success;
 }
 

test/bun/package.json

@@ -20,7 +20,7 @@
 	"scripts": {
 		"test:unit": "bun test test/unit",
 		"lint": "biome check",
-		"benchmark": "bun benchmark:files 'test/perf/index.test.ts'",
+		"benchmark": "bun benchmark:files 'test/perf/*.test.ts'",
 		"benchmark:files": "bun ./node_modules/.bin/benchmark --config .benchrc.yaml --defaultBranch main",
 		"lint:fix": "yarn lint --write"
 	}

test/bun/test/perf/committeeIndices.test.ts

@@ -0,0 +1,57 @@
+import { randomBytes } from "node:crypto";
+import { bench, describe } from "@chainsafe/benchmark";
+import {
+	EFFECTIVE_BALANCE_INCREMENT,
+	MAX_EFFECTIVE_BALANCE_ELECTRA,
+	SHUFFLE_ROUND_COUNT,
+	SYNC_COMMITTEE_SIZE,
+} from "@lodestar/params";
+import { computeSyncCommitteeIndicesElectra } from "../../src/index.js";
+import { naiveComputeSyncCommitteeIndicesElectra } from "../referenceImplementation.js";
+
+describe("computeIndices", () => {
+	for (const listSize of [
+		16384, 250_000,
+		1_000_000,
+		// Don't run 4_000_000 since it's very slow and not testnet has gotten there yet
+		// 4e6,
+	]) {
+		// don't want to generate random seed to investigate performance
+		// each seed may lead to different cached items hence different performance
+		const seed = new Uint8Array(32).fill(1);
+		const vc = listSize;
+		const activeIndices = new Uint32Array(
+			Array.from({ length: vc }, (_, i) => i),
+		);
+		const effectiveBalanceIncrements = new Uint16Array(vc);
+		for (let i = 0; i < vc; i++) {
+			effectiveBalanceIncrements[i] = 32 + 32 * (i % 64);
+		}
+
+		bench({
+			id: `JS   - computeSyncCommitteeIndices - ${listSize} indices`,
+			fn: () => {
+				naiveComputeSyncCommitteeIndicesElectra(
+					seed,
+					activeIndices,
+					effectiveBalanceIncrements,
+				);
+			},
+		});
+
+		bench({
+			id: `Zig   - computeSyncCommitteeIndices - ${listSize} indices`,
+			fn: () => {
+				computeSyncCommitteeIndicesElectra(
+					seed,
+					activeIndices,
+					effectiveBalanceIncrements,
+					SYNC_COMMITTEE_SIZE,
+					MAX_EFFECTIVE_BALANCE_ELECTRA,
+					EFFECTIVE_BALANCE_INCREMENT,
+					SHUFFLE_ROUND_COUNT,
+				);
+			},
+		});
+	}
+});

test/bun/test/perf/shuffle.test.ts

@@ -0,0 +1,50 @@
+import { bench, describe } from "@chainsafe/benchmark";
+import { unshuffleList } from "../../src/index.js";
+import * as referenceImplementation from "../referenceImplementation.js";
+
+const SHUFFLE_ROUNDS_MINIMAL = 10;
+// TODO: below is Rust implementation, double check with Zig
+//          Lighthouse  Lodestar
+// 512      254.04 us  1.6034 ms (x6)
+// 16384    6.2046 ms  18.272 ms (x3)
+// 4000000  1.5617 s   4.9690 s  (x3)
+
+for (const listSize of [
+	16384, 250_000,
+	1_000_000,
+	// Don't run 4_000_000 since it's very slow and not testnet has gotten there yet
+	// 4e6,
+]) {
+	describe(`shuffle list - ${listSize} indices`, () => {
+		const seed = Buffer.alloc(32, 0xac);
+		const input: number[] = [];
+		for (let i = 0; i < listSize; i++) input[i] = i;
+		const indices = new Uint32Array(input);
+
+		bench<Uint32Array, Uint32Array>({
+			id: `JS   - unshuffleList - ${listSize} indices`,
+			fn: () => {
+				referenceImplementation.unshuffleList(
+					indices,
+					seed,
+					SHUFFLE_ROUNDS_MINIMAL,
+				);
+			},
+		});
+
+		bench<Uint32Array, Uint32Array>({
+			id: `Zig - unshuffleList - ${listSize} indices`,
+			fn: () => {
+				unshuffleList(indices, seed, SHUFFLE_ROUNDS_MINIMAL);
+			},
+		});
+
+		// skip this as Bun use polling solution
+		// bench<Uint32Array, Uint32Array>({
+		//   id: `Zig - asyncUnshuffleList - ${listSize} indices`,
+		//   fn: async () => {
+		//     await asyncUnshuffleList(indices, seed, SHUFFLE_ROUNDS_MINIMAL);
+		//   },
+		// });
+	});
+}