From cd40398d63b246187dd980eb1d0feb7b9cc9c757 Mon Sep 17 00:00:00 2001
From: Jose Felix <jfelx@outlook.com>
Date: Fri, 6 Dec 2024 11:53:17 -0400
Subject: [PATCH] (Swap tool) Improve gas amount estimation reliability for
 multi-hop swaps (#3997)

---
 packages/tx/src/__tests__/events.spec.ts | 66 ++++++++++++++++++++++++
 packages/tx/src/events.ts                | 36 +++++++++++++
 2 files changed, 102 insertions(+)

diff --git a/packages/tx/src/__tests__/events.spec.ts b/packages/tx/src/__tests__/events.spec.ts
index 1fc006ca41..fb53d6cd4e 100644
--- a/packages/tx/src/__tests__/events.spec.ts
+++ b/packages/tx/src/__tests__/events.spec.ts
@@ -92,6 +92,18 @@ describe("getSumTotalSpenderCoinsSpent", () => {
 
     expect(coins).toEqual(expectedResult);
   });
+
+  it("should handle token swaps coins spent", () => {
+    const coins = getSumTotalSpenderCoinsSpent(
+      "osmo1gtgx92pxk6hvhc3c3g0xlkrwqq6knymu0e0caw",
+      mockTokenSwapEvents
+    );
+
+    expect(coins).toEqual([
+      { denom: "uosmo", amount: "100000000" },
+      { denom: "uion", amount: "100000000" },
+    ]);
+  });
 });
 
 describe("matchRawCoinValue", () => {
@@ -277,3 +289,57 @@ const mockMultipleEvents = [
     ],
   },
 ];
+
+const mockTokenSwapEvents = [
+  {
+    type: "message",
+    attributes: [
+      {
+        key: "action",
+        value: "/osmosis.gamm.v1beta1.MsgJoinPool",
+        index: true,
+      },
+      {
+        key: "sender",
+        value: "osmo1gtgx92pxk6hvhc3c3g0xlkrwqq6knymu0e0caw",
+        index: true,
+      },
+      {
+        key: "msg_index",
+        value: "0",
+        index: true,
+      },
+    ],
+  },
+  {
+    type: "coin_spent",
+    attributes: [
+      {
+        key: "spender",
+        value: "osmo1gtgx92pxk6hvhc3c3g0xlkrwqq6knymu0e0caw",
+        index: true,
+      },
+      {
+        key: "amount",
+        value:
+          "58573factory/osmo1z0qrq605sjgcqpylfl4aa6s90x738j7m58wyatt0tdzflg2ha26q67k743/wbtc,10541840163500050ibc/EA1D43981D5C9A1C4AAEA9C23BB1D4FA126BA9BC7020A25E0AE4AA841EA25DC5,65605821922677750ibc/F4A070A6D78496D53127EA85C094A9EC87DFC1F36071B8CCDDBD020F933D213D,47899690uosmo",
+        index: true,
+      },
+      {
+        key: "msg_index",
+        value: "0",
+        index: true,
+      },
+    ],
+  },
+  {
+    type: "token_swapped",
+    attributes: [
+      {
+        key: "tokens_in",
+        value: "100000000uosmo,100000000uion",
+        index: true,
+      },
+    ],
+  },
+];
diff --git a/packages/tx/src/events.ts b/packages/tx/src/events.ts
index bc61bed57d..f9e55aa9b4 100644
--- a/packages/tx/src/events.ts
+++ b/packages/tx/src/events.ts
@@ -10,6 +10,42 @@ export function getSumTotalSpenderCoinsSpent(
   // denom => sum amount
   const coinsSpentMap = new Map<string, Int>();
 
+  const tokenSwapEvent = txEvents.find(({ type }) => type === "token_swapped");
+
+  /**
+   * If the transaction is a token swap, we need to extract the coins spent from the
+   * "tokens_in" attribute.
+   */
+  if (tokenSwapEvent) {
+    const tokenSwapEventAttribute = tokenSwapEvent.attributes.find(
+      ({ key }) => key === "tokens_in"
+    );
+    if (tokenSwapEventAttribute) {
+      const coinsSpentRaw = tokenSwapEventAttribute.value.split(",");
+
+      coinsSpentRaw.forEach((coinSpentRaw) => {
+        const coin = matchRawCoinValue(coinSpentRaw);
+
+        if (coin) {
+          const existingCoin = coinsSpentMap.get(coin.denom);
+          if (existingCoin) {
+            coinsSpentMap.set(
+              coin.denom,
+              existingCoin.add(new Int(coin.amount))
+            );
+          } else {
+            coinsSpentMap.set(coin.denom, new Int(coin.amount));
+          }
+        }
+      });
+
+      return Array.from(coinsSpentMap, ([denom, amount]) => ({
+        denom,
+        amount: amount.toString(),
+      }));
+    }
+  }
+
   txEvents.forEach(({ type, attributes }) => {
     // validate that it's a spend event affecting the spender
     if (type !== "coin_spent") return;