ojcmp/
byte_read.rs

1use std::fs::File;
2use std::io::{self, BufRead, Read};
3use std::panic::panic_any;
4use std::ptr;
5use std::slice;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub struct IoByte {
9    byte: u8,
10    eof: bool,
11}
12
13impl IoByte {
14    pub const EOF: IoByte = IoByte { byte: 0, eof: true };
15
16    #[inline(always)]
17    pub fn from_u8(byte: u8) -> Self {
18        Self { byte, eof: false }
19    }
20
21    #[inline(always)]
22    pub fn as_u8(self) -> u8 {
23        self.byte
24    }
25
26    #[inline(always)]
27    pub fn is_eof(self) -> bool {
28        self.eof
29    }
30}
31
32pub trait ByteRead: BufRead {
33    fn next_byte(&mut self) -> IoByte;
34
35    #[allow(clippy::missing_safety_doc)]
36    unsafe fn consume_unchecked(&mut self, amt: usize);
37}
38
39impl ByteRead for &'_ [u8] {
40    fn next_byte(&mut self) -> IoByte {
41        match self {
42            [] => IoByte::EOF,
43            [byte, remain @ ..] => {
44                *self = remain;
45                IoByte::from_u8(*byte)
46            }
47        }
48    }
49    unsafe fn consume_unchecked(&mut self, amt: usize) {
50        *self = &self[amt..];
51    }
52}
53
54pub unsafe trait TrustedRead: Read {
55    #[inline]
56    fn trusted_read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
57        let nread = <Self as Read>::read(self, buf)?;
58        Ok(nread.min(buf.len()))
59    }
60}
61
62unsafe impl TrustedRead for File {}
63
64#[derive(Debug)]
65pub struct ByteReader<R> {
66    inner: R,
67    buf: Box<[u8]>,
68    head: *const u8,
69    tail: *const u8,
70}
71
72impl<R: Read> ByteReader<R> {
73    pub fn with_capacity(capacity: usize, reader: R) -> Self {
74        Self {
75            inner: reader,
76            buf: vec![0; capacity].into(),
77            head: ptr::null(),
78            tail: ptr::null(),
79        }
80    }
81
82    #[allow(clippy::missing_safety_doc)]
83    pub unsafe fn from_raw(buf: *mut [u8], reader: R) -> Self {
84        Self {
85            inner: reader,
86            buf: Box::from_raw(buf),
87            head: ptr::null(),
88            tail: ptr::null(),
89        }
90    }
91
92    #[allow(clippy::missing_safety_doc)]
93    pub unsafe fn into_raw(self) -> (*mut [u8], R) {
94        let buf = Box::into_raw(self.buf);
95        let reader = self.inner;
96        (buf, reader)
97    }
98}
99
100impl<R: Read> Read for ByteReader<R> {
101    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
102        self.inner.read(buf)
103    }
104}
105
106impl<R: TrustedRead> BufRead for ByteReader<R> {
107    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
108        if self.head != self.tail {
109            let len = self.tail as usize - self.head as usize;
110            Ok(unsafe { slice::from_raw_parts(self.head, len) })
111        } else {
112            let nread = self.inner.trusted_read(&mut *self.buf)?;
113            if nread == 0 {
114                self.head = ptr::null();
115                self.tail = ptr::null();
116                Ok(&[])
117            } else {
118                self.head = self.buf.as_ptr();
119                self.tail = unsafe { self.head.add(nread) };
120                Ok(unsafe { slice::from_raw_parts(self.head, nread) })
121            }
122        }
123    }
124
125    fn consume(&mut self, amt: usize) {
126        self.head = (self.head as usize)
127            .saturating_add(amt)
128            .min(self.tail as usize) as *const u8;
129    }
130}
131
132impl<R: TrustedRead> ByteRead for ByteReader<R> {
133    #[inline(always)]
134    fn next_byte(&mut self) -> IoByte {
135        if self.head != self.tail {
136            unsafe {
137                let byte = *self.head;
138                self.head = self.head.add(1);
139                IoByte::from_u8(byte)
140            }
141        } else {
142            match self.inner.trusted_read(&mut *self.buf) {
143                Ok(nread) => {
144                    if nread == 0 {
145                        IoByte::EOF
146                    } else {
147                        unsafe {
148                            let byte = *self.buf.as_ptr();
149                            self.head = self.buf.as_ptr().add(1);
150                            self.tail = self.buf.as_ptr().add(nread);
151                            IoByte::from_u8(byte)
152                        }
153                    }
154                }
155                Err(e) => panic_any(e),
156            }
157        }
158    }
159    unsafe fn consume_unchecked(&mut self, amt: usize) {
160        self.head = self.head.add(amt);
161    }
162}
163
164#[cfg(unix)]
165pub mod unix {
166    use super::TrustedRead;
167
168    use std::fs::File;
169    use std::io::{self, Read};
170    use std::os::raw::c_void;
171    use std::os::unix::io::AsRawFd;
172    use std::panic::panic_any;
173
174    #[derive(Debug)]
175    pub struct UnixFdReader {
176        file: File,
177    }
178
179    impl UnixFdReader {
180        pub fn from_file(file: File) -> Self {
181            Self { file }
182        }
183    }
184
185    impl Read for UnixFdReader {
186        #[inline(always)]
187        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
188            unsafe {
189                let buf_ptr: *mut c_void = buf.as_mut_ptr().cast();
190                let fd = self.file.as_raw_fd();
191                let ret: isize = libc::read(fd, buf_ptr, buf.len());
192                if ret < 0 {
193                    panic_any(io::Error::last_os_error());
194                }
195                assert!(ret as usize <= buf.len());
196                Ok(ret as usize)
197            }
198        }
199    }
200
201    unsafe impl TrustedRead for UnixFdReader {
202        fn trusted_read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
203            <Self as Read>::read(self, buf)
204        }
205    }
206}
ojcmp/byte_read.rs

ojcmp/
byte_read.rs
