splinter_rs/codec/
encoder.rs

1use std::ops::RangeInclusive;
2
3use bitvec::{boxed::BitBox, order::Lsb0};
4use bytes::BufMut;
5use crc64fast_nvme::Digest;
6use zerocopy::{IntoBytes, transmute_ref};
7
8use crate::{
9    PartitionRead,
10    codec::{footer::Footer, runs_ref::EncodedRun, tree_ref::TreeIndexBuilder},
11    level::{Block, Level},
12    partition::Partition,
13    partition_kind::PartitionKind,
14    traits::TruncateFrom,
15};
16
17pub struct Encoder<B: BufMut> {
18    buf: B,
19    bytes_written: usize,
20    checksum: Digest,
21}
22
23impl<B: BufMut> Encoder<B> {
24    pub fn new(buf: B) -> Self {
25        Self {
26            buf,
27            bytes_written: 0,
28            checksum: Digest::new(),
29        }
30    }
31
32    /// Retrieve the wrapped buffer from the `Encoder`
33    pub fn into_inner(self) -> B {
34        self.buf
35    }
36
37    /// Write an entire encoded splinter to the buffer
38    pub(crate) fn write_splinter(&mut self, splinter: &[u8]) {
39        self.buf.put_slice(splinter);
40        self.bytes_written += splinter.len();
41    }
42
43    /// Write the checksum and Splinter Magic value to the buffer
44    pub(crate) fn write_footer(&mut self) {
45        let footer = Footer::from_checksum(self.checksum.sum64());
46        self.put_slice(footer.as_bytes());
47    }
48
49    /// The total number of bytes written to the buffer since this Encoder was
50    /// initialized.
51    pub(crate) fn bytes_written(&self) -> usize {
52        self.bytes_written
53    }
54
55    /// Encode a Bitmap partition into the buffer.
56    pub(crate) fn put_bitmap_partition(&mut self, bitmap: &BitBox<u64, Lsb0>) {
57        self.put_bitmap_raw(bitmap);
58    }
59
60    /// Encode a Vec partition into the buffer.
61    pub(crate) fn put_vec_partition<L: Level>(&mut self, values: &[L::Value]) {
62        self.put_iter::<L>(values.iter().copied());
63        self.put_length::<L>(values.len());
64    }
65
66    /// Encode a Run partition into the buffer.
67    pub(crate) fn put_run_partition<L: Level>(
68        &mut self,
69        runs: impl Iterator<Item = RangeInclusive<L::Value>>,
70    ) {
71        let mut num_runs = 0;
72        for run in runs {
73            let run: EncodedRun<L> = run.into();
74            self.put_slice(run.as_bytes());
75            num_runs += 1;
76        }
77        self.put_length::<L>(num_runs);
78    }
79
80    /// Encode a Tree partition into the buffer.
81    pub(crate) fn put_tree_index<L: Level>(&mut self, tree_index_builder: TreeIndexBuilder<L>) {
82        let (num_children, segments, offsets) = tree_index_builder.build();
83        assert!(
84            num_children > 0 && num_children <= Block::MAX_LEN,
85            "num_children out of range"
86        );
87
88        self.put_iter::<L>(offsets);
89
90        match segments {
91            Partition::Full => {}
92            Partition::Bitmap(p) => self.put_bitmap_raw(p.as_bitbox()),
93            Partition::Vec(p) => self.put_iter::<Block>(p.iter()),
94            Partition::Run(_) | Partition::Tree(_) => unreachable!(),
95        }
96
97        self.put_length::<Block>(num_children);
98    }
99
100    pub(crate) fn put_iter<L: Level>(&mut self, values: impl Iterator<Item = L::Value>) {
101        for value in values {
102            self.put_value::<L>(value);
103        }
104    }
105
106    pub(crate) fn put_kind(&mut self, k: PartitionKind) {
107        let d = [k as u8];
108        self.put_slice(&d)
109    }
110
111    #[inline]
112    fn put_length<L: Level>(&mut self, len: usize) {
113        assert_ne!(len, 0, "Length must be greater than zero");
114        // serialize lengths to len-1 to ensure that they fit in the storage type
115        self.put_value::<L>(L::Value::truncate_from(len - 1));
116    }
117
118    #[inline]
119    fn put_value<L: Level>(&mut self, v: L::Value) {
120        self.put_slice(L::ValueUnaligned::from(v).as_bytes());
121    }
122
123    fn put_bitmap_raw(&mut self, bitmap: &BitBox<u64, Lsb0>) {
124        let raw = bitmap.as_raw_slice();
125        static_assertions::assert_cfg!(target_endian = "little");
126        let raw: &[zerocopy::U64<zerocopy::LittleEndian>] = transmute_ref!(raw);
127        self.put_slice(raw.as_bytes());
128    }
129
130    fn put_slice(&mut self, data: &[u8]) {
131        self.checksum.write(data);
132        self.buf.put_slice(data);
133        self.bytes_written += data.len();
134    }
135}
splinter_rs/codec/encoder.rs

splinter_rs/codec/
encoder.rs
