1use alloc::{sync::Arc, vec::Vec};
2use alloy_eip7928::bal::DecodedBal;
3pub use alloy_eip7928::bal::RawBal;
4use alloy_eips::NumHash;
5use alloy_primitives::{BlockHash, BlockNumber, Bytes};
6use reth_storage_errors::provider::ProviderResult;
7
8#[auto_impl::auto_impl(&, Arc, Box)]
14pub trait BalStore: Send + Sync + 'static {
15 fn insert(&self, num_hash: NumHash, bal: RawBal) -> ProviderResult<()>;
20
21 fn insert_many(&self, entries: Vec<(NumHash, RawBal)>) -> ProviderResult<()> {
25 for (num_hash, bal) in entries {
26 self.insert(num_hash, bal)?;
27 }
28 Ok(())
29 }
30
31 fn flush(&self, _blocks: &[NumHash]) -> ProviderResult<()> {
35 Ok(())
36 }
37
38 fn prune(&self, tip: BlockNumber) -> ProviderResult<usize>;
42
43 fn get_by_hashes(&self, block_hashes: &[BlockHash]) -> ProviderResult<Vec<Option<Bytes>>>;
47
48 fn get_by_hash(&self, block_hash: BlockHash) -> ProviderResult<Option<Bytes>> {
50 Ok(self.get_by_hashes(&[block_hash])?.into_iter().next().flatten())
51 }
52
53 fn get_decoded_by_hash(&self, block_hash: BlockHash) -> ProviderResult<Option<DecodedBal>> {
55 self.get_by_hash(block_hash)?
56 .map(DecodedBal::from_rlp_bytes)
57 .transpose()
58 .map_err(Into::into)
59 }
60
61 fn get_by_hashes_with_limit(
67 &self,
68 block_hashes: &[BlockHash],
69 limit: GetBlockAccessListLimit,
70 ) -> ProviderResult<Vec<Option<Bytes>>> {
71 let mut out = Vec::new();
72 self.append_by_hashes_with_limit(block_hashes, limit, &mut out)?;
73 out.shrink_to_fit();
74 Ok(out)
75 }
76
77 fn append_by_hashes_with_limit(
81 &self,
82 block_hashes: &[BlockHash],
83 limit: GetBlockAccessListLimit,
84 out: &mut Vec<Option<Bytes>>,
85 ) -> ProviderResult<()> {
86 let mut size = 0;
87 for bal in self.get_by_hashes(block_hashes)? {
88 size += bal.as_ref().map_or(1, |bytes| bytes.len());
89 out.push(bal);
90
91 if limit.exceeds(size) {
92 break
93 }
94 }
95 Ok(())
96 }
97}
98
99#[derive(Debug, Clone, Copy, Eq, PartialEq)]
101pub enum GetBlockAccessListLimit {
102 None,
104 ResponseSizeSoftLimit(usize),
106}
107
108impl GetBlockAccessListLimit {
109 #[inline]
111 pub const fn exceeds(&self, size: usize) -> bool {
112 match self {
113 Self::None => false,
114 Self::ResponseSizeSoftLimit(limit) => size > *limit,
115 }
116 }
117}
118
119#[derive(Clone)]
121pub struct BalStoreHandle {
122 inner: Arc<dyn BalStore>,
123}
124
125impl BalStoreHandle {
126 pub fn new(inner: impl BalStore) -> Self {
128 Self { inner: Arc::new(inner) }
129 }
130
131 pub fn noop() -> Self {
133 Self::new(NoopBalStore)
134 }
135
136 #[inline]
138 pub fn insert(&self, num_hash: NumHash, bal: RawBal) -> ProviderResult<()> {
139 self.inner.insert(num_hash, bal)
140 }
141
142 #[inline]
144 pub fn insert_many(&self, entries: Vec<(NumHash, RawBal)>) -> ProviderResult<()> {
145 self.inner.insert_many(entries)
146 }
147
148 #[inline]
150 pub fn flush(&self, blocks: &[NumHash]) -> ProviderResult<()> {
151 self.inner.flush(blocks)
152 }
153
154 #[inline]
156 pub fn prune(&self, tip: BlockNumber) -> ProviderResult<usize> {
157 self.inner.prune(tip)
158 }
159
160 #[inline]
162 pub fn get_by_hashes(&self, block_hashes: &[BlockHash]) -> ProviderResult<Vec<Option<Bytes>>> {
163 self.inner.get_by_hashes(block_hashes)
164 }
165
166 #[inline]
168 pub fn get_by_hash(&self, block_hash: BlockHash) -> ProviderResult<Option<Bytes>> {
169 self.inner.get_by_hash(block_hash)
170 }
171
172 #[inline]
174 pub fn get_decoded_by_hash(&self, block_hash: BlockHash) -> ProviderResult<Option<DecodedBal>> {
175 self.inner.get_decoded_by_hash(block_hash)
176 }
177
178 #[inline]
181 pub fn get_by_hashes_with_limit(
182 &self,
183 block_hashes: &[BlockHash],
184 limit: GetBlockAccessListLimit,
185 ) -> ProviderResult<Vec<Option<Bytes>>> {
186 self.inner.get_by_hashes_with_limit(block_hashes, limit)
187 }
188
189 #[inline]
191 pub fn append_by_hashes_with_limit(
192 &self,
193 block_hashes: &[BlockHash],
194 limit: GetBlockAccessListLimit,
195 out: &mut Vec<Option<Bytes>>,
196 ) -> ProviderResult<()> {
197 self.inner.append_by_hashes_with_limit(block_hashes, limit, out)
198 }
199}
200
201impl Default for BalStoreHandle {
202 fn default() -> Self {
203 Self::noop()
204 }
205}
206
207impl core::fmt::Debug for BalStoreHandle {
208 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
209 f.debug_struct("BalStoreHandle").finish_non_exhaustive()
210 }
211}
212
213#[auto_impl::auto_impl(&, Arc)]
215pub trait BalProvider {
216 fn bal_store(&self) -> &BalStoreHandle;
218
219 fn get_bal_by_hash(&self, block_hash: BlockHash) -> ProviderResult<Option<Bytes>> {
221 self.bal_store().get_by_hash(block_hash)
222 }
223
224 fn get_bals_by_hashes(&self, block_hashes: &[BlockHash]) -> ProviderResult<Vec<Option<Bytes>>> {
226 self.bal_store().get_by_hashes(block_hashes)
227 }
228
229 fn get_bals_by_hashes_with_limit(
232 &self,
233 block_hashes: &[BlockHash],
234 limit: GetBlockAccessListLimit,
235 ) -> ProviderResult<Vec<Option<Bytes>>> {
236 self.bal_store().get_by_hashes_with_limit(block_hashes, limit)
237 }
238}
239
240#[derive(Debug, Default, Clone, Copy)]
242pub struct NoopBalStore;
243
244impl BalStore for NoopBalStore {
245 fn insert(&self, _num_hash: NumHash, _bal: RawBal) -> ProviderResult<()> {
246 Ok(())
247 }
248
249 fn insert_many(&self, _entries: Vec<(NumHash, RawBal)>) -> ProviderResult<()> {
250 Ok(())
251 }
252
253 fn prune(&self, _tip: BlockNumber) -> ProviderResult<usize> {
254 Ok(0)
255 }
256
257 fn get_by_hashes(&self, block_hashes: &[BlockHash]) -> ProviderResult<Vec<Option<Bytes>>> {
258 Ok(block_hashes.iter().map(|_| None).collect())
259 }
260
261 fn append_by_hashes_with_limit(
262 &self,
263 block_hashes: &[BlockHash],
264 limit: GetBlockAccessListLimit,
265 out: &mut Vec<Option<Bytes>>,
266 ) -> ProviderResult<()> {
267 let mut size = 0;
268 for _ in block_hashes {
269 size += 1;
270 out.push(None);
271
272 if limit.exceeds(size) {
273 break
274 }
275 }
276 Ok(())
277 }
278}
279
280#[cfg(test)]
281mod tests {
282 use super::*;
283 use alloy_primitives::B256;
284
285 const EMPTY_LIST_CODE: u8 = 0xc0;
286
287 #[test]
288 fn noop_store_returns_empty_results() {
289 let store = BalStoreHandle::default();
290 let hashes = [B256::random(), B256::random()];
291
292 let by_hash = store.get_by_hashes(&hashes).unwrap();
293
294 assert_eq!(by_hash, vec![None, None]);
295 assert!(store.get_by_hash(B256::random()).unwrap().is_none());
296 assert_eq!(store.prune(10).unwrap(), 0);
297 }
298
299 #[test]
300 fn noop_provider_returns_empty_results() {
301 let provider = crate::noop::NoopProvider::default();
302 let hashes = [B256::random(), B256::random()];
303
304 assert_eq!(provider.get_bals_by_hashes(&hashes).unwrap(), vec![None, None]);
305 assert_eq!(
306 provider
307 .get_bals_by_hashes_with_limit(
308 &hashes,
309 GetBlockAccessListLimit::ResponseSizeSoftLimit(0),
310 )
311 .unwrap(),
312 vec![None]
313 );
314 assert!(provider.get_bal_by_hash(B256::random()).unwrap().is_none());
315 }
316
317 #[test]
318 fn noop_store_flush_is_noop() {
319 let store = BalStoreHandle::default();
320
321 store.flush(&[]).unwrap();
322 }
323
324 #[test]
325 fn noop_store_decoded_lookup_returns_none() {
326 let store = BalStoreHandle::default();
327
328 assert!(store.get_decoded_by_hash(B256::random()).unwrap().is_none());
329 }
330
331 #[test]
332 fn decoded_lookup_decodes_raw_bal() {
333 let hash = B256::random();
334 let raw_bal = Bytes::from_static(&[EMPTY_LIST_CODE]);
335 let store = BalStoreHandle::new(TestBalStore { hash, raw_bal: raw_bal.clone() });
336
337 assert_eq!(store.get_by_hash(hash).unwrap(), Some(raw_bal.clone()));
338
339 let decoded = store.get_decoded_by_hash(hash).unwrap().unwrap();
340
341 assert_eq!(decoded.as_raw(), &raw_bal);
342 }
343
344 #[test]
345 fn noop_store_limited_lookup_returns_prefix() {
346 let store = BalStoreHandle::default();
347 let hashes = [B256::random(), B256::random(), B256::random()];
348
349 let limited = store
350 .get_by_hashes_with_limit(&hashes, GetBlockAccessListLimit::ResponseSizeSoftLimit(1))
351 .unwrap();
352
353 assert_eq!(limited, vec![None, None]);
354 }
355
356 #[test]
357 fn block_access_list_limit() {
358 let limit_none = GetBlockAccessListLimit::None;
359 assert!(!limit_none.exceeds(usize::MAX));
360
361 let size_limit_2mb = GetBlockAccessListLimit::ResponseSizeSoftLimit(2 * 1024 * 1024);
362 assert!(!size_limit_2mb.exceeds(1024 * 1024));
363 assert!(!size_limit_2mb.exceeds(2 * 1024 * 1024));
364 assert!(size_limit_2mb.exceeds(3 * 1024 * 1024));
365 }
366
367 #[derive(Debug)]
368 struct TestBalStore {
369 hash: B256,
370 raw_bal: Bytes,
371 }
372
373 impl BalStore for TestBalStore {
374 fn insert(&self, _num_hash: NumHash, _bal: RawBal) -> ProviderResult<()> {
375 Ok(())
376 }
377
378 fn prune(&self, _tip: BlockNumber) -> ProviderResult<usize> {
379 Ok(0)
380 }
381
382 fn get_by_hashes(&self, block_hashes: &[BlockHash]) -> ProviderResult<Vec<Option<Bytes>>> {
383 Ok(block_hashes
384 .iter()
385 .map(|hash| (*hash == self.hash).then(|| self.raw_bal.clone()))
386 .collect())
387 }
388 }
389}