enable flat portion option on swaps

src/entities/protocols/uniswap.ts

@@ -20,12 +20,18 @@ import { Currency, TradeType, CurrencyAmount, Percent } from '@uniswap/sdk-core'
 import { Command, RouterTradeType, TradeConfig } from '../Command'
 import { SENDER_AS_RECIPIENT, ROUTER_AS_RECIPIENT, CONTRACT_BALANCE } from '../../utils/constants'
 import { encodeFeeBips } from '../../utils/numbers'
-import { BigNumber } from 'ethers'
+import { BigNumber, BigNumberish } from 'ethers'
+
+export type FlatFeeOptions = {
+  amount: BigNumberish
+  recipient: string
+}
 
 // the existing router permit object doesn't include enough data for permit2
 // so we extend swap options with the permit2 permit
 export type SwapOptions = Omit<RouterSwapOptions, 'inputTokenPermit'> & {
   inputTokenPermit?: Permit2Permit
+  flatFee?: FlatFeeOptions
 }
 
 const REFUND_ETH_PRICE_IMPACT_THRESHOLD = new Percent(50, 100)
@@ -40,7 +46,9 @@ interface Swap<TInput extends Currency, TOutput extends Currency> {
 // also translates trade objects from previous (v2, v3) SDKs
 export class UniswapTrade implements Command {
   readonly tradeType: RouterTradeType = RouterTradeType.UniswapTrade
-  constructor(public trade: RouterTrade<Currency, Currency, TradeType>, public options: SwapOptions) {}
+  constructor(public trade: RouterTrade<Currency, Currency, TradeType>, public options: SwapOptions) {
+    if (!!options.fee && !!options.flatFee) throw new Error('Only one fee option permitted')
+  }
 
   encode(planner: RoutePlanner, _config: TradeConfig): void {
     let payerIsUser = true
@@ -66,7 +74,7 @@ export class UniswapTrade implements Command {
       this.trade.tradeType === TradeType.EXACT_INPUT && this.trade.routes.length > 2
     const outputIsNative = this.trade.outputAmount.currency.isNative
     const inputIsNative = this.trade.inputAmount.currency.isNative
-    const routerMustCustody = performAggregatedSlippageCheck || outputIsNative || !!this.options.fee
+    const routerMustCustody = performAggregatedSlippageCheck || outputIsNative || hasFeeOption(this.options)
 
     for (const swap of this.trade.swaps) {
       switch (swap.route.protocol) {
@@ -107,6 +115,25 @@ export class UniswapTrade implements Command {
         }
       }
 
+      // If there is a flat fee, that absolute amount is sent to the fee recipient
+      // In the case where ETH is the output currency, the fee is taken in WETH (for gas reasons)
+      if (!!this.options.flatFee) {
+        const feeAmount = this.options.flatFee.amount
+        if (minimumAmountOut.lt(feeAmount)) throw new Error('Flat fee amount greater than minimumAmountOut')
+
+        planner.addCommand(CommandType.TRANSFER, [
+          this.trade.outputAmount.currency.wrapped.address,
+          this.options.flatFee.recipient,
+          feeAmount,
+        ])
+
+        // If the trade is exact output, and a fee was taken, we must adjust the amount out to be the amount after the fee
+        // Otherwise we continue as expected with the trade's normal expected output
+        if (this.trade.tradeType === TradeType.EXACT_OUTPUT) {
+          minimumAmountOut = minimumAmountOut.sub(feeAmount)
+        }
+      }
+
       // The remaining tokens that need to be sent to the user after the fee is taken will be caught
       // by this if-else clause.
       if (outputIsNative) {
@@ -289,3 +316,7 @@ function addMixedSwap<TInput extends Currency, TOutput extends Currency>(
 function riskOfPartialFill(trade: RouterTrade<Currency, Currency, TradeType>): boolean {
   return trade.priceImpact.greaterThan(REFUND_ETH_PRICE_IMPACT_THRESHOLD)
 }
+
+function hasFeeOption(swapOptions: SwapOptions): boolean {
+  return !!swapOptions.fee || !!swapOptions.flatFee
+}

test/forge/SwapERC20CallParameters.t.sol

@@ -244,6 +244,46 @@ contract SwapERC20CallParametersTest is Test, Interop, DeployRouter {
         assertEq(address(router).balance, 0);
     }
 
+    function testV2ExactOutputSingleNativeInputWithFlatFee() public {
+        MethodParameters memory params = readFixture(json, "._UNISWAP_V2_ETH_FOR_1000_USDC_WITH_FLAT_FEE");
+
+        uint256 outputAmount = 1000 * ONE_USDC;
+        uint256 feeAmount = 50 * ONE_USDC;
+
+        assertEq(from.balance, BALANCE);
+        assertEq(USDC.balanceOf(RECIPIENT), 0);
+        assertEq(USDC.balanceOf(FEE_RECIPIENT), 0);
+
+        (bool success,) = address(router).call{value: params.value}(params.data);
+        require(success, "call failed");
+        assertLe(from.balance, BALANCE - params.value);
+        assertEq(USDC.balanceOf(RECIPIENT), outputAmount);
+        assertEq(USDC.balanceOf(FEE_RECIPIENT), feeAmount);
+        assertEq(WETH.balanceOf(address(router)), 0);
+        assertEq(address(router).balance, 0);
+    }
+
+    function testV2ExactOutputSingleNativeOutputWithFlatFee() public {
+        MethodParameters memory params = readFixture(json, "._UNISWAP_V2_USCD_FOR_10_ETH_WITH_FLAT_FEE");
+
+        deal(address(USDC), from, BALANCE);
+        USDC.approve(address(permit2), BALANCE);
+        permit2.approve(address(USDC), address(router), uint160(BALANCE), uint48(block.timestamp + 1000));
+
+        uint256 outputAmount = 10 ether;
+        uint256 feeAmount = 5 ether;
+
+        assertEq(WETH.balanceOf(FEE_RECIPIENT), 0);
+        uint256 recipientBalanceBefore = RECIPIENT.balance;
+
+        (bool success,) = address(router).call{value: params.value}(params.data);
+        require(success, "call failed");
+        assertGt(RECIPIENT.balance - recipientBalanceBefore, outputAmount); // tiny imprecision with exactOut
+        assertEq(WETH.balanceOf(FEE_RECIPIENT), feeAmount);
+        assertEq(WETH.balanceOf(address(router)), 0);
+        assertEq(address(router).balance, 0);
+    }
+
     function testV2ExactOutputSingleERC20() public {
         MethodParameters memory params = readFixture(json, "._UNISWAP_V2_USDC_FOR_1_ETH");
 
@@ -292,6 +332,26 @@ contract SwapERC20CallParametersTest is Test, Interop, DeployRouter {
         assertGt(totalOut, 1000 * ONE_USDC);
     }
 
+    function testV3ExactInputSingleNativeWithFlatFee() public {
+        MethodParameters memory params = readFixture(json, "._UNISWAP_V3_1_ETH_FOR_USDC_WITH_FLAT_FEE");
+
+        assertEq(from.balance, BALANCE);
+        assertEq(USDC.balanceOf(RECIPIENT), 0);
+        assertEq(USDC.balanceOf(FEE_RECIPIENT), 0);
+
+        (bool success,) = address(router).call{value: params.value}(params.data);
+        require(success, "call failed");
+        assertLe(from.balance, BALANCE - params.value);
+
+        uint256 recipientBalance = USDC.balanceOf(RECIPIENT);
+        uint256 feeRecipientBalance = USDC.balanceOf(FEE_RECIPIENT);
+        uint256 totalOut = recipientBalance + feeRecipientBalance;
+        uint256 expectedFee = 50 * ONE_USDC;
+        assertEq(feeRecipientBalance, expectedFee);
+        assertEq(recipientBalance, totalOut - expectedFee);
+        assertGt(totalOut, 1000 * ONE_USDC);
+    }
+
     function testV3ExactInputSingleERC20() public {
         MethodParameters memory params = readFixture(json, "._UNISWAP_V3_1000_USDC_FOR_ETH");