@indigo-labs/indigo-sdk 0.2.42 → 0.3.0

package/src/contracts/stability-pool/helpers.ts

@@ -1,311 +1,795 @@
-import { UTxO } from '@lucid-evolution/lucid';
-import { getInlineDatumOrThrow } from '../../utils/lucid-utils';
 import {
-  EpochToScaleToSum,
+  addAssets,
+  Address,
+  Assets,
+  credentialToAddress,
+  getInputIndices,
+  LucidEvolution,
+  OutRef,
+  toHex,
+  UTxO,
+  validatorToScriptHash,
+} from '@lucid-evolution/lucid';
+import {
+  AccountContent,
+  AssetSnapshot,
+  AssetState,
+  E2S2SIndex,
+  E2S2SIndicesPerAsset,
+  EpochToScaleKey,
+  EpochToScaleToSumEntry,
   fromSPInteger,
   mkSPInteger,
-  parseSnapshotEpochToScaleToSumDatum,
-  parseStabilityPoolDatum,
+  parseSnapshotEpochToScaleToSumDatumOrThrow,
+  ProcessRequestAccountContent,
+  SnapshotEpochToScaleToSumContent,
   spAdd,
   spDiv,
   SPInteger,
   spMul,
   spSub,
   StabilityPoolContent,
-  StabilityPoolSnapshot,
+  StateSnapshot,
+  SumSnapshot,
 } from './types-new';
-import { match } from 'ts-pattern';
+import {
+  readonlyArray as RA,
+  array as A,
+  function as F,
+  option as O,
+} from 'fp-ts';
+import {
+  AssetClass,
+  getInlineDatumOrThrow,
+  isSameAssetClass,
+  isSameOutRef,
+  mkAssetsOf,
+} from '@3rd-eye-labs/cardano-offchain-common';
+import { repsertWithReadonlyArr } from '../../utils/array-utils';
+import { match, P } from 'ts-pattern';
+import {
+  fromSysParamsCredential,
+  SystemParams,
+} from '../../types/system-params';
+import { mkStabilityPoolValidatorFromSP } from './scripts';
 
+export const BASE_MAX_TX_FEE = 1_000_000n;
+
+export const MAX_E2S2S_ENTRIES_COUNT = 5;
 const newScaleMultiplier = 1_000_000_000n;
 
