From eb1ce58d70a2d5e5cf4567d403bbd2ec57e97e2d Mon Sep 17 00:00:00 2001
From: ctrlc03 <93448202+ctrlc03@users.noreply.github.com>
Date: Fri, 30 Aug 2024 17:24:54 +0100
Subject: [PATCH] chore: update message tree depth programmatically in
 benchmarks (#1799)

---
 .../contracts/tasks/helpers/Deployment.ts     | 16 ++++++++
 .../contracts/tasks/helpers/benchmarks.ts     | 40 +++++++------------
 packages/contracts/tasks/runner/benchmarks.ts |  7 ++++
 3 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/packages/contracts/tasks/helpers/Deployment.ts b/packages/contracts/tasks/helpers/Deployment.ts
index f0619b31a..e5b8607d6 100644
--- a/packages/contracts/tasks/helpers/Deployment.ts
+++ b/packages/contracts/tasks/helpers/Deployment.ts
@@ -396,6 +396,22 @@ export class Deployment {
     return value;
   }
 
+  /**
+   * Update deploy config field (see deploy-config.json)
+   * @param id - contract name
+   * @param field - config field key
+   * @param value - config field value
+   */
+  updateDeployConfig<T = string | number | boolean, ID extends string = EContracts>(
+    id: ID,
+    field: string,
+    value: T,
+  ): void {
+    this.checkHre();
+
+    this.config.set(`${this.hre!.network.name}.${id}.${field}`, value).write();
+  }
+
   /**
    * Get contract by name and group key
    *
diff --git a/packages/contracts/tasks/helpers/benchmarks.ts b/packages/contracts/tasks/helpers/benchmarks.ts
index 8fac6d256..f778fff04 100644
--- a/packages/contracts/tasks/helpers/benchmarks.ts
+++ b/packages/contracts/tasks/helpers/benchmarks.ts
@@ -35,40 +35,30 @@ export async function publishBatch(
   const messageBatch = Array.from({ length: batchSize }, () => message.asContractParam());
   const pubKeyBatch = Array.from({ length: batchSize }, () => keypair.pubKey.asContractParam());
 
-  let low = 1;
-  let high = batchSize;
-  let optimalBatchSize = 0;
+  let optimalBatchSize = batchSize;
 
-  while (low <= high) {
-    const mid = Math.floor((low + high) / 2);
-    const testMessageBatch = messageBatch.slice(0, mid);
-    const testPubKeyBatch = pubKeyBatch.slice(0, mid);
+  while (optimalBatchSize > 0) {
+    const finalMessageBatch = messageBatch.slice(0, optimalBatchSize);
+    const finalPubKeyBatch = pubKeyBatch.slice(0, optimalBatchSize);
 
     try {
       // eslint-disable-next-line no-await-in-loop
-      await pollContract.publishMessageBatch.estimateGas(testMessageBatch, testPubKeyBatch);
-      optimalBatchSize = mid;
-      low = mid + 1;
+      const tx = await pollContract.publishMessageBatch(finalMessageBatch, finalPubKeyBatch);
+      // eslint-disable-next-line no-await-in-loop
+      const receipt = await tx.wait();
+      console.log(`Successfully published batch of ${optimalBatchSize} messages`);
+      console.log(`Gas used: ${receipt?.gasUsed.toString()} wei`);
+      console.log(`Tx: ${tx.hash}`);
+      break; // If successful, we've found the largest working batch size
     } catch (error) {
-      high = mid - 1;
+      // If this size doesn't work, reduce by 1 and try again
+      optimalBatchSize -= 1;
     }
   }
 
   if (optimalBatchSize > 0) {
-    const finalMessageBatch = messageBatch.slice(0, optimalBatchSize);
-    const finalPubKeyBatch = pubKeyBatch.slice(0, optimalBatchSize);
-
-    try {
-      const tx = await pollContract.publishMessageBatch(finalMessageBatch, finalPubKeyBatch);
-      const receipt = await tx.wait();
-      console.log(`Gas used: ${receipt?.gasUsed.toString()} wei\n`);
-      console.log(`Tx: ${tx.hash}\n`);
-
-      console.log(`Submitted a batch of ${optimalBatchSize} messages\n`);
-    } catch (err) {
-      console.error(`Failed to submit a batch of ${optimalBatchSize} messages\n`);
-    }
+    console.log(`Found optimal batch size: ${optimalBatchSize}`);
   } else {
-    console.error("Unable to submit even a single message\n");
+    console.error("Unable to publish even a single message");
   }
 }
diff --git a/packages/contracts/tasks/runner/benchmarks.ts b/packages/contracts/tasks/runner/benchmarks.ts
index fb71b4d2f..db98b2f57 100644
--- a/packages/contracts/tasks/runner/benchmarks.ts
+++ b/packages/contracts/tasks/runner/benchmarks.ts
@@ -3,6 +3,7 @@ import { task } from "hardhat/config";
 import { Keypair, PCommand } from "maci-domainobjs";
 
 import { Deployment } from "../helpers/Deployment";
+import { EContracts } from "../helpers/types";
 
 task("benchmark", "Run benchmarks").setAction(async (_, hre) => {
   const deployment = Deployment.getInstance(hre);
@@ -15,8 +16,14 @@ task("benchmark", "Run benchmarks").setAction(async (_, hre) => {
   await deployment.runSteps(steps, 0);
 
   // deploy a Poll
+  // get original tree depth
+  const messageTreeDepth = deployment.getDeployConfigField<number>(EContracts.VkRegistry, "messageTreeDepth");
+  // update it
+  deployment.updateDeployConfig(EContracts.VkRegistry, "messageTreeDepth", 3);
   const pollSteps = await deployment.start("poll", { incremental: true, verify: false });
   await deployment.runSteps(pollSteps, 0);
+  // restore it
+  deployment.updateDeployConfig(EContracts.VkRegistry, "messageTreeDepth", messageTreeDepth);
 
   try {
     const startBalance = await deployer.provider.getBalance(deployer);