saturn_mempool_oracle/
mempool_entries.rs

1use std::collections::HashMap;
2
3use borsh::{BorshDeserialize, BorshSerialize};
4use serde::{Deserialize, Serialize};
5use serde_with::{serde_as, Same};
6
7use crate::{mempool_entry::MempoolEntry, txid::Txid};
8
9/// Maximum capacity within Solana's 5MB account size limit
10const MAX_ENTRIES: usize = 25_600;
11const BUCKET_COUNT: usize = 256;
12const BUCKET_SIZE: usize = MAX_ENTRIES / BUCKET_COUNT;
13
14#[serde_as]
15#[derive(Debug, BorshSerialize, BorshDeserialize, Serialize, Deserialize, Clone)]
16#[repr(C)]
17pub struct MempoolEntries {
18    #[serde_as(as = "[Same; BUCKET_COUNT]")]
19    pub counts: [u16; BUCKET_COUNT],
20    #[serde_as(as = "[[Same; BUCKET_SIZE]; BUCKET_COUNT]")]
21    pub entries: [[MempoolEntry; BUCKET_SIZE]; BUCKET_COUNT],
22}
23
24impl MempoolEntries {
25    fn bucket_for(&self, txid: &Txid) -> usize {
26        // Use first 2 bytes for better distribution
27        ((txid[0] as usize) << 8 | (txid[1] as usize)) % BUCKET_COUNT
28    }
29
30    pub fn get(&self, txid: &Txid) -> Option<&MempoolEntry> {
31        let bucket = self.bucket_for(txid);
32        let count = self.counts[bucket] as usize;
33
34        // Linear search within the bucket
35        for i in 0..count {
36            if &self.entries[bucket][i].txid == txid {
37                return Some(&self.entries[bucket][i]);
38            }
39        }
40        None
41    }
42
43    pub fn iter(&self) -> impl Iterator<Item = &MempoolEntry> {
44        (0..BUCKET_COUNT).flat_map(move |bucket| {
45            let count = self.counts[bucket] as usize;
46            (0..count).map(move |i| &self.entries[bucket][i])
47        })
48    }
49
50    pub fn keys(&self) -> impl Iterator<Item = &Txid> {
51        (0..BUCKET_COUNT).flat_map(move |bucket| {
52            let count = self.counts[bucket] as usize;
53            (0..count).map(move |i| &self.entries[bucket][i].txid)
54        })
55    }
56
57    pub fn contains_key(&self, txid: &Txid) -> bool {
58        self.get(txid).is_some()
59    }
60
61    pub fn extend(&mut self, other: Self) {
62        for entry in other.iter() {
63            self.add_entry(*entry);
64        }
65    }
66
67    /// Adds a new entry. If there's already an entry with this txid, it's replaced.
68    /// If there's no space left in the bucket, the oldest entry is overwritten.
69    pub fn add_entry(&mut self, entry: MempoolEntry) {
70        let txid = entry.txid;
71        let bucket = self.bucket_for(&txid);
72        let count = self.counts[bucket] as usize;
73
74        // Check if this txid already exists in the bucket
75        for i in 0..count {
76            if &self.entries[bucket][i].txid == &txid {
77                // If entry is different, replace it
78                self.entries[bucket][i] = entry;
79                return;
80            }
81        }
82
83        // If we have space in this bucket, add it
84        if count < BUCKET_SIZE {
85            self.entries[bucket][count] = entry;
86            self.counts[bucket] += 1;
87        } else {
88            // No space in bucket - implement replacement (FIFO)
89            // Shift all entries to make room at the end
90            for i in 0..BUCKET_SIZE - 1 {
91                self.entries[bucket][i] = self.entries[bucket][i + 1];
92            }
93
94            // Insert at the last position
95            let last = BUCKET_SIZE - 1;
96            self.entries[bucket][last] = entry;
97        }
98    }
99
100    /// Removes an entry with the specified txid
101    pub fn remove_entry(&mut self, txid: Txid) {
102        let bucket = self.bucket_for(&txid);
103        let count = self.counts[bucket] as usize;
104
105        for i in 0..count {
106            if &self.entries[bucket][i].txid == &txid {
107                // Shift all following elements down
108                for j in i..count - 1 {
109                    self.entries[bucket][j] = self.entries[bucket][j + 1];
110                }
111                self.counts[bucket] -= 1;
112                return;
113            }
114        }
115    }
116
117    /// Updates an existing entry or adds a new one
118    pub fn update_entry(&mut self, entry: MempoolEntry) {
119        self.add_entry(entry);
120    }
121
122    /// Deserializes a MempoolEntries struct from a slice of bytes
123    pub fn deserialize_from_slice(data: &[u8]) -> std::io::Result<HashMap<Txid, MempoolEntry>> {
124        let mut result = HashMap::new();
125
126        // Read counts array
127        if data.len() < std::mem::size_of::<[u16; BUCKET_COUNT]>() {
128            return Err(std::io::Error::new(
129                std::io::ErrorKind::UnexpectedEof,
130                "Data too small",
131            ));
132        }
133
134        let mut pos = 0;
135
136        // Process each bucket without creating the entire structure
137        for bucket in 0..BUCKET_COUNT {
138            // Read count
139            let count_bytes = [data[pos], data[pos + 1]];
140            let count = u16::from_le_bytes(count_bytes) as usize;
141            pos += 2;
142
143            // Skip to the entries section for this bucket
144            let entries_offset = std::mem::size_of::<[u16; BUCKET_COUNT]>()
145                + bucket * BUCKET_SIZE * std::mem::size_of::<MempoolEntry>();
146
147            // Process each entry in this bucket
148            for i in 0..count {
149                if entries_offset
150                    + i * std::mem::size_of::<MempoolEntry>()
151                    + std::mem::size_of::<MempoolEntry>()
152                    > data.len()
153                {
154                    break;
155                }
156
157                let entry_pos = entries_offset + i * std::mem::size_of::<MempoolEntry>();
158                let entry: MempoolEntry =
159                    borsh::BorshDeserialize::deserialize(&mut &data[entry_pos..])?;
160                result.insert(entry.txid, entry);
161            }
162        }
163
164        Ok(result)
165    }
166}