viem-tx-sim 0.1.1 → 0.2.1

package/src/txSimulator.ts

@@ -1,10 +1,16 @@
+import type { Address } from "viem";
+
 import { DEFAULT_SIMULATION_GAS_LIMIT } from "./constants.js";
 import { InvalidSimulationInputError } from "./errors.js";
-import { estimateAssetRequirements } from "./internal/requirements.js";
+import { discoverErc20s, forUserBalanceQueries } from "./internal/queryDiscovery.js";
+import { estimateTokenOverrideRequirements } from "./internal/requirements.js";
 import { blockOptionsSpread, type ClientArgs } from "./internal/rpc.js";
-import { discoverCandidateAddresses, runSimulator } from "./internal/simulator.js";
-import { prepareAllowanceOverrides, prepareBalanceOverrides } from "./internal/slots.js";
+import { runSimulator } from "./internal/simulator.js";
+import { prepareAllowanceTokenOverrides, prepareBalanceTokenOverrides } from "./internal/slots.js";
 import type {
+  BalanceDelta,
+  BalanceQuery,
+  ForUserBalanceQueriesArgs,
   PreparedAllowanceOverrides,
   PreparedBalanceOverrides,
   PrepareAllowanceOverridesArgs,
@@ -32,15 +38,13 @@ type BoundCallDefaults = {
  */
 export interface TxSimulator {
   /**
-   * Simulates one call or sequential batch and returns raw native/token balance deltas.
+   * Simulates one call or sequential batch and returns requested balance deltas.
    *
-   * This uses `eth_createAccessList` for candidate discovery and one `eth_call` with state
-   * overrides that injects the simulator at `from`. It does not automatically forge balances or
-   * allowances; preparation methods return ready-to-use `tokenSlotOverrides` for view-only or
-   * unfunded accounts. Transaction reverts return `status: "reverted"` instead of throwing.
+   * This performs one `eth_call` with state overrides that inject the simulator at `from`.
+   * Balances are observed only for `balanceQueries`; query the tokens you forge if you want to
+   * observe them. Transaction reverts return `status: "reverted"` instead of throwing.
    *
    * @throws InvalidSimulationInputError when `calls` is empty.
-   * @throws AccessListUnsupportedError when the RPC endpoint cannot provide access lists.
    * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides or
    * returns undecodable simulator output.
    *
@@ -50,50 +54,75 @@ export interface TxSimulator {
    * const result = await sim.simulate({
    *   from,
    *   calls: [{ to, data, value: 0n }],
+   *   balanceQueries: [{ asset: "native", account: from }],
    * });
    * ```
    */
   simulate: (args: SimulateArgs) => Promise<SimulationResult>;
 
-  /**
-   * Prepares ERC-20 balance overrides for `from`.
-   *
-   * Each token is probed with RPC-only access lists and sentinel state overrides. Tokens the
-   * simulator cannot `deal` by verified storage write are returned in `unresolved` rather than
-   * thrown.
-   *
-   * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides.
-   */
-  prepareBalanceOverrides: (args: PrepareBalanceOverridesArgs) => Promise<PreparedBalanceOverrides>;
+  readonly balanceQueries: {
+    /**
+     * Discovers wallet-style balance queries for `from`.
+     *
+     * This runs access-list candidate discovery, then one token-filter `eth_call`, and returns
+     * native plus token balance queries for `from`. Pass the result to `simulate`.
+     *
+     * @throws AccessListUnsupportedError when the RPC endpoint cannot provide access lists.
+     * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides.
+     */
+    forUser: (args: ForUserBalanceQueriesArgs) => Promise<BalanceQuery[]>;
 
-  /**
-   * Prepares ERC-20 allowance overrides for `from` and the requested token/spender pairs.
-   *
-   * Standard Solidity allowance layouts are inferred after one verified probe per token where
-   * possible; non-standard layouts fall back to per-pair probing. Pairs the simulator cannot `deal`
-   * via verified storage write are returned in `unresolved` rather than thrown.
-   *
-   * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides.
-   */
-  prepareAllowanceOverrides: (
-    args: PrepareAllowanceOverridesArgs,
-  ) => Promise<PreparedAllowanceOverrides>;
+    /**
+     * Discovers ERC-20 contracts touched by the calls that answer `balanceOf(from)`.
+     *
+     * This is the discovery half of `forUser`; map the returned addresses yourself when observing a
+     * different account.
+     *
+     * @throws AccessListUnsupportedError when the RPC endpoint cannot provide access lists.
+     * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides.
+     */
+    discoverErc20s: (args: ForUserBalanceQueriesArgs) => Promise<Address[]>;
+  };
 
-  /**
-   * Estimates the balances and approvals needed to execute the observed path.
-   *
-   * Use this when the tokens or spenders are not known ahead of time. Returned amounts are estimated
-   * under forged balances/allowances and should be padded before display or transaction assembly;
-   * unreliable measurements are reported under `unresolved`.
-   *
-   * @throws InvalidSimulationInputError when `calls` is empty.
-   * @throws AccessListUnsupportedError when the RPC endpoint cannot provide access lists.
-   * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides or
-   * returns undecodable simulator output.
-   */
-  estimateAssetRequirements: (
-    args: EstimateAssetRequirementsArgs,
-  ) => Promise<EstimatedAssetRequirements>;
+  readonly tokenOverrides: {
+    /**
+     * Prepares ERC-20 balance overrides for `from`.
+     *
+     * Each token is probed with RPC-only access lists and sentinel state overrides. Tokens the
+     * simulator cannot `deal` by verified storage write are returned in `unresolved` rather than
+     * thrown.
+     *
+     * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides.
+     */
+    forBalances: (args: PrepareBalanceOverridesArgs) => Promise<PreparedBalanceOverrides>;
+
+    /**
+     * Prepares ERC-20 allowance overrides for `from` and the requested token/spender pairs.
+     *
+     * Standard Solidity allowance layouts are inferred after one verified probe per token where
+     * possible; non-standard layouts fall back to per-pair probing. Pairs the simulator cannot
+     * `deal` via verified storage write are returned in `unresolved` rather than thrown.
+     *
+     * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides.
+     */
+    forAllowances: (args: PrepareAllowanceOverridesArgs) => Promise<PreparedAllowanceOverrides>;
+
+    /**
+     * Estimates the balances and approvals needed to execute the observed path.
+     *
+     * Use this when the tokens or spenders are not known ahead of time. Returned amounts are
+     * estimated under forged balances/allowances and should be padded before display or transaction
+     * assembly; unreliable measurements are reported under `unresolved`.
+     *
+     * @throws InvalidSimulationInputError when `calls` is empty.
+     * @throws AccessListUnsupportedError when the RPC endpoint cannot provide access lists.
+     * @throws StateOverrideUnsupportedError when the RPC endpoint cannot execute state overrides or
+     * returns undecodable simulator output.
+     */
+    estimateRequirements: (
+      args: EstimateAssetRequirementsArgs,
+    ) => Promise<EstimatedAssetRequirements>;
+  };
 }
 
 /** Factory for {@link TxSimulator} instances bound to one viem public client. */
@@ -131,12 +160,24 @@ export const TxSimulator = {
 
     return {
       simulate: (args) => runSimulate({ ...args, ...revertDefaults(args), client: bound.client }),
-      prepareBalanceOverrides: (args) =>
-        prepareBalanceOverrides({ ...args, ...defaults(args), client: bound.client }),
-      prepareAllowanceOverrides: (args) =>
-        prepareAllowanceOverrides({ ...args, ...defaults(args), client: bound.client }),
-      estimateAssetRequirements: (args) =>
-        estimateAssetRequirements({ ...args, ...revertDefaults(args), client: bound.client }),
+      balanceQueries: {
+        forUser: (args) =>
+          forUserBalanceQueries({ ...args, ...defaults(args), client: bound.client }),
+        discoverErc20s: (args) =>
+          discoverErc20s({ ...args, ...defaults(args), client: bound.client }),
+      },
+      tokenOverrides: {
+        forBalances: (args) =>
+          prepareBalanceTokenOverrides({ ...args, ...defaults(args), client: bound.client }),
+        forAllowances: (args) =>
+          prepareAllowanceTokenOverrides({ ...args, ...defaults(args), client: bound.client }),
+        estimateRequirements: (args) =>
+          estimateTokenOverrideRequirements({
+            ...args,
+            ...revertDefaults(args),
+            client: bound.client,
+          }),
+      },
     };
   },
 };
@@ -151,26 +192,86 @@ async function runSimulate(args: SimulateArgs & ClientArgs): Promise<SimulationR
     data: call.data,
     value: call.value ?? 0n,
   })) satisfies SimulatedCall[];
