From c46f59e5da424e46cf42b349e69694b89837e9b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Pr=C3=A9vost?=
 <998369+prevostc@users.noreply.github.com>
Date: Mon, 10 Jun 2024 21:58:13 +0200
Subject: [PATCH] Add totalBalance

---
 schema.graphql                 | 8 ++++++++
 src/classic/entity/position.ts | 1 +
 src/classic/interaction.ts     | 2 ++
 src/clm/entity/position.ts     | 1 +
 src/clm/interaction.ts         | 2 ++
 5 files changed, 14 insertions(+)

diff --git a/schema.graphql b/schema.graphql
index 9904657..32cc2c1 100644
--- a/schema.graphql
+++ b/schema.graphql
@@ -288,6 +288,8 @@ type ClassicPosition @entity {
   vaultBalance: BigInt!
   "The amount of vault shares the investor holds in a boost"
   boostBalance: BigInt!
+  "Total amount of vault shares the investor holds. Should always equal vaultBalance + boostBalance. This is mostly used for filtering."
+  totalBalance: BigInt!
 
   "All investor position interactions"
   interactions: [ClassicPositionInteraction!]! @derivedFrom(field: "investorPosition")
@@ -330,6 +332,8 @@ type ClassicPositionInteraction @entity(immutable: true) {
   vaultBalance: BigInt!
   "The amount of vault shares the investor holds in a boost at the time of the interaction"
   boostBalance: BigInt!
+  "Total amount of vault shares the investor holds. Should always equal vaultBalance + boostBalance. This is mostly used for filtering."
+  totalBalance: BigInt!
 
   "Amount of vault shares change in the interaction"
   vaultBalanceDelta: BigInt!
@@ -642,6 +646,8 @@ type ClmPosition @entity {
   managerBalance: BigInt!
   "Amount of reward pool shares the investor holds"
   rewardPoolBalance: BigInt!
+  "Total amount of CLM shares the investor holds. Should always equal to managerBalance + rewardPoolBalance. This is mostly used for filtering."
+  totalBalance: BigInt!
 
   "All investor position interactions"
   interactions: [ClmPositionInteraction!]! @derivedFrom(field: "investorPosition")
@@ -686,6 +692,8 @@ type ClmPositionInteraction @entity(immutable: true) {
   managerBalance: BigInt!
   "The amount of reward pool shares the investor holds at the time of the interaction"
   rewardPoolBalance: BigInt!
+  "Total amount of CLM shares the investor holds. Should always equal to managerBalance + rewardPoolBalance. This is mostly used for filtering."
+  totalBalance: BigInt!
   "The amount of first underlying tokens the investor is entitled to at the time of the interaction"
   underlyingBalance0: BigInt!
   "The amount of second underlying tokens the investor is entitled to at the time of the interaction"
diff --git a/src/classic/entity/position.ts b/src/classic/entity/position.ts
index 60151f0..59552ba 100644
--- a/src/classic/entity/position.ts
+++ b/src/classic/entity/position.ts
@@ -19,6 +19,7 @@ export function getClassicPosition(classic: Classic, investor: Investor): Classi
     position.createdWith = ADDRESS_ZERO
     position.vaultBalance = ZERO_BI
     position.boostBalance = ZERO_BI
+    position.totalBalance = ZERO_BI
   }
   return position
 }
diff --git a/src/classic/interaction.ts b/src/classic/interaction.ts
index 0fb42dd..70bd231 100644
--- a/src/classic/interaction.ts
+++ b/src/classic/interaction.ts
@@ -66,6 +66,7 @@ function updateUserPosition(
   }
   position.vaultBalance = position.vaultBalance.plus(vaultBalanceDelta)
   position.boostBalance = position.boostBalance.plus(boostBalanceDelta)
+  position.totalBalance = position.vaultBalance.plus(position.boostBalance)
   position.save()
 
   ///////
@@ -94,6 +95,7 @@ function updateUserPosition(
   }
   interaction.vaultBalance = position.vaultBalance
   interaction.boostBalance = position.boostBalance
+  interaction.totalBalance = position.totalBalance
   interaction.vaultBalanceDelta = vaultBalanceDelta
   interaction.boostBalanceDelta = boostBalanceDelta
 
diff --git a/src/clm/entity/position.ts b/src/clm/entity/position.ts
index 4b17ce3..1d4b451 100644
--- a/src/clm/entity/position.ts
+++ b/src/clm/entity/position.ts
@@ -19,6 +19,7 @@ export function getClmPosition(clm: CLM, investor: Investor): ClmPosition {
     position.createdWith = ADDRESS_ZERO
     position.managerBalance = ZERO_BI
     position.rewardPoolBalance = ZERO_BI
+    position.totalBalance = ZERO_BI
   }
   return position
 }
diff --git a/src/clm/interaction.ts b/src/clm/interaction.ts
index 6bb65c1..76d0db4 100644
--- a/src/clm/interaction.ts
+++ b/src/clm/interaction.ts
@@ -115,6 +115,7 @@ function updateUserPosition(
   }
   position.managerBalance = position.managerBalance.plus(managerBalanceDelta)
   position.rewardPoolBalance = position.rewardPoolBalance.plus(rewardPoolBalanceDelta)
+  position.totalBalance = position.managerBalance.plus(position.rewardPoolBalance)
   position.save()
 
   ///////
@@ -143,6 +144,7 @@ function updateUserPosition(
   }
   interaction.managerBalance = position.managerBalance
   interaction.rewardPoolBalance = position.rewardPoolBalance
+  interaction.totalBalance = position.totalBalance
   interaction.managerBalanceDelta = managerBalanceDelta
   interaction.rewardPoolBalanceDelta = rewardPoolBalanceDelta