pyra_redis/drift/

positions.rs

//! Typed position data fetching and computation.
//!
//! Fetches Drift position data (DriftUser, SpotMarket, prices) from Redis
//! and computes portfolio values in signed USD cents.

use std::collections::HashMap;

use pyra_tokens::AssetId;

use pyra_margin::get_token_balance;
use pyra_types::{Cache, DriftUser, SpotMarket, Vault};

use crate::{RedisClient, RedisError, RedisKey, RedisResult};

/// Convert a token balance (i128) to a signed USD value in cents using the given price and decimals.
///
/// Returns `MathOverflow` if the result doesn't fit in i64.
fn balance_to_value_cents(token_balance: i128, decimals: u32, price: f64) -> RedisResult<i64> {
    let decimals_pow = 10_f64.powi(i32::try_from(decimals).map_err(|_| RedisError::MathOverflow)?);
    let value = (token_balance as f64) / decimals_pow * price * 100.0;
    let rounded = value.round();
    if rounded.is_finite() && rounded >= i64::MIN as f64 && rounded <= i64::MAX as f64 {
        Ok(rounded as i64)
    } else {
        Err(RedisError::MathOverflow)
    }
}

// ── Data structures ──────────────────────────────────────────────────

/// Parsed Redis position data for a single vault.
pub struct VaultPositionData {
    /// The cached DriftUser account (includes slot information).
    pub drift_user: Cache<DriftUser>,
    /// Spot markets keyed by asset ID.
    pub spot_markets: HashMap<AssetId, SpotMarket>,
    /// Token prices (USD) keyed by asset ID.
    pub prices: HashMap<AssetId, f64>,
}

/// All drift account positions with shared market/price data.
pub struct AllDriftPositionsData {
    /// `(vault_address, drift_user_cache)` pairs for every drift account in Redis.
    pub drift_users: Vec<(String, Cache<DriftUser>)>,
    /// Shared spot markets keyed by asset ID.
    pub spot_markets: HashMap<AssetId, SpotMarket>,
    /// Shared token prices (USD) keyed by asset ID.
    pub prices: HashMap<AssetId, f64>,
    /// `vault_address → owner solana address` (only populated when requested).
    pub vault_owners: HashMap<String, String>,
    /// Number of drift user keys that failed to parse (skipped).
    /// Callers can use this for observability / alerting on high failure rates.
    pub skipped_drift_users: usize,
}

// ── Fetch methods on RedisClient ─────────────────────────────────────

impl RedisClient {
    /// Fetch DriftUser, SpotMarkets, and prices for a vault via a single MGET.
    ///
    /// Builds keys as: `[drift_user, spot_market_0..N, price_0..N]`
    /// and parses the results into typed structures.
    ///
    /// Returns `NotFound` if the DriftUser key is missing.
    pub async fn fetch_vault_position_data(
        &self,
        vault_address: &str,
        asset_ids: &[AssetId],
    ) -> RedisResult<VaultPositionData> {
        let drift_user_key = format!("{}:{vault_address}", RedisKey::DRIFT_USER_PREFIX);
        let mut keys = vec![drift_user_key];
        for &id in asset_ids {
            keys.push(RedisKey::drift_spot_market(id).to_string());
        }
        for &id in asset_ids {
            keys.push(RedisKey::price(id).to_string());
        }

        let values = self.mget(&keys).await?;

        let drift_user_raw = values
            .first()
            .and_then(|v| v.as_ref())
            .ok_or_else(|| RedisError::NotFound("DriftUser not found in Redis".into()))?;
        let drift_user: Cache<DriftUser> = serde_json::from_str(drift_user_raw)?;

        let num_markets = asset_ids.len();
        let mut spot_markets: HashMap<AssetId, SpotMarket> = HashMap::new();
        let mut prices: HashMap<AssetId, f64> = HashMap::new();

        for (i, &id) in asset_ids.iter().enumerate() {
            if let Some(Some(raw)) =
                values.get(1usize.checked_add(i).ok_or(RedisError::MathOverflow)?)
            {
                if let Ok(cache) = serde_json::from_str::<Cache<SpotMarket>>(raw) {
                    spot_markets.insert(id, cache.account);
                }
            }
            if let Some(Some(raw)) = values.get(
                1usize
                    .checked_add(num_markets)
                    .ok_or(RedisError::MathOverflow)?
                    .checked_add(i)
                    .ok_or(RedisError::MathOverflow)?,
            ) {
                if let Ok(price) = serde_json::from_str::<f64>(raw) {
                    prices.insert(id, price);
                }
            }
        }

        Ok(VaultPositionData {
            drift_user,
            spot_markets,
            prices,
        })
    }

