Add state override set to context (#116)

Author: hazim-j

src/builder.ts

@@ -3,7 +3,7 @@ import { OpToJSON } from "./utils";
 import { UserOperationMiddlewareCtx } from "./context";
 import {
   IUserOperation,
-  IUserOperationBuilder,
+  StateOverrideSet,
   UserOperationMiddlewareFn,
 } from "./types";
 
@@ -25,7 +25,7 @@ export const DEFAULT_USER_OP: IUserOperation = {
   signature: ethers.utils.hexlify("0x"),
 };
 
-export class UserOperationBuilder implements IUserOperationBuilder {
+export class UserOperationBuilder {
   private defaultOp: IUserOperation;
   private currOp: IUserOperation;
   private middlewareStack: Array<UserOperationMiddlewareFn>;
@@ -197,11 +197,16 @@ export class UserOperationBuilder implements IUserOperationBuilder {
     return this;
   }
 
-  async buildOp(entryPoint: string, chainId: BigNumberish) {
+  async buildOp(
+    entryPoint: string,
+    chainId: BigNumberish,
+    stateOverrides?: StateOverrideSet
+  ) {
     const ctx = new UserOperationMiddlewareCtx(
       this.currOp,
       entryPoint,
-      chainId
+      chainId,
+      stateOverrides
     );
 
     for (const fn of this.middlewareStack) {

src/client.ts

@@ -1,17 +1,13 @@
 import { BigNumberish, ethers } from "ethers";
-import {
-  IClient,
-  IUserOperationBuilder,
-  ISendUserOperationOpts,
-  IClientOpts,
-} from "./types";
+import { UserOperationBuilder } from "./builder";
+import { ISendUserOperationOpts, IClientOpts, StateOverrideSet } from "./types";
 import { EntryPoint, EntryPoint__factory } from "./typechain";
 import { OpToJSON } from "./utils";
 import { UserOperationMiddlewareCtx } from "./context";
 import { ERC4337 } from "./constants";
 import { BundlerJsonRpcProvider } from "./provider";
 
-export class Client implements IClient {
+export class Client {
   private provider: ethers.providers.JsonRpcProvider;
 
   public entryPoint: EntryPoint;
@@ -41,16 +37,23 @@ export class Client implements IClient {
     return instance;
   }
 
-  async buildUserOperation(builder: IUserOperationBuilder) {
-    return builder.buildOp(this.entryPoint.address, this.chainId);
+  async buildUserOperation(
+    builder: UserOperationBuilder,
+    stateOverrides?: StateOverrideSet
+  ) {
+    return builder.buildOp(
+      this.entryPoint.address,
+      this.chainId,
+      stateOverrides
+    );
   }
 
   async sendUserOperation(
-    builder: IUserOperationBuilder,
+    builder: UserOperationBuilder,
     opts?: ISendUserOperationOpts
   ) {
     const dryRun = Boolean(opts?.dryRun);
-    const op = await this.buildUserOperation(builder);
+    const op = await this.buildUserOperation(builder, opts?.stateOverrides);
     opts?.onBuild?.(op);
 
     const userOpHash = dryRun

src/context.ts

@@ -1,15 +1,26 @@
 import { BigNumberish, ethers } from "ethers";
-import { IUserOperationMiddlewareCtx, IUserOperation } from "./types";
+import {
+  IUserOperationMiddlewareCtx,
+  IUserOperation,
+  StateOverrideSet,
+} from "./types";
 
 export class UserOperationMiddlewareCtx implements IUserOperationMiddlewareCtx {
   public op: IUserOperation;
   readonly entryPoint: string;
   readonly chainId: BigNumberish;
+  readonly stateOverrides?: StateOverrideSet | undefined;
 
-  constructor(op: IUserOperation, entryPoint: string, chainId: BigNumberish) {
+  constructor(
+    op: IUserOperation,
+    entryPoint: string,
+    chainId: BigNumberish,
+    stateOverrides?: StateOverrideSet
+  ) {
     this.op = { ...op };
     this.entryPoint = ethers.utils.getAddress(entryPoint);
     this.chainId = ethers.BigNumber.from(chainId);
+    this.stateOverrides = stateOverrides;
   }
 
   getUserOpHash() {

src/preset/middleware/gasLimit.ts

@@ -14,10 +14,14 @@ interface GasEstimate {
 export const estimateUserOperationGas =
   (provider: ethers.providers.JsonRpcProvider): UserOperationMiddlewareFn =>
   async (ctx) => {
-    const est = (await provider.send("eth_estimateUserOperationGas", [
-      OpToJSON(ctx.op),
-      ctx.entryPoint,
-    ])) as GasEstimate;
+    const params =
+      ctx.stateOverrides !== undefined
+        ? [OpToJSON(ctx.op), ctx.entryPoint, ctx.stateOverrides]
+        : [OpToJSON(ctx.op), ctx.entryPoint];
+    const est = (await provider.send(
+      "eth_estimateUserOperationGas",
+      params
+    )) as GasEstimate;
 
     ctx.op.preVerificationGas = est.preVerificationGas;
     ctx.op.verificationGasLimit =

src/types.ts

@@ -1,6 +1,16 @@
 import { BigNumberish, BytesLike, ethers } from "ethers";
 import { UserOperationEventEvent } from "./typechain/EntryPoint";
 
+export interface ISateOverrideAccount {
+  nonce: BigNumberish;
+  code: BytesLike;
+  balance: BigNumberish;
+  state: Record<string, BytesLike>;
+  stateDiff: Record<string, BytesLike>;
+}
+
+export type StateOverrideSet = Record<string, Partial<ISateOverrideAccount>>;
+
 export interface IUserOperation {
   sender: string;
   nonce: BigNumberish;
@@ -15,55 +25,6 @@ export interface IUserOperation {
   signature: BytesLike;
 }
 
-export interface IUserOperationBuilder {
-  // get methods.
-  getSender: () => string;
-  getNonce: () => BigNumberish;
-  getInitCode: () => BytesLike;
-  getCallData: () => BytesLike;
-  getCallGasLimit: () => BigNumberish;
-  getVerificationGasLimit: () => BigNumberish;
-  getPreVerificationGas: () => BigNumberish;
-  getMaxFeePerGas: () => BigNumberish;
-  getMaxPriorityFeePerGas: () => BigNumberish;
-  getPaymasterAndData: () => BytesLike;
-  getSignature: () => BytesLike;
-  getOp: () => IUserOperation;
-
-  // set methods.
-  setSender: (address: string) => IUserOperationBuilder;
-  setNonce: (nonce: BigNumberish) => IUserOperationBuilder;
-  setInitCode: (code: BytesLike) => IUserOperationBuilder;
-  setCallData: (data: BytesLike) => IUserOperationBuilder;
-  setCallGasLimit: (gas: BigNumberish) => IUserOperationBuilder;
-  setVerificationGasLimit: (gas: BigNumberish) => IUserOperationBuilder;
-  setPreVerificationGas: (gas: BigNumberish) => IUserOperationBuilder;
-  setMaxFeePerGas: (fee: BigNumberish) => IUserOperationBuilder;
-  setMaxPriorityFeePerGas: (fee: BigNumberish) => IUserOperationBuilder;
-  setPaymasterAndData: (data: BytesLike) => IUserOperationBuilder;
-  setSignature: (bytes: BytesLike) => IUserOperationBuilder;
-  setPartial: (partialOp: Partial<IUserOperation>) => IUserOperationBuilder;
-
-  // Sets the default values that won't be wiped on reset.
-  useDefaults: (partialOp: Partial<IUserOperation>) => IUserOperationBuilder;
-  resetDefaults: () => IUserOperationBuilder;
-
-  // Some fields may require arbitrary logic to build an op.
-  // Middleware functions allow you to set custom logic for building op fragments.
-  useMiddleware: (fn: UserOperationMiddlewareFn) => IUserOperationBuilder;
-  resetMiddleware: () => IUserOperationBuilder;
-
-  // This will construct a UserOperation that can be sent to a client.
-  // It will run through your entire middleware stack in the process.
-  buildOp: (
-    entryPoint: string,
-    chainId: BigNumberish
-  ) => Promise<IUserOperation>;
-
-  // Will reset all fields back to default value.
-  resetOp: () => IUserOperationBuilder;
-}
-
 export type UserOperationMiddlewareFn = (
   context: IUserOperationMiddlewareCtx
 ) => Promise<void>;
@@ -72,22 +33,12 @@ export interface IUserOperationMiddlewareCtx {
   op: IUserOperation;
   entryPoint: string;
   chainId: BigNumberish;
+  stateOverrides?: StateOverrideSet;
 
   // A userOpHash is a unique hash of op + entryPoint + chainId.
   getUserOpHash: () => string;
 }
 
-export interface IClient {
-  sendUserOperation: (
-    builder: IUserOperationBuilder,
-    opts?: ISendUserOperationOpts
-  ) => Promise<ISendUserOperationResponse>;
-
-  buildUserOperation: (
-    builder: IUserOperationBuilder
-  ) => Promise<IUserOperation>;
-}
-
 export interface IClientOpts {
   entryPoint?: string;
   overrideBundlerRpc?: string;
@@ -96,6 +47,7 @@ export interface IClientOpts {
 export interface ISendUserOperationOpts {
   dryRun?: boolean;
   onBuild?: (op: IUserOperation) => Promise<any> | any;
+  stateOverrides?: StateOverrideSet;
 }
 
 export interface ISendUserOperationResponse {