shared_buffer/
owned.rs

1use std::{
2    borrow::Cow,
3    fs::File,
4    hash::Hash,
5    ops::{Deref, Index, RangeBounds},
6    path::{Path, PathBuf},
7};
8
9use bytes::{Buf, Bytes};
10
11use crate::{mmap::MmappedSlice, MmapError};
12
13/// A cheaply cloneable owned buffer that may be backed by a memory-mapped file
14/// or [`Bytes`] in memory.
15#[derive(Debug, Clone, Default)]
16pub struct OwnedBuffer {
17    repr: Repr,
18}
19
20impl OwnedBuffer {
21    pub const fn new() -> Self {
22        OwnedBuffer {
23            repr: Repr::InMemory(Bytes::new()),
24        }
25    }
26
27    pub const fn from_static(bytes: &'static [u8]) -> Self {
28        OwnedBuffer {
29            repr: Repr::InMemory(Bytes::from_static(bytes)),
30        }
31    }
32
33    pub fn from_bytes(bytes: impl Into<Bytes>) -> Self {
34        bytes.into().into()
35    }
36
37    pub fn mmap(path: impl AsRef<Path>) -> Result<Self, MmapError> {
38        let mmap = MmappedSlice::from_path(path.as_ref())?;
39        Ok(mmap.into())
40    }
41
42    pub fn from_file(file: &File) -> Result<Self, MmapError> {
43        let mmap = MmappedSlice::from_file(file)?;
44        Ok(mmap.into())
45    }
46
47    pub fn into_bytes(self) -> Bytes {
48        match self.repr {
49            Repr::InMemory(b) => b,
50            Repr::Mmap(m) => Bytes::from(m.as_slice().to_vec()),
51        }
52    }
53
54    pub fn as_bytes(&self) -> Option<&Bytes> {
55        self.repr.as_bytes()
56    }
57
58    #[inline]
59    pub fn as_slice(&self) -> &[u8] {
60        match &self.repr {
61            Repr::InMemory(b) => b,
62            Repr::Mmap(m) => m.as_slice(),
63        }
64    }
65
66    #[inline]
67    pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
68        self.repr.slice(range).into()
69    }
70
71    pub fn is_mmapped(&self) -> bool {
72        matches!(self.repr, Repr::Mmap(_))
73    }
74}
75
76impl Deref for OwnedBuffer {
77    type Target = [u8];
78
79    #[inline]
80    fn deref(&self) -> &Self::Target {
81        self.as_slice()
82    }
83}
84
85impl AsRef<[u8]> for OwnedBuffer {
86    fn as_ref(&self) -> &[u8] {
87        self.as_slice()
88    }
89}
90
91impl PartialEq for OwnedBuffer {
92    #[inline]
93    fn eq(&self, other: &Self) -> bool {
94        self.as_slice() == other.as_slice()
95    }
96}
97
98impl PartialEq<[u8]> for OwnedBuffer {
99    fn eq(&self, other: &[u8]) -> bool {
100        self.as_slice() == other
101    }
102}
103
104impl PartialEq<&[u8]> for OwnedBuffer {
105    fn eq(&self, other: &&[u8]) -> bool {
106        self.as_slice() == *other
107    }
108}
109
110impl<const N: usize> PartialEq<[u8; N]> for OwnedBuffer {
111    fn eq(&self, other: &[u8; N]) -> bool {
112        self.as_slice() == other
113    }
114}
115
116impl<const N: usize> PartialEq<&[u8; N]> for OwnedBuffer {
117    fn eq(&self, other: &&[u8; N]) -> bool {
118        self.as_slice() == *other
119    }
120}
121
122impl PartialEq<Vec<u8>> for OwnedBuffer {
123    fn eq(&self, other: &Vec<u8>) -> bool {
124        self.as_slice() == other
125    }
126}
127
128impl PartialEq<&Vec<u8>> for OwnedBuffer {
129    fn eq(&self, other: &&Vec<u8>) -> bool {
130        self.as_slice() == *other
131    }
132}
133
134impl PartialEq<OwnedBuffer> for [u8] {
135    fn eq(&self, other: &OwnedBuffer) -> bool {
136        other == self
137    }
138}
139
140impl PartialEq<OwnedBuffer> for &[u8] {
141    fn eq(&self, other: &OwnedBuffer) -> bool {
142        other == self
143    }
144}
145
146impl<const N: usize> PartialEq<OwnedBuffer> for [u8; N] {
147    fn eq(&self, other: &OwnedBuffer) -> bool {
148        other == self
149    }
150}
151
152impl<const N: usize> PartialEq<OwnedBuffer> for &[u8; N] {
153    fn eq(&self, other: &OwnedBuffer) -> bool {
154        other == self
155    }
156}
157
158impl PartialEq<OwnedBuffer> for Vec<u8> {
159    fn eq(&self, other: &OwnedBuffer) -> bool {
160        other == self
161    }
162}
163
164impl PartialEq<OwnedBuffer> for &Vec<u8> {
165    fn eq(&self, other: &OwnedBuffer) -> bool {
166        other == self
167    }
168}
169
170impl<'a> IntoIterator for &'a OwnedBuffer {
171    type Item = &'a u8;
172    type IntoIter = std::slice::Iter<'a, u8>;
173
174    fn into_iter(self) -> Self::IntoIter {
175        self.as_slice().iter()
176    }
177}
178
179impl IntoIterator for OwnedBuffer {
180    type Item = u8;
181    type IntoIter = OwnedIntoIter;
182
183    fn into_iter(self) -> Self::IntoIter {
184        OwnedIntoIter {
185            buffer: self,
186            next_index: 0,
187        }
188    }
189}
190
191impl<I> Index<I> for OwnedBuffer
192where
193    [u8]: Index<I>,
194{
195    type Output = <[u8] as Index<I>>::Output;
196
197    fn index(&self, index: I) -> &Self::Output {
198        &self.as_slice()[index]
199    }
200}
201
202impl Eq for OwnedBuffer {}
203
204impl Ord for OwnedBuffer {
205    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
206        self.as_slice().cmp(other.as_slice())
207    }
208}
209
210impl PartialOrd for OwnedBuffer {
211    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
212        Some(self.cmp(other))
213    }
214}
215
216impl Hash for OwnedBuffer {
217    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
218        self.as_slice().hash(state);
219    }
220}
221
222impl Buf for OwnedBuffer {
223    fn remaining(&self) -> usize {
224        self.len()
225    }
226
227    fn chunk(&self) -> &[u8] {
228        self.as_slice()
229    }
230
231    fn advance(&mut self, cnt: usize) {
232        match &mut self.repr {
233            Repr::InMemory(b) => b.advance(cnt),
234            Repr::Mmap(m) => m.advance(cnt),
235        }
236    }
237}
238
239impl From<Bytes> for OwnedBuffer {
240    fn from(value: Bytes) -> Self {
241        OwnedBuffer {
242            repr: Repr::InMemory(value),
243        }
244    }
245}
246
247impl From<&[u8]> for OwnedBuffer {
248    fn from(value: &[u8]) -> Self {
249        value.to_vec().into()
250    }
251}
252
253impl From<Cow<'_, [u8]>> for OwnedBuffer {
254    fn from(value: Cow<'_, [u8]>) -> Self {
255        match value {
256            Cow::Borrowed(value) => value.into(),
257            Cow::Owned(value) => value.into(),
258        }
259    }
260}
261
262impl From<Vec<u8>> for OwnedBuffer {
263    fn from(value: Vec<u8>) -> Self {
264        Bytes::from(value).into()
265    }
266}
267
268impl From<&Vec<u8>> for OwnedBuffer {
269    fn from(value: &Vec<u8>) -> Self {
270        value.as_slice().into()
271    }
272}
273
274impl From<OwnedBuffer> for Vec<u8> {
275    fn from(value: OwnedBuffer) -> Self {
276        value.to_vec()
277    }
278}
279
280impl From<Repr> for OwnedBuffer {
281    fn from(repr: Repr) -> Self {
282        OwnedBuffer { repr }
283    }
284}
285
286impl From<MmappedSlice> for OwnedBuffer {
287    fn from(mmap: MmappedSlice) -> Self {
288        Repr::Mmap(mmap).into()
289    }
290}
291
292impl TryFrom<&Path> for OwnedBuffer {
293    type Error = MmapError;
294
295    fn try_from(value: &Path) -> Result<Self, Self::Error> {
296        OwnedBuffer::mmap(value)
297    }
298}
299
300impl TryFrom<PathBuf> for OwnedBuffer {
301    type Error = MmapError;
302
303    fn try_from(value: PathBuf) -> Result<Self, Self::Error> {
304        OwnedBuffer::mmap(value)
305    }
306}
307
308impl TryFrom<&PathBuf> for OwnedBuffer {
309    type Error = MmapError;
310
311    fn try_from(value: &PathBuf) -> Result<Self, Self::Error> {
312        OwnedBuffer::mmap(value)
313    }
314}
315
316impl TryFrom<&std::fs::File> for OwnedBuffer {
317    type Error = MmapError;
318
319    fn try_from(file: &std::fs::File) -> Result<Self, Self::Error> {
320        OwnedBuffer::from_file(file)
321    }
322}
323
324impl TryFrom<std::fs::File> for OwnedBuffer {
325    type Error = MmapError;
326
327    fn try_from(file: std::fs::File) -> Result<Self, Self::Error> {
328        OwnedBuffer::from_file(&file)
329    }
330}
331
332impl TryFrom<OwnedBuffer> for Bytes {
333    type Error = OwnedBuffer;
334
335    fn try_from(value: OwnedBuffer) -> Result<Self, Self::Error> {
336        if let Repr::InMemory(bytes) = value.repr {
337            Ok(bytes)
338        } else {
339            Err(value)
340        }
341    }
342}
343
344#[derive(Debug, Clone)]
345enum Repr {
346    InMemory(Bytes),
347    Mmap(MmappedSlice),
348}
349
350impl Repr {
351    #[inline]
352    fn slice(&self, range: impl RangeBounds<usize>) -> Self {
353        match self {
354            Repr::InMemory(b) => Repr::InMemory(b.slice(range)),
355            Repr::Mmap(m) => Repr::Mmap(m.slice(range)),
356        }
357    }
358
359    #[inline]
360    fn as_bytes(&self) -> Option<&Bytes> {
361        match self {
362            Repr::InMemory(b) => Some(b),
363            _ => None,
364        }
365    }
366}
367
368impl Default for Repr {
369    fn default() -> Self {
370        Repr::InMemory(Bytes::new())
371    }
372}
373
374#[derive(Debug, Clone)]
375pub struct OwnedIntoIter {
376    buffer: OwnedBuffer,
377    next_index: usize,
378}
379
380impl Iterator for OwnedIntoIter {
381    type Item = u8;
382
383    #[inline]
384    fn next(&mut self) -> Option<Self::Item> {
385        let &byte = self.buffer.as_slice().get(self.next_index)?;
386        self.next_index += 1;
387        Some(byte)
388    }
389}