    /// Fetch ALL drift user positions from Redis via SCAN + MGET.
    ///
    /// Scans every `account:drift:user:*` key and fetches all position data
    /// in a single MGET round-trip.
    ///
    /// When `include_vault_owners` is true, vault accounts are also fetched
    /// to map `vault_address → owner solana address`.
    pub async fn fetch_all_drift_positions(
        &self,
        asset_ids: &[AssetId],
        include_vault_owners: bool,
    ) -> RedisResult<AllDriftPositionsData> {
        // 1. Scan for all drift user keys
        let drift_keys = self
            .scan_keys(&RedisKey::pattern(RedisKey::DRIFT_USER_PREFIX))
            .await?;

        // Extract vault addresses from key names
        let prefix_with_colon = format!("{}:", RedisKey::DRIFT_USER_PREFIX);
        let vault_addresses: Vec<&str> = drift_keys
            .iter()
            .filter_map(|k| k.strip_prefix(prefix_with_colon.as_str()))
            .collect();

        // 2. Build a single MGET request for everything
        //    Layout: [drift_users...] [vaults... (optional)] [spot_markets...] [prices...]
        let num_drift = drift_keys.len();
        let mut all_keys: Vec<String> = drift_keys.clone();

        if include_vault_owners {
            for vault_addr in &vault_addresses {
                all_keys.push(format!("{}:{vault_addr}", RedisKey::VAULT_PREFIX));
            }
        }
        for &id in asset_ids {
            all_keys.push(RedisKey::drift_spot_market(id).to_string());
        }
        for &id in asset_ids {
            all_keys.push(RedisKey::price(id).to_string());
        }

        let values = self.mget(&all_keys).await?;

        // 3. Parse drift users (track failures for observability)
        let mut drift_users: Vec<(String, Cache<DriftUser>)> = Vec::new();
        let mut skipped_drift_users: usize = 0;
        for (i, vault_addr) in vault_addresses.iter().enumerate() {
            if let Some(Some(raw)) = values.get(i) {
                match serde_json::from_str::<Cache<DriftUser>>(raw) {
                    Ok(du) => drift_users.push(((*vault_addr).to_string(), du)),
                    Err(_) => {
                        skipped_drift_users = skipped_drift_users.saturating_add(1);
                    }
                }
            }
        }

        // 4. Parse vault owners if requested
        let mut vault_owners: HashMap<String, String> = HashMap::new();
        if include_vault_owners {
            for (i, vault_addr) in vault_addresses.iter().enumerate() {
                let offset = num_drift.checked_add(i).ok_or(RedisError::MathOverflow)?;
                if let Some(Some(raw)) = values.get(offset) {
                    if let Ok(vault_cache) = serde_json::from_str::<Cache<Vault>>(raw) {
                        vault_owners.insert(
                            (*vault_addr).to_string(),
                            vault_cache.account.owner.to_string(),
                        );
                    }
                }
            }
        }

        // 5. Parse shared spot markets and prices
        let vault_count = if include_vault_owners { num_drift } else { 0 };
        let market_base = num_drift
            .checked_add(vault_count)
            .ok_or(RedisError::MathOverflow)?;
        let num_markets = asset_ids.len();

        let mut spot_markets: HashMap<AssetId, SpotMarket> = HashMap::new();
        let mut prices: HashMap<AssetId, f64> = HashMap::new();

        for (i, &id) in asset_ids.iter().enumerate() {
            let market_offset = market_base.checked_add(i).ok_or(RedisError::MathOverflow)?;
            if let Some(Some(raw)) = values.get(market_offset) {
                if let Ok(cache) = serde_json::from_str::<Cache<SpotMarket>>(raw) {
                    spot_markets.insert(id, cache.account);
                }
            }

            let price_offset = market_base
                .checked_add(num_markets)
                .ok_or(RedisError::MathOverflow)?
                .checked_add(i)
                .ok_or(RedisError::MathOverflow)?;
            if let Some(Some(raw)) = values.get(price_offset) {
                if let Ok(price) = serde_json::from_str::<f64>(raw) {
                    prices.insert(id, price);
                }
            }
        }

        Ok(AllDriftPositionsData {
            drift_users,
            spot_markets,
            prices,
            vault_owners,
            skipped_drift_users,
        })
    }
}

// ── Computation helpers ──────────────────────────────────────────────

