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, cardinalities) = tree_index_builder.build();
83        assert!(
84            num_children > 0 && num_children <= Block::MAX_LEN,
85            "num_children out of range"
86        );
87
88        // Encoding order: [offsets][cardinalities][segments][num_children]
89        self.put_iter::<L>(offsets);
90        self.put_iter::<L>(cardinalities);
91
92        match segments {
93            Partition::Full => {}
94            Partition::Bitmap(p) => self.put_bitmap_raw(p.as_bitbox()),
95            Partition::Vec(p) => self.put_iter::<Block>(p.iter()),
96            Partition::Run(_) | Partition::Tree(_) => unreachable!(),
97        }
98
99        self.put_length::<Block>(num_children);
100    }
101
102    pub(crate) fn put_iter<L: Level>(&mut self, values: impl Iterator<Item = L::Value>) {
103        for value in values {
104            self.put_value::<L>(value);
105        }
106    }
107
108    pub(crate) fn put_kind(&mut self, k: PartitionKind) {
109        let d = [k as u8];
110        self.put_slice(&d)
111    }
112
113    #[inline]
114    fn put_length<L: Level>(&mut self, len: usize) {
115        assert_ne!(len, 0, "Length must be greater than zero");
116        // serialize lengths to len-1 to ensure that they fit in the storage type
117        self.put_value::<L>(L::Value::truncate_from(len - 1));
118    }
119
120    #[inline]
121    fn put_value<L: Level>(&mut self, v: L::Value) {
122        self.put_slice(L::ValueUnaligned::from(v).as_bytes());
123    }
124
125    fn put_bitmap_raw(&mut self, bitmap: &BitBox<u64, Lsb0>) {
126        let raw = bitmap.as_raw_slice();
127        static_assertions::assert_cfg!(target_endian = "little");
128        let raw: &[zerocopy::U64<zerocopy::LittleEndian>] = transmute_ref!(raw);
129        self.put_slice(raw.as_bytes());
130    }
131
132    fn put_slice(&mut self, data: &[u8]) {
133        self.checksum.write(data);
134        self.buf.put_slice(data);
135        self.bytes_written += data.len();
136    }
137}
splinter_rs/codec/encoder.rs

splinter_rs/codec/
encoder.rs
