feat: running balance calculations for optimistic proposals

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@across-protocol/sdk",
   "author": "UMA Team",
-  "version": "4.3.50",
+  "version": "4.3.51",
   "license": "AGPL-3.0",
   "homepage": "https://docs.across.to/reference/sdk",
   "files": [

src/clients/BundleDataClient/utils/DataworkerUtils.ts

@@ -10,6 +10,7 @@ import {
   PoolRebalanceLeaf,
   Refund,
   RunningBalances,
+  SpokePoolClientsByChain,
 } from "../../../interfaces";
 import {
   bnZero,
@@ -18,6 +19,9 @@ import {
   count2DDictionaryValues,
   count3DDictionaryValues,
   toAddressType,
+  getImpliedBundleBlockRanges,
+  isDefined,
+  EvmAddress,
 } from "../../../utils";
 import {
   addLastRunningBalance,
@@ -28,6 +32,7 @@ import {
 } from "./PoolRebalanceUtils";
 import { AcrossConfigStoreClient } from "../../AcrossConfigStoreClient";
 import { HubPoolClient } from "../../HubPoolClient";
+import { BundleDataClient } from "../../BundleDataClient";
 import { buildPoolRebalanceLeafTree } from "./MerkleTreeUtils";
 
 // and expired deposits.
@@ -114,17 +119,258 @@ export function getEndBlockBuffers(
   return chainIdListForBundleEvaluationBlockNumbers.map((chainId: number) => blockRangeEndBlockBuffer[chainId] ?? 0);
 }
 
-export function _buildPoolRebalanceRoot(
+/*
+ * @notice Constructs a new pool rebalance root given the input bundle data.
+ * @dev It is assumed that the input bundle data corresponds to the block ranges of the input mainnetBundleEndBlock.
+ * If the mainnetBundleEndBlock does not correspond to any historical or pending root bundle, then the output pool rebalance
+ * root will be constructed under the assumption that the pending root bundle is valid and passes liveness.
+ * @param latestMainnetBlock The latest mainnet block number.
+ * @param mainnetBundleEndBlock The end block number of the block range corresponding to the bundle data.
+ * @param bundleV3Deposits Deposit bundle data for the implied block range given by the mainnetBundleEndBlock.
+ * @param bundleFillsV3 Fill bundle data.
+ * @param bundleSlowFillsV3 Slow fill bundle data.
+ * @param unexecutableSlowFills Expired slow fill bundle data.
+ * @param expiredDepositsToRefundV3 Expired deposit bundle data.
+ * @param clients Clients required to construct a new pool rebalance root.
+ * @maxL1TokenCountOverride Optional parameter to cap the number of tokens in a single pool rebalance leaf.
+ */
+export async function _buildPoolRebalanceRoot(
   latestMainnetBlock: number,
   mainnetBundleEndBlock: number,
   bundleV3Deposits: BundleDepositsV3,
   bundleFillsV3: BundleFillsV3,
   bundleSlowFillsV3: BundleSlowFills,
   unexecutableSlowFills: BundleExcessSlowFills,
   expiredDepositsToRefundV3: ExpiredDepositsToRefundV3,
-  clients: { hubPoolClient: HubPoolClient; configStoreClient: AcrossConfigStoreClient },
+  clients: {
+    hubPoolClient: HubPoolClient;
+    configStoreClient: AcrossConfigStoreClient;
+    bundleDataClient: BundleDataClient;
+    spokePoolClients: SpokePoolClientsByChain;
+  },
+  maxL1TokenCountOverride?: number
+): Promise<PoolRebalanceRoot> {
+  // If there is a pending proposal and the mainnet bundle end block is greater than the pending proposal's mainnet end block, then this pool rebalance root is being built during the liveness
+  // of a different root bundle, so running balance calculations will be slightly different.
+  if (
+    clients.hubPoolClient.hasPendingProposal() &&
+    clients.hubPoolClient.getPendingRootBundle()!.bundleEvaluationBlockNumbers[0] < mainnetBundleEndBlock
+  ) {
+    return await _buildOptimisticPoolRebalanceRoot(
+      latestMainnetBlock,
+      mainnetBundleEndBlock,
+      bundleV3Deposits,
+      bundleFillsV3,
+      bundleSlowFillsV3,
+      unexecutableSlowFills,
+      expiredDepositsToRefundV3,
+      clients,
+      maxL1TokenCountOverride
+    );
+  }
+  // Otherwise, we can synchronously reconstruct a historical pool rebalance root from the input data.
+  return _buildHistoricalPoolRebalanceRoot(
+    latestMainnetBlock,
+    mainnetBundleEndBlock,
+    bundleV3Deposits,
+    bundleFillsV3,
+    bundleSlowFillsV3,
+    unexecutableSlowFills,
+    expiredDepositsToRefundV3,
+    clients,
+    maxL1TokenCountOverride
+  );
+}
+
+/*
+ * @notice Constructs a new pool rebalance root given historical bundle data.
+ * @param latestMainnetBlock The latest mainnet block number.
+ * @param mainnetBundleEndBlock The end block number of the block range corresponding to the bundle data.
+ * @param bundleV3Deposits Deposit bundle data for the implied block range given by the mainnetBundleEndBlock.
+ * @param bundleFillsV3 Fill bundle data.
+ * @param bundleSlowFillsV3 Slow fill bundle data.
+ * @param unexecutableSlowFills Expired slow fill bundle data.
+ * @param expiredDepositsToRefundV3 Expired deposit bundle data.
+ * @param clients Clients required to construct a new pool rebalance root.
+ * @maxL1TokenCountOverride Optional parameter to cap the number of tokens in a single pool rebalance leaf.
+ */
+export function _buildHistoricalPoolRebalanceRoot(
+  latestMainnetBlock: number,
+  mainnetBundleEndBlock: number,
+  bundleV3Deposits: BundleDepositsV3,
+  bundleFillsV3: BundleFillsV3,
+  bundleSlowFillsV3: BundleSlowFills,
+  unexecutableSlowFills: BundleExcessSlowFills,
+  expiredDepositsToRefundV3: ExpiredDepositsToRefundV3,
+  clients: {
+    hubPoolClient: HubPoolClient;
+    configStoreClient: AcrossConfigStoreClient;
+  },
   maxL1TokenCountOverride?: number
 ): PoolRebalanceRoot {
+  const { runningBalances, realizedLpFees, chainWithRefundsOnly } = _getMarginalRunningBalances(
+    mainnetBundleEndBlock,
+    bundleV3Deposits,
+    bundleFillsV3,
+    bundleSlowFillsV3,
+    unexecutableSlowFills,
+    expiredDepositsToRefundV3,
+    clients
+  );
+  addLastRunningBalance(latestMainnetBlock, runningBalances, clients.hubPoolClient);
+  const leaves: PoolRebalanceLeaf[] = constructPoolRebalanceLeaves(
+    mainnetBundleEndBlock,
+    runningBalances,
+    realizedLpFees,
+    Array.from(chainWithRefundsOnly).filter((chainId) => !Object.keys(runningBalances).includes(chainId.toString())),
+    clients.configStoreClient,
+    maxL1TokenCountOverride
+  );
+  return {
+    runningBalances,
+    realizedLpFees,
+    leaves,
+    tree: buildPoolRebalanceLeafTree(leaves),
+  };
+}
+
+/*
+ * @notice Constructs a new pool rebalance root given the input bundle data. This function assumes there is a pending root bundle which will eventually clear liveness.
+ * @param latestMainnetBlock The latest mainnet block number.
+ * @param mainnetBundleEndBlock The end block number of the block range corresponding to the bundle data.
+ * @param bundleV3Deposits Deposit bundle data for the implied block range given by the mainnetBundleEndBlock.
+ * @param bundleFillsV3 Fill bundle data.
+ * @param bundleSlowFillsV3 Slow fill bundle data.
+ * @param unexecutableSlowFills Expired slow fill bundle data.
+ * @param expiredDepositsToRefundV3 Expired deposit bundle data.
+ * @param clients Clients required to construct a new pool rebalance root.
+ * @maxL1TokenCountOverride Optional parameter to cap the number of tokens in a single pool rebalance leaf.
+ */
+export async function _buildOptimisticPoolRebalanceRoot(
+  latestMainnetBlock: number,
+  mainnetBundleEndBlock: number,
+  bundleV3Deposits: BundleDepositsV3,
+  bundleFillsV3: BundleFillsV3,
+  bundleSlowFillsV3: BundleSlowFills,
+  unexecutableSlowFills: BundleExcessSlowFills,
+  expiredDepositsToRefundV3: ExpiredDepositsToRefundV3,
+  clients: {
+    hubPoolClient: HubPoolClient;
+    configStoreClient: AcrossConfigStoreClient;
+    bundleDataClient: BundleDataClient;
+    spokePoolClients: SpokePoolClientsByChain;
+  },
+  maxL1TokenCountOverride?: number
+): Promise<PoolRebalanceRoot> {
+  const { runningBalances, realizedLpFees, chainWithRefundsOnly } = _getMarginalRunningBalances(
+    mainnetBundleEndBlock,
+    bundleV3Deposits,
+    bundleFillsV3,
+    bundleSlowFillsV3,
+    unexecutableSlowFills,
+    expiredDepositsToRefundV3,
+    clients
+  );
+  // Get the pool rebalance root for the pending bundle so that we may account for its calculated running balances.
+  // @dev It is safe to index the hub pool client's proposed root bundles here since there is guaranteed to be a pending proposal in this code block.
+  const mostRecentProposedRootBundle = clients.hubPoolClient.getLatestProposedRootBundle();
+  const blockRangesForChains = getImpliedBundleBlockRanges(
+    clients.hubPoolClient,
+    clients.configStoreClient,
+    mostRecentProposedRootBundle
+  );
+  // We are loading data from a pending root bundle which should be well into liveness, so we want to use arweave if possible.
+  const pendingRootBundleData = await clients.bundleDataClient.loadData(
+    blockRangesForChains,
+    clients.spokePoolClients,
+    true
+  );
+  // Build the pool rebalance root for the pending root bundle.
+  const { leaves, tree } = _buildHistoricalPoolRebalanceRoot(
+    latestMainnetBlock,
+    blockRangesForChains[0][1],
+    pendingRootBundleData.bundleDepositsV3,
+    pendingRootBundleData.bundleFillsV3,
+    pendingRootBundleData.bundleSlowFillsV3,
+    pendingRootBundleData.unexecutableSlowFills,
+    pendingRootBundleData.expiredDepositsToRefundV3,
+    clients,
+    maxL1TokenCountOverride
+  );
+  // Assert that the rebuilt pool rebalance root matches the pending root bundle's value. If it does not, then we likely misconstructed the pending root bundle and should throw.
+  assert(tree.getHexRoot() === mostRecentProposedRootBundle.poolRebalanceRoot);
+
+  // Only add marginal pending running balances if there is already an entry in `runningBalances`. If there is no entry in `runningBalances`, then
+  // The running balance for this entry was unchanged since the last root bundle.
+  Object.keys(runningBalances).forEach((_repaymentChainId) => {
+    Object.keys(runningBalances[Number(_repaymentChainId)]).forEach((_l1TokenAddress) => {
+      const repaymentChainId = Number(_repaymentChainId);
+      const l1TokenAddress = EvmAddress.from(_l1TokenAddress);
+      const pendingPoolRebalanceLeaf = leaves.find(
+        (leaf) => leaf.chainId === repaymentChainId && leaf.l1Tokens.some((l1Token) => l1Token.eq(l1TokenAddress))
+      );
+      // If the pending pool rebalance root has running balances defined, then add it to `runningBalances`.
+      if (isDefined(pendingPoolRebalanceLeaf)) {
+        const pendingLeafTokenIdx = pendingPoolRebalanceLeaf.l1Tokens.findIndex((l1Token) =>
+          l1Token.eq(l1TokenAddress)
+        );
+        assert(pendingLeafTokenIdx !== -1);
+        const pendingRunningBalanceAmount = pendingPoolRebalanceLeaf.runningBalances[pendingLeafTokenIdx];
+        if (!pendingRunningBalanceAmount.eq(bnZero)) {
+          updateRunningBalance(runningBalances, repaymentChainId, _l1TokenAddress, pendingRunningBalanceAmount);
+        }
+      } else {
+        // Otherwise, add the last running balance for this token.
+        const { runningBalance: lastExecutedBundleRunningBalance } =
+          clients.hubPoolClient.getRunningBalanceBeforeBlockForChain(
+            latestMainnetBlock,
+            repaymentChainId,
+            l1TokenAddress
+          );
+        if (!lastExecutedBundleRunningBalance.eq(bnZero)) {
+          updateRunningBalance(runningBalances, repaymentChainId, _l1TokenAddress, lastExecutedBundleRunningBalance);
+        }
+      }
+    });
+  });
+  const poolRebalanceLeaves: PoolRebalanceLeaf[] = constructPoolRebalanceLeaves(
+    mainnetBundleEndBlock,
+    runningBalances,
+    realizedLpFees,
+    Array.from(chainWithRefundsOnly).filter((chainId) => !Object.keys(runningBalances).includes(chainId.toString())),
+    clients.configStoreClient,
+    maxL1TokenCountOverride
+  );
+  return {
+    runningBalances,
+    realizedLpFees,
+    leaves: poolRebalanceLeaves,
+    tree: buildPoolRebalanceLeafTree(poolRebalanceLeaves),
+  };
+}
+
+/*
+ * @notice Gets the running balance amounts derived from the input bundle data.
+ * @param mainnetBundleEndBlock The end block number of the block range corresponding to the bundle data.
+ * @param bundleV3Deposits Deposit bundle data for the implied block range given by the mainnetBundleEndBlock.
+ * @param bundleFillsV3 Fill bundle data.
+ * @param bundleSlowFillsV3 Slow fill bundle data.
+ * @param unexecutableSlowFills Expired slow fill bundle data.
+ * @param expiredDepositsToRefundV3 Expired deposit bundle data.
+ * @param clients Clients required to construct a new pool rebalance root.
+ */
+export function _getMarginalRunningBalances(
+  mainnetBundleEndBlock: number,
+  bundleV3Deposits: BundleDepositsV3,
+  bundleFillsV3: BundleFillsV3,
+  bundleSlowFillsV3: BundleSlowFills,
+  unexecutableSlowFills: BundleExcessSlowFills,
+  expiredDepositsToRefundV3: ExpiredDepositsToRefundV3,
+  clients: {
+    hubPoolClient: HubPoolClient;
+    configStoreClient: AcrossConfigStoreClient;
+  }
+): { chainWithRefundsOnly: Set<number>; realizedLpFees: RunningBalances; runningBalances: RunningBalances } {
   // Running balances are the amount of tokens that we need to send to each SpokePool to pay for all instant and
   // slow relay refunds. They are decreased by the amount of funds already held by the SpokePool. Balances are keyed
   // by the SpokePool's network and L1 token equivalent of the L2 token to refund.
@@ -296,24 +542,5 @@ export function _buildPoolRebalanceRoot(
       });
     });
   });
-
-  // Add to the running balance value from the last valid root bundle proposal for {chainId, l1Token}
-  // combination if found.
-  addLastRunningBalance(latestMainnetBlock, runningBalances, clients.hubPoolClient);
-
-  const leaves: PoolRebalanceLeaf[] = constructPoolRebalanceLeaves(
-    mainnetBundleEndBlock,
-    runningBalances,
-    realizedLpFees,
-    Array.from(chainWithRefundsOnly).filter((chainId) => !Object.keys(runningBalances).includes(chainId.toString())),
-    clients.configStoreClient,
-    maxL1TokenCountOverride
-  );
-
-  return {
-    runningBalances,
-    realizedLpFees,
-    leaves,
-    tree: buildPoolRebalanceLeafTree(leaves),
-  };
+  return { runningBalances, chainWithRefundsOnly, realizedLpFees };
 }