Rebalance WETH and USDC to a target share

// Rebalances the Pot's WETH and USDC toward `targetBps` of their combined value in WETH, priced by
// Uniswap V3's 0.05% WETH/USDC pool. Once WETH's share is more than `bandBps` from the target, it
// sells the excess side on that pool, with a price limit `maxSlippageBps` away from the pool's
// price: a pool that reaches the limit fills part of the trade and the rest waits for the next
// run. Inside the band it trades nothing and returns normally. The cook reverts when the pool does
// not exist, the Pot holds neither token, or the trade's average price is past the limit plus the
// pool fee. Both tokens stay in the Pot.
//
//   sauce simulate rebalance.sauce.ts --chain <chain> --owner <address> \
//     --assume-balance WETH:<whole-tokens> --assume-balance USDC:<whole-tokens> \
//     --args <targetBps> <bandBps> <maxSlippageBps>
//   sauce run rebalance.sauce.ts --chain <chain> --wallet <name> \
//     --args <targetBps> <bandBps> <maxSlippageBps>
//
//   targetBps: Uint256       WETH's share of the two by value, in hundredths of a percent:
//                            5000 is half
//   bandBps: Uint256         how far the share may drift before a trade, same units: 500 trades
//                            below 45% or above 55% at a 5000 target
//   maxSlippageBps: Uint256  how far the trade may move the pool's price, in hundredths of a
//                            percent: 30 is 0.3%
//
// On Base, half and half, trading once the share leaves 45-55%, moving the price at most 0.3%:
//
//   sauce run rebalance.sauce.ts --chain base --wallet me --args 5000 500 30
import { uniswapV3 } from "@sauce/router";

// USDC base units per WETH at a pool sqrt price (token1 per token0, x 2^96).
function priceAt(sqrtPriceX96: Uint256): Uint256 {
  const ratio = Math.mulDiv(sqrtPriceX96, sqrtPriceX96, 1 << 96);
  return WETH < USDC
    ? Math.mulDiv(ratio, 10 ** 18, 1 << 96)
    : Math.mulDiv(10 ** 18, 1 << 96, ratio);
}

// The inverse: the pool sqrt price at `price` USDC base units per WETH.
function sqrtPriceAt(price: Uint256): Uint256 {
  return WETH < USDC
    ? Math.sqrt(Math.mulDiv(price, 1 << 192, 10 ** 18))
    : Math.sqrt(Math.mulDiv(1 << 192, 10 ** 18, price));
}

// WETH's share of the two balances' value at `price`, in hundredths of a percent.
function wethShare(weth: Uint256, usdc: Uint256, price: Uint256): Uint256 {
  const wethValue = Math.mulDiv(weth, price, 10 ** 18);
  return Math.mulDiv(wethValue, 10000, wethValue + usdc);
}

function main(
  targetBps: Uint256,
  bandBps: Uint256,
  maxSlippageBps: Uint256,
): { before: Uint256; after: Uint256; soldWeth: bool; amountIn: Uint256; amountOut: Uint256 } {
  require(targetBps <= 10000 && maxSlippageBps < 10000, "a share above 100%");
  const pool = UniswapV3.Factory.getPool(WETH, USDC, 500);
  require(pool != 0, "no WETH/USDC pool at 0.05%");
  const price = priceAt(UniswapV3.Pool(pool).slot0().sqrtPriceX96);
  const weth = pot.WETH;
  const usdc = pot.USDC;
  require(weth > 0 || usdc > 0, "the Pot holds neither");

  const before = wethShare(weth, usdc, price);
  const soldWeth = before > targetBps;
  if ((soldWeth ? before - targetBps : targetBps - before) <= bandBps) {
    return { before, after: before, soldWeth, amountIn: 0, amountOut: 0 };
  }
  // The value to move, in USDC, to bring WETH to its target share.
  const wethValue = Math.mulDiv(weth, price, 10 ** 18);
  const wanted = Math.mulDiv(wethValue + usdc, targetBps, 10000);
  const fee = UniswapV3.Pool(pool).fee();
  let amountIn: Uint256 = 0;
  let amountOut: Uint256 = 0;
  if (soldWeth) {
    // Sell WETH no lower than `floor`: every unit fetches at least that, less the pool fee.
    const floor = Math.mulDiv(price, 10000 - maxSlippageBps, 10000);
    const sell = Math.mulDiv(wethValue - wanted, 10 ** 18, price);
    uniswapV3(pool, WETH, USDC, sell, sqrtPriceAt(floor), 0, pot);
    amountIn = weth - pot.WETH;
    amountOut = pot.USDC - usdc;
    const least = Math.mulDiv(Math.mulDiv(amountIn, floor, 10 ** 18), 1000000 - fee, 1000000);
    require(amountOut >= least, "sold below the floor");
  } else {
    // Buy WETH no higher than `ceiling`: no unit costs more than that, plus the pool fee.
    const ceiling = Math.mulDiv(price, 10000 + maxSlippageBps, 10000);
    uniswapV3(pool, USDC, WETH, wanted - wethValue, sqrtPriceAt(ceiling), 0, pot);
    amountIn = usdc - pot.USDC;
    amountOut = pot.WETH - weth;
    const most = Math.mulDiv(Math.mulDiv(amountOut, ceiling, 10 ** 18), 1000000, 1000000 - fee);
    require(amountIn <= most + 1, "bought above the ceiling");
  }
  const after = wethShare(pot.WETH, pot.USDC, priceAt(UniswapV3.Pool(pool).slot0().sqrtPriceX96));
  return { before, after, soldWeth, amountIn, amountOut };
}