/// Compute the signed USD-cent value for each non-zero spot position.
///
/// Positive values are deposits; negative values are borrows.
pub fn compute_position_values(data: &VaultPositionData) -> RedisResult<Vec<i64>> {
    compute_user_position_values(&data.drift_user.account, &data.spot_markets, &data.prices)
}

/// Compute per-asset data: `(asset_id, signed_token_balance_i64, signed_value_cents)`.
pub fn compute_asset_data(data: &VaultPositionData) -> RedisResult<Vec<(AssetId, i64, i64)>> {
    compute_user_asset_data(&data.drift_user.account, &data.spot_markets, &data.prices)
}

/// Compute signed USD-cent position values for a single drift user,
/// using shared spot-market and price maps (keyed by asset ID).
pub fn compute_user_position_values(
    drift_user: &DriftUser,
    spot_markets: &HashMap<AssetId, SpotMarket>,
    prices: &HashMap<AssetId, f64>,
) -> RedisResult<Vec<i64>> {
    let mut results = Vec::new();
    for position in &drift_user.spot_positions {
        if position.scaled_balance == 0 {
            continue;
        }
        let Some(token) = pyra_tokens::Token::find_by_drift_market_index(position.market_index)
        else {
            continue;
        };
        let asset_id = token.asset_id;
        let Some(market) = spot_markets.get(&asset_id) else {
            continue;
        };
        let Some(&price) = prices.get(&asset_id) else {
            continue;
        };
        let token_balance = get_token_balance(position, market)?;
        let value_cents = balance_to_value_cents(token_balance, market.decimals, price)?;
        results.push(value_cents);
    }
    Ok(results)
}

/// Compute per-asset data for a single drift user using shared market/price maps.
///
/// Returns `(asset_id, signed_token_balance_i64, signed_value_cents)` tuples.
pub fn compute_user_asset_data(
    drift_user: &DriftUser,
    spot_markets: &HashMap<AssetId, SpotMarket>,
    prices: &HashMap<AssetId, f64>,
) -> RedisResult<Vec<(AssetId, i64, i64)>> {
    let mut results = Vec::new();
    for position in &drift_user.spot_positions {
        if position.scaled_balance == 0 {
            continue;
        }
        let Some(token) = pyra_tokens::Token::find_by_drift_market_index(position.market_index)
        else {
            continue;
        };
        let asset_id = token.asset_id;
        let Some(market) = spot_markets.get(&asset_id) else {
            continue;
        };
        let Some(&price) = prices.get(&asset_id) else {
            continue;
        };
        let token_balance_i128 = get_token_balance(position, market)?;
        let token_balance =
            i64::try_from(token_balance_i128).map_err(|_| RedisError::MathOverflow)?;
        let value_cents = balance_to_value_cents(token_balance_i128, market.decimals, price)?;
        results.push((asset_id, token_balance, value_cents));
    }
    Ok(results)
}

#[cfg(test)]
#[allow(
    clippy::allow_attributes,
    clippy::allow_attributes_without_reason,
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::arithmetic_side_effects
)]
mod tests {
    use super::*;
    use pyra_types::SpotBalanceType;

    fn make_spot_position(market_index: u16, scaled_balance: u64) -> pyra_types::SpotPosition {
        pyra_types::SpotPosition {
            market_index,
            scaled_balance,
            balance_type: SpotBalanceType::Deposit,
            ..Default::default()
        }
    }

    fn make_spot_market(market_index: u16, decimals: u32) -> SpotMarket {
        // Precision = 10^(19 - decimals). At 1.0x interest, cumulative = precision.
        let precision = 10u128.pow(19u32.saturating_sub(decimals));
        SpotMarket {
            pubkey: vec![],
            market_index,
            initial_asset_weight: 0,
            initial_liability_weight: 0,
            imf_factor: 0,
            scale_initial_asset_weight_start: 0,
            decimals,
            cumulative_deposit_interest: precision,
            cumulative_borrow_interest: precision,
            deposit_balance: 0,
            borrow_balance: 0,
            optimal_utilization: 0,
            optimal_borrow_rate: 0,
            max_borrow_rate: 0,
            min_borrow_rate: 0,
            insurance_fund: Default::default(),
            historical_oracle_data: Default::default(),
            oracle: None,
        }
    }

    fn make_drift_user(positions: Vec<pyra_types::SpotPosition>) -> DriftUser {
        DriftUser {
            authority: Default::default(),
            spot_positions: positions,
        }
    }

