Skip to main content

glycin_utils/
img_buf.rs

1use std::os::fd::{AsRawFd, RawFd};
2
3use crate::{DimensionTooLargerError, editing};
4
5pub enum ImgBuf {
6    MMap {
7        mmap: memmap::MmapMut,
8        raw_fd: RawFd,
9    },
10    Vec(Vec<u8>),
11}
12
13impl ImgBuf {
14    pub unsafe fn from_raw_fd(raw_fd: impl AsRawFd) -> std::io::Result<Self> {
15        let mmap = unsafe { memmap::MmapMut::map_mut(&raw_fd)? };
16        Ok(Self::MMap {
17            mmap,
18            raw_fd: raw_fd.as_raw_fd(),
19        })
20    }
21
22    pub fn as_slice(&self) -> &[u8] {
23        match self {
24            Self::MMap { mmap, .. } => mmap.as_ref(),
25            Self::Vec(v) => v.as_slice(),
26        }
27    }
28
29    pub fn as_mut_slice(&mut self) -> &mut [u8] {
30        match self {
31            Self::MMap { mmap, .. } => mmap.as_mut(),
32            Self::Vec(v) => v.as_mut_slice(),
33        }
34    }
35
36    pub fn into_vec(self) -> Vec<u8> {
37        match self {
38            Self::Vec(vec) => vec,
39            Self::MMap { .. } => self.to_vec(),
40        }
41    }
42
43    pub fn resize(self, new_len: i64) -> Result<Self, editing::Error> {
44        if self.len() == new_len as usize {
45            return Ok(self);
46        }
47
48        match self {
49            ImgBuf::MMap { mmap, raw_fd, .. } => {
50                let borrowed_fd = unsafe { std::os::fd::BorrowedFd::borrow_raw(raw_fd) };
51
52                // This mmap would have the wrong size after ftruncate
53                drop(mmap);
54
55                nix::unistd::ftruncate(
56                    borrowed_fd,
57                    libc::off_t::try_from(new_len).map_err(|_| DimensionTooLargerError)?,
58                )
59                .map_err(std::io::Error::from)?;
60
61                // Need a new mmap with correct size
62                let mmap = unsafe { memmap::MmapMut::map_mut(raw_fd) }?;
63
64                Ok(ImgBuf::MMap { mmap, raw_fd })
65            }
66            Self::Vec(mut vec) => {
67                vec.resize(new_len as usize, 0);
68                Ok(Self::Vec(vec))
69            }
70        }
71    }
72}
73
74impl std::ops::Deref for ImgBuf {
75    type Target = [u8];
76    fn deref(&self) -> &Self::Target {
77        self.as_slice()
78    }
79}
80
81impl std::ops::DerefMut for ImgBuf {
82    fn deref_mut(&mut self) -> &mut Self::Target {
83        self.as_mut_slice()
84    }
85}