Skip to main content

parallel_processor/buckets/
bucket_writer.rs

1use std::{
2    io::{Read, Write},
3    marker::PhantomData,
4};
5
6use bincode::{Decode, Encode};
7
8pub trait BucketItemSerializer {
9    type InputElementType<'a>: ?Sized;
10    type ExtraData;
11    type ReadBuffer: Default;
12    type ExtraDataBuffer: Default;
13    type ReadType<'a>;
14    type InitData;
15
16    type CheckpointData: Encode + Decode<()> + 'static;
17
18    fn clear_buffer(buffer: &mut Self::ReadBuffer);
19
20    /// Creates a new instance
21    fn new(init_data: Self::InitData) -> Self;
22    /// Reset on non continuous data
23    fn reset(&mut self);
24
25    fn write_to(
26        &mut self,
27        element: &Self::InputElementType<'_>,
28        bucket: &mut Vec<u8>,
29        extra_data: &Self::ExtraData,
30        extra_read_buffer: &Self::ExtraDataBuffer,
31    );
32    fn read_from<'a, S: Read>(
33        &mut self,
34        stream: S,
35        read_buffer: &'a mut Self::ReadBuffer,
36        extra_read_buffer: &mut Self::ExtraDataBuffer,
37    ) -> Option<Self::ReadType<'a>>;
38
39    fn get_size(&self, element: &Self::InputElementType<'_>, extra: &Self::ExtraData) -> usize;
40}
41
42#[repr(transparent)]
43pub struct BytesArrayBuffer<const SIZE: usize>([u8; SIZE]);
44
45impl<const SIZE: usize> Default for BytesArrayBuffer<SIZE> {
46    fn default() -> Self {
47        Self([0; SIZE])
48    }
49}
50
51pub struct BytesArraySerializer<const SIZE: usize>(PhantomData<[(); SIZE]>);
52impl<const SIZE: usize> BucketItemSerializer for BytesArraySerializer<SIZE> {
53    type InputElementType<'a> = [u8; SIZE];
54    type ExtraData = ();
55    type ExtraDataBuffer = ();
56    type ReadBuffer = BytesArrayBuffer<SIZE>;
57    type ReadType<'a> = &'a [u8; SIZE];
58    type InitData = ();
59
60    type CheckpointData = ();
61
62    fn clear_buffer(_buffer: &mut Self::ReadBuffer) {}
63
64    #[inline(always)]
65    fn new(_: ()) -> Self {
66        Self(PhantomData)
67    }
68
69    #[inline(always)]
70    fn reset(&mut self) {}
71
72    #[inline(always)]
73    fn write_to(
74        &mut self,
75        element: &Self::InputElementType<'_>,
76        bucket: &mut Vec<u8>,
77        _: &Self::ExtraData,
78        _: &Self::ExtraDataBuffer,
79    ) {
80        bucket.write(element).unwrap();
81    }
82
83    fn read_from<'a, S: Read>(
84        &mut self,
85        mut stream: S,
86        read_buffer: &'a mut Self::ReadBuffer,
87        _: &mut Self::ExtraDataBuffer,
88    ) -> Option<Self::ReadType<'a>> {
89        stream.read_exact(&mut read_buffer.0).ok()?;
90        Some(&read_buffer.0)
91    }
92
93    #[inline(always)]
94    fn get_size(&self, element: &Self::InputElementType<'_>, _: &()) -> usize {
95        element.len()
96    }
97}
98
99pub struct BytesSliceSerializer;
100impl BucketItemSerializer for BytesSliceSerializer {
101    type InputElementType<'a> = [u8];
102    type ExtraData = ();
103    type ExtraDataBuffer = ();
104    type ReadBuffer = ();
105    type ReadType<'a> = ();
106    type InitData = ();
107
108    type CheckpointData = ();
109
110    fn clear_buffer(_buffer: &mut Self::ReadBuffer) {}
111
112    #[inline(always)]
113    fn new(_: ()) -> Self {
114        Self
115    }
116
117    #[inline(always)]
118    fn reset(&mut self) {}
119
120    #[inline(always)]
121    fn write_to(
122        &mut self,
123        element: &Self::InputElementType<'_>,
124        bucket: &mut Vec<u8>,
125        _extra_data: &Self::ExtraData,
126        _: &Self::ExtraDataBuffer,
127    ) {
128        bucket.write(element).unwrap();
129    }
130
131    fn read_from<'a, S: Read>(
132        &mut self,
133        _stream: S,
134        _read_buffer: &'a mut Self::ReadBuffer,
135        _: &mut Self::ExtraDataBuffer,
136    ) -> Option<Self::ReadType<'a>> {
137        unimplemented!("Cannot read slices of unknown size!")
138    }
139
140    #[inline(always)]
141    fn get_size(&self, element: &Self::InputElementType<'_>, _: &()) -> usize {
142        element.len()
143    }
144}