From a3a5a61876abec56a31290af0b43bb7ecfe400ff Mon Sep 17 00:00:00 2001
From: 0xdavinchee <0xdavinchee@gmail.com>
Date: Mon, 11 Dec 2023 12:03:39 +0200
Subject: [PATCH] add total units

---
 packages/subgraph/schema.graphql                 |  3 +++
 packages/subgraph/src/mappings/gdav1.ts          | 12 ++++++++----
 packages/subgraph/src/mappings/superfluidPool.ts |  9 +++++++--
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/packages/subgraph/schema.graphql b/packages/subgraph/schema.graphql
index 2aaed85a02..8fd3b4a12d 100644
--- a/packages/subgraph/schema.graphql
+++ b/packages/subgraph/schema.graphql
@@ -722,6 +722,7 @@ type InstantDistributionUpdatedEvent implements Event @entity(immutable: true) {
     operator: Bytes!
     requestedAmount: BigInt!
     actualAmount: BigInt!
+    totalUnits: BigInt!
     userData: Bytes!
 
     pool: Pool!
@@ -755,6 +756,7 @@ type FlowDistributionUpdatedEvent implements Event @entity(immutable: true) {
     newTotalDistributionFlowRate: BigInt!
     adjustmentFlowRecipient: Bytes!
     adjustmentFlowRate: BigInt!
+    totalUnits: BigInt!
     userData: Bytes!
 
     pool: Pool!
@@ -811,6 +813,7 @@ type MemberUnitsUpdatedEvent implements Event @entity(immutable: true) {
     token: Bytes!
     oldUnits: BigInt!
     units: BigInt!
+    totalUnits: BigInt!
 
     pool: Pool!
     poolMember: PoolMember!
diff --git a/packages/subgraph/src/mappings/gdav1.ts b/packages/subgraph/src/mappings/gdav1.ts
index 1e554a3e44..8de55c57f8 100644
--- a/packages/subgraph/src/mappings/gdav1.ts
+++ b/packages/subgraph/src/mappings/gdav1.ts
@@ -266,7 +266,7 @@ export function handleFlowDistributionUpdated(
     );
 
     // Create Event Entity
-    _createFlowDistributionUpdatedEntity(event, poolDistributor.id);
+    _createFlowDistributionUpdatedEntity(event, poolDistributor.id, pool.totalUnits);
 }
 
 export function handleInstantDistributionUpdated(
@@ -344,7 +344,7 @@ export function handleInstantDistributionUpdated(
     );
 
     // Create Event Entity
-    _createInstantDistributionUpdatedEntity(event, poolDistributor.id);
+    _createInstantDistributionUpdatedEntity(event, poolDistributor.id, pool.totalUnits);
 }
 
 // Event Entity Creation Functions
@@ -417,7 +417,8 @@ function _createBufferAdjustedEntity(
 
 function _createInstantDistributionUpdatedEntity(
     event: InstantDistributionUpdated,
-    poolDistributorId: string
+    poolDistributorId: string,
+    totalUnits: BigInt
 ): InstantDistributionUpdatedEvent {
     const ev = new InstantDistributionUpdatedEvent(
         createEventID("InstantDistributionUpdated", event)
@@ -435,6 +436,7 @@ function _createInstantDistributionUpdatedEntity(
     ev.actualAmount = event.params.actualAmount;
     ev.pool = event.params.pool.toHex();
     ev.poolDistributor = poolDistributorId;
+    ev.totalUnits = totalUnits;
     ev.userData = event.params.userData;
 
     ev.save();
@@ -444,7 +446,8 @@ function _createInstantDistributionUpdatedEntity(
 
 function _createFlowDistributionUpdatedEntity(
     event: FlowDistributionUpdated,
-    poolDistributorId: string
+    poolDistributorId: string,
+    totalUnits: BigInt
 ): FlowDistributionUpdatedEvent {
     const ev = new FlowDistributionUpdatedEvent(
         createEventID("FlowDistributionUpdated", event)
@@ -465,6 +468,7 @@ function _createFlowDistributionUpdatedEntity(
     ev.adjustmentFlowRate = event.params.adjustmentFlowRate;
     ev.pool = event.params.pool.toHex();
     ev.poolDistributor = poolDistributorId;
+    ev.totalUnits = totalUnits;
     ev.userData = event.params.userData;
 
     ev.save();
diff --git a/packages/subgraph/src/mappings/superfluidPool.ts b/packages/subgraph/src/mappings/superfluidPool.ts
index 847947b049..e4e86008b9 100644
--- a/packages/subgraph/src/mappings/superfluidPool.ts
+++ b/packages/subgraph/src/mappings/superfluidPool.ts
@@ -124,7 +124,7 @@ export function handleMemberUnitsUpdated(event: MemberUnitsUpdated): void {
     }
 
     // Create Event Entity
-    _createMemberUnitsUpdatedEntity(event, poolMember.id);
+    _createMemberUnitsUpdatedEntity(event, poolMember.id, pool.totalUnits);
 }
 
 function _createDistributionClaimedEntity(event: DistributionClaimed, poolMemberId: string): DistributionClaimedEvent {
@@ -141,13 +141,18 @@ function _createDistributionClaimedEntity(event: DistributionClaimed, poolMember
     return ev;
 }
 
-function _createMemberUnitsUpdatedEntity(event: MemberUnitsUpdated, poolMemberId: string): MemberUnitsUpdatedEvent {
+function _createMemberUnitsUpdatedEntity(
+    event: MemberUnitsUpdated,
+    poolMemberId: string,
+    totalUnits: BigInt
+): MemberUnitsUpdatedEvent {
     const ev = new MemberUnitsUpdatedEvent(createEventID("MemberUnitsUpdated", event));
     initializeEventEntity(ev, event, [event.params.token, event.address, event.params.member]);
 
     ev.token = event.params.token;
     ev.oldUnits = event.params.oldUnits;
     ev.units = event.params.newUnits;
+    ev.totalUnits = totalUnits;
     ev.pool = event.address.toHex();
     ev.poolMember = poolMemberId;
     ev.save();