hypercall/rsm/

ledger.rs

//! Balance ledger types.
//!
//! `BalanceLedger` is the engine-owned runtime source of truth for account
//! balances. Runtime cash movement must be reduced through `BalanceUpdate`
//! effects so replay, snapshots, and followers share one ordered contract.

use async_trait::async_trait;
use hypercall_types::{BalanceUpdate, WalletAddress};
use rust_decimal::Decimal;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::sync::Arc;

/// Error type for ledger operations.
#[derive(Debug, Clone)]
pub enum LedgerError {
    /// Database error
    Db(String),
    /// Invalid operation
    InvalidOperation(String),
}

impl fmt::Display for LedgerError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            LedgerError::Db(msg) => write!(f, "Ledger DB error: {}", msg),
            LedgerError::InvalidOperation(msg) => write!(f, "Invalid ledger operation: {}", msg),
        }
    }
}

impl std::error::Error for LedgerError {}

/// Durable account balance persistence interface.
///
/// The engine-owned `BalanceLedger` is authoritative for runtime reads. This
/// trait exists for DB persistence adapters and tests that exercise durable
/// account balance writes:
/// - Position closes (realized PnL)
/// - Explicit deposits/withdrawals
/// - Premium settlement for Standard margin mode
///
/// Opening positions does NOT touch the ledger in Portfolio mode.
#[async_trait]
pub trait Ledger: Send + Sync {
    /// Get the current collateral balance for an account.
    ///
    /// Returns 0 if the account doesn't exist (lazy account creation).
    async fn get_balance(&self, wallet: &WalletAddress) -> Result<Decimal, LedgerError>;

    /// Apply realized PnL to an account's collateral balance.
    ///
    /// This is called when positions are closed:
    /// - Positive realized_pnl: balance increases (profitable close)
    /// - Negative realized_pnl: balance decreases (losing close)
    ///
    /// Creates the account if it doesn't exist.
    async fn apply_pnl(
        &self,
        wallet: &WalletAddress,
        realized_pnl: Decimal,
    ) -> Result<(), LedgerError>;

    /// Set the balance directly (for deposits, withdrawals, and tests).
    ///
    /// Creates the account if it doesn't exist.
    async fn set_balance(
        &self,
        wallet: &WalletAddress,
        balance: Decimal,
    ) -> Result<(), LedgerError>;

    /// Apply option premium settlement to an account's balance.
    ///
    /// This is called for Standard margin mode accounts on option fills:
    /// - Negative premium_delta: balance decreases (buying options)
    /// - Positive premium_delta: balance increases (selling options)
    ///
    /// Creates the account if it doesn't exist.
    ///
    /// Note: For Portfolio margin mode, this is NOT called - premium is
    /// financed via margin instead of directly debiting the balance.
    async fn apply_premium(
        &self,
        wallet: &WalletAddress,
        premium_delta: Decimal,
    ) -> Result<(), LedgerError> {
        // Default implementation delegates to apply_pnl since they have the same semantics
        self.apply_pnl(wallet, premium_delta).await
    }
}

/// Read-only source for runtime account balances.
#[async_trait]
pub trait BalanceProvider: Send + Sync {
    async fn get_balance(&self, wallet: &WalletAddress) -> Result<Decimal, LedgerError>;
}

pub struct LedgerBalanceProvider {
    ledger: Arc<dyn Ledger + Send + Sync>,
}

impl LedgerBalanceProvider {
    pub fn new(ledger: Arc<dyn Ledger + Send + Sync>) -> Self {
        Self { ledger }
    }
}

#[async_trait]
impl BalanceProvider for LedgerBalanceProvider {
    async fn get_balance(&self, wallet: &WalletAddress) -> Result<Decimal, LedgerError> {
        self.ledger.get_balance(wallet).await
    }
}

/// Engine-owned runtime balance ledger.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct BalanceLedger {
    balances: HashMap<WalletAddress, Decimal>,
    last_balance_update_seq: u64,
    last_applied_balance_update: Option<BalanceUpdate>,
}

