Skip to main content

satellite_bitcoin_transactions/
mempool.rs

1use arch_program::rune::RuneAmount;
2use satellite_collections::generic::fixed_set::FixedCapacitySet;
3
4use crate::UtxoInfo;
5
6#[derive(Clone, Copy, Debug, Default)]
7pub struct MempoolInfo {
8    pub total_fee: u64,
9    pub total_size: u64,
10}
11
12#[derive(Clone, Debug, Default)]
13pub enum TxStatus {
14    Pending(MempoolInfo),
15    #[default]
16    Confirmed,
17}
18
19#[derive(Debug, Clone)]
20pub struct AccountMempoolInfo {
21    pub ancestors_count: u16,
22    pub descendants_count: u16,
23    // Smaller value means higher similarity priority. 0 is the closest match.
24    pub priority_index: u16,
25}
26
27impl Default for AccountMempoolInfo {
28    fn default() -> Self {
29        Self {
30            ancestors_count: 0,
31            descendants_count: 0,
32            priority_index: u16::MAX,
33        }
34    }
35}
36
37#[derive(Debug)]
38pub struct MempoolData<const MAX_UTXOS: usize, const MAX_ACCOUNTS: usize> {
39    // Up to N user-provided UTXOs
40    utxo_mempool_info: [Option<([u8; 32], MempoolInfo)>; MAX_UTXOS],
41    accounts_utxo_mempool_info: [AccountMempoolInfo; MAX_ACCOUNTS],
42}
43
44impl<const MAX_UTXOS: usize, const MAX_ACCOUNTS: usize> MempoolData<MAX_UTXOS, MAX_ACCOUNTS> {
45    pub fn new(
46        utxo_mempool_info: [Option<([u8; 32], MempoolInfo)>; MAX_UTXOS],
47        accounts_utxo_mempool_info: [AccountMempoolInfo; MAX_ACCOUNTS],
48    ) -> Self {
49        Self {
50            utxo_mempool_info,
51            accounts_utxo_mempool_info,
52        }
53    }
54
55    pub fn get_utxo_status(&self, txid: [u8; 32]) -> TxStatus {
56        // Linear search is efficient for small collections and cache-friendly
57        for (stored_txid, info) in self.utxo_mempool_info.iter().flatten() {
58            if *stored_txid == txid {
59                return TxStatus::Pending(*info);
60            }
61        }
62        TxStatus::Confirmed
63    }
64
65    pub fn get_mempool_info_for_accounts(&self, n_accounts: usize) -> &[AccountMempoolInfo] {
66        &self.accounts_utxo_mempool_info.split_at(n_accounts).0
67    }
68
69    /// Compute a stable ordering of account indices [0, n_accounts) by similarity.
70    /// Lower `priority_index` comes first. Ties preserve input order.
71    pub fn account_order_by_similarity(&self, n_accounts: usize) -> [usize; MAX_ACCOUNTS] {
72        // Initialize indices 0..n_accounts
73        let mut indices: [usize; MAX_ACCOUNTS] = std::array::from_fn(|i| i);
74
75        // Simple stable insertion sort on the first n_accounts elements
76        let slice = &mut indices[..n_accounts];
77        for i in 1..slice.len() {
78            let key = slice[i];
79            let key_prio = self.accounts_utxo_mempool_info[key].priority_index;
80            let mut j = i;
81            while j > 0 {
82                let prev = slice[j - 1];
83                let prev_prio = self.accounts_utxo_mempool_info[prev].priority_index;
84                if prev_prio <= key_prio {
85                    break;
86                }
87                slice[j] = prev;
88                j -= 1;
89            }
90            slice[j] = key;
91        }
92
93        indices
94    }
95}
96
97impl<const MAX_UTXOS: usize, const MAX_ACCOUNTS: usize> Default
98    for MempoolData<MAX_UTXOS, MAX_ACCOUNTS>
99{
100    fn default() -> Self {
101        MempoolData::new(
102            [None; MAX_UTXOS],
103            std::array::from_fn(|_| AccountMempoolInfo::default()),
104        )
105    }
106}
107
108pub trait MempoolDataView {
109    fn get_utxo_status(&self, txid: [u8; 32]) -> TxStatus;
110    fn get_mempool_info_for_accounts(&self, n_accounts: usize) -> &[AccountMempoolInfo];
111}
112
113impl<const MAX_UTXOS: usize, const MAX_ACCOUNTS: usize> MempoolDataView
114    for MempoolData<MAX_UTXOS, MAX_ACCOUNTS>
115{
116    fn get_utxo_status(&self, txid: [u8; 32]) -> TxStatus {
117        self.get_utxo_status(txid)
118    }
119
120    fn get_mempool_info_for_accounts(&self, n_accounts: usize) -> &[AccountMempoolInfo] {
121        self.get_mempool_info_for_accounts(n_accounts)
122    }
123}
124
125pub(crate) fn generate_mempool_info<RuneSet: FixedCapacitySet<Item = RuneAmount>>(
126    user_utxos: &[UtxoInfo<RuneSet>],
127    mempool_data: &impl MempoolDataView,
128) -> MempoolInfo {
129    let mut mempool_info = MempoolInfo::default();
130    for (i, utxo) in user_utxos.iter().enumerate() {
131        let txid: [u8; 32] = utxo.meta.txid_big_endian();
132
133        // Check if we've already processed this txid in previous UTXOs
134        let already_processed = user_utxos[..i].iter().any(|prev_utxo| {
135            let prev_txid: [u8; 32] = prev_utxo.meta.txid_big_endian();
136            prev_txid == txid
137        });
138
139        if already_processed {
140            continue;
141        }
142
143        let status = mempool_data.get_utxo_status(txid);
144
145        match status {
146            TxStatus::Pending(info) => {
147                mempool_info.total_fee += info.total_fee;
148                mempool_info.total_size += info.total_size;
149            }
150            TxStatus::Confirmed => {}
151        }
152    }
153
154    mempool_info
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160    use crate::utxo_info::SingleRuneSet;
161    use arch_program::utxo::UtxoMeta;
162
163    /// Convenience helper to construct a mock `UtxoInfo` with the desired txid/vout.
164    fn make_utxo(txid: [u8; 32], vout: u32) -> UtxoInfo<SingleRuneSet> {
165        UtxoInfo::new(UtxoMeta::from(txid, vout), 0)
166    }
167
168    #[test]
169    fn aggregates_fees_and_sizes_for_unique_pending_txids() {
170        const MAX_UTXOS: usize = 4;
171        const MAX_ACCOUNTS: usize = 1;
172
173        // Two distinct pending transactions in the mempool.
174        let txid1 = [1u8; 32];
175        let txid2 = [2u8; 32];
176
177        let mut utxo_mempool_info: [Option<([u8; 32], MempoolInfo)>; MAX_UTXOS] = [None; MAX_UTXOS];
178        utxo_mempool_info[0] = Some((
179            txid1,
180            MempoolInfo {
181                total_fee: 100,
182                total_size: 200,
183            },
184        ));
185        utxo_mempool_info[1] = Some((
186            txid2,
187            MempoolInfo {
188                total_fee: 50,
189                total_size: 75,
190            },
191        ));
192
193        let accounts_info: [AccountMempoolInfo; MAX_ACCOUNTS] =
194            std::array::from_fn(|_| AccountMempoolInfo::default());
195
196        let mempool_data =
197            MempoolData::<MAX_UTXOS, MAX_ACCOUNTS>::new(utxo_mempool_info, accounts_info);
198
199        let user_utxos = vec![make_utxo(txid1, 0), make_utxo(txid2, 1)];
200
201        let info = generate_mempool_info::<SingleRuneSet>(&user_utxos, &mempool_data);
202
203        assert_eq!(info.total_fee, 150);
204        assert_eq!(info.total_size, 275);
205    }
206
207    #[test]
208    fn ignores_duplicate_txids() {
209        const MAX_UTXOS: usize = 4;
210        const MAX_ACCOUNTS: usize = 1;
211
212        let txid = [42u8; 32];
213
214        let mut utxo_mempool_info: [Option<([u8; 32], MempoolInfo)>; MAX_UTXOS] = [None; MAX_UTXOS];
215        utxo_mempool_info[0] = Some((
216            txid,
217            MempoolInfo {
218                total_fee: 80,
219                total_size: 160,
220            },
221        ));
222
223        let accounts_info: [AccountMempoolInfo; MAX_ACCOUNTS] =
224            std::array::from_fn(|_| AccountMempoolInfo::default());
225        let mempool_data =
226            MempoolData::<MAX_UTXOS, MAX_ACCOUNTS>::new(utxo_mempool_info, accounts_info);
227
228        // Two UTXOs share the same transaction id but different vouts.
229        let user_utxos = vec![make_utxo(txid, 0), make_utxo(txid, 1)];
230
231        let info = generate_mempool_info::<SingleRuneSet>(&user_utxos, &mempool_data);
232
233        assert_eq!(info.total_fee, 80);
234        assert_eq!(info.total_size, 160);
235    }
236
237    #[test]
238    fn ignores_confirmed_transactions() {
239        const MAX_UTXOS: usize = 4;
240        const MAX_ACCOUNTS: usize = 1;
241
242        let pending_txid = [7u8; 32];
243        let confirmed_txid = [8u8; 32];
244
245        let mut utxo_mempool_info: [Option<([u8; 32], MempoolInfo)>; MAX_UTXOS] = [None; MAX_UTXOS];
246        utxo_mempool_info[0] = Some((
247            pending_txid,
248            MempoolInfo {
249                total_fee: 30,
250                total_size: 60,
251            },
252        ));
253
254        let accounts_info: [AccountMempoolInfo; MAX_ACCOUNTS] =
255            std::array::from_fn(|_| AccountMempoolInfo::default());
256        let mempool_data =
257            MempoolData::<MAX_UTXOS, MAX_ACCOUNTS>::new(utxo_mempool_info, accounts_info);
258
259        // Only one UTXO is pending, the other txid is not in the mempool (confirmed).
260        let user_utxos = vec![make_utxo(pending_txid, 0), make_utxo(confirmed_txid, 0)];
261
262        let info = generate_mempool_info::<SingleRuneSet>(&user_utxos, &mempool_data);
263
264        assert_eq!(info.total_fee, 30);
265        assert_eq!(info.total_size, 60);
266    }
267}