feat: Use dataloader in multicall wrapper

package.json

@@ -64,6 +64,7 @@
     "class-validator": "^0.13.2",
     "conventional-changelog-conventionalcommits": "^4.6.3",
     "copyfiles": "^2.4.1",
+    "dataloader": "^2.1.0",
     "dedent": "^0.7.0",
     "dotenv": "^16.0.0",
     "eslint": "^8.11.0",
@@ -105,6 +106,7 @@
     "cache-manager": "^3.4.1",
     "class-transformer": "^0.4.0",
     "class-validator": "^0.13.2",
+    "dataloader": "^2.1.0",
     "ethers": "^5.5.1",
     "graphql": "14 || 15 || 16",
     "graphql-request": "^3",

src/multicall/multicall.contract.ts

@@ -0,0 +1,36 @@
+import { Fragment, FunctionFragment, JsonFragment } from '@ethersproject/abi';
+
+import { ContractCall } from './multicall.ethers';
+
+export class MulticallContract {
+  private _address: string;
+  private _abi: Fragment[];
+  private _functions: FunctionFragment[];
+
+  get address() {
+    return this._address;
+  }
+
+  get abi() {
+    return this._abi;
+  }
+
+  get functions() {
+    return this._functions;
+  }
+
+  constructor(address: string, abi: JsonFragment[] | string[] | Fragment[]) {
+    this._address = address;
+
+    this._abi = abi.map((item: JsonFragment | string | Fragment) => Fragment.from(item));
+    this._functions = this._abi.filter(x => x.type === 'function').map(x => FunctionFragment.from(x));
+    const fragments = this._functions.filter(x => x.stateMutability === 'pure' || x.stateMutability === 'view');
+
+    for (const frag of fragments) {
+      const fn = (...params: any[]): ContractCall => ({ fragment: frag, address, params });
+      if (!this[frag.name]) Object.defineProperty(this, frag.name, { enumerable: true, writable: false, value: fn });
+    }
+  }
+
+  [method: string]: any;
+}

src/multicall/multicall.ethers.ts

@@ -1,206 +1,81 @@
-import { FunctionFragment } from '@ethersproject/abi';
-import { Contract } from '@ethersproject/contracts';
-import { ethers } from 'ethers';
+import DataLoader from 'dataloader';
+import { Contract } from 'ethers';
+import { FunctionFragment, Interface } from 'ethers/lib/utils';
 
-import { Multicall as MulticallContract } from '~contract/contracts';
+import { Multicall } from '~contract/contracts';
 
-export type EthersMulticallConfig = {
-  batchInterval?: number;
-  batchMaxSize?: number;
-};
-
-type MulticallAggregateReturnData = Awaited<ReturnType<MulticallContract['callStatic']['aggregate']>>['returnData'];
-type TargetContract = Pick<Contract, 'functions' | 'interface' | 'callStatic' | 'address'>;
-type HasContractFunctions = { functions: Record<string, unknown> };
-
-// Extracts the contract functions in `functions`
-type AllContractFunctions<T extends HasContractFunctions> = keyof {
-  [P in keyof T['functions']]: T[P];
-};
+import { MulticallContract } from './multicall.contract';
 
