chore: simplify rln zerokit to prefer direct calls over events

packages/rln/src/contract/rln_contract.ts

@@ -1,11 +1,11 @@
 import { Logger } from "@waku/utils";
 import { hexToBytes } from "@waku/utils/bytes";
-import { ethers } from "ethers";
 
 import type { RLNInstance } from "../rln.js";
 import { MerkleRootTracker } from "../root_tracker.js";
 import { zeroPadLE } from "../utils/bytes.js";
 
+import { ContractStateError } from "./errors.js";
 import { RLNBaseContract } from "./rln_base_contract.js";
 import { RLNContractInitOptions } from "./types.js";
 
@@ -14,6 +14,7 @@ const log = new Logger("waku:rln:contract");
 export class RLNContract extends RLNBaseContract {
   private instance: RLNInstance;
   private merkleRootTracker: MerkleRootTracker;
+  private lastSyncedBlock: number = 0;
 
   /**
    * Asynchronous initializer for RLNContract.
@@ -24,121 +25,145 @@ export class RLNContract extends RLNBaseContract {
     options: RLNContractInitOptions
   ): Promise<RLNContract> {
     const rlnContract = new RLNContract(rlnInstance, options);
-
+    await rlnContract.syncState();
     return rlnContract;
   }
 
+  /**
+   * Override base contract method to keep Merkle tree in sync
+   * Registers a new membership with the given commitment and rate limit
+   */
+  public override async registerMembership(
+    idCommitment: string,
+    rateLimit: number = this.getRateLimit()
+  ): Promise<void> {
+    await super.registerMembership(idCommitment, rateLimit);
+    await this.syncState();
+  }
+
+  /**
+   * Override base contract method to keep Merkle tree in sync
+   * Erases an existing membership from the contract
+   */
+  public override async eraseMembership(
+    idCommitment: string,
+    eraseFromMembershipSet: boolean = true
+  ): Promise<void> {
+    await super.eraseMembership(idCommitment, eraseFromMembershipSet);
+    await this.syncState();
+  }
+
+  /**
+   * Gets the current Merkle root
+   * Returns the latest valid root or empty array if no roots exist
+   */
+  public async getMerkleRoot(): Promise<Uint8Array> {
+    await this.syncState();
+    const roots = this.merkleRootTracker.roots();
+    return roots.length > 0 ? roots[0] : new Uint8Array();
+  }
+
   private constructor(
     rlnInstance: RLNInstance,
     options: RLNContractInitOptions
   ) {
     super(options);
-
     this.instance = rlnInstance;
-
     const initialRoot = rlnInstance.zerokit.getMerkleRoot();
     this.merkleRootTracker = new MerkleRootTracker(5, initialRoot);
   }
 
-  public override processEvents(events: ethers.Event[]): void {
+  /**
-    const toRemoveTable = new Map<number, number[]>();
+   * Syncs the local Merkle tree with the current contract state
-    const toInsertTable = new Map<number, ethers.Event[]>();
+   */
+  private async syncState(): Promise<void> {
+    try {
+      const currentBlock = await this.provider.getBlockNumber();
 
-    events.forEach((evt) => {
+      // If we're already synced, just get new members
-      if (!evt.args) {
+      if (this.lastSyncedBlock > 0) {
+        await this.syncNewMembers(this.lastSyncedBlock, currentBlock);
+        this.lastSyncedBlock = currentBlock;
         return;
       }
 
-      if (
+      // First time sync - get all members
-        evt.event === "MembershipErased" ||
+      const nextIndex = await this.contract.nextFreeIndex();
-        evt.event === "MembershipExpired"
+      const members = await this.getMembersInRange(0, nextIndex.toNumber());
-      ) {
-        let index = evt.args.index;
 
-        if (!index) {
+      // Clear existing members by deleting them one by one
-          return;
+      // This effectively resets the tree without needing resetTree()
+      for (let i = 0; i < nextIndex.toNumber(); i++) {
+        try {
+          this.instance.zerokit.deleteMember(i);
+        } catch (error) {
+          // Ignore errors for non-existent members
+          continue;
         }
-
-        if (typeof index === "number" || typeof index === "string") {
-          index = ethers.BigNumber.from(index);
-        } else {
-          log.error("Index is not a number or string", {
-            index,
-            event: evt
-          });
-          return;
-        }
-
-        const toRemoveVal = toRemoveTable.get(evt.blockNumber);
-        if (toRemoveVal != undefined) {
-          toRemoveVal.push(index.toNumber());
-          toRemoveTable.set(evt.blockNumber, toRemoveVal);
-        } else {
-          toRemoveTable.set(evt.blockNumber, [index.toNumber()]);
-        }
-      } else if (evt.event === "MembershipRegistered") {
-        let eventsPerBlock = toInsertTable.get(evt.blockNumber);
-        if (eventsPerBlock == undefined) {
-          eventsPerBlock = [];
-        }
-
-        eventsPerBlock.push(evt);
-        toInsertTable.set(evt.blockNumber, eventsPerBlock);
       }
-    });
 
-    this.removeMembers(this.instance, toRemoveTable);
+      // Insert all members
-    this.insertMembers(this.instance, toInsertTable);
+      for (const member of members) {
+        const idCommitment = zeroPadLE(hexToBytes(member.idCommitment), 32);
+        this.instance.zerokit.insertMember(idCommitment);
+      }
+
+      // Update root tracker
+      const currentRoot = this.instance.zerokit.getMerkleRoot();
+      this.merkleRootTracker.pushRoot(currentBlock, currentRoot);
+      this.lastSyncedBlock = currentBlock;
+
+      log.info(
+        `Synced ${members.length} members to current block ${currentBlock}`
+      );
+    } catch (error) {
+      log.error("Failed to sync state", error);
+      throw new ContractStateError("Failed to sync contract state");
+    }
   }
 
-  private insertMembers(
+  /**
-    rlnInstance: RLNInstance,
+   * Syncs new members added between fromBlock and toBlock
-    toInsert: Map<number, ethers.Event[]>
+   */
-  ): void {
+  private async syncNewMembers(
-    toInsert.forEach((events: ethers.Event[], blockNumber: number) => {
+    fromBlock: number,
-      events.forEach((evt) => {
+    toBlock: number
-        if (!evt.args) return;
+  ): Promise<void> {
+    // Get members that were added
+    const filter = this.contract.filters.MembershipRegistered();
+    const addEvents = await this.contract.queryFilter(
+      filter,
+      fromBlock,
+      toBlock
+    );
 
-        const _idCommitment = evt.args.idCommitment as string;
+    // Get members that were removed
-        let index = evt.args.index;
+    const removeFilter = this.contract.filters.MembershipErased();
+    const removeEvents = await this.contract.queryFilter(
+      removeFilter,
+      fromBlock,
+      toBlock
+    );
 
-        if (!_idCommitment || !index) {
+    // Process removals first (in reverse block order)
-          return;
+    for (const evt of removeEvents.sort(
-        }
+      (a, b) => b.blockNumber - a.blockNumber
+    )) {
+      if (!evt.args) continue;
+      const index = evt.args.index.toNumber();
+      this.instance.zerokit.deleteMember(index);
+      this.merkleRootTracker.backFill(evt.blockNumber);
+    }
 
-        if (typeof index === "number" || typeof index === "string") {
+    // Then process additions
-          index = ethers.BigNumber.from(index);
+    for (const evt of addEvents) {
-        }
+      if (!evt.args) continue;
-
+      const idCommitment = zeroPadLE(hexToBytes(evt.args.idCommitment), 32);
-        const idCommitment = zeroPadLE(hexToBytes(_idCommitment), 32);
+      this.instance.zerokit.insertMember(idCommitment);
-        rlnInstance.zerokit.insertMember(idCommitment);
+      this.merkleRootTracker.pushRoot(
-
+        evt.blockNumber,
-        const numericIndex = index.toNumber();
+        this.instance.zerokit.getMerkleRoot()
-        this._members.set(numericIndex, {
+      );
-          index,
+    }
-          idCommitment: _idCommitment
-        });
-      });
-
-      const currentRoot = rlnInstance.zerokit.getMerkleRoot();
-      this.merkleRootTracker.pushRoot(blockNumber, currentRoot);
-    });
-  }
-
-  private removeMembers(
-    rlnInstance: RLNInstance,
-    toRemove: Map<number, number[]>
-  ): void {
-    const removeDescending = new Map([...toRemove].reverse());
-    removeDescending.forEach((indexes: number[], blockNumber: number) => {
-      indexes.forEach((index) => {
-        if (this._members.has(index)) {
-          this._members.delete(index);
-          rlnInstance.zerokit.deleteMember(index);
-        }
-      });
-
-      this.merkleRootTracker.backFill(blockNumber);
-    });
   }
 }