persy/util/
io.rs

1use crate::{error::PERes, index::string_wrapper::StringWrapper, ByteVec};
2use std::io::{Read, Write};
3use std::sync::Arc;
4use unsigned_varint::{encode, io};
5use zigzag::ZigZag;
6
7pub trait InfallibleRead {
8    fn read_exact(&mut self, buf: &mut [u8]);
9    fn read_slice(&mut self, len: usize) -> ArcSliceRead;
10}
11
12pub trait InfallibleWrite {
13    fn write_all(&mut self, buf: &[u8]);
14}
15
16struct InfallibleReadWrapper<'a, T: InfallibleRead + ?Sized> {
17    reader: &'a mut T,
18}
19
20#[allow(dead_code)]
21impl<'a, T: InfallibleRead + ?Sized> InfallibleReadWrapper<'a, T> {
22    fn new(reader: &'a mut T) -> Self {
23        InfallibleReadWrapper { reader }
24    }
25}
26impl<'a, T: InfallibleRead + ?Sized> Read for InfallibleReadWrapper<'a, T> {
27    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
28        InfallibleRead::read_exact(self.reader, buf);
29        Ok(buf.len())
30    }
31}
32
33// End of snippet taken from integer-encoding specialized for InfallibleRead
34
35pub(crate) trait InfallibleReadVarInt: InfallibleRead {
36    #[inline]
37    fn read_varint_u8(&mut self) -> u8 {
38        io::read_u8(InfallibleReadWrapper::new(self)).expect("infallible")
39    }
40
41    #[inline]
42    fn read_varint_u16(&mut self) -> u16 {
43        io::read_u16(InfallibleReadWrapper::new(self)).expect("infallible")
44    }
45
46    #[inline]
47    fn read_varint_u32(&mut self) -> u32 {
48        io::read_u32(InfallibleReadWrapper::new(self)).expect("infallible")
49    }
50
51    #[inline]
52    fn read_varint_u64(&mut self) -> u64 {
53        io::read_u64(InfallibleReadWrapper::new(self)).expect("infallible")
54    }
55
56    #[inline]
57    fn read_varint_u128(&mut self) -> u128 {
58        io::read_u128(InfallibleReadWrapper::new(self)).expect("infallible")
59    }
60
61    #[inline]
62    fn read_varint_i8(&mut self) -> i8 {
63        ZigZag::decode(self.read_varint_u8())
64    }
65
66    #[inline]
67    fn read_varint_i16(&mut self) -> i16 {
68        ZigZag::decode(self.read_varint_u16())
69    }
70
71    #[inline]
72    fn read_varint_i32(&mut self) -> i32 {
73        ZigZag::decode(self.read_varint_u32())
74    }
75
76    #[inline]
77    fn read_varint_i64(&mut self) -> i64 {
78        ZigZag::decode(self.read_varint_u64())
79    }
80
81    #[inline]
82    fn read_varint_i128(&mut self) -> i128 {
83        ZigZag::decode(self.read_varint_u128())
84    }
85}
86impl<R: InfallibleRead + ?Sized> InfallibleReadVarInt for R {}
87
88pub(crate) trait InfallibleWriteVarInt: InfallibleWrite {
89    #[inline]
90    fn write_varint_u8(&mut self, value: u8) {
91        self.write_all(encode::u8(value, &mut encode::u8_buffer()));
92    }
93
94    #[inline]
95    fn write_varint_u16(&mut self, value: u16) {
96        self.write_all(encode::u16(value, &mut encode::u16_buffer()));
97    }
98
99    #[inline]
100    fn write_varint_u32(&mut self, value: u32) {
101        self.write_all(encode::u32(value, &mut encode::u32_buffer()));
102    }
103
104    #[inline]
105    fn write_varint_u64(&mut self, value: u64) {
106        self.write_all(encode::u64(value, &mut encode::u64_buffer()));
107    }
108
109    #[inline]
110    fn write_varint_u128(&mut self, value: u128) {
111        self.write_all(encode::u128(value, &mut encode::u128_buffer()));
112    }
113
114    #[inline]
115    fn write_varint_i8(&mut self, value: i8) {
116        self.write_varint_u8(ZigZag::encode(value))
117    }
118
119    #[inline]
120    fn write_varint_i16(&mut self, value: i16) {
121        self.write_varint_u16(ZigZag::encode(value))
122    }
123
124    #[inline]
125    fn write_varint_i32(&mut self, value: i32) {
126        self.write_varint_u32(ZigZag::encode(value))
127    }
128
129    #[inline]
130    fn write_varint_i64(&mut self, value: i64) {
131        self.write_varint_u64(ZigZag::encode(value))
132    }
133
134    #[inline]
135    fn write_varint_i128(&mut self, value: i128) {
136        self.write_varint_u128(ZigZag::encode(value))
137    }
138}
139
140impl<W: InfallibleWrite + ?Sized> InfallibleWriteVarInt for W {}
141
142#[allow(unused)]
143pub(crate) trait InfallibleWriteFormat: InfallibleWrite {
144    #[inline]
145    fn write_u8(&mut self, value: u8) {
146        self.write_all(&[value])
147    }
148
149    #[inline]
150    fn write_u16(&mut self, value: u16) {
151        self.write_all(&value.to_be_bytes())
152    }
153
154    #[inline]
155    fn write_u32(&mut self, value: u32) {
156        self.write_all(&value.to_be_bytes())
157    }
158
159    #[inline]
160    fn write_u64(&mut self, value: u64) {
161        self.write_all(&value.to_be_bytes())
162    }
163
164    #[inline]
165    fn write_u128(&mut self, value: u128) {
166        self.write_all(&value.to_be_bytes())
167    }
168
169    #[inline]
170    fn write_i8(&mut self, value: i8) {
171        self.write_all(&value.to_be_bytes())
172    }
173
174    #[inline]
175    fn write_i16(&mut self, value: i16) {
176        self.write_all(&value.to_be_bytes())
177    }
178
179    #[inline]
180    fn write_i32(&mut self, value: i32) {
181        self.write_all(&value.to_be_bytes())
182    }
183
184    #[inline]
185    fn write_i64(&mut self, value: i64) {
186        self.write_all(&value.to_be_bytes())
187    }
188
189    #[inline]
190    fn write_i128(&mut self, value: i128) {
191        self.write_all(&value.to_be_bytes())
192    }
193
194    #[inline]
195    fn write_f32(&mut self, value: f32) {
196        self.write_all(&value.to_be_bytes())
197    }
198
199    #[inline]
200    fn write_f64(&mut self, value: f64) {
201        self.write_all(&value.to_be_bytes())
202    }
203}
204
205impl<W: InfallibleWrite + ?Sized> InfallibleWriteFormat for W {}
206
207impl InfallibleWrite for Vec<u8> {
208    fn write_all(&mut self, buf: &[u8]) {
209        self.extend_from_slice(buf);
210    }
211}
212
213#[allow(unused)]
214pub(crate) trait InfallibleReadFormat: InfallibleRead {
215    #[inline]
216    fn read_u8(&mut self) -> u8 {
217        let mut value = [0; 1];
218        self.read_exact(&mut value);
219        value[0]
220    }
221
222    #[inline]
223    fn read_u16(&mut self) -> u16 {
224        let mut value = [0; 2];
225        self.read_exact(&mut value);
226        u16::from_be_bytes(value)
227    }
228
229    #[inline]
230    fn read_u32(&mut self) -> u32 {
231        let mut value = [0; 4];
232        self.read_exact(&mut value);
233        u32::from_be_bytes(value)
234    }
235
236    #[inline]
237    fn read_u64(&mut self) -> u64 {
238        let mut value = [0; 8];
239        self.read_exact(&mut value);
240        u64::from_be_bytes(value)
241    }
242
243    #[inline]
244    fn read_u128(&mut self) -> u128 {
245        let mut value = [0; 16];
246        self.read_exact(&mut value);
247        u128::from_be_bytes(value)
248    }
249
250    #[inline]
251    fn read_i8(&mut self) -> i8 {
252        let mut value = [0; 1];
253        self.read_exact(&mut value);
254        i8::from_be_bytes(value)
255    }
256
257    #[inline]
258    fn read_i16(&mut self) -> i16 {
259        let mut value = [0; 2];
260        self.read_exact(&mut value);
261        i16::from_be_bytes(value)
262    }
263
264    #[inline]
265    fn read_i32(&mut self) -> i32 {
266        let mut value = [0; 4];
267        self.read_exact(&mut value);
268        i32::from_be_bytes(value)
269    }
270
271    #[inline]
272    fn read_i64(&mut self) -> i64 {
273        let mut value = [0; 8];
274        self.read_exact(&mut value);
275        i64::from_be_bytes(value)
276    }
277
278    #[inline]
279    fn read_i128(&mut self) -> i128 {
280        let mut value = [0; 16];
281        self.read_exact(&mut value);
282        i128::from_be_bytes(value)
283    }
284
285    #[inline]
286    fn read_f32(&mut self) -> f32 {
287        let mut value = [0; 4];
288        self.read_exact(&mut value);
289        f32::from_be_bytes(value)
290    }
291
292    #[inline]
293    fn read_f64(&mut self) -> f64 {
294        let mut value = [0; 8];
295        self.read_exact(&mut value);
296        f64::from_be_bytes(value)
297    }
298}
299
300impl<R: InfallibleRead + ?Sized> InfallibleReadFormat for R {}
301
302pub struct ArcSliceRead {
303    buff: Arc<Vec<u8>>,
304    cursor: usize,
305    limit: usize,
306}
307
308impl ArcSliceRead {
309    pub(crate) fn new(buff: Arc<Vec<u8>>, cursor: usize, limit: usize) -> Self {
310        Self { buff, cursor, limit }
311    }
312
313    pub(crate) fn new_vec(buff: Vec<u8>) -> Self {
314        let limit = buff.len();
315        Self {
316            buff: Arc::new(buff),
317            cursor: 0,
318            limit,
319        }
320    }
321
322    pub(crate) fn to_vec(self) -> Vec<u8> {
323        let mut buf = vec![0; self.limit - self.cursor];
324        Read::read(&mut &self.buff[self.cursor..self.limit], &mut buf).unwrap();
325        buf
326    }
327
328    #[allow(unused)]
329    pub(crate) fn cursor(&self) -> usize {
330        self.cursor
331    }
332
333    pub(crate) fn to_byte_vec(self) -> ByteVec {
334        ByteVec::new_slice(self.buff, self.cursor, self.limit)
335    }
336
337    pub(crate) fn to_string_wrapper(self) -> StringWrapper {
338        StringWrapper::new_slice(self.buff, self.cursor, self.limit)
339    }
340}
341
342impl Read for ArcSliceRead {
343    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
344        let len = self.limit;
345        let amt = std::cmp::min(self.cursor, len);
346        let read = Read::read(&mut &self.buff[amt..len], buf)?;
347        self.cursor += read;
348        Ok(read)
349    }
350}
351
352impl InfallibleRead for ArcSliceRead {
353    fn read_exact(&mut self, buf: &mut [u8]) {
354        assert!(self.limit >= self.cursor + buf.len());
355        buf.copy_from_slice(&self.buff[self.cursor..(self.cursor + buf.len())]);
356        self.cursor += buf.len();
357    }
358
359    fn read_slice(&mut self, size: usize) -> ArcSliceRead {
360        assert!(self.limit >= self.cursor + size);
361        let read = ArcSliceRead::new(self.buff.clone(), self.cursor, size);
362        self.cursor += size;
363        read
364    }
365}
366
367pub(crate) fn read_u16(buf: &[u8]) -> u16 {
368    let mut value = [0; 2];
369    value[0] = buf[0];
370    value[1] = buf[1];
371    u16::from_be_bytes(value)
372}
373pub(crate) fn write_u16(buf: &mut [u8], val: u16) {
374    let bv = val.to_be_bytes();
375    buf[0] = bv[0];
376    buf[1] = bv[1];
377}
378pub(crate) fn read_u64(buf: &[u8]) -> u64 {
379    let mut value = [0; 8];
380    value[..8].clone_from_slice(&buf[..8]);
381    u64::from_be_bytes(value)
382}
383pub(crate) fn write_u64(buf: &mut [u8], val: u64) {
384    let value = val.to_be_bytes();
385    buf[..8].clone_from_slice(&value[..8]);
386}
387
388#[allow(unused)]
389pub(crate) trait WriteFormat: Write {
390    #[inline]
391    fn write_u8(&mut self, value: u8) -> PERes<()> {
392        Ok(self.write_all(&[value])?)
393    }
394
395    #[inline]
396    fn write_u16(&mut self, value: u16) -> PERes<()> {
397        Ok(self.write_all(&value.to_be_bytes())?)
398    }
399
400    #[inline]
401    fn write_u32(&mut self, value: u32) -> PERes<()> {
402        Ok(self.write_all(&value.to_be_bytes())?)
403    }
404
405    #[inline]
406    fn write_u64(&mut self, value: u64) -> PERes<()> {
407        Ok(self.write_all(&value.to_be_bytes())?)
408    }
409
410    #[inline]
411    fn write_u128(&mut self, value: u128) -> PERes<()> {
412        Ok(self.write_all(&value.to_be_bytes())?)
413    }
414
415    #[inline]
416    fn write_i8(&mut self, value: i8) -> PERes<()> {
417        Ok(self.write_all(&value.to_be_bytes())?)
418    }
419
420    #[inline]
421    fn write_i16(&mut self, value: i16) -> PERes<()> {
422        Ok(self.write_all(&value.to_be_bytes())?)
423    }
424
425    #[inline]
426    fn write_i32(&mut self, value: i32) -> PERes<()> {
427        Ok(self.write_all(&value.to_be_bytes())?)
428    }
429
430    #[inline]
431    fn write_i64(&mut self, value: i64) -> PERes<()> {
432        Ok(self.write_all(&value.to_be_bytes())?)
433    }
434
435    #[inline]
436    fn write_i128(&mut self, value: i128) -> PERes<()> {
437        Ok(self.write_all(&value.to_be_bytes())?)
438    }
439
440    #[inline]
441    fn write_f32(&mut self, value: f32) -> PERes<()> {
442        Ok(self.write_all(&value.to_be_bytes())?)
443    }
444
445    #[inline]
446    fn write_f64(&mut self, value: f64) -> PERes<()> {
447        Ok(self.write_all(&value.to_be_bytes())?)
448    }
449}
450
451impl<W: Write + ?Sized> WriteFormat for W {}
452
453#[allow(unused)]
454pub(crate) trait ReadFormat: Read {
455    #[inline]
456    fn read_u8(&mut self) -> PERes<u8> {
457        let mut buf = [0; 1];
458        self.read_exact(&mut buf)?;
459        Ok(buf[0])
460    }
461
462    #[inline]
463    fn read_u16(&mut self) -> PERes<u16> {
464        let mut buf = [0; 2];
465        self.read_exact(&mut buf)?;
466        Ok(u16::from_be_bytes(buf))
467    }
468
469    #[inline]
470    fn read_u32(&mut self) -> PERes<u32> {
471        let mut buf = [0; 4];
472        self.read_exact(&mut buf)?;
473        Ok(u32::from_be_bytes(buf))
474    }
475
476    #[inline]
477    fn read_u64(&mut self) -> PERes<u64> {
478        let mut buf = [0; 8];
479        self.read_exact(&mut buf)?;
480        Ok(u64::from_be_bytes(buf))
481    }
482
483    #[inline]
484    fn read_u128(&mut self) -> PERes<u128> {
485        let mut buf = [0; 16];
486        self.read_exact(&mut buf)?;
487        Ok(u128::from_be_bytes(buf))
488    }
489
490    #[inline]
491    fn read_i8(&mut self) -> PERes<i8> {
492        let mut buf = [0; 1];
493        self.read_exact(&mut buf)?;
494        Ok(buf[0] as i8)
495    }
496
497    #[inline]
498    fn read_i16(&mut self) -> PERes<i16> {
499        let mut buf = [0; 2];
500        self.read_exact(&mut buf)?;
501        Ok(i16::from_be_bytes(buf))
502    }
503
504    #[inline]
505    fn read_i32(&mut self) -> PERes<i32> {
506        let mut buf = [0; 4];
507        self.read_exact(&mut buf)?;
508        Ok(i32::from_be_bytes(buf))
509    }
510
511    #[inline]
512    fn read_i64(&mut self) -> PERes<i64> {
513        let mut buf = [0; 8];
514        self.read_exact(&mut buf)?;
515        Ok(i64::from_be_bytes(buf))
516    }
517
518    #[inline]
519    fn read_i128(&mut self) -> PERes<i128> {
520        let mut buf = [0; 16];
521        self.read_exact(&mut buf)?;
522        Ok(i128::from_be_bytes(buf))
523    }
524
525    #[inline]
526    fn read_f32(&mut self) -> PERes<f32> {
527        let mut buf = [0; 4];
528        self.read_exact(&mut buf)?;
529        Ok(f32::from_be_bytes(buf))
530    }
531
532    #[inline]
533    fn read_f64(&mut self) -> PERes<f64> {
534        let mut buf = [0; 8];
535        self.read_exact(&mut buf)?;
536        Ok(f64::from_be_bytes(buf))
537    }
538}
539
540impl<R: Read + ?Sized> ReadFormat for R {}
541
542#[allow(unused)]
543pub(crate) trait ReadVarInt: Read {
544    #[inline]
545    fn read_varint_u8(&mut self) -> PERes<u8> {
546        Ok(io::read_u8(self)?)
547    }
548
549    #[inline]
550    fn read_varint_u16(&mut self) -> PERes<u16> {
551        Ok(io::read_u16(self)?)
552    }
553
554    #[inline]
555    fn read_varint_u32(&mut self) -> PERes<u32> {
556        Ok(io::read_u32(self)?)
557    }
558
559    #[inline]
560    fn read_varint_u64(&mut self) -> PERes<u64> {
561        Ok(io::read_u64(self)?)
562    }
563
564    #[inline]
565    fn read_varint_u128(&mut self) -> PERes<u128> {
566        Ok(io::read_u128(self)?)
567    }
568
569    #[inline]
570    fn read_varint_i8(&mut self) -> PERes<i8> {
571        Ok(ZigZag::decode(self.read_varint_u8()?))
572    }
573
574    #[inline]
575    fn read_varint_i16(&mut self) -> PERes<i16> {
576        Ok(ZigZag::decode(self.read_varint_u16()?))
577    }
578
579    #[inline]
580    fn read_varint_i32(&mut self) -> PERes<i32> {
581        Ok(ZigZag::decode(self.read_varint_u32()?))
582    }
583
584    #[inline]
585    fn read_varint_i64(&mut self) -> PERes<i64> {
586        Ok(ZigZag::decode(self.read_varint_u64()?))
587    }
588
589    #[inline]
590    fn read_varint_i128(&mut self) -> PERes<i128> {
591        Ok(ZigZag::decode(self.read_varint_u128()?))
592    }
593}
594impl<R: Read + ?Sized> ReadVarInt for R {}
595
596#[cfg(test)]
597mod tests {
598    use std::io::Cursor;
599    #[test]
600    pub fn test_format() {
601        use super::{ReadFormat, WriteFormat};
602        let mut buffer = Vec::new();
603        buffer.write_i8(10).unwrap();
604        buffer.write_i16(11).unwrap();
605        buffer.write_i32(12).unwrap();
606        buffer.write_i64(13).unwrap();
607        buffer.write_i128(14).unwrap();
608        buffer.write_u8(15).unwrap();
609        buffer.write_u16(16).unwrap();
610        buffer.write_u32(17).unwrap();
611        buffer.write_u64(18).unwrap();
612        buffer.write_u128(19).unwrap();
613        buffer.write_f32(10.0).unwrap();
614        buffer.write_f64(11.0).unwrap();
615        let mut read = Cursor::new(buffer);
616        assert_eq!(read.read_i8().unwrap(), 10);
617        assert_eq!(read.read_i16().unwrap(), 11);
618        assert_eq!(read.read_i32().unwrap(), 12);
619        assert_eq!(read.read_i64().unwrap(), 13);
620        assert_eq!(read.read_i128().unwrap(), 14);
621        assert_eq!(read.read_u8().unwrap(), 15);
622        assert_eq!(read.read_u16().unwrap(), 16);
623        assert_eq!(read.read_u32().unwrap(), 17);
624        assert_eq!(read.read_u64().unwrap(), 18);
625        assert_eq!(read.read_u128().unwrap(), 19);
626        assert_eq!(read.read_f32().unwrap(), 10.0);
627        assert_eq!(read.read_f64().unwrap(), 11.0);
628    }
629
630    #[test]
631    pub fn test_infallible_format() {
632        use super::{ArcSliceRead, InfallibleRead, InfallibleReadFormat, InfallibleWriteFormat, Read};
633
634        impl InfallibleRead for Cursor<Vec<u8>> {
635            fn read_exact(&mut self, buf: &mut [u8]) {
636                Read::read_exact(self, buf).expect("in memory buff never fail");
637            }
638
639            fn read_slice(&mut self, size: usize) -> ArcSliceRead {
640                let mut slice: Vec<u8> = vec![0; size];
641                InfallibleRead::read_exact(self, &mut slice);
642                ArcSliceRead::new_vec(slice)
643            }
644        }
645        let mut buffer = Vec::new();
646        buffer.write_i8(10);
647        buffer.write_i16(11);
648        buffer.write_i32(12);
649        buffer.write_i64(13);
650        buffer.write_i128(14);
651        buffer.write_u8(15);
652        buffer.write_u16(16);
653        buffer.write_u32(17);
654        buffer.write_u64(18);
655        buffer.write_u128(19);
656        buffer.write_f32(10.0);
657        buffer.write_f64(11.0);
658        let mut read = Cursor::new(buffer);
659        assert_eq!(read.read_i8(), 10);
660        assert_eq!(read.read_i16(), 11);
661        assert_eq!(read.read_i32(), 12);
662        assert_eq!(read.read_i64(), 13);
663        assert_eq!(read.read_i128(), 14);
664        assert_eq!(read.read_u8(), 15);
665        assert_eq!(read.read_u16(), 16);
666        assert_eq!(read.read_u32(), 17);
667        assert_eq!(read.read_u64(), 18);
668        assert_eq!(read.read_u128(), 19);
669        assert_eq!(read.read_f32(), 10.0);
670        assert_eq!(read.read_f64(), 11.0);
671    }
672}