-export const initSpSnapshot: StabilityPoolSnapshot = {
-  productVal: { value: 1_000_000_000_000_000_000n },
+export const initSumVal: SPInteger = { value: 0n };
+
+export const initSpState: StateSnapshot = {
+  productVal: mkSPInteger(1n),
   depositVal: { value: 0n },
-  sumVal: { value: 0n },
   epoch: 0n,
   scale: 0n,
 };
 
-export const initEpochToScaleToSumMap = (): EpochToScaleToSum =>
-  new Map([[{ epoch: 0n, scale: 0n }, { value: 0n }]]);
-
-export function getSumFromEpochToScaleToSum(
-  e2s2s: EpochToScaleToSum,
-  epoch: bigint,
-  scale: bigint,
-): SPInteger | undefined {
-  for (const [key, value] of e2s2s.entries()) {
-    if (key.epoch === epoch && key.scale === scale) {
-      return value;
-    }
-  }
-  return undefined;
+export function mkStabilityPoolAddr(
+  lucid: LucidEvolution,
+  sysParams: SystemParams,
+): Address {
+  return credentialToAddress(
+    lucid.config().network!,
+    {
+      hash: validatorToScriptHash(
+        mkStabilityPoolValidatorFromSP(sysParams.stabilityPoolParams),
+      ),
+      type: 'Script',
+    },
+    sysParams.stabilityPoolParams.stakeCredential != null
+      ? fromSysParamsCredential(sysParams.stabilityPoolParams.stakeCredential)
+      : undefined,
+  );
 }
 
-/**
- * It's necessary to use this to update map entries because typescript can't compare keys as objects.
- */
-export function setSumInEpochToScaleToSum(
-  e2s2s: EpochToScaleToSum,
-  epoch: bigint,
-  scale: bigint,
-  sum: SPInteger,
-): EpochToScaleToSum {
-  const map = new Map<{ epoch: bigint; scale: bigint }, SPInteger>();
-  for (const [key, value] of e2s2s.entries()) {
-    if (!(key.epoch === epoch && key.scale === scale)) {
-      map.set(key, value);
-    }
-  }
+export function isSameEpochToScaleKey(
+  a: EpochToScaleKey,
+  b: EpochToScaleKey,
+): boolean {
+  return a.epoch === b.epoch && a.scale === b.scale;
+}
 
-  map.set({ epoch, scale }, sum);
-  return map;
+type StabilityPoolListIdx = {
+  spListIdx: number;
+  sumSnapshot: readonly [EpochToScaleKey, SumSnapshot];
+};
+type SnapshotIdx = {
+  snapshotUtxo: UTxO;
+  snapshotDatum: SnapshotEpochToScaleToSumContent;
+  snapshotListIdx: number;
+  sumSnapshot: readonly [EpochToScaleKey, SumSnapshot];
+};
+export type FindE2S2SIdxResult = StabilityPoolListIdx | SnapshotIdx;
+
+function findE2s2sIdx(
+  expectedKey: EpochToScaleKey,
+  collateralAsset: AssetClass,
+  spAssetState: AssetSnapshot,
+  e2s2sUtxos: [UTxO, SnapshotEpochToScaleToSumContent][],
+): O.Option<FindE2S2SIdxResult> {
+  return F.pipe(
+    // Try to find the s1 in sp e2s2s.
+    F.pipe(
+      spAssetState.epoch2scale2sum,
+      RA.findIndex(([key, _]) => isSameEpochToScaleKey(key, expectedKey)),
+    ),
+    O.match<number, O.Option<FindE2S2SIdxResult>>(
+      // When such e2sKey non existent in sp e2s2s, find it in e2s2s snapshots
+      () =>
+        F.pipe(
+          e2s2sUtxos,
+          A.findFirst(
+            ([_, datum]) =>
+              isSameAssetClass(datum.collateralAsset, collateralAsset) &&
+              F.pipe(
+                datum.snapshot,
+                RA.exists(([key, _]) =>
+                  isSameEpochToScaleKey(expectedKey, key),
+                ),
+              ),
+          ),
+          O.map((res) => {
+            const listIdx = F.pipe(
+              res[1].snapshot,
+              RA.findIndex(([key, _]) =>
+                isSameEpochToScaleKey(key, expectedKey),
+              ),
+              O.getOrElse<number>(() => {
+                throw new Error(
+                  'It was supposed to be there. Some logic error.',
+                );
+              }),
+            );
+
+            return {
+              snapshotUtxo: res[0],
+              snapshotDatum: res[1],
+              snapshotListIdx: listIdx,
+              sumSnapshot: res[1].snapshot[listIdx],
+            } satisfies SnapshotIdx;
+          }),
+        ),
+      (poolIdx) =>
+        O.some({
+          spListIdx: poolIdx,
+          sumSnapshot: spAssetState.epoch2scale2sum[poolIdx],
+        } satisfies StabilityPoolListIdx),
+    ),
+  );
 }
 
-export function getAccountReward(
-  account: StabilityPoolSnapshot,
-  e2s2s: EpochToScaleToSum,
-): SPInteger {
-  const s1 = getSumFromEpochToScaleToSum(e2s2s, account.epoch, account.scale);
-  if (!s1) throw new Error('No scale found for epoch and scale');
-  const s2 =
-    getSumFromEpochToScaleToSum(e2s2s, account.epoch, account.scale + 1n) ?? s1;
-  const a1 = spSub(s1, account.sumVal);
-  const a2 = spDiv(spSub(s2, s1), mkSPInteger(newScaleMultiplier));
+export async function findRelevantE2s2sIdxs(
+  lucid: LucidEvolution,
+  stabilityPool: StabilityPoolContent,
+  accountState: StateSnapshot,
+  allSnapshotsOutRefs: OutRef[],
+): Promise<[FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>][]> {
+  const e2s2sUtxos = (await lucid.utxosByOutRef(allSnapshotsOutRefs))
+    .map((utxo) => {
+      return [
+        utxo,
+        parseSnapshotEpochToScaleToSumDatumOrThrow(getInlineDatumOrThrow(utxo)),
+      ] satisfies [UTxO, SnapshotEpochToScaleToSumContent];
+    })
+    .filter(([_, d]) => toHex(d.iasset) === toHex(stabilityPool.iasset));
+
+  const s1E2sKey: EpochToScaleKey = {
+    epoch: accountState.epoch,
+    scale: accountState.scale,
+  };
+  const s2E2sKey: EpochToScaleKey = {
+    epoch: accountState.epoch,
+    scale: accountState.scale + 1n,
+  };
+
+  return stabilityPool.assetStates.map(([collateralAsset, spAssetState]) => {
+    const s1Res = F.pipe(
+      findE2s2sIdx(s1E2sKey, collateralAsset, spAssetState, e2s2sUtxos),
+      // When it's non-existent, we need to find proof it doesn't exist.
+      O.getOrElse(() =>
+        F.pipe(
+          e2s2sUtxos,
+          A.findFirst(
+            ([_, datum]) =>
+              isSameAssetClass(datum.collateralAsset, collateralAsset) &&
+              F.pipe(
+                datum.snapshot,
+                RA.exists(([_, val]) => val.isFirstSnapshot),
+              ),
+          ),
+          O.match<[UTxO, SnapshotEpochToScaleToSumContent], FindE2S2SIdxResult>(
+            // Try to find the proof in pool's list.
+            () =>
+              F.pipe(
+                spAssetState.epoch2scale2sum,
+                RA.findIndex(([_, dat]) => dat.isFirstSnapshot),
+                O.match(
+                  () => {
+                    throw new Error(
+                      "Couldn't find relevant proof for s1 non-existence.",
+                    );
+                  },
+                  (poolIdx) =>
+                    ({
+                      spListIdx: poolIdx,
+                      sumSnapshot: spAssetState.epoch2scale2sum[poolIdx],
+                    }) satisfies StabilityPoolListIdx,
+                ),
+              ),
+            (snapshotProof) => {
+              const listIdx = F.pipe(
+                snapshotProof[1].snapshot,
+                RA.findIndex(([_, val]) => val.isFirstSnapshot),
+                O.getOrElse<number>(() => {
+                  throw new Error(
+                    'It was supposed to be there. Some logic error.',
+                  );
+                }),
+              );
+
+              return {
+                snapshotUtxo: snapshotProof[0],
+                snapshotDatum: snapshotProof[1],
+                snapshotListIdx: listIdx,
+                sumSnapshot: snapshotProof[1].snapshot[listIdx],
+              } satisfies SnapshotIdx;
+            },
+          ),
+        ),
+      ),
+    );
+
+    const s2Res = findE2s2sIdx(
+      s2E2sKey,
+      collateralAsset,
+      spAssetState,
+      e2s2sUtxos,
+    );
 
-  return spDiv(spMul(spAdd(a1, a2), account.depositVal), account.productVal);
+    // Is actual s1 index?
+    if (isSameEpochToScaleKey(s1Res.sumSnapshot[0], s1E2sKey)) {
+      return [
+        s1Res,
+        F.pipe(
+          s2Res,
+          O.match(
+            () => {
+              // When s1 is not just a proof and is last in epoch, it proves s2 non-existant.
+              if (s1Res.sumSnapshot[1].isLastInEpoch) {
+                return O.none;
+              } else {
+                throw new Error('Expected s2 to be existent.');
+              }
+            },
+            (res) => O.some(res),
+          ),
+        ),
+      ];
+    } else {
+      // When s1 is just a proof.
+      return [
+        s1Res,
+        F.pipe(
+          s2Res,
+          O.match(
+            () => {
+              // When the non-existance proof works for s2 as well
+              if (
+                s1Res.sumSnapshot[1].isFirstSnapshot &&
+                (s1Res.sumSnapshot[0].epoch > s2E2sKey.epoch ||
+                  (s1Res.sumSnapshot[0].epoch === s2E2sKey.epoch &&
+                    s1Res.sumSnapshot[0].scale > s2E2sKey.scale))
+              ) {
+                return O.none;
+              }
+
+              throw new Error("S1 proof doesn't work for s2.");
+            },
+            // s2 exists.
+            (res) => O.some(res),
+          ),
+        ),
+      ];
+    }
+  });
 }
 
-export function getAccountFund(
-  pool: StabilityPoolSnapshot,
-  account: StabilityPoolSnapshot,
-): SPInteger {
-  let fund = spDiv(
-    spMul(account.depositVal, pool.productVal),
-    account.productVal,
+export function createProcessRequestAccountRedeemer(
+  e2s2sIdxs: [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>][],
+  otherRefInputs: UTxO[],
+  currentTime: bigint,
+): {
+  e2s2sRefInputs: UTxO[];
+  mkProcessRequestAccountRedeemerContent: (
+    poolInputIdx: bigint,
+    accountInputIdx: bigint,
+  ) => ProcessRequestAccountContent;
+} {
+  const allE2s2sSnapshotRefInputs = F.pipe(
+    e2s2sIdxs,
+    A.flatMap(([s1, s2]) => {
+      const s1RefInput = match(s1)
+        .with({ snapshotUtxo: P.select() }, (utxo) => O.some(utxo))
+        .otherwise(() => O.none);
+
+      const s2RefInput = F.pipe(
+        s2,
+        O.match(
+          () => [],
+          (s) =>
+            match(s)
+              .with({ snapshotUtxo: P.select() }, (utxo) =>
+                F.pipe(
+                  s1RefInput,
+                  O.match(
+                    () => [utxo],
+                    // when s1 ref input is same as s2 ref input, don't reference it again
+                    (s1In) => (isSameOutRef(s1In, utxo) ? [] : [utxo]),
+                  ),
+                ),
+              )
+              .otherwise(() => []),
+        ),
+      );
+
+      return [...F.pipe([s1RefInput], A.compact), ...s2RefInput];
+    }),
   );
-  if (pool.epoch > account.epoch) fund = mkSPInteger(0n);
-  if (pool.scale - account.scale > 1n) fund = mkSPInteger(0n);
-  if (pool.scale > account.scale)
-    fund = spDiv(spMul(account.depositVal, pool.productVal), {
-      value: account.productVal.value * newScaleMultiplier,
-    });
 
-  if (fund.value < spDiv(account.depositVal, mkSPInteger(1000000000n)).value)
-    return mkSPInteger(0n);
-  return fund;
-}
+  const e2s2sInputsIndices = getInputIndices(allE2s2sSnapshotRefInputs, [
+    ...otherRefInputs,
+    ...allE2s2sSnapshotRefInputs,
+  ]);
+
+  const createE2s2sIdx = (s: FindE2S2SIdxResult): E2S2SIndex =>
+    match(s)
+      .returnType<E2S2SIndex>()
+      .with({ spListIdx: P.select() }, (spListIdx) => ({
+        StabilityPoolListIdx: BigInt(spListIdx),
+      }))
+      .with({ snapshotUtxo: P.any }, (obj) => {
+        const idxForIdx = F.pipe(
+          allE2s2sSnapshotRefInputs,
+          A.findIndex((input) => isSameOutRef(input, obj.snapshotUtxo)),
+          O.getOrElse<number>(() => {
+            throw new Error('Expected to find the index.');
+          }),
+        );
 
-export function adjust(
-  pool: StabilityPoolSnapshot,
-  account: StabilityPoolSnapshot,
-  e2s2s: EpochToScaleToSum,
-): [StabilityPoolSnapshot, bigint] {
-  const newAccountSnapshot: StabilityPoolSnapshot = {
-    ...pool,
-    depositVal: getAccountFund(pool, account),
+        return {
+          RefInputIdx: {
+            refInputIdx: e2s2sInputsIndices[idxForIdx],
+            snapshotListIdx: BigInt(obj.snapshotListIdx),
+          },
+        };
+      })
+      .exhaustive();
+
+  return {
+    e2s2sRefInputs: allE2s2sSnapshotRefInputs,
+    mkProcessRequestAccountRedeemerContent: (
+      poolInputIdx: bigint,
+      accountInputIdx: bigint,
+    ) => ({
+      poolInputIdx: poolInputIdx,
+      accountInputIdx: accountInputIdx,
+      e2s2sIdxs: F.pipe(
+        e2s2sIdxs,
+        A.reduce<
+          [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>],
+          E2S2SIndicesPerAsset
+        >([], (acc, [s1, s2]) => {
+          return [
+            ...acc,
+            [
+              createE2s2sIdx(s1),
+              F.pipe(
+                s2,
+                O.match(
+                  () => null,
+                  (s) => createE2s2sIdx(s),
+                ),
+              ),
+            ],
+          ];
+        }),
+      ),
+      currentTime: currentTime,
+    }),
   };
+}
 
-  const accountReward = fromSPInteger(getAccountReward(account, e2s2s));
+function calculateReward(
+  s1: SPInteger,
+  s2: SPInteger,
+  accountSumVal: SPInteger,
+  accountState: StateSnapshot,
+): bigint {
+  const a1 = spSub(s1, accountSumVal);
+  const a2 = spDiv(spSub(s2, s1), mkSPInteger(newScaleMultiplier));
 
-  return [newAccountSnapshot, accountReward];
+  return F.pipe(
+    spDiv(
+      spMul(spAdd(a1, a2), accountState.depositVal),
+      accountState.productVal,
+    ),
+    fromSPInteger,
+  );
 }
 
-type SnapshotESSSearchResult = { utxo: UTxO; e2s2s: EpochToScaleToSum };
-
-function findEpochToScaleToSum(
-  snapshotEpochToScaleToSumTokenRef1: UTxO,
-  snapshotEpochToScaleToSumTokenRef2: UTxO | undefined,
-): [SnapshotESSSearchResult, SnapshotESSSearchResult | undefined] {
-  let ess1: EpochToScaleToSum;
-  try {
-    ess1 = parseSnapshotEpochToScaleToSumDatum(
-      getInlineDatumOrThrow(snapshotEpochToScaleToSumTokenRef1),
-    ).snapshot;
-  } catch (_) {
-    ess1 = parseStabilityPoolDatum(
-      getInlineDatumOrThrow(snapshotEpochToScaleToSumTokenRef1),
-    ).epochToScaleToSum;
-  }
+function getE2s2sEntry(
+  idx: FindE2S2SIdxResult,
+  assetState: AssetSnapshot,
+): EpochToScaleToSumEntry {
+  return match(idx)
+    .returnType<EpochToScaleToSumEntry>()
+    .with(
+      { spListIdx: P.select() },
+      (spIdx) => assetState.epoch2scale2sum[spIdx],
+    )
+    .with(
+      { snapshotUtxo: P.any },
+      (obj) => obj.snapshotDatum.snapshot[obj.snapshotListIdx],
+    )
+    .exhaustive();
+}
 
-  const ess1Ref: SnapshotESSSearchResult = {
-    utxo: snapshotEpochToScaleToSumTokenRef1,
-    e2s2s: ess1,
+function rewardsPerAsset(
+  poolAssetStates: readonly AssetState[],
+  e2s2sIdxs: [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>][],
+  accountAssetSums: readonly (readonly [AssetClass, SPInteger])[],
+  accountState: StateSnapshot,
+): [AssetClass, bigint][] {
+  const expectedS1Key: EpochToScaleKey = {
+    epoch: accountState.epoch,
+    scale: accountState.scale,
   };
 
-  if (snapshotEpochToScaleToSumTokenRef2) {
-    const ess2 = parseSnapshotEpochToScaleToSumDatum(
-      getInlineDatumOrThrow(snapshotEpochToScaleToSumTokenRef2),
-    );
-    const ess2Ref: SnapshotESSSearchResult = {
-      utxo: snapshotEpochToScaleToSumTokenRef2,
-      e2s2s: ess2.snapshot,
-    };
-    return [ess1Ref, ess2Ref];
-  }
-
-  return [ess1Ref, undefined];
-}
+  return F.pipe(
+    RA.zip(e2s2sIdxs)(poolAssetStates),
+    RA.reduce<
+      readonly [AssetState, [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>]],
+      [AssetClass, bigint][]
+    >([], (acc, [[poolAsset, poolAssetState], [s1Idx, s2Idx]]) => {
+      const s1Res = getE2s2sEntry(s1Idx, poolAssetState);
+
+      const s1 = isSameEpochToScaleKey(s1Res[0], expectedS1Key)
+        ? s1Res[1].sumVal
+        : initSumVal;
+
+      const s2 = F.pipe(
+        s2Idx,
+        O.match(
+          () => s1,
+          (s2Res) => s2Res.sumSnapshot[1].sumVal,
+        ),
+      );
+
+      const reward = calculateReward(
+        s1,
+        s2,
+        F.pipe(
+          accountAssetSums,
+          RA.findFirst(([accountAsset, _]) =>
+            isSameAssetClass(accountAsset, poolAsset),
+          ),
+          O.match(
+            // When account doesn't have this asset in snapshots
+            () => initSumVal,
+            (accountSumVal) => accountSumVal[1],
+          ),
+        ),
+        accountState,
+      );
 
-export function adjustmentHelper(
-  spESTSTokenRef1: UTxO,
-  spESTSTokenRef2: UTxO | undefined,
-  pool: StabilityPoolSnapshot,
-  e2s2s: EpochToScaleToSum,
-  account: StabilityPoolSnapshot,
-): [StabilityPoolSnapshot, bigint, UTxO[]] {
-  let accumulatedEpochToScaleToSum: EpochToScaleToSum;
-  let refInputs: UTxO[];
-
-  const value = getSumFromEpochToScaleToSum(
-    e2s2s,
-    account.epoch,
-    account.scale,
+      return [[poolAsset, reward], ...acc];
+    }),
   );
+}
 
-  if (value) {
-    accumulatedEpochToScaleToSum = e2s2s;
-    refInputs = [];
+export function getUpdatedAccountDeposit(
+  poolState: StateSnapshot,
+  accountState: StateSnapshot,
+): SPInteger {
+  if (poolState.epoch > accountState.epoch) {
+    return mkSPInteger(0n);
+  } else if (poolState.scale - accountState.scale > 1) {
+    return mkSPInteger(0n);
+  } else if (poolState.scale > accountState.scale) {
+    return spDiv(spMul(accountState.depositVal, poolState.productVal), {
+      value: accountState.productVal.value * newScaleMultiplier,
+    });
   } else {
-    const [ess1, ess2] = findEpochToScaleToSum(
-      spESTSTokenRef1,
-      spESTSTokenRef2,
+    return spDiv(
+      spMul(accountState.depositVal, poolState.productVal),
+      accountState.productVal,
     );
-    if (ess2) {
-      accumulatedEpochToScaleToSum = new Map<
-        { epoch: bigint; scale: bigint },
-        SPInteger
-      >([
-        ...Array.from(ess1.e2s2s.entries()),
-        ...Array.from(ess2.e2s2s.entries()),
-        ...Array.from(e2s2s.entries()),
-      ]);
-      refInputs = [ess1.utxo, ess2.utxo];
-    } else {
-      accumulatedEpochToScaleToSum = new Map<
-        { epoch: bigint; scale: bigint },
-        SPInteger
-      >([...Array.from(ess1.e2s2s.entries()), ...Array.from(e2s2s.entries())]);
-      refInputs = [ess1.utxo];
-    }
   }
-
-  const [newAccountSnapshot, accountReward] = adjust(
-    pool,
-    account,
-    accumulatedEpochToScaleToSum,
-  );
-
-  return [newAccountSnapshot, accountReward, refInputs];
 }
 
-export function updatePoolSnapshotWithdrawalFee(
-  withdrawalFeeAmount: SPInteger,
-  newPoolDepositExcludingFee: SPInteger,
-  pool: StabilityPoolSnapshot,
-): [SPInteger, SPInteger] {
-  const newPoolDepositVal = spAdd(
-    newPoolDepositExcludingFee,
-    withdrawalFeeAmount,
+export function updateAccount(
+  pool: StabilityPoolContent,
+  account: AccountContent,
+  e2s2sIdxs: [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>][],
+): {
+  updatedAccountContent: AccountContent;
+  reward: Assets;
+} {
+  const accountState = account.state;
+  const fund = getUpdatedAccountDeposit(pool.state, accountState);
+
+  const rewards = rewardsPerAsset(
+    pool.assetStates,
+    e2s2sIdxs,
+    account.assetSums,
+    accountState,
   );
-  const newPoolProduct =
-    withdrawalFeeAmount.value === 0n
-      ? pool.productVal
-      : spMul(
-          pool.productVal,
-          spAdd(
-            mkSPInteger(1n),
-            spDiv(withdrawalFeeAmount, newPoolDepositExcludingFee),
-          ),
-        );
 
-  return [newPoolDepositVal, newPoolProduct];
-}
+  const newDepositVal =
+    fund.value <
+    spDiv(accountState.depositVal, mkSPInteger(1_000_000_000n)).value
+      ? mkSPInteger(0n)
+      : fund;
 
+  return {
+    updatedAccountContent: {
+      ...account,
+      state: { ...pool.state, depositVal: newDepositVal },
+      assetSums: F.pipe(
+        pool.assetStates,
+        RA.map(([key, assetState]) => [key, assetState.currentSumVal]),
+      ),
+    },
+    reward: F.pipe(
+      rewards,
+      A.reduce<[AssetClass, bigint], Assets>({}, (acc, [asset, amt]) =>
+        addAssets(acc, mkAssetsOf(asset, amt)),
+      ),
+    ),
+  };
+}
 export function liquidationHelper(
-  spContent: StabilityPoolContent,
+  poolContent: StabilityPoolContent,
+  collateralAsset: AssetClass,
   iassetBurnAmt: bigint,
   /**
    * The collateral absorbed
    */
   reward: bigint,
-): { newSpContent: StabilityPoolContent } {
+): StabilityPoolContent {
   const lossPerUnitStaked = spDiv(
     mkSPInteger(iassetBurnAmt),
-    spContent.poolSnapshot.depositVal,
+    poolContent.state.depositVal,
   );
+
   const productFactor = spSub(mkSPInteger(1n), lossPerUnitStaked);
 
   const isScaleIncrease =
-    spMul(spContent.poolSnapshot.productVal, productFactor).value <
+    spMul(poolContent.state.productVal, productFactor).value <
     newScaleMultiplier;
 
-  const newSumSnapshot = spAdd(
-    spContent.poolSnapshot.sumVal,
-    spDiv(
-      spMul(mkSPInteger(reward), spContent.poolSnapshot.productVal),
-      spContent.poolSnapshot.depositVal,
-    ),
-  );
-  const newProductSnapshot = spMul(
+  const newSnapshotP = spMul(
     {
       value:
-        spContent.poolSnapshot.productVal.value *
+        poolContent.state.productVal.value *
         (isScaleIncrease ? newScaleMultiplier : 1n),
     },
     productFactor,
   );
 
-  const isEpochIncrease = newProductSnapshot.value <= 0n;
+  const currentS = F.pipe(
+    poolContent.assetStates,
+    RA.findFirstMap(([ac, assetSnap]) =>
+      isSameAssetClass(ac, collateralAsset)
+        ? O.some(assetSnap.currentSumVal)
+        : O.none,
+    ),
+    O.getOrElse(() => initSumVal),
+  );
+
+  const newSnapshotS = spAdd(
+    currentS,
+    spDiv(
+      spMul(mkSPInteger(reward), poolContent.state.productVal),
+      poolContent.state.depositVal,
+    ),
+  );
+
+  const isEpochIncrease = newSnapshotP.value <= 0;
 
-  const newSnapshot: StabilityPoolSnapshot = isEpochIncrease
-    ? { ...initSpSnapshot, epoch: spContent.poolSnapshot.epoch + 1n }
+  const newState: StateSnapshot = isEpochIncrease
+    ? { ...initSpState, epoch: poolContent.state.epoch + 1n }
     : {
-        productVal: newProductSnapshot,
+        productVal: newSnapshotP,
         depositVal: spSub(
-          spContent.poolSnapshot.depositVal,
+          poolContent.state.depositVal,
           mkSPInteger(iassetBurnAmt),
         ),
-        sumVal: newSumSnapshot,
-        epoch: spContent.poolSnapshot.epoch,
-        scale: spContent.poolSnapshot.scale + (isScaleIncrease ? 1n : 0n),
+        epoch: poolContent.state.epoch,
+        scale: poolContent.state.scale + (isScaleIncrease ? 1n : 0n),
       };
 
-  const newMap = setSumInEpochToScaleToSum(
-    spContent.epochToScaleToSum,
-    spContent.poolSnapshot.epoch,
-    spContent.poolSnapshot.scale,
-    newSumSnapshot,
-  );
+  const currentE2S2SKey: EpochToScaleKey = {
+    epoch: poolContent.state.epoch,
+    scale: poolContent.state.scale,
+  };
 
-  const newEpochToScaleToSum = match(true)
-    .when(
-      () => isEpochIncrease,
-      () =>
-        setSumInEpochToScaleToSum(
-          newMap,
-          spContent.poolSnapshot.epoch + 1n,
-          spContent.poolSnapshot.scale,
-          { value: 0n },
-        ),
-    )
-    .when(
-      () => isScaleIncrease,
-      () =>
-        setSumInEpochToScaleToSum(
-          newMap,
-          spContent.poolSnapshot.epoch,
-          spContent.poolSnapshot.scale + 1n,
-          newSumSnapshot,
-        ),
-    )
-    .otherwise(() => newMap);
+  const newCollateralAssetSumVal = isEpochIncrease ? initSumVal : newSnapshotS;
+
+  const newAssetStates = (() => {
+    const updatedAssetStates = F.pipe(
+      poolContent.assetStates,
+      repsertWithReadonlyArr<AssetClass, AssetSnapshot>(
+        (key) => isSameAssetClass(key, collateralAsset),
+        (assetSnap) => ({
+          currentSumVal: newCollateralAssetSumVal,
+          epoch2scale2sum: F.pipe(
+            assetSnap.epoch2scale2sum,
+            RA.modifyAt(0, ([key, val]) => [
+              key,
+              { ...val, sumVal: newSnapshotS } satisfies SumSnapshot,
+            ]),
+            O.getOrElse<readonly EpochToScaleToSumEntry[]>(() => {
+              throw new Error('There has to be first entry');
+            }),
+          ),
+        }),
+        () => ({
+          currentSumVal: newCollateralAssetSumVal,
+          epoch2scale2sum: [
+            [
+              currentE2S2SKey,
+              {
+                sumVal: newSnapshotS,
+                isFirstSnapshot: true,
+                isLastInEpoch: true,
+              },
+            ],
+          ],
+        }),
+        () => collateralAsset,
+      ),
+    );
+
+    if (isEpochIncrease) {
+      return F.pipe(
+        updatedAssetStates,
+        RA.map<AssetState, AssetState>(([key, assetSnap]) => [
+          key,
+          {
+            currentSumVal: initSumVal,
+            epoch2scale2sum: [
+              [
+                {
+                  epoch: poolContent.state.epoch + 1n,
+                  scale: 0n,
+                },
+                {
+                  sumVal: initSumVal,
+                  isLastInEpoch: true,
+                  isFirstSnapshot: false,
+                },
+              ],
+              ...assetSnap.epoch2scale2sum,
+            ],
+          },
+        ]),
+      );
+    } else if (isScaleIncrease) {
+      return F.pipe(
+        updatedAssetStates,
+        RA.map<AssetState, AssetState>(([key, assetSnap]) => [
+          key,
+          {
+            ...assetSnap,
+            epoch2scale2sum: match(assetSnap.epoch2scale2sum)
+              .returnType<readonly EpochToScaleToSumEntry[]>()
+              .with([[P.any, P.any], ...P.array()], ([[key, val], ...rest]) => {
+                const newScaleEntry: EpochToScaleToSumEntry = [
+                  {
+                    epoch: poolContent.state.epoch,
+                    scale: poolContent.state.scale + 1n,
+                  },
+                  {
+                    sumVal: val.sumVal,
+                    isLastInEpoch: true,
+                    isFirstSnapshot: false,
+                  },
+                ];
+
+                return [
+                  newScaleEntry,
+                  [key, { ...val, isLastInEpoch: false } satisfies SumSnapshot],
+                  ...rest,
+                ];
+              })
+              .otherwise(() => {
+                throw new Error('There has to be at least 1 entry');
+              }),
+          },
+        ]),
+      );
+    } else {
+      return updatedAssetStates;
+    }
+  })();
 
   return {
-    newSpContent: {
-      asset: spContent.asset,
-      epochToScaleToSum: newEpochToScaleToSum,
-      poolSnapshot: newSnapshot,
-    },
+    ...poolContent,
+    assetStates: newAssetStates,
+    state: newState,
   };
 }
+
+export function updatePoolStateWhenWithdrawalFee(
+  withdrawalFeeAmt: bigint,
+  updatedPoolState: StateSnapshot,
+): StateSnapshot {
+  if (withdrawalFeeAmt === 0n) {
+    return updatedPoolState;
+  } else {
+    const withdrawalFeeSpInt = mkSPInteger(withdrawalFeeAmt);
+    const newDepositVal = spAdd(
+      updatedPoolState.depositVal,
+      withdrawalFeeSpInt,
+    );
+
+    const productFactor = spAdd(
+      mkSPInteger(1n),
+      spDiv(withdrawalFeeSpInt, updatedPoolState.depositVal),
+    );
+
+    const newProductVal = spMul(updatedPoolState.productVal, productFactor);
+
+    return {
+      ...updatedPoolState,
+      productVal: newProductVal,
+      depositVal: newDepositVal,
+    };
+  }
+}
+
+export function partitionEpochToScaleToSums(
+  spContent: StabilityPoolContent,
+): readonly [
+  readonly SnapshotEpochToScaleToSumContent[],
+  readonly AssetState[],
+] {
+  const res = F.pipe(
+    spContent.assetStates,
+    RA.map<AssetState, [SnapshotEpochToScaleToSumContent[], AssetState]>(
+      ([collateralAsset, assetState]) => {
+        if (assetState.epoch2scale2sum.length >= MAX_E2S2S_ENTRIES_COUNT) {
+          const { right: remaining, left: snapshotMapItems } = F.pipe(
+            assetState.epoch2scale2sum,
+            RA.partition(
+              ([e2sKey, _]) =>
+                e2sKey.epoch === spContent.state.epoch &&
+                e2sKey.scale === spContent.state.scale,
+            ),
+          );
+
+          return [
+            [
+              {
+                iasset: spContent.iasset,
+                collateralAsset: collateralAsset,
+                snapshot: snapshotMapItems,
+              },
+            ],
+            [collateralAsset, { ...assetState, epoch2scale2sum: remaining }],
+          ];
+        } else {
+          return [[], [collateralAsset, assetState]];
+        }
+      },
+    ),
+  );
+
+  const [newSnapshots, newAssetStates] = RA.unzip(res);
+
+  return [RA.flatten(newSnapshots), newAssetStates];
+}