impl BalanceLedger {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn from_map(balances: HashMap<WalletAddress, Decimal>) -> Self {
        Self::from_map_with_sequence(balances, 0)
    }

    pub fn from_map_with_sequence(
        balances: HashMap<WalletAddress, Decimal>,
        last_balance_update_seq: u64,
    ) -> Self {
        let mut ledger = Self::new();
        for (wallet, balance) in balances {
            ledger.set(wallet, balance);
        }
        ledger.last_balance_update_seq = last_balance_update_seq;
        ledger
    }

    pub fn balance(&self, wallet: &WalletAddress) -> Decimal {
        self.balances.get(wallet).copied().unwrap_or(Decimal::ZERO)
    }

    pub fn balances(&self) -> &HashMap<WalletAddress, Decimal> {
        &self.balances
    }

    pub fn last_balance_update_seq(&self) -> u64 {
        self.last_balance_update_seq
    }

    pub fn next_balance_update_seq(&self) -> u64 {
        self.last_balance_update_seq
            .checked_add(1)
            .expect("STATE_CORRUPTION: balance_update_seq overflow")
    }

    pub fn apply_balance_update(&mut self, update: &BalanceUpdate) -> Result<bool, LedgerError> {
        if update.balance_update_seq == self.last_balance_update_seq {
            match self.last_applied_balance_update.as_ref() {
                Some(last_update) if last_update == update => return Ok(false),
                Some(_) => {
                    return Err(LedgerError::InvalidOperation(format!(
                        "balance update seq {} duplicate payload mismatch",
                        update.balance_update_seq
                    )));
                }
                None => return Ok(false),
            }
        }

        if update.balance_update_seq < self.last_balance_update_seq {
            return Err(LedgerError::InvalidOperation(format!(
                "balance update seq {} is older than last applied seq {}",
                update.balance_update_seq, self.last_balance_update_seq
            )));
        }

        let expected_seq = self.next_balance_update_seq();
        if update.balance_update_seq != expected_seq {
            return Err(LedgerError::InvalidOperation(format!(
                "balance update seq gap: expected {}, got {}",
                expected_seq, update.balance_update_seq
            )));
        }

        let current = self.balance(&update.wallet);
        let expected_balance_after = current + update.delta;
        if expected_balance_after != update.balance_after {
            return Err(LedgerError::InvalidOperation(format!(
                "balance update seq {} for {} has invalid balance_after: current {} + delta {} = {}, got {}",
                update.balance_update_seq,
                update.wallet,
                current,
                update.delta,
                expected_balance_after,
                update.balance_after
            )));
        }

        self.set(update.wallet, update.balance_after);
        self.last_balance_update_seq = update.balance_update_seq;
        self.last_applied_balance_update = Some(update.clone());
        Ok(true)
    }

    pub fn set(&mut self, wallet: WalletAddress, balance: Decimal) {
        if balance == Decimal::ZERO {
            self.balances.remove(&wallet);
        } else {
            self.balances.insert(wallet, balance);
        }
    }

    pub fn insert(&mut self, wallet: WalletAddress, balance: Decimal) -> Option<Decimal> {
        let previous = self.balances.get(&wallet).copied();
        self.set(wallet, balance);
        previous
    }

    pub fn add_delta(&mut self, wallet: WalletAddress, delta: Decimal) {
        if delta == Decimal::ZERO {
            return;
        }

        let next = self.balance(&wallet) + delta;
        self.set(wallet, next);
    }

    pub fn remove_zero(&mut self, wallet: &WalletAddress) {
        if self.balance(wallet) == Decimal::ZERO {
            self.balances.remove(wallet);
        }
    }

    pub fn snapshot_map(&self) -> HashMap<WalletAddress, Decimal> {
        self.balances.clone()
    }

    pub fn sorted_nonzero_entries(&self) -> Vec<(WalletAddress, Decimal)> {
        let mut entries: Vec<_> = self
            .balances
            .iter()
            .filter_map(|(wallet, balance)| {
                if *balance == Decimal::ZERO {
                    None
                } else {
                    Some((*wallet, *balance))
                }
            })
            .collect();
        entries.sort_by_key(|(wallet, _)| *wallet);
        entries
    }

