sim-cli 0.7.0

//! Per-partition join and P&L diff scoring.

use std::collections::{BTreeSet, HashMap};

use super::model::{
    CompareTransaction, ImprovementEntry, LogDiffEntry, PnlDiff, RegressionEntry, SolDiff, SplDiff,
};

impl PnlDiff {
    fn new(base: &CompareTransaction, exp: &CompareTransaction) -> Option<Self> {
        let base_sol: HashMap<&str, i64> = base
            .sol_changes
            .iter()
            .map(|c| (c.pubkey.as_str(), c.delta()))
            .collect();
        let exp_sol: HashMap<&str, i64> = exp
            .sol_changes
            .iter()
            .map(|c| (c.pubkey.as_str(), c.delta()))
            .collect();
        let sol_keys: BTreeSet<&str> = base_sol.keys().chain(exp_sol.keys()).copied().collect();
        let sol_diffs: Vec<SolDiff> = sol_keys
            .into_iter()
            .filter_map(|pk| {
                let b = base_sol.get(pk).copied().unwrap_or(0);
                let e = exp_sol.get(pk).copied().unwrap_or(0);
                if b == e {
                    return None;
                }
                Some(SolDiff {
                    pubkey: pk.to_string(),
                    baseline_delta: b,
                    experiment_delta: e,
                })
            })
            .collect();

        let base_spl: HashMap<(&str, &str), (i64, u8)> = base
            .token_changes
            .iter()
            .map(|c| {
                (
                    (c.pubkey.as_str(), c.mint.as_str()),
                    (c.delta(), c.decimals),
                )
            })
            .collect();
        let exp_spl: HashMap<(&str, &str), (i64, u8)> = exp
            .token_changes
            .iter()
            .map(|c| {
                (
                    (c.pubkey.as_str(), c.mint.as_str()),
                    (c.delta(), c.decimals),
                )
            })
            .collect();
        let spl_keys: BTreeSet<(&str, &str)> =
            base_spl.keys().chain(exp_spl.keys()).copied().collect();
        let token_diffs: Vec<SplDiff> = spl_keys
            .into_iter()
            .filter_map(|(pk, mint)| {
                let (b, decimals) = base_spl.get(&(pk, mint)).copied().unwrap_or_else(|| {
                    exp_spl
                        .get(&(pk, mint))
                        .map(|(_, dec)| (0, *dec))
                        .unwrap_or((0, 0))
                });
                let (e, _) = exp_spl.get(&(pk, mint)).copied().unwrap_or((0, decimals));
                if b == e {
                    return None;
                }
                Some(SplDiff {
                    pubkey: pk.to_string(),
                    mint: mint.to_string(),
                    symbol: String::new(),
                    decimals,
                    baseline_delta: b,
                    experiment_delta: e,
                })
            })
            .collect();

        let normalized = |base: i64, exp: i64| -> f64 {
            let diff = (exp - base).unsigned_abs() as f64;
            let base_abs = base.unsigned_abs() as f64;
            // geometric mean of absolute diff and base magnitude; weights by position size
            // so small changes on tiny positions don't dominate the ranking.
            // fall back to raw diff when base is zero (still a valid change, just no base to weight against)
            if base == 0 {
                diff
            } else {
                (diff * base_abs).sqrt()
            }
        };
        let discrepancy: f64 = sol_diffs
            .iter()
            .map(|d| normalized(d.baseline_delta, d.experiment_delta))
            .sum::<f64>()
            + token_diffs
                .iter()
                .map(|d| normalized(d.baseline_delta, d.experiment_delta))
                .sum::<f64>();

        if discrepancy == 0.0 {
            return None;
        }
        Some(PnlDiff {
            slot: base.slot,
            signature: base.signature.clone(),
            sol: sol_diffs,
            tokens: token_diffs,
            score: discrepancy,
        })
    }
}

#[derive(Default)]
pub(super) struct CompareAccumulator {
    pub(super) regressions: Vec<RegressionEntry>,
    pub(super) improvements: Vec<ImprovementEntry>,
    pub(super) log_diffs: Vec<LogDiffEntry>,
    pub(super) missing: Vec<String>,
    pub(super) pnl_diffs: Vec<PnlDiff>,
    pub(super) common: usize,
}

impl CompareAccumulator {
    /// Fold another partition's accumulator into this one. Ordering across
    /// partitions is irrelevant: the caller sorts every section afterwards.
    pub(super) fn merge(&mut self, other: Self) {
        self.regressions.extend(other.regressions);
        self.improvements.extend(other.improvements);
        self.log_diffs.extend(other.log_diffs);
        self.missing.extend(other.missing);
        self.pnl_diffs.extend(other.pnl_diffs);
        self.common += other.common;
    }

    /// Join two signature-keyed partition maps (already deduped by
    /// `read_partition`); every section is sorted afterwards, so map
    /// iteration order is irrelevant.
    pub(super) fn join_partition(
        &mut self,
        base_map: HashMap<String, CompareTransaction>,
        exp_map: HashMap<String, CompareTransaction>,
    ) {
        for (sig, base_tx) in base_map {
            let Some(exp_tx) = exp_map.get(&sig) else {
                self.missing.push(sig);
                continue;
            };
            self.common += 1;
            if base_tx.success
                && exp_tx.success
                && let Some(diff) = PnlDiff::new(&base_tx, exp_tx)
            {
                self.pnl_diffs.push(diff);
            }
            match (base_tx.success, exp_tx.success) {
                (true, false) => self.regressions.push(RegressionEntry {
                    slot: base_tx.slot,
                    signature: base_tx.signature,
                    error: exp_tx.error.clone(),
                }),
                (false, true) => self.improvements.push(ImprovementEntry {
                    slot: base_tx.slot,
                    signature: base_tx.signature,
                    was: base_tx.error,
                }),
                _ => {
                    if base_tx.log_count != exp_tx.log_count {
                        self.log_diffs.push(LogDiffEntry {
                            slot: base_tx.slot,
                            signature: base_tx.signature,
                            baseline_log_count: base_tx.log_count,
                            experiment_log_count: exp_tx.log_count,
                        });
                    }
                }
            }
        }
    }
}