From a68c386573bcd582d691608f42f294833479687f Mon Sep 17 00:00:00 2001
From: johngrantuk <johngrantuk@googlemail.com>
Date: Fri, 26 Jul 2024 11:09:14 +0100
Subject: [PATCH] feat: Add max add/remove methods.

---
 typescript/src/constants.ts             |  2 ++
 typescript/src/stable/stablePool.ts     | 22 ++++++++++++++++++++++
 typescript/src/vault/types.ts           |  8 ++++++++
 typescript/src/weighted/weightedPool.ts | 22 ++++++++++++++++++++++
 4 files changed, 54 insertions(+)
 create mode 100644 typescript/src/constants.ts

diff --git a/typescript/src/constants.ts b/typescript/src/constants.ts
new file mode 100644
index 0000000..f76553a
--- /dev/null
+++ b/typescript/src/constants.ts
@@ -0,0 +1,2 @@
+export const MAX_UINT256 =
+    115792089237316195423570985008687907853269984665640564039457584007913129639935n;
diff --git a/typescript/src/stable/stablePool.ts b/typescript/src/stable/stablePool.ts
index c2a0b18..47bb8f9 100644
--- a/typescript/src/stable/stablePool.ts
+++ b/typescript/src/stable/stablePool.ts
@@ -1,3 +1,4 @@
+import { MAX_UINT256 } from '../constants';
 import { MathSol } from '../utils/math';
 import {
     MaxSwapParams,
@@ -45,6 +46,27 @@ export class Stable implements PoolBase {
         );
     }
 
+    getMaxSingleTokenAddAmount(): bigint {
+        return MAX_UINT256;
+    }
+
+    getMaxSingleTokenExitAmount(
+        isExactIn: boolean,
+        totalSupply: bigint,
+        tokenOutBalance: bigint,
+        tokenOutScalingFactor: bigint,
+        tokenOutRate: bigint,
+    ): bigint {
+        return this.getMaxSwapAmount({
+            swapKind: isExactIn ? SwapKind.GivenIn : SwapKind.GivenOut,
+            balancesLiveScaled18: [totalSupply, tokenOutBalance],
+            tokenRates: [1000000000000000000n, tokenOutRate],
+            scalingFactors: [1000000000000000000n, tokenOutScalingFactor],
+            indexIn: 0,
+            indexOut: 1,
+        });
+    }
+
     onSwap(swapParams: SwapParams): bigint {
         const {
             swapKind,
diff --git a/typescript/src/vault/types.ts b/typescript/src/vault/types.ts
index 8f581a4..5582c13 100644
--- a/typescript/src/vault/types.ts
+++ b/typescript/src/vault/types.ts
@@ -17,6 +17,14 @@ export enum SwapKind {
 
 export interface PoolBase {
     getMaxSwapAmount(maxSwapParams: MaxSwapParams): bigint;
+    getMaxSingleTokenExitAmount(
+        isExactIn: boolean,
+        totalSupply: bigint,
+        tokenOutBalance: bigint,
+        tokenOutScalingFactor: bigint,
+        tokenOutRate: bigint,
+    ): bigint;
+    getMaxSingleTokenAddAmount(): bigint;
     onSwap(swapParams: SwapParams): bigint;
     computeInvariant(balancesLiveScaled18: bigint[]): bigint;
     computeBalance(
diff --git a/typescript/src/weighted/weightedPool.ts b/typescript/src/weighted/weightedPool.ts
index a622d5f..264a50d 100644
--- a/typescript/src/weighted/weightedPool.ts
+++ b/typescript/src/weighted/weightedPool.ts
@@ -1,3 +1,4 @@
+import { MAX_UINT256 } from '../constants';
 import { MathSol } from '../utils/math';
 import {
     MaxSwapParams,
@@ -39,6 +40,27 @@ export class Weighted implements PoolBase {
         );
     }
 
+    getMaxSingleTokenAddAmount(): bigint {
+        return MAX_UINT256;
+    }
+
+    getMaxSingleTokenExitAmount(
+        isExactIn: boolean,
+        totalSupply: bigint,
+        tokenOutBalance: bigint,
+        tokenOutScalingFactor: bigint,
+        tokenOutRate: bigint,
+    ): bigint {
+        return this.getMaxSwapAmount({
+            swapKind: isExactIn ? SwapKind.GivenIn : SwapKind.GivenOut,
+            balancesLiveScaled18: [totalSupply, tokenOutBalance],
+            tokenRates: [1000000000000000000n, tokenOutRate],
+            scalingFactors: [1000000000000000000n, tokenOutScalingFactor],
+            indexIn: 0,
+            indexOut: 1,
+        });
+    }
+
     onSwap(swapParams: SwapParams): bigint {
         const {
             swapKind,