From 4d49b3f0c1d98da8667d62821fb776f796eb1571 Mon Sep 17 00:00:00 2001
From: Bruno Eidam Guerios <bruno@balancerlabs.dev>
Date: Wed, 13 Sep 2023 16:16:27 -0300
Subject: [PATCH] Use index provided by the api to sort pool tokens

---
 src/entities/join/index.ts                 | 3 ++-
 src/entities/join/weighted/weightedJoin.ts | 6 +++---
 test/weightedJoin.integration.test.ts      | 7 ++++++-
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/entities/join/index.ts b/src/entities/join/index.ts
index cafa2217..b448a52f 100644
--- a/src/entities/join/index.ts
+++ b/src/entities/join/index.ts
@@ -17,7 +17,8 @@ export type PoolState = {
     tokens: {
         address: Address;
         decimals: number;
-    }[]; // already properly sorted in case different versions sort them differently
+        index: number;
+    }[];
 };
 
 // This will be extended for each pools specific input requirements
diff --git a/src/entities/join/weighted/weightedJoin.ts b/src/entities/join/weighted/weightedJoin.ts
index 7f446f46..7b34c7f3 100644
--- a/src/entities/join/weighted/weightedJoin.ts
+++ b/src/entities/join/weighted/weightedJoin.ts
@@ -30,9 +30,9 @@ export class WeightedJoin implements BaseJoin {
         // TODO: Extend input validation for cases we'd like to check
         checkInputs(input, poolState);
 
-        const poolTokens = poolState.tokens.map(
-            (t) => new Token(input.chainId, t.address, t.decimals),
-        );
+        const poolTokens = poolState.tokens
+            .sort((a, b) => a.index - b.index)
+            .map((t) => new Token(input.chainId, t.address, t.decimals));
         let maxAmountsIn = Array(poolTokens.length).fill(MAX_UINT256);
         let userData: Address;
 
diff --git a/test/weightedJoin.integration.test.ts b/test/weightedJoin.integration.test.ts
index 1267bb7c..b8b07453 100644
--- a/test/weightedJoin.integration.test.ts
+++ b/test/weightedJoin.integration.test.ts
@@ -328,7 +328,8 @@ describe('weighted join test', () => {
 
 export class MockApi {
     public async getPool(id: Address): Promise<PoolState> {
-        let tokens: { address: Address; decimals: number }[] = [];
+        let tokens: { address: Address; decimals: number; index: number }[] =
+            [];
         if (
             id ===
             '0x5c6ee304399dbdb9c8ef030ab642b10820db8f56000200000000000000000014'
@@ -337,10 +338,12 @@ export class MockApi {
                 {
                     address: '0xba100000625a3754423978a60c9317c58a424e3d', // BAL
                     decimals: 18,
+                    index: 0,
                 },
                 {
                     address: '0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2', // wETH
                     decimals: 18,
+                    index: 1,
                 },
             ];
         } else if (
@@ -351,10 +354,12 @@ export class MockApi {
                 {
                     address: '0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0', // wstETH slot 0
                     decimals: 18,
+                    index: 0,
                 },
                 {
                     address: '0xc00e94cb662c3520282e6f5717214004a7f26888', // COMP slot 1
                     decimals: 18,
+                    index: 1,
                 },
             ];
         }