    pub fn nonzero_wallets(&self) -> Vec<WalletAddress> {
        self.sorted_nonzero_entries()
            .into_iter()
            .map(|(wallet, _)| wallet)
            .collect()
    }
}

impl Deref for BalanceLedger {
    type Target = HashMap<WalletAddress, Decimal>;

    fn deref(&self) -> &Self::Target {
        &self.balances
    }
}

impl DerefMut for BalanceLedger {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.balances
    }
}

/// In-memory test double for code that still exercises the async durable ledger
/// trait.
#[derive(Default)]
pub struct InMemoryLedger {
    balances: std::sync::Arc<tokio::sync::RwLock<BalanceLedger>>,
}

impl InMemoryLedger {
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn nonzero_wallets(&self) -> Vec<WalletAddress> {
        let balances = self.balances.read().await;
        balances.nonzero_wallets()
    }
}

#[async_trait]
impl Ledger for InMemoryLedger {
    async fn get_balance(&self, wallet: &WalletAddress) -> Result<Decimal, LedgerError> {
        let balances = self.balances.read().await;
        Ok(balances.balance(wallet))
    }

    async fn apply_pnl(
        &self,
        wallet: &WalletAddress,
        realized_pnl: Decimal,
    ) -> Result<(), LedgerError> {
        let mut balances = self.balances.write().await;
        balances.add_delta(*wallet, realized_pnl);
        Ok(())
    }

    async fn set_balance(
        &self,
        wallet: &WalletAddress,
        balance: Decimal,
    ) -> Result<(), LedgerError> {
        let mut balances = self.balances.write().await;
        balances.set(*wallet, balance);
        Ok(())
    }
}

/// Mock balance provider for tests -- always returns zero balance.
#[cfg(any(test, feature = "test-utils"))]
pub struct MockBalanceProvider;

