@waku/rln 0.1.5-053bb95.0 → 0.1.5-4adf870.0

package/src/contract/rln_base_contract.ts

@@ -7,15 +7,8 @@ import { DecryptedCredentials } from "../keystore/types.js";
 import { RLN_ABI } from "./abi.js";
 import { DEFAULT_RATE_LIMIT, RATE_LIMIT_PARAMS } from "./constants.js";
 import {
-  InvalidMembershipError,
-  InvalidRateLimitError,
-  MembershipExistsError,
-  MembershipNotFoundError,
-  RateLimitExceededError,
-  RLNContractError,
-  TransactionError
-} from "./errors.js";
-import {
+  CustomQueryOptions,
+  FetchMembersOptions,
   Member,
   MembershipInfo,
   MembershipRegisteredEvent,
@@ -27,9 +20,28 @@ const log = new Logger("waku:rln:contract:base");
 
 export class RLNBaseContract {
   public contract: ethers.Contract;
+  private deployBlock: undefined | number;
   private rateLimit: number;
 
+  protected _members: Map<number, Member> = new Map();
+  private _membersFilter: ethers.EventFilter;
+  private _membershipErasedFilter: ethers.EventFilter;
+  private _membersExpiredFilter: ethers.EventFilter;
+
+  /**
+   * Constructor for RLNBaseContract.
+   * Allows injecting a mocked contract for testing purposes.
+   */
   public constructor(options: RLNContractInitOptions) {
+    // Initialize members and subscriptions
+    this.fetchMembers()
+      .then(() => {
+        this.subscribeToMembers();
+      })
+      .catch((error) => {
+        log.error("Failed to initialize members", { error });
+      });
+
     const {
       address,
       signer,
@@ -38,8 +50,14 @@ export class RLNBaseContract {
     } = options;
 
     this.validateRateLimit(rateLimit);
+
     this.contract = contract || new ethers.Contract(address, RLN_ABI, signer);
     this.rateLimit = rateLimit;
+
+    // Initialize event filters
+    this._membersFilter = this.contract.filters.MembershipRegistered();
+    this._membershipErasedFilter = this.contract.filters.MembershipErased();
+    this._membersExpiredFilter = this.contract.filters.MembershipExpired();
   }
 
   /**
@@ -120,85 +138,266 @@ export class RLNBaseContract {
     this.rateLimit = newRateLimit;
   }
 
-  /**
-   * Gets all members in the given range
-   * @param startIndex Start index (inclusive)
-   * @param endIndex End index (exclusive)
-   */
-  public async getMembersInRange(
-    startIndex: number,
-    endIndex: number
-  ): Promise<Member[]> {
-    try {
-      // Get all commitments in one call
-      const idCommitments =
-        await this.contract.getRateCommitmentsInRangeBoundsInclusive(
-          startIndex,
-          endIndex - 1 // -1 because getRateCommitmentsInRangeBoundsInclusive is inclusive
-        );
+  public get members(): Member[] {
+    const sortedMembers = Array.from(this._members.values()).sort(
+      (left, right) => left.index.toNumber() - right.index.toNumber()
+    );
+    return sortedMembers;
+  }
 
-      // Get membership info for each commitment in a single batch
-      const membershipPromises = idCommitments.map(
-        (idCommitment: ethers.BigNumber) =>
-          this.contract
-            .memberships(idCommitment)
-            .then((info: { index: number | ethers.BigNumber }) => ({
-              idCommitment: idCommitment.toString(),
-              index: ethers.BigNumber.from(info.index)
-            }))
-            .catch(() => null) // Skip invalid members
-      );
+  public async fetchMembers(options: FetchMembersOptions = {}): Promise<void> {
+    const registeredMemberEvents = await RLNBaseContract.queryFilter(
+      this.contract,
+      {
+        fromBlock: this.deployBlock,
+        ...options,
+        membersFilter: this.membersFilter
+      }
+    );
+    const removedMemberEvents = await RLNBaseContract.queryFilter(
+      this.contract,
+      {
+        fromBlock: this.deployBlock,
+        ...options,
+        membersFilter: this.membershipErasedFilter
+      }
+    );
+    const expiredMemberEvents = await RLNBaseContract.queryFilter(
+      this.contract,
+      {
+        fromBlock: this.deployBlock,
+        ...options,
+        membersFilter: this.membersExpiredFilter
+      }
+    );
+
+    const events = [
+      ...registeredMemberEvents,
+      ...removedMemberEvents,
+      ...expiredMemberEvents
+    ];
+    this.processEvents(events);
+  }
+
+  public static async queryFilter(
+    contract: ethers.Contract,
+    options: CustomQueryOptions
+  ): Promise<ethers.Event[]> {
+    const FETCH_CHUNK = 5;
+    const BLOCK_RANGE = 3000;
+
+    const {
+      fromBlock,
+      membersFilter,
+      fetchRange = BLOCK_RANGE,
+      fetchChunks = FETCH_CHUNK
+    } = options;
+
+    if (fromBlock === undefined) {
+      return contract.queryFilter(membersFilter);
+    }
+
+    if (!contract.provider) {
+      throw Error("No provider found on the contract.");
+    }
 
-      // Wait for all promises to resolve
-      const members = (await Promise.all(membershipPromises)).filter(
-        (m): m is Member => m !== null
+    const toBlock = await contract.provider.getBlockNumber();
+
+    if (toBlock - fromBlock < fetchRange) {
+      return contract.queryFilter(membersFilter, fromBlock, toBlock);
+    }
+
+    const events: ethers.Event[][] = [];
+    const chunks = RLNBaseContract.splitToChunks(
+      fromBlock,
+      toBlock,
+      fetchRange
+    );
+
+    for (const portion of RLNBaseContract.takeN<[number, number]>(
+      chunks,
+      fetchChunks
+    )) {
+      const promises = portion.map(([left, right]) =>
+        RLNBaseContract.ignoreErrors(
+          contract.queryFilter(membersFilter, left, right),
+          []
+        )
       );
-      return members;
-    } catch (error) {
+      const fetchedEvents = await Promise.all(promises);
+      events.push(fetchedEvents.flatMap((v) => v));
+    }
+
+    return events.flatMap((v) => v);
+  }
+
+  public processEvents(events: ethers.Event[]): void {
+    const toRemoveTable = new Map<number, number[]>();
+    const toInsertTable = new Map<number, ethers.Event[]>();
+
+    events.forEach((evt) => {
+      if (!evt.args) {
+        return;
+      }
+
       if (
-        error instanceof Error &&
-        error.message.includes("InvalidPaginationQuery")
+        evt.event === "MembershipErased" ||
+        evt.event === "MembershipExpired"
       ) {
-        throw new RLNContractError(
-          `Invalid pagination range: start=${startIndex}, end=${endIndex}`
-        );
+        let index = evt.args.index;
+
+        if (!index) {
+          return;
+        }
+
+        if (typeof index === "number" || typeof index === "string") {
+          index = ethers.BigNumber.from(index);
+        }
+
+        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);
       }
-      throw error;
+    });
+  }
+
+  public static splitToChunks(
+    from: number,
+    to: number,
+    step: number
+  ): Array<[number, number]> {
+    const chunks: Array<[number, number]> = [];
+
+    let left = from;
+    while (left < to) {
+      const right = left + step < to ? left + step : to;
+
+      chunks.push([left, right] as [number, number]);
+
+      left = right;
     }
+
+    return chunks;
   }
 
-  /**
-   * Gets all current members
-   */
-  public async getAllMembers(): Promise<Member[]> {
-    const nextIndex = (await this.contract.nextFreeIndex()).toNumber();
-    return this.getMembersInRange(0, nextIndex);
+  public static *takeN<T>(array: T[], size: number): Iterable<T[]> {
+    let start = 0;
+
+    while (start < array.length) {
+      const portion = array.slice(start, start + size);
+
+      yield portion;
+
+      start += size;
+    }
+  }
+
+  public static async ignoreErrors<T>(
+    promise: Promise<T>,
+    defaultValue: T
+  ): Promise<T> {
+    try {
+      return await promise;
+    } catch (err: unknown) {
+      if (err instanceof Error) {
+        log.info(`Ignoring an error during query: ${err.message}`);
+      } else {
+        log.info(`Ignoring an unknown error during query`);
+      }
+      return defaultValue;
+    }
+  }
+
+  public subscribeToMembers(): void {
+    this.contract.on(
+      this.membersFilter,
+      (
+        _idCommitment: string,
+        _membershipRateLimit: ethers.BigNumber,
+        _index: ethers.BigNumber,
+        event: ethers.Event
+      ) => {
+        this.processEvents([event]);
+      }
+    );
+
+    this.contract.on(
+      this.membershipErasedFilter,
+      (
+        _idCommitment: string,
+        _membershipRateLimit: ethers.BigNumber,
+        _index: ethers.BigNumber,
+        event: ethers.Event
+      ) => {
+        this.processEvents([event]);
+      }
+    );
+
+    this.contract.on(
+      this.membersExpiredFilter,
+      (
+        _idCommitment: string,
+        _membershipRateLimit: ethers.BigNumber,
+        _index: ethers.BigNumber,
+        event: ethers.Event
+      ) => {
+        this.processEvents([event]);
+      }
+    );
   }
 
   /**
-   * Gets the member index if it exists, or null if it doesn't
-   * Throws only on actual errors (invalid input, network issues, etc)
+   * Helper method to get remaining messages in current epoch
+   * @param membershipId The ID of the membership to check
+   * @returns number of remaining messages allowed in current epoch
    */
-  private async getMemberIndex(
-    idCommitment: string
-  ): Promise<ethers.BigNumber | null> {
+  public async getRemainingMessages(membershipId: number): Promise<number> {
     try {
-      const isValid = await this.contract.isInMembershipSet(idCommitment);
-      if (!isValid) {
-        return null;
-      }
+      const [startTime, , rateLimit] =
+        await this.contract.getMembershipInfo(membershipId);
 
-      const membershipInfo = await this.contract.memberships(idCommitment);
-      return ethers.BigNumber.from(membershipInfo.index);
+      // Calculate current epoch
+      const currentTime = Math.floor(Date.now() / 1000);
+      const epochsPassed = Math.floor(
+        (currentTime - startTime) / RATE_LIMIT_PARAMS.EPOCH_LENGTH
+      );
+      const currentEpochStart =
+        startTime + epochsPassed * RATE_LIMIT_PARAMS.EPOCH_LENGTH;
+
+      // Get message count in current epoch using contract's function
+      const messageCount = await this.contract.getMessageCount(
+        membershipId,
+        currentEpochStart
+      );
+      return Math.max(
+        0,
+        ethers.BigNumber.from(rateLimit)
+          .sub(ethers.BigNumber.from(messageCount))
+          .toNumber()
+      );
     } catch (error) {
-      log.error(`Error getting member index: ${(error as Error).message}`);
-      throw new InvalidMembershipError(idCommitment);
+      log.error(
+        `Error getting remaining messages: ${(error as Error).message}`
+      );
+      return 0; // Fail safe: assume no messages remaining on error
     }
   }
 
   public async getMembershipInfo(
     idCommitment: string
-  ): Promise<MembershipInfo> {
+  ): Promise<MembershipInfo | undefined> {
     try {
       const [startBlock, endBlock, rateLimit] =
         await this.contract.getMembershipInfo(idCommitment);
@@ -214,9 +413,7 @@ export class RLNBaseContract {
       }
 
       const index = await this.getMemberIndex(idCommitment);
-      if (index === null) {
-        throw new MembershipNotFoundError(idCommitment);
-      }
+      if (!index) return undefined;
 
       return {
         index,
@@ -227,40 +424,30 @@ export class RLNBaseContract {
         state
       };
     } catch (error) {
-      if (error instanceof RLNContractError) {
-        throw error;
-      }
-      log.error(`Error getting membership info: ${(error as Error).message}`);
-      throw new InvalidMembershipError(idCommitment);
+      return undefined;
     }
   }
 
-  public async extendMembership(idCommitment: string): Promise<void> {
-    const tx = await this.contract.extendMemberships([idCommitment]);
-    await this.confirmTransaction(tx, "MembershipExtended", (event) => ({
-      idCommitment: event.args!.idCommitment,
-      endBlock: event.args!.endBlock
-    }));
+  public async extendMembership(
+    idCommitment: string
+  ): Promise<ethers.ContractTransaction> {
+    return this.contract.extendMemberships([idCommitment]);
   }
 
   public async eraseMembership(
     idCommitment: string,
     eraseFromMembershipSet: boolean = true
-  ): Promise<void> {
-    const tx = await this.contract.eraseMemberships(
+  ): Promise<ethers.ContractTransaction> {
+    return this.contract.eraseMemberships(
       [idCommitment],
       eraseFromMembershipSet
     );
-    await this.confirmTransaction(tx, "MembershipErased", (event) => ({
-      idCommitment: event.args!.idCommitment,
-      index: event.args!.index
-    }));
   }
 
   public async registerMembership(
     idCommitment: string,
     rateLimit: number = DEFAULT_RATE_LIMIT
-  ): Promise<void> {
+  ): Promise<ethers.ContractTransaction> {
     if (
       rateLimit < RATE_LIMIT_PARAMS.MIN_RATE ||
       rateLimit > RATE_LIMIT_PARAMS.MAX_RATE
@@ -269,25 +456,15 @@ export class RLNBaseContract {
         `Rate limit must be between ${RATE_LIMIT_PARAMS.MIN_RATE} and ${RATE_LIMIT_PARAMS.MAX_RATE}`
       );
     }
-    const tx = await this.contract.register(idCommitment, rateLimit, []);
-    await this.confirmTransaction(tx, "MembershipRegistered", (event) => ({
-      idCommitment: event.args!.idCommitment,
-      membershipRateLimit: event.args!.membershipRateLimit,
-      index: event.args!.index
-    }));
+    return this.contract.register(idCommitment, rateLimit, []);
   }
 
   public async withdraw(token: string, holder: string): Promise<void> {
     try {
       const tx = await this.contract.withdraw(token, { from: holder });
-      await this.confirmTransaction(tx, "TokenWithdrawn", (event) => ({
-        token: event.args!.token,
-        holder: event.args!.holder,
-        amount: event.args!.amount
-      }));
+      await tx.wait();
     } catch (error) {
       log.error(`Error in withdraw: ${(error as Error).message}`);
-      throw error;
     }
   }
 
@@ -303,18 +480,18 @@ export class RLNBaseContract {
       const existingIndex = await this.getMemberIndex(
         identity.IDCommitmentBigInt.toString()
       );
-
-      if (existingIndex !== null) {
-        throw new MembershipExistsError(
-          identity.IDCommitmentBigInt.toString(),
-          existingIndex.toString()
+      if (existingIndex) {
+        throw new Error(
+          `ID commitment is already registered with index ${existingIndex}`
         );
       }
 
       // Check if there's enough remaining rate limit
       const remainingRateLimit = await this.getRemainingTotalRateLimit();
       if (remainingRateLimit < this.rateLimit) {
-        throw new RateLimitExceededError(this.rateLimit, remainingRateLimit);
+        throw new Error(
+          `Not enough remaining rate limit. Requested: ${this.rateLimit}, Available: ${remainingRateLimit}`
+        );
       }
 
       const estimatedGas = await this.contract.estimateGas.register(
@@ -324,23 +501,37 @@ export class RLNBaseContract {
       );
       const gasLimit = estimatedGas.add(10000);
 
-      const tx = await this.contract.register(
-        identity.IDCommitmentBigInt,
-        this.rateLimit,
-        [],
-        { gasLimit }
-      );
+      const txRegisterResponse: ethers.ContractTransaction =
+        await this.contract.register(
+          identity.IDCommitmentBigInt,
+          this.rateLimit,
+          [],
+          { gasLimit }
+        );
+
+      const txRegisterReceipt = await txRegisterResponse.wait();
+
+      if (txRegisterReceipt.status === 0) {
+        throw new Error("Transaction failed on-chain");
+      }
 
-      const decodedData = await this.confirmTransaction(
-        tx,
-        "MembershipRegistered",
-        (event): MembershipRegisteredEvent => ({
-          idCommitment: event.args!.idCommitment,
-          membershipRateLimit: event.args!.membershipRateLimit,
-          index: event.args!.index
-        })
+      const memberRegistered = txRegisterReceipt.events?.find(
+        (event) => event.event === "MembershipRegistered"
       );
 
+      if (!memberRegistered || !memberRegistered.args) {
+        log.error(
+          "Failed to register membership: No MembershipRegistered event found"
+        );
+        return undefined;
+      }
+
+      const decodedData: MembershipRegisteredEvent = {
+        idCommitment: memberRegistered.args.idCommitment,
+        membershipRateLimit: memberRegistered.args.membershipRateLimit,
+        index: memberRegistered.args.index
+      };
+
       log.info(
         `Successfully registered membership with index ${decodedData.index} ` +
           `and rate limit ${decodedData.membershipRateLimit}`
@@ -348,55 +539,44 @@ export class RLNBaseContract {
 
       const network = await this.contract.provider.getNetwork();
       const address = this.contract.address;
-      const membershipId = decodedData.index.toString();
+      const membershipId = Number(decodedData.index);
 
       return {
         identity,
         membership: {
           address,
-          treeIndex: parseInt(membershipId),
-          chainId: network.chainId.toString(),
+          treeIndex: membershipId,
+          chainId: network.chainId,
           rateLimit: decodedData.membershipRateLimit.toNumber()
         }
       };
     } catch (error) {
-      if (error instanceof RLNContractError) {
-        throw error;
-      }
-
       if (error instanceof Error) {
         const errorMessage = error.message;
         log.error("registerWithIdentity - error message:", errorMessage);
         log.error("registerWithIdentity - error stack:", error.stack);
 
-        // Map contract errors to our custom errors
+        // Try to extract more specific error information
         if (errorMessage.includes("CannotExceedMaxTotalRateLimit")) {
-          throw new RateLimitExceededError(
-            this.rateLimit,
-            await this.getRemainingTotalRateLimit()
+          throw new Error(
+            "Registration failed: Cannot exceed maximum total rate limit"
           );
         } else if (errorMessage.includes("InvalidIdCommitment")) {
-          throw new InvalidMembershipError(
-            identity.IDCommitmentBigInt.toString()
-          );
+          throw new Error("Registration failed: Invalid ID commitment");
         } else if (errorMessage.includes("InvalidMembershipRateLimit")) {
-          throw new InvalidRateLimitError(
-            this.rateLimit,
-            RATE_LIMIT_PARAMS.MIN_RATE,
-            RATE_LIMIT_PARAMS.MAX_RATE
-          );
+          throw new Error("Registration failed: Invalid membership rate limit");
         } else if (errorMessage.includes("execution reverted")) {
-          throw new TransactionError(
+          throw new Error(
             "Contract execution reverted. Check contract requirements."
           );
+        } else {
+          throw new Error(`Error in registerWithIdentity: ${errorMessage}`);
         }
-
-        throw new RLNContractError(
-          `Error in registerWithIdentity: ${errorMessage}`
-        );
+      } else {
+        throw new Error("Unknown error in registerWithIdentity", {
+          cause: error
+        });
       }
-
-      throw new RLNContractError("Unknown error in registerWithIdentity");
     }
   }
 
@@ -416,27 +596,36 @@ export class RLNBaseContract {
         `Registering identity with permit and rate limit: ${this.rateLimit} messages/epoch`
       );
 
-      const tx = await this.contract.registerWithPermit(
-        permit.owner,
-        permit.deadline,
-        permit.v,
-        permit.r,
-        permit.s,
-        identity.IDCommitmentBigInt,
-        this.rateLimit,
-        idCommitmentsToErase.map((id) => ethers.BigNumber.from(id))
-      );
+      const txRegisterResponse: ethers.ContractTransaction =
+        await this.contract.registerWithPermit(
+          permit.owner,
+          permit.deadline,
+          permit.v,
+          permit.r,
+          permit.s,
+          identity.IDCommitmentBigInt,
+          this.rateLimit,
+          idCommitmentsToErase.map((id) => ethers.BigNumber.from(id))
+        );
+      const txRegisterReceipt = await txRegisterResponse.wait();
 
-      const decodedData = await this.confirmTransaction(
-        tx,
-        "MembershipRegistered",
-        (event): MembershipRegisteredEvent => ({
-          idCommitment: event.args!.idCommitment,
-          membershipRateLimit: event.args!.membershipRateLimit,
-          index: event.args!.index
-        })
+      const memberRegistered = txRegisterReceipt.events?.find(
+        (event) => event.event === "MembershipRegistered"
       );
 
+      if (!memberRegistered || !memberRegistered.args) {
+        log.error(
+          "Failed to register membership with permit: No MembershipRegistered event found"
+        );
+        return undefined;
+      }
+
+      const decodedData: MembershipRegisteredEvent = {
+        idCommitment: memberRegistered.args.idCommitment,
+        membershipRateLimit: memberRegistered.args.membershipRateLimit,
+        index: memberRegistered.args.index
+      };
+
       log.info(
         `Successfully registered membership with permit. Index: ${decodedData.index}, ` +
           `Rate limit: ${decodedData.membershipRateLimit}, Erased ${idCommitmentsToErase.length} commitments`
@@ -444,14 +633,14 @@ export class RLNBaseContract {
 
       const network = await this.contract.provider.getNetwork();
       const address = this.contract.address;
-      const membershipId = decodedData.index.toString();
+      const membershipId = Number(decodedData.index);
 
       return {
         identity,
         membership: {
           address,
-          treeIndex: parseInt(membershipId),
-          chainId: network.chainId.toString(),
+          treeIndex: membershipId,
+          chainId: network.chainId,
           rateLimit: decodedData.membershipRateLimit.toNumber()
         }
       };
@@ -459,7 +648,7 @@ export class RLNBaseContract {
       log.error(
         `Error in registerWithPermitAndErase: ${(error as Error).message}`
       );
-      throw error;
+      return undefined;
     }
   }
 
@@ -472,29 +661,47 @@ export class RLNBaseContract {
       rateLimit < RATE_LIMIT_PARAMS.MIN_RATE ||
       rateLimit > RATE_LIMIT_PARAMS.MAX_RATE
     ) {
-      throw new InvalidRateLimitError(
-        rateLimit,
-        RATE_LIMIT_PARAMS.MIN_RATE,
-        RATE_LIMIT_PARAMS.MAX_RATE
+      throw new Error(
+        `Rate limit must be between ${RATE_LIMIT_PARAMS.MIN_RATE} and ${RATE_LIMIT_PARAMS.MAX_RATE} messages per epoch`
       );
     }
   }
 
-  /**
-   * Helper to confirm a transaction and extract event data
-   */
-  private async confirmTransaction<T>(
-    tx: ethers.ContractTransaction,
-    expectedEvent: string,
-    transform: (event: ethers.Event) => T
-  ): Promise<T> {
-    const receipt = await tx.wait();
-    const event = receipt.events?.find((e) => e.event === expectedEvent);
+  private get membersFilter(): ethers.EventFilter {
+    if (!this._membersFilter) {
+      throw Error("Members filter was not initialized.");
+    }
+    return this._membersFilter;
+  }
 
-    if (!event || !event.args) {
-      throw new TransactionError(`Expected event ${expectedEvent} not found`);
+  private get membershipErasedFilter(): ethers.EventFilter {
+    if (!this._membershipErasedFilter) {
+      throw Error("MembershipErased filter was not initialized.");
     }
+    return this._membershipErasedFilter;
+  }
 
-    return transform(event);
+  private get membersExpiredFilter(): ethers.EventFilter {
+    if (!this._membersExpiredFilter) {
+      throw Error("MembersExpired filter was not initialized.");
+    }
+    return this._membersExpiredFilter;
+  }
+
+  private async getMemberIndex(
+    idCommitment: string
+  ): Promise<ethers.BigNumber | undefined> {
+    try {
+      const events = await this.contract.queryFilter(
+        this.contract.filters.MembershipRegistered(idCommitment)
+      );
+      if (events.length === 0) return undefined;
+
+      // Get the most recent registration event
+      const event = events[events.length - 1];
+      return event.args?.index;
+    } catch (error) {
+      return undefined;
+    }
   }
 }