Supply referral initial module

packages/margin_trading/sources/helper/margin_constants.move

@@ -11,6 +11,7 @@ const MIN_LEVERAGE: u64 = 1_000_000_000; // 1x
 const MAX_LEVERAGE: u64 = 20_000_000_000; // 20x
 const YEAR_MS: u64 = 365 * 24 * 60 * 60 * 1000;
 const MIN_MIN_BORROW: u64 = 1000;
+const DEFAULT_REFERRAL: address = @0x0;
 
 public fun margin_version(): u64 {
     MARGIN_VERSION
@@ -43,3 +44,7 @@ public fun year_ms(): u64 {
 public fun min_min_borrow(): u64 {
     MIN_MIN_BORROW
 }
+
+public fun default_referral(): address {
+    DEFAULT_REFERRAL
+}

packages/margin_trading/sources/margin_pool.move

@@ -8,7 +8,8 @@ use margin_trading::{
     margin_registry::{MarginRegistry, MaintainerCap, MarginPoolCap},
     margin_state::{Self, State},
     position_manager::{Self, PositionManager},
-    protocol_config::{InterestConfig, MarginPoolConfig, ProtocolConfig}
+    protocol_config::{InterestConfig, MarginPoolConfig, ProtocolConfig},
+    protocol_fees::{Self, ProtocolFees, Referral}
 };
 use std::type_name::{Self, TypeName};
 use sui::{balance::{Self, Balance}, clock::Clock, coin::Coin, event, vec_set::{Self, VecSet}};
@@ -29,6 +30,7 @@ public struct MarginPool<phantom Asset> has key, store {
     vault: Balance<Asset>,
     state: State,
     config: ProtocolConfig,
+    protocol_fees: ProtocolFees,
     positions: PositionManager,
     allowed_deepbook_pools: VecSet<ID>,
 }
@@ -99,6 +101,7 @@ public fun create_margin_pool<Asset>(
         vault: balance::zero<Asset>(),
         state: margin_state::default(clock),
         config,
+        protocol_fees: protocol_fees::default_protocol_fees(ctx, clock),
         positions: position_manager::create_position_manager(ctx),
         allowed_deepbook_pools: vec_set::empty(),
     };
@@ -204,20 +207,26 @@ public fun update_margin_pool_config<Asset>(
 }
 
 // === Public Functions * LENDING * ===
-/// Allows anyone to supply the margin pool. Returns the new user supply amount.
+/// Supply to the margin pool. Returns the new user supply amount.
 public fun supply<Asset>(
     self: &mut MarginPool<Asset>,
     registry: &MarginRegistry,
     coin: Coin<Asset>,
+    referral: Option<address>,
     clock: &Clock,
     ctx: &TxContext,
-) {
+): u64 {
     registry.load_inner();
-    let supplier = ctx.sender();
     let supply_amount = coin.value();
-
-    let supply_shares = self.state.increase_supply(&self.config, supply_amount, clock);
-    self.positions.increase_user_supply(supplier, supply_shares);
+    let (supply_shares, protocol_fees) = self
+        .state
+        .increase_supply(&self.config, supply_amount, clock);
+    self.protocol_fees.increase_fees_per_share(self.state.supply_shares(), protocol_fees);
+    let (total_user_supply, previous_referral) = self
+        .positions
+        .increase_user_supply(referral, supply_shares, ctx);
+    self.protocol_fees.decrease_shares(previous_referral, total_user_supply - supply_shares, clock);
+    self.protocol_fees.increase_shares(referral, total_user_supply, clock);
 
     let balance = coin.into_balance();
     self.vault.join(balance);
@@ -227,14 +236,16 @@ public fun supply<Asset>(
     event::emit(AssetSupplied {
         margin_pool_id: self.id(),
         asset_type: type_name::with_defining_ids<Asset>(),
-        supplier,
+        supplier: ctx.sender(),
         supply_amount,
         supply_shares,
         timestamp: clock.timestamp_ms(),
     });
+
+    total_user_supply
 }
 
-/// Allows withdrawal from the margin pool. Returns the withdrawn coin.
+/// Withdraw from the margin pool. Returns the withdrawn coin.
 public fun withdraw<Asset>(
     self: &mut MarginPool<Asset>,
     registry: &MarginRegistry,
@@ -243,20 +254,25 @@ public fun withdraw<Asset>(
     ctx: &mut TxContext,
 ): Coin<Asset> {
     registry.load_inner();
-    let supplier = ctx.sender();
-    let supplied_shares = self.positions.user_supply_shares(supplier);
+    let supplied_shares = self.positions.user_supply_shares(ctx);
     let supplied_amount = self.state.supply_shares_to_amount(supplied_shares, &self.config, clock);
     let withdraw_amount = amount.destroy_with_default(supplied_amount);
     let withdraw_shares = math::mul(supplied_shares, math::div(withdraw_amount, supplied_amount));
 
-    self.positions.decrease_user_supply(supplier, withdraw_shares);
+    let (_, protocol_fees) = self
+        .state
+        .decrease_supply_shares(&self.config, withdraw_shares, clock);
+    self.protocol_fees.increase_fees_per_share(self.state.supply_shares(), protocol_fees);
+
+    let (_, previous_referral) = self.positions.decrease_user_supply(withdraw_shares, ctx);
+    self.protocol_fees.decrease_shares(previous_referral, withdraw_shares, clock);
     assert!(withdraw_amount <= self.vault.value(), ENotEnoughAssetInPool);
     let coin = self.vault.split(withdraw_amount).into_coin(ctx);
 
     event::emit(AssetWithdrawn {
         margin_pool_id: self.id(),
         asset_type: type_name::with_defining_ids<Asset>(),
-        supplier,
+        supplier: ctx.sender(),
         withdraw_amount,
         withdraw_shares,
         timestamp: clock.timestamp_ms(),
@@ -265,6 +281,19 @@ public fun withdraw<Asset>(
     coin
 }
 
+/// Withdraw the referral fees.
+public fun withdraw_referral_fees<Asset>(
+    self: &mut MarginPool<Asset>,
+    referral: &mut Referral,
+    clock: &Clock,
+    ctx: &mut TxContext,
+): Coin<Asset> {
+    let referral_fees = self.protocol_fees.calculate_and_claim(referral, clock);
+    let coin = self.vault.split(referral_fees).into_coin(ctx);
+
+    coin
+}
+
 // === Public-View Functions ===
 public fun deepbook_pool_allowed<Asset>(self: &MarginPool<Asset>, deepbook_pool_id: ID): bool {
     self.allowed_deepbook_pools.contains(&deepbook_pool_id)
@@ -280,9 +309,10 @@ public(package) fun borrow<Asset>(
 ): (Coin<Asset>, u64, u64) {
     assert!(amount <= self.vault.value(), ENotEnoughAssetInPool);
     assert!(amount >= self.config.min_borrow(), EBorrowAmountTooLow);
-    let (total_borrow, total_borrow_shares) = self
+    let (total_borrow, total_borrow_shares, protocol_fees) = self
         .state
         .increase_borrow(&self.config, amount, clock);
+    self.protocol_fees.increase_fees_per_share(self.state.supply_shares(), protocol_fees);
     assert!(
         self.state.utilization_rate() <= self.config.max_utilization_rate(),
         EMaxPoolBorrowPercentageExceeded,
@@ -297,7 +327,8 @@ public(package) fun repay<Asset>(
     coin: Coin<Asset>,
     clock: &Clock,
 ) {
-    let amount = self.state.decrease_borrow_shares(&self.config, shares, clock);
+    let (amount, protocol_fees) = self.state.decrease_borrow_shares(&self.config, shares, clock);
+    self.protocol_fees.increase_fees_per_share(self.state.supply_shares(), protocol_fees);
     assert!(coin.value() == amount, EInvalidRepayQuantity);
     self.vault.join(coin.into_balance());
 }
@@ -311,7 +342,8 @@ public(package) fun repay_liquidation<Asset>(
     coin: Coin<Asset>,
     clock: &Clock,
 ): (u64, u64, u64) {
-    let amount = self.state.decrease_borrow_shares(&self.config, shares, clock); // decreased 48.545 shares, 97.087 USDC
+    let (amount, protocol_fees) = self.state.decrease_borrow_shares(&self.config, shares, clock); // decreased 48.545 shares, 97.087 USDC
+    self.protocol_fees.increase_fees_per_share(self.state.supply_shares(), protocol_fees);
     let coin_value = coin.value(); // 100 USDC
     let (reward, default) = if (coin_value > amount) {
         self.state.increase_supply_absolute(coin_value - amount);

packages/margin_trading/sources/margin_pool/margin_state.move

@@ -29,29 +29,29 @@ public(package) fun increase_supply(
     config: &ProtocolConfig,
     amount: u64,
     clock: &Clock,
-): u64 {
-    self.update(config, clock);
+): (u64, u64) {
+    let protocol_fees = self.update(config, clock);
     let ratio = self.supply_ratio();
     let shares = math::div(amount, ratio);
     self.supply_shares = self.supply_shares + shares;
     self.supply = self.supply + amount;
 
-    shares
+    (shares, protocol_fees)
 }
 
 public(package) fun decrease_supply_shares(
     self: &mut State,
     config: &ProtocolConfig,
     shares: u64,
     clock: &Clock,
-): u64 {
-    self.update(config, clock);
+): (u64, u64) {
+    let protocol_fees = self.update(config, clock);
     let ratio = self.supply_ratio();
     let amount = math::mul(shares, ratio);
     self.supply_shares = self.supply_shares - shares;
     self.supply = self.supply - amount;
 
-    amount
+    (amount, protocol_fees)
 }
 
 public(package) fun increase_supply_absolute(self: &mut State, amount: u64) {
@@ -67,14 +67,14 @@ public(package) fun increase_borrow(
     config: &ProtocolConfig,
     amount: u64,
     clock: &Clock,
-): (u64, u64) {
-    self.update(config, clock);
+): (u64, u64, u64) {
+    let protocol_fees = self.update(config, clock);
     let ratio = self.borrow_ratio();
     let shares = math::div(amount, ratio);
     self.borrow_shares = self.borrow_shares + shares;
     self.borrow = self.borrow + amount;
 
-    (self.borrow, self.borrow_shares)
+    (self.borrow, self.borrow_shares, protocol_fees)
 }
 
 /// Decrease borrowed shares and return the corresponding amount
@@ -83,14 +83,14 @@ public(package) fun decrease_borrow_shares(
     config: &ProtocolConfig,
     shares: u64,
     clock: &Clock,
-): u64 {
-    self.update(config, clock);
+): (u64, u64) {
+    let protocol_fees = self.update(config, clock);
     let ratio = self.borrow_ratio();
     let amount = math::mul(shares, ratio);
     self.borrow_shares = self.borrow_shares - shares;
     self.borrow = self.borrow - amount;
 
-    amount
+    (amount, protocol_fees)
 }
 
 public(package) fun utilization_rate(self: &State): u64 {
@@ -105,6 +105,10 @@ public(package) fun supply(self: &State): u64 {
     self.supply
 }
 
+public(package) fun supply_shares(self: &State): u64 {
+    self.supply_shares
+}
+
 public(package) fun borrow_shares_to_amount(
     self: &State,
     shares: u64,
@@ -115,7 +119,8 @@ public(package) fun borrow_shares_to_amount(
     let elapsed = now - self.last_update_timestamp;
 
     let time_adjusted_rate = config.time_adjusted_rate(self.utilization_rate(), elapsed);
-    let borrow = self.borrow + math::mul(self.borrow, time_adjusted_rate);
+    let interest = math::mul(self.borrow, time_adjusted_rate);
+    let borrow = self.borrow + interest;
     let ratio = if (self.borrow_shares == 0) {
         constants::float_scaling()
     } else {
@@ -135,7 +140,9 @@ public(package) fun supply_shares_to_amount(
     let elapsed = now - self.last_update_timestamp;
 
     let time_adjusted_rate = config.time_adjusted_rate(self.utilization_rate(), elapsed);
-    let supply = self.supply + math::mul(self.borrow, time_adjusted_rate);
+    let interest = math::mul(self.borrow, time_adjusted_rate);
+    let protocol_fees = math::mul(interest, config.protocol_spread());
+    let supply = self.supply + interest - protocol_fees;
     let ratio = if (self.supply_shares == 0) {
         constants::float_scaling()
     } else {
@@ -145,14 +152,18 @@ public(package) fun supply_shares_to_amount(
     math::div(shares, ratio)
 }
 
-fun update(self: &mut State, config: &ProtocolConfig, clock: &Clock) {
+fun update(self: &mut State, config: &ProtocolConfig, clock: &Clock): u64 {
     let now = clock.timestamp_ms();
     let elapsed = now - self.last_update_timestamp;
 
     let time_adjusted_rate = config.time_adjusted_rate(self.utilization_rate(), elapsed);
-    self.supply = self.supply + math::mul(self.borrow, time_adjusted_rate);
-    self.borrow = self.borrow + math::mul(self.borrow, time_adjusted_rate);
+    let interest = math::mul(self.borrow, time_adjusted_rate);
+    let protocol_fees = math::mul(interest, config.protocol_spread());
+    self.supply = self.supply + interest - protocol_fees;
+    self.borrow = self.borrow + interest;
     self.last_update_timestamp = now;
+
+    protocol_fees
 }
 
 fun supply_ratio(self: &State): u64 {

packages/margin_trading/sources/margin_pool/position_manager.move

@@ -8,54 +8,73 @@ module margin_trading::position_manager;
 use sui::table::{Self, Table};
 
 public struct PositionManager has store {
-    supply_shares: Table<address, u64>,
+    positions: Table<address, Position>,
+}
+
+public struct Position has store {
+    shares: u64,
+    referral: Option<address>,
 }
 
 public(package) fun create_position_manager(ctx: &mut TxContext): PositionManager {
     PositionManager {
-        supply_shares: table::new(ctx),
+        positions: table::new(ctx),
     }
 }
 
 /// Increase the supply shares of the user and return outstanding supply shares.
 public(package) fun increase_user_supply(
     self: &mut PositionManager,
-    user: address,
+    referral: Option<address>,
     supply_shares: u64,
-): u64 {
-    self.add_supply_entry(user);
-    let user_supply_shares = self.supply_shares.borrow_mut(user);
-    *user_supply_shares = *user_supply_shares + supply_shares;
+    ctx: &TxContext,
+): (u64, Option<address>) {
+    let user = ctx.sender();
+    self.add_supply_entry(referral, ctx);
+    let user_position = self.positions.borrow_mut(user);
+    let current_referral = user_position.referral;
+    user_position.shares = user_position.shares + supply_shares;
+    user_position.referral = referral;
 
-    *user_supply_shares
+    (user_position.shares, current_referral)
 }
 
 /// Decrease the supply shares of the user and return outstanding supply shares.
 public(package) fun decrease_user_supply(
     self: &mut PositionManager,
-    user: address,
     supply_shares: u64,
-): u64 {
-    let user_supply_shares = self.supply_shares.borrow_mut(user);
-    *user_supply_shares = *user_supply_shares - supply_shares;
+    ctx: &TxContext,
+): (u64, Option<address>) {
+    let user = ctx.sender();
+    let user_position = self.positions.borrow_mut(user);
+    user_position.shares = user_position.shares - supply_shares;
 
-    *user_supply_shares
+    (user_position.shares, user_position.referral)
 }
 
-public(package) fun add_supply_entry(self: &mut PositionManager, user: address) {
-    if (!self.supply_shares.contains(user)) {
+public(package) fun add_supply_entry(
+    self: &mut PositionManager,
+    referral: Option<address>,
+    ctx: &TxContext,
+) {
+    let user = ctx.sender();
+    if (!self.positions.contains(user)) {
         self
-            .supply_shares
+            .positions
             .add(
                 user,
-                0,
+                Position {
+                    shares: 0,
+                    referral,
+                },
             );
     }
 }
 
-public(package) fun user_supply_shares(self: &PositionManager, user: address): u64 {
-    if (self.supply_shares.contains(user)) {
-        *self.supply_shares.borrow(user)
+public(package) fun user_supply_shares(self: &PositionManager, ctx: &TxContext): u64 {
+    let user = ctx.sender();
+    if (self.positions.contains(user)) {
+        self.positions.borrow(user).shares
     } else {
         0
     }