    #[test]
    fn compute_position_values_basic() {
        let drift_user = make_drift_user(vec![make_spot_position(0, 1_000_000)]);
        let mut spot_markets = HashMap::new();
        spot_markets.insert(AssetId::new(0).unwrap(), make_spot_market(0, 6));
        let mut prices = HashMap::new();
        prices.insert(AssetId::new(0).unwrap(), 1.0);

        let values = compute_user_position_values(&drift_user, &spot_markets, &prices).unwrap();
        assert_eq!(values.len(), 1);
        assert_eq!(values[0], 100); // 1 USDC = 100 cents
    }

    #[test]
    fn compute_position_values_multiple_markets() {
        let drift_user = make_drift_user(vec![
            make_spot_position(0, 2_000_000),   // 2 USDC
            make_spot_position(1, 100_000_000), // 0.1 SOL (9 decimals)
        ]);
        let mut spot_markets = HashMap::new();
        spot_markets.insert(AssetId::new(0).unwrap(), make_spot_market(0, 6));
        spot_markets.insert(AssetId::new(1).unwrap(), make_spot_market(1, 9));
        let mut prices = HashMap::new();
        prices.insert(AssetId::new(0).unwrap(), 1.0);
        prices.insert(AssetId::new(1).unwrap(), 150.0);

        let values = compute_user_position_values(&drift_user, &spot_markets, &prices).unwrap();
        assert_eq!(values.len(), 2);
        assert_eq!(values[0], 200); // 2 USDC = 200 cents
        assert_eq!(values[1], 1500); // 0.1 SOL * $150 = $15 = 1500 cents
    }

    #[test]
    fn compute_position_values_skips_zero_balance() {
        let drift_user = make_drift_user(vec![
            make_spot_position(0, 0),         // zero — skip
            make_spot_position(1, 1_000_000), // 1 USDC
        ]);
        let mut spot_markets = HashMap::new();
        spot_markets.insert(AssetId::new(0).unwrap(), make_spot_market(0, 6));
        spot_markets.insert(AssetId::new(1).unwrap(), make_spot_market(1, 6));
        let mut prices = HashMap::new();
        prices.insert(AssetId::new(0).unwrap(), 1.0);
        prices.insert(AssetId::new(1).unwrap(), 1.0);

        let values = compute_user_position_values(&drift_user, &spot_markets, &prices).unwrap();
        assert_eq!(values.len(), 1);
        assert_eq!(values[0], 100);
    }

    #[test]
    fn compute_position_values_skips_missing_market() {
        let drift_user = make_drift_user(vec![make_spot_position(99, 1_000_000)]);
        let spot_markets = HashMap::new();
        let prices = HashMap::new();

        let values = compute_user_position_values(&drift_user, &spot_markets, &prices).unwrap();
        assert!(values.is_empty());
    }

    #[test]
    fn compute_position_values_skips_missing_price() {
        let drift_user = make_drift_user(vec![make_spot_position(0, 1_000_000)]);
        let mut spot_markets = HashMap::new();
        spot_markets.insert(AssetId::new(0).unwrap(), make_spot_market(0, 6));
        let prices = HashMap::new();

        let values = compute_user_position_values(&drift_user, &spot_markets, &prices).unwrap();
        assert!(values.is_empty());
    }

    #[test]
    fn compute_asset_data_returns_tuples() {
        let drift_user = make_drift_user(vec![make_spot_position(0, 5_000_000)]);
        let mut spot_markets = HashMap::new();
        spot_markets.insert(AssetId::new(0).unwrap(), make_spot_market(0, 6));
        let mut prices = HashMap::new();
        prices.insert(AssetId::new(0).unwrap(), 1.0);

        let data = compute_user_asset_data(&drift_user, &spot_markets, &prices).unwrap();
        assert_eq!(data.len(), 1);
        let (asset_id, token_balance, value_cents) = data[0];
        assert_eq!(asset_id, AssetId::new(0).unwrap());
        assert_eq!(token_balance, 5_000_000);
        assert_eq!(value_cents, 500);
    }

    #[test]
    fn compute_position_values_delegates_to_user_variant() {
        let drift_user = Cache {
            account: make_drift_user(vec![make_spot_position(0, 1_000_000)]),
            last_updated_slot: 12345,
        };
        let mut spot_markets = HashMap::new();
        spot_markets.insert(AssetId::new(0).unwrap(), make_spot_market(0, 6));
        let mut prices = HashMap::new();
        prices.insert(AssetId::new(0).unwrap(), 1.0);

        let vpd = VaultPositionData {
            drift_user,
            spot_markets,
            prices,
        };
        let values = compute_position_values(&vpd).unwrap();
        assert_eq!(values.len(), 1);
        assert_eq!(values[0], 100);
    }
}