Skip to main content

ruzstd/decoding/
literals_section_decoder.rs

1//! This module contains the decompress_literals function, used to take a
2//! parsed literals header and a source and decompress it.
3
4use super::super::blocks::literals_section::{LiteralsSection, LiteralsSectionType};
5use super::scratch::HuffmanScratch;
6use crate::bit_io::BitReaderReversed;
7use crate::decoding::errors::DecompressLiteralsError;
8use crate::huff0::HuffmanDecoder;
9use alloc::vec::Vec;
10
11/// Decode and decompress the provided literals section into `target`, returning the number of bytes read.
12pub fn decode_literals(
13    section: &LiteralsSection,
14    scratch: &mut HuffmanScratch,
15    source: &[u8],
16    target: &mut Vec<u8>,
17) -> Result<u32, DecompressLiteralsError> {
18    match section.ls_type {
19        LiteralsSectionType::Raw => {
20            target.extend(&source[0..section.regenerated_size as usize]);
21            Ok(section.regenerated_size)
22        }
23        LiteralsSectionType::RLE => {
24            target.resize(target.len() + section.regenerated_size as usize, source[0]);
25            Ok(1)
26        }
27        LiteralsSectionType::Compressed | LiteralsSectionType::Treeless => {
28            let bytes_read = decompress_literals(section, scratch, source, target)?;
29
30            //return sum of used bytes
31            Ok(bytes_read)
32        }
33    }
34}
35
36/// Decompress the provided literals section and source into the provided `target`.
37/// This function is used when the literals section is `Compressed` or `Treeless`
38///
39/// Returns the number of bytes read.
40fn decompress_literals(
41    section: &LiteralsSection,
42    scratch: &mut HuffmanScratch,
43    source: &[u8],
44    target: &mut Vec<u8>,
45) -> Result<u32, DecompressLiteralsError> {
46    use DecompressLiteralsError as err;
47
48    let compressed_size = section.compressed_size.ok_or(err::MissingCompressedSize)? as usize;
49    let num_streams = section.num_streams.ok_or(err::MissingNumStreams)?;
50
51    target.reserve(section.regenerated_size as usize);
52    let source = &source[0..compressed_size];
53    let mut bytes_read = 0;
54
55    match section.ls_type {
56        LiteralsSectionType::Compressed => {
57            //read Huffman tree description
58            bytes_read += scratch.table.build_decoder(source)?;
59            vprintln!("Built huffman table using {} bytes", bytes_read);
60        }
61        LiteralsSectionType::Treeless if scratch.table.max_num_bits == 0 => {
62            return Err(err::UninitializedHuffmanTable);
63        }
64
65        _ => { /* nothing to do, huffman tree has been provided by previous block */ }
66    }
67
68    let source = &source[bytes_read as usize..];
69
70    if num_streams == 4 {
71        //build jumptable
72        if source.len() < 6 {
73            return Err(err::MissingBytesForJumpHeader { got: source.len() });
74        }
75        let jump1 = source[0] as usize + ((source[1] as usize) << 8);
76        let jump2 = jump1 + source[2] as usize + ((source[3] as usize) << 8);
77        let jump3 = jump2 + source[4] as usize + ((source[5] as usize) << 8);
78        bytes_read += 6;
79        let source = &source[6..];
80
81        if source.len() < jump3 {
82            return Err(err::MissingBytesForLiterals {
83                got: source.len(),
84                needed: jump3,
85            });
86        }
87
88        //decode 4 streams
89        let stream1 = &source[..jump1];
90        let stream2 = &source[jump1..jump2];
91        let stream3 = &source[jump2..jump3];
92        let stream4 = &source[jump3..];
93
94        for stream in &[stream1, stream2, stream3, stream4] {
95            let mut decoder = HuffmanDecoder::new(&scratch.table);
96            let mut br = BitReaderReversed::new(stream);
97            //skip the 0 padding at the end of the last byte of the bit stream and throw away the first 1 found
98            let mut skipped_bits = 0;
99            loop {
100                let val = br.get_bits(1);
101                skipped_bits += 1;
102                if val == 1 || skipped_bits > 8 {
103                    break;
104                }
105            }
106            if skipped_bits > 8 {
107                //if more than 7 bits are 0, this is not the correct end of the bitstream. Either a bug or corrupted data
108                return Err(DecompressLiteralsError::ExtraPadding { skipped_bits });
109            }
110            decoder.init_state(&mut br);
111
112            while br.bits_remaining() > -(scratch.table.max_num_bits as isize) {
113                target.push(decoder.decode_symbol());
114                decoder.next_state(&mut br);
115            }
116            if br.bits_remaining() != -(scratch.table.max_num_bits as isize) {
117                return Err(DecompressLiteralsError::BitstreamReadMismatch {
118                    read_til: br.bits_remaining(),
119                    expected: -(scratch.table.max_num_bits as isize),
120                });
121            }
122        }
123
124        bytes_read += source.len() as u32;
125    } else {
126        //just decode the one stream
127        assert!(num_streams == 1);
128        let mut decoder = HuffmanDecoder::new(&scratch.table);
129        let mut br = BitReaderReversed::new(source);
130        let mut skipped_bits = 0;
131        loop {
132            let val = br.get_bits(1);
133            skipped_bits += 1;
134            if val == 1 || skipped_bits > 8 {
135                break;
136            }
137        }
138        if skipped_bits > 8 {
139            //if more than 7 bits are 0, this is not the correct end of the bitstream. Either a bug or corrupted data
140            return Err(DecompressLiteralsError::ExtraPadding { skipped_bits });
141        }
142        decoder.init_state(&mut br);
143        while br.bits_remaining() > -(scratch.table.max_num_bits as isize) {
144            target.push(decoder.decode_symbol());
145            decoder.next_state(&mut br);
146        }
147        bytes_read += source.len() as u32;
148    }
149
150    if target.len() != section.regenerated_size as usize {
151        return Err(DecompressLiteralsError::DecodedLiteralCountMismatch {
152            decoded: target.len(),
153            expected: section.regenerated_size as usize,
154        });
155    }
156
157    Ok(bytes_read)
158}