-// Extracts all methods from type
-type Methods<T> = { [P in keyof T as T[P] extends (...args: any[]) => any ? P : never]: T[P] };
-
-// Picks only the methods on the contract that are defined in `functions`
-type ExtractContractFunctions<T extends HasContractFunctions> = Pick<Methods<T>, AllContractFunctions<T>>;
-
-type BatchItem = {
-  callData: string;
-  callTarget: string;
-  functionFragment: FunctionFragment;
-  resolve: (value?: any) => void;
-  reject: (reason?: any) => void;
+export type ContractCall = {
+  fragment: FunctionFragment;
+  address: string;
+  params: any[];
 };
 
 export class EthersMulticall {
-  private readonly batchInterval: number;
-  private readonly batchMaxSize: number;
-  private readonly multicallContract: MulticallContract;
-  private batchRequests: BatchItem[] | null;
-
-  constructor(multicallContact: MulticallContract, opts: EthersMulticallConfig = {}) {
-    const { batchInterval = 10, batchMaxSize = 250 } = opts;
-
-    this.multicallContract = multicallContact;
-    this.batchInterval = batchInterval;
-    this.batchMaxSize = batchMaxSize;
-    this.batchRequests = null;
+  private multicall: Multicall;
+  private dataLoader: DataLoader<ContractCall, any>;
+
+  constructor(
+    multicall: Multicall,
+    dataLoaderOptions: DataLoader.Options<ContractCall, any> = { cache: false, maxBatchSize: 250 },
+  ) {
+    this.multicall = multicall;
+    this.dataLoader = new DataLoader(this.doCalls.bind(this), dataLoaderOptions);
   }
 
-  get contract() {
-    return this.multicallContract;
-  }
+  private async doCalls(calls: readonly ContractCall[]) {
+    const callRequests = calls.map(call => ({
+      target: call.address,
+      callData: new Interface([]).encodeFunctionData(call.fragment, call.params),
+    }));
 
-  /**
-   * Intercepts provider calls to functions / static functions on the contract,
-   * replaces execution with encoding the function data of the target & args, then
-   * adds the function data to the multicall consumption queue
-   */
-  private hijackExecution(contract: TargetContract, method: string) {
-    return (...args: any[]) => {
-      const functionFragment = contract.interface.getFunction(method as string);
-      if (!functionFragment) throw new Error('Cannot find function on the given contract');
-
-      if (!this.batchRequests) {
-        this.batchRequests = [];
-      }
+    const response = await this.multicall.callStatic.aggregate(callRequests, false);
 
-      return new Promise((resolve, reject) => {
-        this.batchRequests?.push({
-          functionFragment,
-          callData: contract.interface.encodeFunctionData(functionFragment, args),
-          callTarget: contract.address,
-          resolve,
-          reject,
-        });
-
-        if (this.batchRequests?.length === 1) {
-          this.scheduleConsumeQueue();
-        }
-
-        if (this.batchRequests?.length === this.batchMaxSize) {
-          this.consumeQueue();
-        }
-      });
-    };
-  }
-
-  wrap<T extends TargetContract>(contract: T): ExtractContractFunctions<T> & Pick<T, 'callStatic'> {
-    // Removes readonly contraints on the contract properties
-    const configurableContract = Object.create(contract);
-
-    return new Proxy(configurableContract, {
-      get: (target, key) => {
-        const functionName = key as string;
-        const isFunctionCall = functionName in configurableContract.functions;
-        const isStaticFunctionCall = functionName === 'callStatic';
-        // Disregard calls other than provider method invocations
-        if (!isFunctionCall && !isStaticFunctionCall) throw new Error('Invalid multicall operation');
-
-        if (isStaticFunctionCall) {
-          // Removes readonly contraints on the `functions` / `callStatic` properties
-          const configurableTarget = Object.create(target[functionName]);
-
-          return new Proxy(configurableTarget, {
-            get: (_target, staticFunctionName) => this.hijackExecution(contract, staticFunctionName as string),
-          });
-        }
-
-        return this.hijackExecution(contract, functionName);
-      },
-    });
-  }
-
-  private async aggregate(batchItems: BatchItem[]) {
-    // Prepare batch items into payloads
-    const calls = batchItems.map(({ callData, callTarget }) => {
-      return {
-        target: callTarget,
-        callData,
-      };
-    });
+    const result = calls.map((call, i) => {
+      const signature = FunctionFragment.from(call.fragment).format();
+      const callIdentifier = [call.address, signature].join(':');
+      const [success, data] = response.returnData[i];
 
-    // Actual call to multicall's aggregate
-    try {
-      const responses = await this.multicallContract.callStatic.aggregate(calls, false);
-      const returnData = responses.returnData;
-
-      if (returnData.length !== batchItems.length) {
-        throw new Error(`Unexpected response length: received ${returnData.length}; expected ${batchItems.length}`);
+      if (!success) {
+        return new Error(`Multicall call failed for ${callIdentifier}`);
       }
 
-      return returnData;
-    } catch (err) {
-      const exception = err as Error;
-      exception.message = `Multicall aggregate request failed: ${exception.message}`;
-      throw err;
-    }
-  }
-
-  private decodeFunctionData(batchItem: BatchItem, batchReturnData: MulticallAggregateReturnData[number]) {
-    const { callTarget, functionFragment } = batchItem;
-    const [success, data] = batchReturnData;
-    const functionSignature = functionFragment.format();
-
-    // Multicall's response for the batch item failed
-    if (!success) {
-      const callIdentifier = [callTarget, functionSignature].join(':');
-      throw new Error(`Multicall call failed for ${callIdentifier}`);
-    }
-
-    try {
-      const decoder = ethers.utils.defaultAbiCoder;
-      if (!functionFragment.outputs) throw new Error('no outputs received');
-      const decoded = decoder.decode(functionFragment.outputs, data);
-      if (functionFragment.outputs?.length > 1) return decoded;
-      return decoded[0];
-    } catch (err) {
-      const exception = err as Error;
-      const callIdentifier = [callTarget, functionSignature].join(':');
-      exception.message = `Multicall call failed for ${callIdentifier}: ${exception.message}`;
-      throw err;
-    }
-  }
-
-  private async consumeQueue() {
-    const batch = this.batchRequests;
-    if (!batch) return;
-
-    // Clear the batch queue
-    this.batchRequests = null;
-
-    // Call multicall's aggregate with all the batch requests
-    let returnData: MulticallAggregateReturnData;
-    try {
-      returnData = await this.aggregate(batch);
-    } catch (err) {
-      // Each batch rejects with the error reason
-      batch.forEach(({ reject }) => reject(err));
-      return;
-    }
-
-    // Decode the data for each batch item
-    batch.forEach((batchItem, i) => {
-      const { resolve, reject } = batchItem;
       try {
-        const decodedBatchData = this.decodeFunctionData(batchItem, returnData[i]);
-        resolve(decodedBatchData);
+        const outputs = call.fragment.outputs!;
+        const result = new Interface([]).decodeFunctionResult(call.fragment, data);
+        return outputs.length === 1 ? result[0] : result;
       } catch (err) {
-        reject(err);
+        return new Error(`Multicall call failed for ${callIdentifier}`);
       }
     });
+
+    return result;
   }
 
-  private scheduleConsumeQueue(): void {
-    setTimeout(async () => {
-      if (this.batchRequests?.length) {
-        try {
-          await this.consumeQueue();
-        } catch (err) {
-          const exception = err as Error;
-          exception.message = `Multicall unexpected error occurred: ${exception.message}`;
-        }
-      }
-    }, this.batchInterval);
+  wrap<T extends Contract>(contract: T) {
+    const abi = contract.interface.fragments;
+    const multicallContract = new MulticallContract(contract.address, abi as any);
+
+    const funcs = abi.reduce((memo, frag) => {
+      if (frag.type !== 'function') return memo;
+
+      const funcFrag = frag as FunctionFragment;
+      if (!['pure', 'view'].includes(funcFrag.stateMutability)) return memo;
+
+      // Overwrite the function with a dataloader batched call
+      const multicallFunc = multicallContract[funcFrag.name].bind(multicallContract);
+      const newFunc = (...args: any) => {
+        const contractCall = multicallFunc(...args);
+        return this.dataLoader.load(contractCall);
+      };
+
+      memo[funcFrag.name] = newFunc;
+      return memo;
+    }, {} as Record<string, (...args: any) => any>);
+
+    return Object.setPrototypeOf({ ...contract, ...funcs }, Contract.prototype) as any as T;
   }
 }
+
+export default EthersMulticall;