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 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 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}