From 0eb47834878ad701d2422ddb8da8c28a6834eb2a Mon Sep 17 00:00:00 2001
From: Oumar Fall <o_fall@orange.fr>
Date: Thu, 14 Nov 2024 17:52:32 +0100
Subject: [PATCH] feat(simulation-state): clean up unlimited balance result

---
 .../simulation-sdk/src/SimulationState.ts     | 29 ++++++++++++++-----
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/packages/simulation-sdk/src/SimulationState.ts b/packages/simulation-sdk/src/SimulationState.ts
index 38f7af0..58cd98d 100644
--- a/packages/simulation-sdk/src/SimulationState.ts
+++ b/packages/simulation-sdk/src/SimulationState.ts
@@ -47,6 +47,7 @@ import {
 } from "./errors.js";
 import {
   type MaybeDraft,
+  type SimulationResult,
   produceImmutable,
   simulateOperation,
   simulateOperations,
@@ -659,20 +660,20 @@ export class SimulationState implements InputSimulationState {
     operations: Operation[],
     holders: Address[],
   ) {
-    const virtualBundlerData = produceImmutable(this, (draft) => {
+    const virtualHoldersData = produceImmutable(this, (draft) => {
       holders
         .flatMap((holder) => Object.values(draft.holdings[holder] ?? {}))
-        .forEach((bundlerTokenData) => {
+        .forEach((holderTokenData) => {
           // Virtual balance to calculate the amount required.
-          bundlerTokenData.balance += MathLib.MAX_UINT_160;
+          holderTokenData.balance += MathLib.MAX_UINT_160;
         });
     });
 
-    const steps = simulateOperations(operations, virtualBundlerData);
+    const steps = simulateOperations(operations, virtualHoldersData);
 
-    const bundlerTokenDiffs = holders
+    const holdersTokenDiffs = holders
       .flatMap((holder) =>
-        keys(virtualBundlerData.holdings[holder]).map(
+        keys(virtualHoldersData.holdings[holder]).map(
           (token) => [holder, token] as const,
         ),
       )
@@ -700,10 +701,22 @@ export class SimulationState implements InputSimulationState {
       }));
 
     return {
-      requiredTokenAmounts: bundlerTokenDiffs.filter(
+      requiredTokenAmounts: holdersTokenDiffs.filter(
         ({ required }) => required > 0n,
       ),
-      steps,
+      steps: steps.map((step) =>
+        produceImmutable(step, (draft) =>
+          holders
+            .flatMap((holder) => Object.values(draft.holdings[holder] ?? {}))
+            .forEach((holderTokenData) => {
+              // Change virtual balance back to real value
+              holderTokenData.balance = MathLib.max(
+                0n,
+                holderTokenData.balance - MathLib.MAX_UINT_160,
+              );
+            }),
+        ),
+      ) as SimulationResult,
     };
   }
 }