-  const candidateAddresses = await discoverCandidateAddresses({
-    client: args.client,
-    from: args.from,
-    calls,
-    ...blockOptionsSpread(args),
-    gas: args.gas,
-    ...(args.debug !== undefined ? { debug: args.debug } : {}),
-  });
-
   const tokenSlotOverrides = args.tokenSlotOverrides ?? [];
 
-  return runSimulator({
+  const result = await runSimulator({
     client: args.client,
     from: args.from,
     calls,
-    candidates: [...candidateAddresses, ...tokenSlotOverrides.map((slot) => slot.token)],
+    candidates: [],
     tokenSlotOverrides,
+    extraStateOverrides: (args.nativeBalanceOverrides ?? []).map((override) => ({
+      address: override.account,
+      balance: override.amount,
+    })),
+    balanceProbes: args.balanceQueries.map((query) => ({
+      token: query.asset,
+      account: query.account,
+    })),
     debug: args.debug,
     ...blockOptionsSpread(args),
     gas: args.gas,
     ...(args.errorAbi !== undefined ? { errorAbi: args.errorAbi } : {}),
   });
+  const balances = buildBalanceResults(args.balanceQueries, result.probeData, calls.length);
+
+  if (result.status === "reverted") {
+    return {
+      status: "reverted",
+      ...balances,
+      revertData: result.revertData,
+      ...(result.revertReason !== undefined ? { revertReason: result.revertReason } : {}),
+      ...(result.revertError !== undefined ? { revertError: result.revertError } : {}),
+      ...(result.revertSelector !== undefined ? { revertSelector: result.revertSelector } : {}),
+      failingCallIndex: result.failingCallIndex,
+    };
+  }
+
+  return { status: "success", ...balances };
+}
+
+type BalanceResultFields = {
+  balanceDeltas: BalanceDelta[];
+  unresolved: BalanceQuery[];
+};
+
+function buildBalanceResults(
+  queries: readonly BalanceQuery[],
+  probeData: {
+    balanceCheckpoints: readonly bigint[];
+    balanceProbeOk: readonly boolean[];
+  },
+  callsLength: number,
+): BalanceResultFields {
+  const balanceDeltas: BalanceDelta[] = [];
+  const unresolved: BalanceQuery[] = [];
+  const stride = callsLength + 1;
+
+  for (let i = 0; i < queries.length; ++i) {
+    const query = queries[i];
+    if (query === undefined) continue;
+    if (probeData.balanceProbeOk[i] !== true) {
+      unresolved.push(query);
+      continue;
+    }
+    const base = i * stride;
+    const before = probeData.balanceCheckpoints[base] ?? 0n;
+    const after = probeData.balanceCheckpoints[base + callsLength] ?? 0n;
+    const byCall = Array.from(
+      { length: callsLength },
+      (_, callIndex) =>
+        (probeData.balanceCheckpoints[base + callIndex + 1] ?? 0n) -
+        (probeData.balanceCheckpoints[base + callIndex] ?? 0n),
+    );
+    balanceDeltas.push({
+      asset: query.asset,
+      account: query.account,
+      before,
+      after,
+      delta: after - before,
+      byCall,
+    });
+  }
+
+  return { balanceDeltas, unresolved };
 }

package/src/types.ts

@@ -10,6 +10,38 @@ export type SimulatedCall = {
   value?: bigint;
 };
 
+/** One balance to observe during simulation. `asset` is `"native"` or an ERC-20 address. */
+export type BalanceQuery = {
+  asset: "native" | Address;
+  /** Account whose balance is observed; this can be any address, not just `from`. */
+  account: Address;
+};
+
+/** Protocol-level native balance to set before simulating; no slot discovery needed. */
+export type NativeBalanceOverride = {
+  /** Account to fund: `from` or any other address such as a plugin or router. */
+  account: Address;
+  /** Balance in wei to set. */
+  amount: bigint;
+};
+
+/**
+ * Balance observation for one query. `before` is after `tokenSlotOverrides` are applied, so deltas
+ * describe simulated changes under the supplied state assumptions.
+ */
+export type BalanceDelta = {
+  asset: "native" | Address;
+  account: Address;
+  before: bigint;
+  after: bigint;
+  delta: bigint;
+  /**
+   * Signed change per call, index-aligned with `calls`. Sums to `delta`; on a
+   * revert, entries from the failing call onward are 0n.
+   */
+  byCall: readonly bigint[];
+};
+
 /** Structured event emitted before and after each RPC call when debug logging is enabled. */
 export type SimulationDebugEvent = {
   /** Lifecycle phase for the RPC operation. */
@@ -51,7 +83,7 @@ type SimulationOptions = {
 export type TokenSlotOverride = {
   /** Token contract whose storage should be overridden. */
   token: Address;
-  /** Storage slot to write. Usually prepared by `prepareBalanceOverrides` or `prepareAllowanceOverrides`. */
+  /** Storage slot to write. Usually prepared by `tokenOverrides.forBalances` or `tokenOverrides.forAllowances`. */
   slot: Hex;
   /** Value written to the slot. Must be below uint256 max. */
   amount: bigint;
@@ -100,13 +132,31 @@ export type SimulateArgs = SimulationOptions & {
   from: Address;
   /** One call or an ERC-5792-style sequential batch. Must contain at least one call. */
   calls: readonly SimulatedCall[];
-  /** Storage-slot overrides applied before simulating. Usually from override preparation. */
+  /** Balances to observe. Use `[]` to execute without balance observations. */
+  balanceQueries: readonly BalanceQuery[];
+  /**
+   * Storage-slot overrides applied before simulating. Query the tokens you forge if you want to
+   * observe them.
+   */
   tokenSlotOverrides?: readonly TokenSlotOverride[];
+  /**
+   * Native balance overrides applied before simulating. Duplicate accounts use the last amount.
+   * Query forged accounts if you want to observe them.
+   */
+  nativeBalanceOverrides?: readonly NativeBalanceOverride[];
   /** Additional error definitions for decoding this call's reverts; merged after the bound errorAbi. */
   errorAbi?: Abi;
 };
 
-/** Arguments for `TxSimulator.prepareBalanceOverrides`. */
+/** Arguments for `TxSimulator.balanceQueries.forUser`. */
+export type ForUserBalanceQueriesArgs = SimulationOptions & {
+  /** Account whose wallet-style balance queries should be discovered. */
+  from: Address;
+  /** Calls whose access lists should be searched for token candidates. */
+  calls: readonly SimulatedCall[];
+};
+
+/** Arguments for `TxSimulator.tokenOverrides.forBalances`. */
 export type PrepareBalanceOverridesArgs = SimulationOptions & {
   /** Account whose token balance overrides should be prepared. */
   from: Address;
@@ -114,7 +164,7 @@ export type PrepareBalanceOverridesArgs = SimulationOptions & {
   tokens: readonly Address[];
 };
 
-/** Arguments for `TxSimulator.prepareAllowanceOverrides`. */
+/** Arguments for `TxSimulator.tokenOverrides.forAllowances`. */
 export type PrepareAllowanceOverridesArgs = SimulationOptions & {
   /** Account whose allowance overrides should be prepared. */
   from: Address;
@@ -122,7 +172,7 @@ export type PrepareAllowanceOverridesArgs = SimulationOptions & {
   pairs: readonly AllowanceSlotPair[];
 };
 
-/** Arguments for `TxSimulator.estimateAssetRequirements`. */
+/** Arguments for `TxSimulator.tokenOverrides.estimateRequirements`. */
 export type EstimateAssetRequirementsArgs = SimulationOptions & {
   /** Account whose balance and approval needs should be estimated. */
   from: Address;
@@ -204,26 +254,22 @@ export type EstimatedAssetRequirements =
   | EstimatedAssetRequirementsSuccess
   | EstimatedAssetRequirementsReverted;
 
-/** Raw balance delta for native ETH or an ERC-20-style `balanceOf(address)` asset. */
-export type AssetBalanceDelta = {
-  /** `"native"` for ETH, otherwise the token contract address. */
-  asset: "native" | Address;
-  /** Signed raw-unit balance change for `from`; negative means the account lost assets. */
-  delta: bigint;
-};
-
 /** Successful simulation result. */
 export type SimulationSuccess = {
   status: "success";
-  /** Non-zero raw balance deltas observed during the simulated execution. */
-  assetBalanceDeltas: AssetBalanceDelta[];
+  /** Balance observations mirrored 1:1 from successful queries, including zero deltas. */
+  balanceDeltas: BalanceDelta[];
+  /** Queries that could not be read, usually because an ERC-20 `balanceOf` staticcall failed. */
+  unresolved: BalanceQuery[];
 };
 
 /** Simulation result for a transaction revert; infrastructure failures throw typed errors instead. */
 export type SimulationReverted = {
   status: "reverted";
-  /** Non-zero raw balance deltas observed before the failing call. */
-  assetBalanceDeltas: AssetBalanceDelta[];
+  /** Balance observations mirrored 1:1 from successful queries, including zero deltas. */
+  balanceDeltas: BalanceDelta[];
+  /** Queries that could not be read, usually because an ERC-20 `balanceOf` staticcall failed. */
+  unresolved: BalanceQuery[];
   /** Raw EVM revert data from the failing simulated call. */
   revertData: Hex;
   /** Human-readable decoded revert; present when revertData decodes via supplied error definitions or as built-in Error/Panic. */