Move input validation to a higher level to simplify the individual implementations

src/entities/exit/index.ts

@@ -1 +1,2 @@
+export * from './poolExit';
 export * from './types';

src/entities/exit/poolExit.ts

@@ -0,0 +1,52 @@
+import {
+    BaseExit,
+    ExitBuildOutput,
+    ExitCallInput,
+    ExitConfig,
+    ExitInput,
+    ExitQueryResult,
+} from './types';
+import { WeightedExit } from './weighted/weightedExit';
+import { PoolStateInput } from '../types';
+import { validateInputs } from './weighted/validateInputs';
+import { getSortedTokens } from '../utils/getSortedTokens';
+
+export class PoolExit {
+    private readonly poolExits: Record<string, BaseExit> = {};
+
+    constructor(config?: ExitConfig) {
+        const { customPoolExits } = config || {};
+        this.poolExits = {
+            Weighted: new WeightedExit(),
+            // custom pool Exits take precedence over base Exits
+            ...customPoolExits,
+        };
+    }
+
+    public getExit(poolType: string): BaseExit {
+        if (!this.poolExits[poolType]) {
+            throw new Error('Unsupported pool type');
+        }
+
+        return this.poolExits[poolType];
+    }
+
+    public async query(
+        input: ExitInput,
+        poolState: PoolStateInput,
+    ): Promise<ExitQueryResult> {
+        validateInputs(input, poolState);
+
+        const sortedTokens = getSortedTokens(poolState.tokens, input.chainId);
+        const mappedPoolState = {
+            ...poolState,
+            tokens: sortedTokens,
+        };
+
+        return this.getExit(poolState.type).query(input, mappedPoolState);
+    }
+
+    public buildCall(input: ExitCallInput): ExitBuildOutput {
+        return this.getExit(input.poolType).buildCall(input);
+    }
+}

src/entities/exit/types.ts

