Merge pull request #166 from morpho-org/refactor/simulate-required-to…

…ken-amounts refactor(simulation-state): extract simulateRequiredTokenAmounts
morpho-org · Nov 15, 2024 · 72d2da0 · 72d2da0
2 parents 0db203c + 91e2fa9
commit 72d2da0
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 79 deletions.
diff --git a/packages/bundler-sdk-viem/src/operations.ts b/packages/bundler-sdk-viem/src/operations.ts
@@ -12,7 +12,13 @@ import {
   permissionedBackedTokens,
   permissionedWrapperTokens,
 } from "@morpho-org/blue-sdk";
-import { entries, getLast, getValue, keys } from "@morpho-org/morpho-ts";
+import {
+  bigIntComparator,
+  entries,
+  getLast,
+  getValue,
+  keys,
+} from "@morpho-org/morpho-ts";
 import {
   type Erc20Operations,
   type MaybeDraft,
@@ -25,6 +31,7 @@ import {
   handleOperations,
   produceImmutable,
   simulateOperation,
+  simulateOperations,
 } from "@morpho-org/simulation-sdk";
 
 import { maxUint256 } from "viem";
@@ -403,9 +410,9 @@ export const populateSubBundle = (
   } as Operation;
 
   // Operations with callbacks are populated recursively as a side-effect of the simulation, within the callback itself.
-  let { requiredTokenAmounts } = data.simulateWithUnlimitedBalances(
+  let requiredTokenAmounts = simulateRequiredTokenAmounts(
     (operations as Operation[]).concat([simulatedOperation]),
-    [bundler],
+    data,
   );
 
   const allOperations = (operations as BundlerOperation[]).concat([
@@ -428,12 +435,12 @@ export const populateSubBundle = (
 
   const requirementOperations =
     getRequirementOperations?.(requiredTokenAmounts) ?? [];
-  ({ requiredTokenAmounts } = data.simulateWithUnlimitedBalances(
+  requiredTokenAmounts = simulateRequiredTokenAmounts(
     requirementOperations
       .concat(allOperations)
       .map((operation) => getSimulatedBundlerOperation(operation)),
-    [bundler],
-  ));
+    data,
+  );
 
   // Append required input transfers.
   requiredTokenAmounts.forEach(({ token, required }) => {
@@ -847,6 +854,48 @@ export const populateBundle = (
   return { operations, steps };
 };
 
+export const simulateRequiredTokenAmounts = (
+  operations: Operation[],
+  data: MaybeDraft<SimulationState>,
+) => {
+  const { bundler } = getChainAddresses(data.chainId);
+
+  const virtualBundlerData = produceImmutable(data, (draft) => {
+    Object.values(draft.holdings[bundler] ?? {}).forEach((bundlerTokenData) => {
+      // Virtual balance to calculate the amount required.
+      bundlerTokenData.balance += MathLib.MAX_UINT_160;
+    });
+  });
+
+  const steps = simulateOperations(operations, virtualBundlerData);
+
+  const bundlerTokenDiffs = keys(virtualBundlerData.holdings[bundler]).map(
+    (token) => ({
+      token,
+      required: steps
+        .map(
+          (step) =>
+            // When recursively simulated, this will cause tokens to be required at the highest recursion level.
+            // For example: supplyCollateral(x, supplyCollateral(y, borrow(z)))   [provided x, y, z < MAX_UINT_160]
+            //              |                   |                   |=> MAX_UINT_160 - (3 * MAX_UINT_160 + z) < 0
+            //              |                   |=> MAX_UINT_160 - (2 * MAX_UINT_160 - y) < 0
+            //              |=> MAX_UINT_160 - (MAX_UINT_160 - y - x) > 0
+            MathLib.MAX_UINT_160 -
+            (step.holdings[bundler]?.[token]?.balance ?? 0n),
+        )
+        .sort(
+          bigIntComparator(
+            (required) => required,
+            // Take the highest required amount among all operations.
+            "desc",
+          ),
+        )[0]!,
+    }),
+  );
+
+  return bundlerTokenDiffs.filter(({ required }) => required > 0n);
+};
+
 export const getSimulatedBundlerOperation = (
   operation: BundlerOperation,
   { slippage }: { slippage?: bigint } = {},

diff --git a/packages/simulation-sdk/src/SimulationState.ts b/packages/simulation-sdk/src/SimulationState.ts
@@ -45,14 +45,7 @@ import {
   UnknownVaultUserError,
   UnknownWrappedTokenError,
 } from "./errors.js";
-import {
-  type MaybeDraft,
-  type SimulationResult,
-  produceImmutable,
-  simulateOperation,
-  simulateOperations,
-} from "./handlers/index.js";
-import type { Operation } from "./operations.js";
+import { type MaybeDraft, simulateOperation } from "./handlers/index.js";
 
 export interface PublicAllocatorOptions {
   /* The array of vaults to reallocate. Must all have enabled the PublicAllocator. Defaults to all the vaults that have enabled the PublicAllocator. */
@@ -655,68 +648,4 @@ export class SimulationState implements InputSimulationState {
 
     return _getMarketPublicReallocations();
   }
-
-  public simulateWithUnlimitedBalances(
-    operations: Operation[],
-    holders: Address[],
-  ) {
-    const virtualHoldersData = produceImmutable(this, (draft) => {
-      holders
-        .flatMap((holder) => Object.values(draft.holdings[holder] ?? {}))
-        .forEach((holderTokenData) => {
-          // Virtual balance to calculate the amount required.
-          holderTokenData.balance += MathLib.MAX_UINT_160;
-        });
-    });
-
-    const steps = simulateOperations(operations, virtualHoldersData);
-
-    const holdersTokenDiffs = holders
-      .flatMap((holder) =>
-        keys(virtualHoldersData.holdings[holder]).map(
-          (token) => [holder, token] as const,
-        ),
-      )
-      .map(([holder, token]) => ({
-        holder,
-        token,
-        required: steps
-          .map(
-            (step) =>
-              // When recursively simulated, this will cause tokens to be required at the highest recursion level.
-              // For example: supplyCollateral(x, supplyCollateral(y, borrow(z)))   [provided x, y, z < MAX_UINT_160]
-              //              |                   |                   |=> MAX_UINT_160 - (3 * MAX_UINT_160 + z) < 0
-              //              |                   |=> MAX_UINT_160 - (2 * MAX_UINT_160 - y) < 0
-              //              |=> MAX_UINT_160 - (MAX_UINT_160 - y - x) > 0
-              MathLib.MAX_UINT_160 -
-              (step.holdings[holder]?.[token]?.balance ?? 0n),
-          )
-          .sort(
-            bigIntComparator(
-              (required) => required,
-              // Take the highest required amount among all operations.
-              "desc",
-            ),
-          )[0]!,
-      }));
-
-    return {
-      requiredTokenAmounts: holdersTokenDiffs.filter(
-        ({ required }) => required > 0n,
-      ),
-      steps: steps.map((step) =>
-        produceImmutable(step, (draft) =>
-          holders
-            .flatMap((holder) => Object.values(draft.holdings[holder] ?? {}))
-            .forEach((holderTokenData) => {
-              // Change virtual balance back to real value
-              holderTokenData.balance = MathLib.max(
-                0n,
-                holderTokenData.balance - MathLib.MAX_UINT_160,
-              );
-            }),
-        ),
-      ) as SimulationResult,
-    };
-  }
 }
diff --git a/packages/simulation-sdk/src/handlers/dispatchers.ts b/packages/simulation-sdk/src/handlers/dispatchers.ts
@@ -88,7 +88,7 @@ export const simulateOperation = (
 
 export const simulateOperations = (
   operations: Operation[],
-  startData: SimulationState,
+  startData: MaybeDraft<SimulationState>,
 ) =>
   handleOperations(
     operations,