iroh_io/
mem.rs

1use std::io;
2
3use bytes::{Bytes, BytesMut};
4
5use super::{AsyncSliceReader, AsyncSliceWriter};
6
7impl AsyncSliceReader for bytes::Bytes {
8    async fn read_at(&mut self, offset: u64, len: usize) -> io::Result<Bytes> {
9        Ok(get_limited_slice(self, offset, len))
10    }
11
12    async fn size(&mut self) -> io::Result<u64> {
13        Ok(Bytes::len(self) as u64)
14    }
15}
16
17impl AsyncSliceReader for bytes::BytesMut {
18    async fn read_at(&mut self, offset: u64, len: usize) -> io::Result<Bytes> {
19        Ok(copy_limited_slice(self, offset, len))
20    }
21
22    async fn size(&mut self) -> io::Result<u64> {
23        Ok(BytesMut::len(self) as u64)
24    }
25}
26
27impl AsyncSliceReader for &[u8] {
28    async fn read_at(&mut self, offset: u64, len: usize) -> io::Result<Bytes> {
29        Ok(copy_limited_slice(self, offset, len))
30    }
31
32    async fn size(&mut self) -> io::Result<u64> {
33        Ok(self.len() as u64)
34    }
35}
36
37impl AsyncSliceWriter for bytes::BytesMut {
38    async fn write_bytes_at(&mut self, offset: u64, data: Bytes) -> io::Result<()> {
39        write_extend(self, offset, &data)
40    }
41
42    async fn write_at(&mut self, offset: u64, data: &[u8]) -> io::Result<()> {
43        write_extend(self, offset, data)
44    }
45
46    async fn set_len(&mut self, len: u64) -> io::Result<()> {
47        let len = len.try_into().unwrap_or(usize::MAX);
48        self.resize(len, 0);
49        Ok(())
50    }
51
52    async fn sync(&mut self) -> io::Result<()> {
53        Ok(())
54    }
55}
56
57impl AsyncSliceWriter for Vec<u8> {
58    async fn write_bytes_at(&mut self, offset: u64, data: Bytes) -> io::Result<()> {
59        write_extend_vec(self, offset, &data)
60    }
61
62    async fn write_at(&mut self, offset: u64, data: &[u8]) -> io::Result<()> {
63        write_extend_vec(self, offset, data)
64    }
65
66    async fn set_len(&mut self, len: u64) -> io::Result<()> {
67        let len = len.try_into().unwrap_or(usize::MAX);
68        self.resize(len, 0);
69        Ok(())
70    }
71
72    async fn sync(&mut self) -> io::Result<()> {
73        Ok(())
74    }
75}
76
77pub(crate) fn limited_range(offset: u64, len: usize, buf_len: usize) -> std::ops::Range<usize> {
78    if offset < buf_len as u64 {
79        let start = offset as usize;
80        let end = start.saturating_add(len).min(buf_len);
81        start..end
82    } else {
83        0..0
84    }
85}
86
87fn get_limited_slice(bytes: &Bytes, offset: u64, len: usize) -> Bytes {
88    bytes.slice(limited_range(offset, len, bytes.len()))
89}
90
91fn copy_limited_slice(bytes: &[u8], offset: u64, len: usize) -> Bytes {
92    bytes[limited_range(offset, len, bytes.len())]
93        .to_vec()
94        .into()
95}
96
97fn write_extend(bytes: &mut BytesMut, offset: u64, data: &[u8]) -> io::Result<()> {
98    let start = usize::try_from(offset).map_err(|_| {
99        io::Error::new(
100            io::ErrorKind::InvalidInput,
101            "start is too large to fit in usize",
102        )
103    })?;
104    let end = start.checked_add(data.len()).ok_or_else(|| {
105        io::Error::new(
106            io::ErrorKind::InvalidInput,
107            "offset + data.len() is too large to fit in usize",
108        )
109    })?;
110    if data.is_empty() {
111        return Ok(());
112    }
113    if end > BytesMut::len(bytes) {
114        bytes.resize(start, 0);
115        bytes.extend_from_slice(data);
116    } else {
117        bytes[start..end].copy_from_slice(data);
118    }
119
120    Ok(())
121}
122
123fn write_extend_vec(bytes: &mut Vec<u8>, offset: u64, data: &[u8]) -> io::Result<()> {
124    let start = usize::try_from(offset).map_err(|_| {
125        io::Error::new(
126            io::ErrorKind::InvalidInput,
127            "start is too large to fit in usize",
128        )
129    })?;
130    let end = start.checked_add(data.len()).ok_or_else(|| {
131        io::Error::new(
132            io::ErrorKind::InvalidInput,
133            "offset + data.len() is too large to fit in usize",
134        )
135    })?;
136    if data.is_empty() {
137        return Ok(());
138    }
139    if end > Vec::len(bytes) {
140        bytes.resize(start, 0);
141        bytes.extend_from_slice(data);
142    } else {
143        bytes[start..end].copy_from_slice(data);
144    }
145
146    Ok(())
147}