From e6002032eb5a8e7e503c399cc5ffbe4cc45040bd Mon Sep 17 00:00:00 2001
From: Alvaro Revuelta <alvrevuelta@gmail.com>
Date: Fri, 9 Feb 2024 17:06:25 +0100
Subject: [PATCH] Benchmark RLN proof generation/verification (#2410)

---
 Makefile                       |  4 +++
 apps/benchmarks/benchmarks.nim | 47 ++++++++++++++++++++++++++++++++++
 waku.nimble                    |  4 +++
 3 files changed, 55 insertions(+)
 create mode 100644 apps/benchmarks/benchmarks.nim

diff --git a/Makefile b/Makefile
index f18aba2fb..c578efc28 100644
--- a/Makefile
+++ b/Makefile
@@ -174,6 +174,10 @@ wakunode2: | build deps librln
 	echo -e $(BUILD_MSG) "build/$@" && \
 		$(ENV_SCRIPT) nim wakunode2 $(NIM_PARAMS) waku.nims
 
+benchmarks: | build deps librln
+	echo -e $(BUILD_MSG) "build/$@" && \
+		$(ENV_SCRIPT) nim benchmarks $(NIM_PARAMS) waku.nims
+
 testwakunode2: | build deps librln
 	echo -e $(BUILD_MSG) "build/$@" && \
 		$(ENV_SCRIPT) nim testwakunode2 $(NIM_PARAMS) waku.nims
diff --git a/apps/benchmarks/benchmarks.nim b/apps/benchmarks/benchmarks.nim
new file mode 100644
index 000000000..24c011970
--- /dev/null
+++ b/apps/benchmarks/benchmarks.nim
@@ -0,0 +1,47 @@
+import
+  math,
+  std/sequtils,
+  stew/results,
+  options,
+  ../../waku/waku_rln_relay/protocol_types,
+  ../../waku/waku_rln_relay/rln,
+  ../../waku/waku_rln_relay,
+  ../../waku/waku_rln_relay/conversion_utils,
+  ../../waku/waku_rln_relay/group_manager/static/group_manager
+
+import std/[times, os]
+
+proc main(): Future[string] {.async, gcsafe.} =
+  let rlnIns = createRLNInstance(20).get()
+  let credentials = toSeq(0 .. 1000).mapIt(membershipKeyGen(rlnIns).get())
+
+  let manager = StaticGroupManager(
+    rlnInstance: rlnIns,
+    groupSize: 1000,
+    membershipIndex: some(MembershipIndex(900)),
+    groupKeys: credentials,
+  )
+
+  await manager.init()
+
+  let data: seq[byte] = newSeq[byte](1024)
+
+  var proofGenTimes: seq[times.Duration] = @[]
+  var proofVerTimes: seq[times.Duration] = @[]
+  for i in 0 .. 50:
+    var time = getTime()
+    let proof = manager.generateProof(data, getCurrentEpoch()).get()
+    proofGenTimes.add(getTime() - time)
+
+    time = getTime()
+    let res = manager.verifyProof(data, proof).get()
+    proofVerTimes.add(getTime() - time)
+
+  echo "Proof generation times: ", sum(proofGenTimes) div len(proofGenTimes)
+  echo "Proof verification times: ", sum(proofVerTimes) div len(proofVerTimes)
+
+when isMainModule:
+  try:
+    waitFor(main())
+  except CatchableError as e:
+    raise e
diff --git a/waku.nimble b/waku.nimble
index f9337fceb..f1169423f 100644
--- a/waku.nimble
+++ b/waku.nimble
@@ -62,6 +62,10 @@ task wakunode2, "Build Waku v2 cli node":
   let name = "wakunode2"
   buildBinary name, "apps/wakunode2/"
 
+task benchmarks, "Some benchmarks":
+  let name = "benchmarks"
+  buildBinary name, "apps/benchmarks/"
+
 task wakucanary, "Build waku-canary tool":
   let name = "wakucanary"
   buildBinary name, "apps/wakucanary/"