From 7f2c4f3f3211c671463f600eff6b2a44b35c2966 Mon Sep 17 00:00:00 2001
From: Alex Forshtat <forshtat1@gmail.com>
Date: Sun, 4 Aug 2024 17:56:47 +0200
Subject: [PATCH] Move 'userOpMaxGas' calculation into the MempoolEntry class

---
 packages/bundler/src/modules/BundleManager.ts | 15 ++++---------
 packages/bundler/src/modules/MempoolEntry.ts  | 21 +++++++++++++++++++
 .../bundler/src/modules/MempoolManager.ts     | 19 ++++++-----------
 3 files changed, 31 insertions(+), 24 deletions(-)
 create mode 100644 packages/bundler/src/modules/MempoolEntry.ts

diff --git a/packages/bundler/src/modules/BundleManager.ts b/packages/bundler/src/modules/BundleManager.ts
index be4e2e0e..f79c0e83 100644
--- a/packages/bundler/src/modules/BundleManager.ts
+++ b/packages/bundler/src/modules/BundleManager.ts
@@ -23,7 +23,8 @@ import {
 import { EventsManager } from './EventsManager'
 import { GetUserOpHashes__factory } from '../types'
 import { IBundleManager } from './IBundleManager'
-import { MempoolEntry, MempoolManager } from './MempoolManager'
+import { MempoolEntry } from './MempoolEntry'
+import { MempoolManager } from './MempoolManager'
 import { ReputationManager, ReputationStatus } from './ReputationManager'
 
 const debug = Debug('aa.exec.cron')
@@ -313,20 +314,12 @@ export class BundleManager implements IBundleManager {
       }
       mergeStorageMap(storageMap, validationResult.storageMap)
 
-      const userOpGas =
-        BigNumber
-          .from((entry.userOp as UserOperation).preVerificationGas ?? 0)
-          .add(entry.userOp.callGasLimit)
-          .add(entry.userOp.verificationGasLimit)
-          .add(entry.userOp.paymasterVerificationGasLimit ?? 0)
-          .add(entry.userOp.paymasterPostOpGasLimit ?? 0)
-
-      const newBundleGas = userOpGas.add(bundleGas)
+      const newBundleGas = entry.userOpMaxGas.add(bundleGas)
       if (
         maxBundleGas != null &&
         !BigNumber.from(maxBundleGas).eq(0) &&
         newBundleGas.gte(maxBundleGas)) {
-        debug('exiting after maxBundleGas is reached', maxBundleGas, bundleGas, userOpGas)
+        debug('exiting after maxBundleGas is reached', maxBundleGas, bundleGas, entry.userOpMaxGas)
         break
       }
       bundleGas = newBundleGas
diff --git a/packages/bundler/src/modules/MempoolEntry.ts b/packages/bundler/src/modules/MempoolEntry.ts
new file mode 100644
index 00000000..11807496
--- /dev/null
+++ b/packages/bundler/src/modules/MempoolEntry.ts
@@ -0,0 +1,21 @@
+import { BigNumber, BigNumberish } from 'ethers'
+import { OperationBase, ReferencedCodeHashes, UserOperation } from '@account-abstraction/utils'
+
+export class MempoolEntry {
+  userOpMaxGas: BigNumber
+
+  constructor (
+    readonly userOp: OperationBase,
+    readonly userOpHash: string,
+    readonly prefund: BigNumberish,
+    readonly referencedContracts: ReferencedCodeHashes,
+    readonly aggregator?: string
+  ) {
+    this.userOpMaxGas = BigNumber
+      .from((this.userOp as UserOperation).preVerificationGas ?? 0)
+      .add(this.userOp.callGasLimit)
+      .add(this.userOp.verificationGasLimit)
+      .add(this.userOp.paymasterVerificationGasLimit ?? 0)
+      .add(this.userOp.paymasterPostOpGasLimit ?? 0)
+  }
+}
diff --git a/packages/bundler/src/modules/MempoolManager.ts b/packages/bundler/src/modules/MempoolManager.ts
index e5ee3ff1..40f7f94f 100644
--- a/packages/bundler/src/modules/MempoolManager.ts
+++ b/packages/bundler/src/modules/MempoolManager.ts
@@ -1,4 +1,6 @@
 import { BigNumber, BigNumberish } from 'ethers'
+import Debug from 'debug'
+
 import {
   ReferencedCodeHashes,
   RpcError,
@@ -6,20 +8,11 @@ import {
   ValidationErrors,
   requireCond, OperationBase
 } from '@account-abstraction/utils'
+import { MempoolEntry } from './MempoolEntry'
 import { ReputationManager } from './ReputationManager'
-import Debug from 'debug'
 
 const debug = Debug('aa.mempool')
 
-export interface MempoolEntry {
-  userOp: OperationBase
-  userOpHash: string
-  prefund: BigNumberish
-  referencedContracts: ReferencedCodeHashes
-  // aggregator, if one was found during simulation
-  aggregator?: string
-}
-
 type MempoolDump = OperationBase[]
 
 const THROTTLED_ENTITY_MEMPOOL_COUNT = 4
@@ -75,13 +68,13 @@ export class MempoolManager {
     factoryInfo?: StakeInfo,
     aggregatorInfo?: StakeInfo
   ): void {
-    const entry: MempoolEntry = {
+    const entry = new MempoolEntry(
       userOp,
       userOpHash,
       prefund,
       referencedContracts,
-      aggregator: aggregatorInfo?.addr
-    }
+      aggregatorInfo?.addr
+    )
     const index = this._findBySenderNonce(userOp.sender, userOp.nonce)
     if (index !== -1) {
       const oldEntry = this.mempool[index]