@@ -40,6 +40,7 @@ export type ExitInput =
 
 // Returned from a exit query
 export type ExitQueryResult = {
+    poolType: string;
     id: Address;
     exitKind: ExitKind;
     bptIn: TokenAmount;
@@ -54,7 +55,7 @@ export type ExitCallInput = ExitQueryResult & {
     recipient: Address;
 };
 
-export type BuildOutput = {
+export type ExitBuildOutput = {
     call: Address;
     to: Address;
     value: bigint | undefined;
@@ -64,7 +65,7 @@ export type BuildOutput = {
 
 export interface BaseExit {
     query(input: ExitInput, poolState: PoolState): Promise<ExitQueryResult>;
-    buildCall(input: ExitCallInput): BuildOutput;
+    buildCall(input: ExitCallInput): ExitBuildOutput;
 }
 
 export type ExitConfig = {

src/entities/exit/weighted/validateInputs.ts

@@ -1,8 +1,8 @@
 import { ExitInput, ExitKind } from '..';
-import { PoolState } from '../../types';
+import { PoolStateInput } from '../../types';
 import { areTokensInArray } from '../../utils/areTokensInArray';
 
-export function validateInputs(input: ExitInput, poolState: PoolState) {
+export function validateInputs(input: ExitInput, poolState: PoolStateInput) {
     switch (input.kind) {
         case ExitKind.UNBALANCED:
             areTokensInArray(

src/entities/exit/weighted/weightedExit.ts

@@ -1,32 +1,32 @@
 import { encodeFunctionData } from 'viem';
-import { Token, TokenAmount, WeightedEncoder } from '../../..';
+import { Token } from '../../token';
+import { TokenAmount } from '../../tokenAmount';
+import { WeightedEncoder } from '../../encoders/weighted';
 import { Address } from '../../../types';
-import { BALANCER_VAULT, MAX_UINT256, ZERO_ADDRESS } from '../../../utils';
+import {
+    BALANCER_VAULT,
+    MAX_UINT256,
+    ZERO_ADDRESS,
+} from '../../../utils/constants';
 import { vaultAbi } from '../../../abi';
 import { parseExitArgs } from '../../utils/parseExitArgs';
 import {
     BaseExit,
-    BuildOutput,
+    ExitBuildOutput,
     ExitCallInput,
     ExitInput,
     ExitKind,
     ExitQueryResult,
 } from '../types';
-import { getSortedTokens } from '../../utils';
-import { PoolState, AmountsExit } from '../../types';
+import { AmountsExit, PoolState } from '../../types';
 import { doQueryExit } from '../../utils/doQueryExit';
-import { validateInputs } from './validateInputs';
 
 export class WeightedExit implements BaseExit {
     public async query(
         input: ExitInput,
         poolState: PoolState,
     ): Promise<ExitQueryResult> {
-        validateInputs(input, poolState);
-
-        const sortedTokens = getSortedTokens(poolState.tokens, input.chainId);
-
-        const amounts = this.getAmountsQuery(sortedTokens, input);
+        const amounts = this.getAmountsQuery(poolState.tokens, input);
 
         const userData = this.encodeUserData(input.kind, amounts);
 
@@ -35,7 +35,7 @@ export class WeightedExit implements BaseExit {
             chainId: input.chainId,
             exitWithNativeAsset: !!input.exitWithNativeAsset,
             poolId: poolState.id,
-            sortedTokens,
+            sortedTokens: poolState.tokens,
             sender: ZERO_ADDRESS,
             recipient: ZERO_ADDRESS,
             minAmountsOut: amounts.minAmountsOut,
@@ -57,6 +57,7 @@ export class WeightedExit implements BaseExit {
         );
 
         return {
+            poolType: poolState.type,
             exitKind: input.kind,
             id: poolState.id,
             bptIn,
@@ -94,7 +95,7 @@ export class WeightedExit implements BaseExit {
         }
     }
 
-    public buildCall(input: ExitCallInput): BuildOutput {
+    public buildCall(input: ExitCallInput): ExitBuildOutput {
         const amounts = this.getAmountsCall(input);
 
         const userData = this.encodeUserData(input.exitKind, amounts);

src/entities/join/index.ts

@@ -1,73 +1,2 @@
-import { TokenAmount } from '../tokenAmount';
-import { Slippage } from '../slippage';
-import { PoolState } from '../types';
-import { Address, Hex } from '../../types';
-
-export enum JoinKind {
-    Init = 'Init',
-    Unbalanced = 'Unbalanced',
-    SingleAsset = 'SingleAsset',
-    Proportional = 'Proportional',
-}
-
-// This will be extended for each pools specific input requirements
-export type BaseJoinInput = {
-    chainId: number;
-    rpcUrl: string;
-    useNativeAssetAsWrappedAmountIn?: boolean;
-    fromInternalBalance?: boolean;
-};
-
-export type InitJoinInput = BaseJoinInput & {
-    amountsIn: TokenAmount[];
-    kind: JoinKind.Init;
-};
-
-export type UnbalancedJoinInput = BaseJoinInput & {
-    amountsIn: TokenAmount[];
-    kind: JoinKind.Unbalanced;
-};
-
-export type SingleAssetJoinInput = BaseJoinInput & {
-    bptOut: TokenAmount;
-    tokenIn: Address;
-    kind: JoinKind.SingleAsset;
-};
-
-export type ProportionalJoinInput = BaseJoinInput & {
-    bptOut: TokenAmount;
-    kind: JoinKind.Proportional;
-};
-
-export type JoinInput =
-    | InitJoinInput
-    | UnbalancedJoinInput
-    | SingleAssetJoinInput
-    | ProportionalJoinInput;
-
-// Returned from a join query
-export type JoinQueryResult = {
-    poolId: Hex;
-    joinKind: JoinKind;
-    bptOut: TokenAmount;
-    amountsIn: TokenAmount[];
-    fromInternalBalance: boolean;
-    tokenInIndex?: number;
-};
-
-export type JoinCallInput = JoinQueryResult & {
-    slippage: Slippage;
-    sender: Address;
-    recipient: Address;
-};
-
-export interface BaseJoin {
-    query(input: JoinInput, poolState: PoolState): Promise<JoinQueryResult>;
-    buildCall(input: JoinCallInput): {
-        call: Hex;
-        to: Address;
-        value: bigint | undefined;
-        minBptOut: bigint;
-        maxAmountsIn: bigint[];
-    };
-}
+export * from './poolJoin';
+export * from './types';

src/entities/join/poolJoin.ts

@@ -0,0 +1,52 @@
+import {
+    BaseJoin,
+    JoinBuildOutput,
+    JoinConfig,
+    JoinInput,
+    JoinQueryResult,
+    JoinCallInput,
+} from './types';
+import { WeightedJoin } from './weighted/weightedJoin';
+import { PoolStateInput } from '../types';
+import { validateInputs } from './weighted/validateInputs';
+import { getSortedTokens } from '../utils/getSortedTokens';
+
+export class PoolJoin {
+    private readonly poolJoins: Record<string, BaseJoin> = {};
+
+    constructor(config?: JoinConfig) {
+        const { customPoolJoins } = config || {};
+        this.poolJoins = {
+            Weighted: new WeightedJoin(),
+            // custom pool Joins take precedence over base Joins
+            ...customPoolJoins,
+        };
+    }
+
+    public getJoin(poolType: string): BaseJoin {
+        if (!this.poolJoins[poolType]) {
+            throw new Error('Unsupported pool type');
+        }
+
+        return this.poolJoins[poolType];
+    }
+
+    public async query(
+        input: JoinInput,
+        poolState: PoolStateInput,
+    ): Promise<JoinQueryResult> {
+        validateInputs(input, poolState);
+
+        const sortedTokens = getSortedTokens(poolState.tokens, input.chainId);
+        const mappedPoolState = {
+            ...poolState,
+            tokens: sortedTokens,
+        };
+
+        return this.getJoin(poolState.type).query(input, mappedPoolState);
+    }
+
+    public buildCall(input: JoinCallInput): JoinBuildOutput {
+        return this.getJoin(input.poolType).buildCall(input);
+    }
+}