1use std::sync::atomic::{AtomicU64, Ordering};
10
11pub trait RangeReader {
13 fn len(&self) -> u64;
15
16 fn is_empty(&self) -> bool {
18 self.len() == 0
19 }
20
21 fn read_at(&self, offset: u64, len: u64) -> std::io::Result<Vec<u8>>;
24
25 fn read_many(&self, ranges: &[(u64, u64)]) -> std::io::Result<Vec<Vec<u8>>> {
31 ranges
32 .iter()
33 .map(|&(offset, len)| self.read_at(offset, len))
34 .collect()
35 }
36
37 fn concurrency(&self) -> usize {
43 1
44 }
45}
46
47impl<R: RangeReader + ?Sized> RangeReader for std::sync::Arc<R> {
50 fn len(&self) -> u64 {
51 (**self).len()
52 }
53
54 fn read_at(&self, offset: u64, len: u64) -> std::io::Result<Vec<u8>> {
55 (**self).read_at(offset, len)
56 }
57
58 fn read_many(&self, ranges: &[(u64, u64)]) -> std::io::Result<Vec<Vec<u8>>> {
59 (**self).read_many(ranges)
60 }
61
62 fn concurrency(&self) -> usize {
63 (**self).concurrency()
64 }
65}
66
67pub struct SliceReader<'a> {
69 data: &'a [u8],
70}
71
72impl<'a> SliceReader<'a> {
73 pub fn new(data: &'a [u8]) -> Self {
74 Self { data }
75 }
76}
77
78impl RangeReader for SliceReader<'_> {
79 fn len(&self) -> u64 {
80 self.data.len() as u64
81 }
82
83 fn read_at(&self, offset: u64, len: u64) -> std::io::Result<Vec<u8>> {
84 let range = usize::try_from(offset).ok().and_then(|start| {
87 let end = start.checked_add(usize::try_from(len).ok()?)?;
88 (end <= self.data.len()).then_some(start..end)
89 });
90 let range = range.ok_or_else(|| {
91 std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "range out of bounds")
92 })?;
93 Ok(self.data[range].to_vec())
94 }
95}
96
97pub struct CountingReader<R> {
102 inner: R,
103 requests: AtomicU64,
104 bytes: AtomicU64,
105}
106
107impl<R: RangeReader> CountingReader<R> {
108 pub fn new(inner: R) -> Self {
109 Self {
110 inner,
111 requests: AtomicU64::new(0),
112 bytes: AtomicU64::new(0),
113 }
114 }
115
116 pub fn requests(&self) -> u64 {
118 self.requests.load(Ordering::Relaxed)
119 }
120
121 pub fn bytes_read(&self) -> u64 {
123 self.bytes.load(Ordering::Relaxed)
124 }
125}
126
127impl<R: RangeReader> RangeReader for CountingReader<R> {
128 fn len(&self) -> u64 {
129 self.inner.len()
130 }
131
132 fn read_at(&self, offset: u64, len: u64) -> std::io::Result<Vec<u8>> {
133 let out = self.inner.read_at(offset, len)?;
134 self.requests.fetch_add(1, Ordering::Relaxed);
135 self.bytes.fetch_add(out.len() as u64, Ordering::Relaxed);
136 Ok(out)
137 }
138
139 fn read_many(&self, ranges: &[(u64, u64)]) -> std::io::Result<Vec<Vec<u8>>> {
143 let out = self.inner.read_many(ranges)?;
144 self.requests.fetch_add(out.len() as u64, Ordering::Relaxed);
145 self.bytes
146 .fetch_add(out.iter().map(|b| b.len() as u64).sum(), Ordering::Relaxed);
147 Ok(out)
148 }
149
150 fn concurrency(&self) -> usize {
151 self.inner.concurrency()
152 }
153}
154
155pub struct OffsetReader<R> {
161 inner: R,
162 base: u64,
163}
164
165impl<R> OffsetReader<R> {
166 pub fn new(inner: R, base: u64) -> Self {
167 Self { inner, base }
168 }
169 pub fn base(&self) -> u64 {
170 self.base
171 }
172 pub fn into_inner(self) -> R {
173 self.inner
174 }
175}
176
177impl<R: RangeReader> RangeReader for OffsetReader<R> {
178 fn len(&self) -> u64 {
179 self.inner.len().saturating_sub(self.base)
180 }
181
182 fn read_at(&self, offset: u64, len: u64) -> std::io::Result<Vec<u8>> {
183 self.inner.read_at(self.base + offset, len)
184 }
185
186 fn read_many(&self, ranges: &[(u64, u64)]) -> std::io::Result<Vec<Vec<u8>>> {
187 let shifted: Vec<(u64, u64)> = ranges.iter().map(|&(o, l)| (self.base + o, l)).collect();
188 self.inner.read_many(&shifted)
189 }
190
191 fn concurrency(&self) -> usize {
192 self.inner.concurrency()
193 }
194}
195
196pub const POLYGLOT_MARKER: &[u8] = b"RETE-BASE:";
202pub const POLYGLOT_DIGITS: usize = 16;
204
205pub fn detect_polyglot_base(head: &[u8]) -> Option<u64> {
208 let pos = head
209 .windows(POLYGLOT_MARKER.len())
210 .position(|w| w == POLYGLOT_MARKER)?;
211 let start = pos + POLYGLOT_MARKER.len();
212 let digits = head.get(start..start + POLYGLOT_DIGITS)?;
213 std::str::from_utf8(digits).ok()?.parse::<u64>().ok()
214}
215
216pub struct BytesReader(pub std::sync::Arc<[u8]>);
220
221impl RangeReader for BytesReader {
222 fn len(&self) -> u64 {
223 self.0.len() as u64
224 }
225
226 fn read_at(&self, offset: u64, len: u64) -> std::io::Result<Vec<u8>> {
227 let end = offset
228 .checked_add(len)
229 .filter(|&e| e <= self.0.len() as u64)
230 .ok_or_else(|| {
231 std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "range out of bounds")
232 })?;
233 Ok(self.0[offset as usize..end as usize].to_vec())
234 }
235}
236
237pub struct PrefixReader<R> {
242 inner: R,
243 prefix: std::sync::Arc<[u8]>,
244}
245
246impl<R: RangeReader> PrefixReader<R> {
247 pub fn new(inner: R, prefix: std::sync::Arc<[u8]>) -> Self {
249 PrefixReader { inner, prefix }
250 }
251
252 fn cut(&self, offset: u64, len: u64) -> Option<Vec<u8>> {
253 let end = offset.checked_add(len)?;
254 (end <= self.prefix.len() as u64)
255 .then(|| self.prefix[offset as usize..end as usize].to_vec())
256 }
257}
258
259impl<R: RangeReader> RangeReader for PrefixReader<R> {
260 fn len(&self) -> u64 {
261 self.inner.len()
262 }
263
264 fn read_at(&self, offset: u64, len: u64) -> std::io::Result<Vec<u8>> {
265 match self.cut(offset, len) {
266 Some(b) => Ok(b),
267 None => self.inner.read_at(offset, len),
268 }
269 }
270
271 fn read_many(&self, ranges: &[(u64, u64)]) -> std::io::Result<Vec<Vec<u8>>> {
272 let mut out: Vec<Option<Vec<u8>>> = ranges.iter().map(|&(o, l)| self.cut(o, l)).collect();
273 let rest: Vec<(u64, u64)> = ranges
274 .iter()
275 .zip(&out)
276 .filter(|(_, b)| b.is_none())
277 .map(|(r, _)| *r)
278 .collect();
279 if !rest.is_empty() {
280 let mut fetched = self.inner.read_many(&rest)?.into_iter();
281 for slot in out.iter_mut().filter(|b| b.is_none()) {
282 *slot = fetched.next();
283 }
284 }
285 out.into_iter()
286 .map(|b| b.ok_or_else(|| std::io::Error::other("read_many left a range unfetched")))
287 .collect()
288 }
289
290 fn concurrency(&self) -> usize {
291 self.inner.concurrency()
292 }
293}
294
295#[cfg(test)]
296mod tests {
297 use super::*;
298
299 #[test]
300 fn slice_reader_serves_ranges_and_bounds_check() {
301 let data = (0u8..=255).collect::<Vec<_>>();
302 let r = SliceReader::new(&data);
303 assert_eq!(r.len(), 256);
304 assert_eq!(r.read_at(10, 4).unwrap(), vec![10, 11, 12, 13]);
305 assert!(r.read_at(254, 10).is_err()); }
307
308 #[test]
309 fn counting_reader_tallies() {
310 let data = vec![0u8; 100];
311 let r = CountingReader::new(SliceReader::new(&data));
312 r.read_at(0, 10).unwrap();
313 r.read_at(50, 20).unwrap();
314 assert_eq!(r.requests(), 2);
315 assert_eq!(r.bytes_read(), 30);
316 }
317}