ruzstd/decoding/
block_decoder.rs

1use super::super::blocks::block::BlockHeader;
2use super::super::blocks::block::BlockType;
3use super::super::blocks::literals_section::LiteralsSection;
4use super::super::blocks::literals_section::LiteralsSectionType;
5use super::super::blocks::sequence_section::SequencesHeader;
6use super::literals_section_decoder::decode_literals;
7use super::sequence_section_decoder::decode_sequences;
8use crate::common::MAX_BLOCK_SIZE;
9use crate::decoding::errors::DecodeSequenceError;
10use crate::decoding::errors::{
11    BlockHeaderReadError, BlockSizeError, BlockTypeError, DecodeBlockContentError,
12    DecompressBlockError,
13};
14use crate::decoding::scratch::DecoderScratch;
15use crate::decoding::sequence_execution::execute_sequences;
16use crate::io::Read;
17
18pub struct BlockDecoder {
19    header_buffer: [u8; 3],
20    internal_state: DecoderState,
21}
22
23enum DecoderState {
24    ReadyToDecodeNextHeader,
25    ReadyToDecodeNextBody,
26    #[allow(dead_code)]
27    Failed, //TODO put "self.internal_state = DecoderState::Failed;" everywhere an unresolvable error occurs
28}
29
30/// Create a new [BlockDecoder].
31pub fn new() -> BlockDecoder {
32    BlockDecoder {
33        internal_state: DecoderState::ReadyToDecodeNextHeader,
34        header_buffer: [0u8; 3],
35    }
36}
37
38impl BlockDecoder {
39    pub fn decode_block_content(
40        &mut self,
41        header: &BlockHeader,
42        workspace: &mut DecoderScratch, //reuse this as often as possible. Not only if the trees are reused but also reuse the allocations when building new trees
43        mut source: impl Read,
44    ) -> Result<u64, DecodeBlockContentError> {
45        match self.internal_state {
46            DecoderState::ReadyToDecodeNextBody => { /* Happy :) */ }
47            DecoderState::Failed => return Err(DecodeBlockContentError::DecoderStateIsFailed),
48            DecoderState::ReadyToDecodeNextHeader => {
49                return Err(DecodeBlockContentError::ExpectedHeaderOfPreviousBlock)
50            }
51        }
52
53        let block_type = header.block_type;
54        match block_type {
55            BlockType::RLE => {
56                let mut buf = [0u8; 1];
57                source.read_exact(&mut buf[..]).map_err(|err| {
58                    DecodeBlockContentError::ReadError {
59                        step: block_type,
60                        source: err,
61                    }
62                })?;
63                workspace
64                    .buffer
65                    .extend_and_fill(buf[0], header.decompressed_size as usize);
66
67                self.internal_state = DecoderState::ReadyToDecodeNextHeader;
68
69                Ok(1)
70            }
71            BlockType::Raw => {
72                workspace
73                    .buffer
74                    .extend_from_reader(&mut source, header.decompressed_size as usize)
75                    .map_err(|err| DecodeBlockContentError::ReadError {
76                        step: block_type,
77                        source: err,
78                    })?;
79
80                self.internal_state = DecoderState::ReadyToDecodeNextHeader;
81                Ok(u64::from(header.decompressed_size))
82            }
83
84            BlockType::Reserved => {
85                panic!("How did you even get this. The decoder should error out if it detects a reserved-type block");
86            }
87
88            BlockType::Compressed => {
89                self.decompress_block(header, workspace, source)?;
90
91                self.internal_state = DecoderState::ReadyToDecodeNextHeader;
92                Ok(u64::from(header.content_size))
93            }
94        }
95    }
96
97    fn decompress_block(
98        &mut self,
99        header: &BlockHeader,
100        workspace: &mut DecoderScratch, //reuse this as often as possible. Not only if the trees are reused but also reuse the allocations when building new trees
101        mut source: impl Read,
102    ) -> Result<(), DecompressBlockError> {
103        workspace
104            .block_content_buffer
105            .resize(header.content_size as usize, 0);
106
107        source.read_exact(workspace.block_content_buffer.as_mut_slice())?;
108        let raw = workspace.block_content_buffer.as_slice();
109
110        let mut section = LiteralsSection::new();
111        let bytes_in_literals_header = section.parse_from_header(raw)?;
112        let raw = &raw[bytes_in_literals_header as usize..];
113        vprintln!(
114            "Found {} literalssection with regenerated size: {}, and compressed size: {:?}",
115            section.ls_type,
116            section.regenerated_size,
117            section.compressed_size
118        );
119
120        let upper_limit_for_literals = match section.compressed_size {
121            Some(x) => x as usize,
122            None => match section.ls_type {
123                LiteralsSectionType::RLE => 1,
124                LiteralsSectionType::Raw => section.regenerated_size as usize,
125                _ => panic!("Bug in this library"),
126            },
127        };
128
129        if raw.len() < upper_limit_for_literals {
130            return Err(DecompressBlockError::MalformedSectionHeader {
131                expected_len: upper_limit_for_literals,
132                remaining_bytes: raw.len(),
133            });
134        }
135
136        let raw_literals = &raw[..upper_limit_for_literals];
137        vprintln!("Slice for literals: {}", raw_literals.len());
138
139        workspace.literals_buffer.clear(); //all literals of the previous block must have been used in the sequence execution anyways. just be defensive here
140        let bytes_used_in_literals_section = decode_literals(
141            &section,
142            &mut workspace.huf,
143            raw_literals,
144            &mut workspace.literals_buffer,
145        )?;
146        assert!(
147            section.regenerated_size == workspace.literals_buffer.len() as u32,
148            "Wrong number of literals: {}, Should have been: {}",
149            workspace.literals_buffer.len(),
150            section.regenerated_size
151        );
152        assert!(bytes_used_in_literals_section == upper_limit_for_literals as u32);
153
154        let raw = &raw[upper_limit_for_literals..];
155        vprintln!("Slice for sequences with headers: {}", raw.len());
156
157        let mut seq_section = SequencesHeader::new();
158        let bytes_in_sequence_header = seq_section.parse_from_header(raw)?;
159        let raw = &raw[bytes_in_sequence_header as usize..];
160        vprintln!(
161            "Found sequencessection with sequences: {} and size: {}",
162            seq_section.num_sequences,
163            raw.len()
164        );
165
166        assert!(
167            u32::from(bytes_in_literals_header)
168                + bytes_used_in_literals_section
169                + u32::from(bytes_in_sequence_header)
170                + raw.len() as u32
171                == header.content_size
172        );
173        vprintln!("Slice for sequences: {}", raw.len());
174
175        if seq_section.num_sequences != 0 {
176            decode_sequences(
177                &seq_section,
178                raw,
179                &mut workspace.fse,
180                &mut workspace.sequences,
181            )?;
182            vprintln!("Executing sequences");
183            execute_sequences(workspace)?;
184        } else {
185            if !raw.is_empty() {
186                return Err(DecompressBlockError::DecodeSequenceError(
187                    DecodeSequenceError::ExtraBits {
188                        bits_remaining: raw.len() as isize * 8,
189                    },
190                ));
191            }
192            workspace.buffer.push(&workspace.literals_buffer);
193            workspace.sequences.clear();
194        }
195
196        Ok(())
197    }
198
199    /// Reads 3 bytes from the provided reader and returns
200    /// the deserialized header and the number of bytes read.
201    pub fn read_block_header(
202        &mut self,
203        mut r: impl Read,
204    ) -> Result<(BlockHeader, u8), BlockHeaderReadError> {
205        //match self.internal_state {
206        //    DecoderState::ReadyToDecodeNextHeader => {/* Happy :) */},
207        //    DecoderState::Failed => return Err(format!("Cant decode next block if failed along the way. Results will be nonsense")),
208        //    DecoderState::ReadyToDecodeNextBody => return Err(format!("Cant decode next block header, while expecting to decode the body of the previous block. Results will be nonsense")),
209        //}
210
211        r.read_exact(&mut self.header_buffer[0..3])?;
212
213        let btype = self.block_type()?;
214        if let BlockType::Reserved = btype {
215            return Err(BlockHeaderReadError::FoundReservedBlock);
216        }
217
218        let block_size = self.block_content_size()?;
219        let decompressed_size = match btype {
220            BlockType::Raw => block_size,
221            BlockType::RLE => block_size,
222            BlockType::Reserved => 0, //should be caught above, this is an error state
223            BlockType::Compressed => 0, //unknown but will be smaller than 128kb (or window_size if that is smaller than 128kb)
224        };
225        let content_size = match btype {
226            BlockType::Raw => block_size,
227            BlockType::Compressed => block_size,
228            BlockType::RLE => 1,
229            BlockType::Reserved => 0, //should be caught above, this is an error state
230        };
231
232        let last_block = self.is_last();
233
234        self.reset_buffer();
235        self.internal_state = DecoderState::ReadyToDecodeNextBody;
236
237        //just return 3. Blockheaders always take 3 bytes
238        Ok((
239            BlockHeader {
240                last_block,
241                block_type: btype,
242                decompressed_size,
243                content_size,
244            },
245            3,
246        ))
247    }
248
249    fn reset_buffer(&mut self) {
250        self.header_buffer[0] = 0;
251        self.header_buffer[1] = 0;
252        self.header_buffer[2] = 0;
253    }
254
255    fn is_last(&self) -> bool {
256        self.header_buffer[0] & 0x1 == 1
257    }
258
259    fn block_type(&self) -> Result<BlockType, BlockTypeError> {
260        let t = (self.header_buffer[0] >> 1) & 0x3;
261        match t {
262            0 => Ok(BlockType::Raw),
263            1 => Ok(BlockType::RLE),
264            2 => Ok(BlockType::Compressed),
265            3 => Ok(BlockType::Reserved),
266            other => Err(BlockTypeError::InvalidBlocktypeNumber { num: other }),
267        }
268    }
269
270    fn block_content_size(&self) -> Result<u32, BlockSizeError> {
271        let val = self.block_content_size_unchecked();
272        if val > MAX_BLOCK_SIZE {
273            Err(BlockSizeError::BlockSizeTooLarge { size: val })
274        } else {
275            Ok(val)
276        }
277    }
278
279    fn block_content_size_unchecked(&self) -> u32 {
280        u32::from(self.header_buffer[0] >> 3) //push out type and last_block flags. Retain 5 bit
281            | (u32::from(self.header_buffer[1]) << 5)
282            | (u32::from(self.header_buffer[2]) << 13)
283    }
284}
ruzstd/decoding/block_decoder.rs

ruzstd/decoding/
block_decoder.rs