#[cfg(any(test, feature = "test-utils"))]
#[async_trait]
impl BalanceProvider for MockBalanceProvider {
    async fn get_balance(&self, _wallet: &WalletAddress) -> Result<Decimal, LedgerError> {
        Ok(Decimal::ZERO)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use hypercall_types::{wallet_address::test_wallet, BalanceUpdateReason};
    use rust_decimal_macros::dec;

    fn balance_update(
        ledger: &BalanceLedger,
        wallet: WalletAddress,
        delta: Decimal,
    ) -> BalanceUpdate {
        BalanceUpdate {
            balance_update_seq: ledger.next_balance_update_seq(),
            wallet,
            delta,
            balance_after: ledger.balance(&wallet) + delta,
            reason: BalanceUpdateReason::Deposit,
            reference_id: Some("test-reference".to_string()),
            source_command_id: Some(1),
            timestamp_ms: 1234,
        }
    }

    fn sequenced_balance_update(
        seq: u64,
        wallet: WalletAddress,
        delta: Decimal,
        balance_after: Decimal,
    ) -> BalanceUpdate {
        BalanceUpdate {
            balance_update_seq: seq,
            wallet,
            delta,
            balance_after,
            reason: BalanceUpdateReason::Deposit,
            reference_id: Some(format!("test-reference-{seq}")),
            source_command_id: Some(seq as i64),
            timestamp_ms: 1234 + seq,
        }
    }

    #[tokio::test]
    async fn test_in_memory_ledger_get_balance_nonexistent() {
        let ledger = InMemoryLedger::new();
        let balance = ledger.get_balance(&test_wallet(1)).await.unwrap();
        assert_eq!(balance, dec!(0));
    }

    #[tokio::test]
    async fn test_in_memory_ledger_set_and_get_balance() {
        let ledger = InMemoryLedger::new();
        ledger
            .set_balance(&test_wallet(1), dec!(10000))
            .await
            .unwrap();
        let balance = ledger.get_balance(&test_wallet(1)).await.unwrap();
        assert_eq!(balance, dec!(10000));
    }

    #[tokio::test]
    async fn test_in_memory_ledger_nonzero_wallets() {
        let ledger = InMemoryLedger::new();
        ledger.set_balance(&test_wallet(1), dec!(25)).await.unwrap();
        ledger
            .set_balance(&test_wallet(2), rust_decimal::Decimal::ZERO)
            .await
            .unwrap();
        ledger.set_balance(&test_wallet(3), dec!(-5)).await.unwrap();

        let mut wallets = ledger.nonzero_wallets().await;
        wallets.sort_unstable();

        assert_eq!(wallets, vec![test_wallet(1), test_wallet(3)]);
    }

    #[test]
    fn apply_balance_update_validates_and_applies_delta() {
        let wallet = test_wallet(5);
        let mut ledger = BalanceLedger::new();
        ledger.set(wallet, dec!(100));

        let update = balance_update(&ledger, wallet, dec!(25));
        assert!(ledger.apply_balance_update(&update).unwrap());

        assert_eq!(ledger.balance(&wallet), dec!(125));
        assert_eq!(ledger.last_balance_update_seq(), update.balance_update_seq);
    }

    #[test]
    fn apply_balance_update_removes_zero_balance() {
        let wallet = test_wallet(6);
        let mut ledger = BalanceLedger::new();
        ledger.set(wallet, dec!(100));

        let update = balance_update(&ledger, wallet, dec!(-100));
        ledger.apply_balance_update(&update).unwrap();

        assert_eq!(ledger.balance(&wallet), Decimal::ZERO);
        assert!(!ledger.snapshot_map().contains_key(&wallet));
    }

    #[test]
    fn apply_balance_update_rejects_invalid_balance_after() {
        let wallet = test_wallet(7);
        let mut ledger = BalanceLedger::new();
        ledger.set(wallet, dec!(100));

        let mut update = balance_update(&ledger, wallet, dec!(25));
        update.balance_after = dec!(126);

        let error = ledger.apply_balance_update(&update).unwrap_err();
        assert!(
            error.to_string().contains("has invalid balance_after"),
            "unexpected error: {error}"
        );
        assert_eq!(ledger.balance(&wallet), dec!(100));
    }

    #[test]
    fn apply_balance_update_rejects_sequence_gap() {
        let wallet = test_wallet(8);
        let mut ledger = BalanceLedger::new();

        let mut update = balance_update(&ledger, wallet, dec!(10));
        update.balance_update_seq += 1;

        let error = ledger.apply_balance_update(&update).unwrap_err();
        assert!(
            error.to_string().contains("balance update seq gap"),
            "unexpected error: {error}"
        );
    }

    #[test]
    fn apply_balance_update_accepts_exact_duplicate() {
        let wallet = test_wallet(9);
        let mut ledger = BalanceLedger::new();

        let update = balance_update(&ledger, wallet, dec!(10));
        assert!(ledger.apply_balance_update(&update).unwrap());
        assert!(!ledger.apply_balance_update(&update).unwrap());
        assert_eq!(ledger.balance(&wallet), dec!(10));
    }

    #[test]
    fn apply_balance_update_rejects_duplicate_payload_mismatch() {
        let wallet = test_wallet(10);
        let mut ledger = BalanceLedger::new();

        let update = balance_update(&ledger, wallet, dec!(10));
        ledger.apply_balance_update(&update).unwrap();

        let mut duplicate = update.clone();
        duplicate.reference_id = Some("different".to_string());

        let error = ledger.apply_balance_update(&duplicate).unwrap_err();
        assert!(
            error.to_string().contains("duplicate payload mismatch"),
            "unexpected error: {error}"
        );
    }

    #[test]
    fn follower_replay_from_snapshot_converges_to_engine_ledger() {
        let wallet_a = test_wallet(11);
        let wallet_b = test_wallet(12);
        let mut snapshot_balances = HashMap::new();
        snapshot_balances.insert(wallet_a, dec!(100));
        snapshot_balances.insert(wallet_b, dec!(50));
        let mut follower = BalanceLedger::from_map_with_sequence(snapshot_balances, 10);

        let updates = [
            sequenced_balance_update(11, wallet_a, dec!(25), dec!(125)),
            sequenced_balance_update(12, wallet_b, dec!(-50), dec!(0)),
            sequenced_balance_update(13, wallet_a, dec!(-5), dec!(120)),
        ];
        for update in &updates {
            assert!(follower.apply_balance_update(update).unwrap());
        }

        let mut expected = BalanceLedger::new();
        expected.set(wallet_a, dec!(120));
        expected.last_balance_update_seq = 13;
        expected.last_applied_balance_update = Some(updates[2].clone());
        assert_eq!(follower, expected);
        assert!(!follower.snapshot_map().contains_key(&wallet_b));
    }

    #[test]
    fn follower_replay_duplicate_after_snapshot_is_noop() {
        let wallet = test_wallet(13);
        let mut snapshot_balances = HashMap::new();
        snapshot_balances.insert(wallet, dec!(100));
        let mut follower = BalanceLedger::from_map_with_sequence(snapshot_balances, 2);
        let update = sequenced_balance_update(3, wallet, dec!(10), dec!(110));

        assert!(follower.apply_balance_update(&update).unwrap());
        assert!(!follower.apply_balance_update(&update).unwrap());
        assert_eq!(follower.balance(&wallet), dec!(110));
        assert_eq!(follower.last_balance_update_seq(), 3);
    }

    #[test]
    fn follower_replay_duplicate_snapshot_boundary_without_payload_is_noop() {
        let wallet = test_wallet(15);
        let mut snapshot_balances = HashMap::new();
        snapshot_balances.insert(wallet, dec!(100));
        let mut follower = BalanceLedger::from_map_with_sequence(snapshot_balances, 3);
        let update = sequenced_balance_update(3, wallet, dec!(10), dec!(110));

        assert!(!follower.apply_balance_update(&update).unwrap());
        assert_eq!(follower.balance(&wallet), dec!(100));
        assert_eq!(follower.last_balance_update_seq(), 3);
    }

    #[test]
    fn follower_replay_gap_after_snapshot_fails_closed() {
        let wallet = test_wallet(14);
        let mut snapshot_balances = HashMap::new();
        snapshot_balances.insert(wallet, dec!(100));
        let mut follower = BalanceLedger::from_map_with_sequence(snapshot_balances, 7);
        let update = sequenced_balance_update(9, wallet, dec!(10), dec!(110));

        let error = follower.apply_balance_update(&update).unwrap_err();
        assert!(
            error.to_string().contains("seq gap"),
            "unexpected error: {error}"
        );
        assert_eq!(follower.balance(&wallet), dec!(100));
        assert_eq!(follower.last_balance_update_seq(), 7);
    }

    #[tokio::test]
    async fn test_in_memory_ledger_apply_pnl_positive() {
        let ledger = InMemoryLedger::new();
        ledger
            .set_balance(&test_wallet(1), dec!(10000))
            .await
            .unwrap();
        ledger.apply_pnl(&test_wallet(1), dec!(500)).await.unwrap();
        let balance = ledger.get_balance(&test_wallet(1)).await.unwrap();
        assert_eq!(balance, dec!(10500));
    }

    #[tokio::test]
    async fn test_in_memory_ledger_apply_pnl_negative() {
        let ledger = InMemoryLedger::new();
        ledger
            .set_balance(&test_wallet(1), dec!(10000))
            .await
            .unwrap();
        ledger.apply_pnl(&test_wallet(1), dec!(-300)).await.unwrap();
        let balance = ledger.get_balance(&test_wallet(1)).await.unwrap();
        assert_eq!(balance, dec!(9700));
    }

    #[tokio::test]
    async fn test_in_memory_ledger_apply_pnl_creates_account() {
        let ledger = InMemoryLedger::new();
        // Apply PnL to non-existent account
        ledger.apply_pnl(&test_wallet(1), dec!(500)).await.unwrap();
        let balance = ledger.get_balance(&test_wallet(1)).await.unwrap();
        assert_eq!(balance, dec!(500));
    }
}