#![forbid(unsafe_code)]
#[cfg(feature = "alloc")]
use alloc::{vec, vec::Vec};
use crate::decode::BlockDecodeWorkspace;
use crate::error::DecompressError;
use crate::huffman::HuffmanDecodeEntry;
use crate::huffman::decode::{
decode_4_streams, decode_4_streams_into, decode_single_stream, decode_single_stream_vec,
};
use crate::huffman::weights::{
build_huffman_decode_table, build_huffman_decode_table_into, parse_huffman_weights,
};
#[derive(Debug, Clone, Copy)]
pub enum LiteralsBlockType {
Raw,
Rle,
Compressed,
Treeless,
}
pub struct LiteralsSection {
pub data: Vec<u8>,
pub huffman_table: Option<(Vec<HuffmanDecodeEntry>, u8)>,
}
pub fn decode_literals(
data: &[u8],
prev_huffman: Option<&(Vec<HuffmanDecodeEntry>, u8)>,
) -> Result<(LiteralsSection, usize), DecompressError> {
if data.is_empty() {
return Err(DecompressError::CorruptLiterals);
}
let header_byte = data[0];
let block_type = match header_byte & 0x03 {
0 => LiteralsBlockType::Raw,
1 => LiteralsBlockType::Rle,
2 => LiteralsBlockType::Compressed,
3 => LiteralsBlockType::Treeless,
_ => unreachable!(),
};
match block_type {
LiteralsBlockType::Raw => decode_raw_literals(data),
LiteralsBlockType::Rle => decode_rle_literals(data),
LiteralsBlockType::Compressed => decode_compressed_literals(data, false, None),
LiteralsBlockType::Treeless => decode_compressed_literals(data, true, prev_huffman),
}
}
fn decode_raw_literals(data: &[u8]) -> Result<(LiteralsSection, usize), DecompressError> {
let header_byte = data[0];
let size_format = (header_byte >> 2) & 0x03;
let (regen_size, header_size) = match size_format {
0 | 2 => ((header_byte >> 3) as usize, 1),
1 => {
if data.len() < 2 {
return Err(DecompressError::CorruptLiterals);
}
let size = ((header_byte >> 4) as usize) | ((data[1] as usize) << 4);
(size, 2)
}
3 => {
if data.len() < 3 {
return Err(DecompressError::CorruptLiterals);
}
let size = ((header_byte >> 4) as usize)
| ((data[1] as usize) << 4)
| ((data[2] as usize) << 12);
(size, 3)
}
_ => unreachable!(),
};
if data.len() < header_size + regen_size {
return Err(DecompressError::CorruptLiterals);
}
let literals = data[header_size..header_size + regen_size].to_vec();
Ok((
LiteralsSection {
data: literals,
huffman_table: None,
},
header_size + regen_size,
))
}
fn decode_rle_literals(data: &[u8]) -> Result<(LiteralsSection, usize), DecompressError> {
let header_byte = data[0];
let size_format = (header_byte >> 2) & 0x03;
let (regen_size, header_size) = match size_format {
0 | 2 => ((header_byte >> 3) as usize, 1),
1 => {
if data.len() < 2 {
return Err(DecompressError::CorruptLiterals);
}
let size = ((header_byte >> 4) as usize) | ((data[1] as usize) << 4);
(size, 2)
}
3 => {
if data.len() < 3 {
return Err(DecompressError::CorruptLiterals);
}
let size = ((header_byte >> 4) as usize)
| ((data[1] as usize) << 4)
| ((data[2] as usize) << 12);
(size, 3)
}
_ => unreachable!(),
};
if data.len() < header_size + 1 {
return Err(DecompressError::CorruptLiterals);
}
let byte = data[header_size];
let literals = vec![byte; regen_size];
Ok((
LiteralsSection {
data: literals,
huffman_table: None,
},
header_size + 1,
))
}
fn decode_compressed_literals(
data: &[u8],
treeless: bool,
prev_huffman: Option<&(Vec<HuffmanDecodeEntry>, u8)>,
) -> Result<(LiteralsSection, usize), DecompressError> {
let header_byte = data[0];
let size_format = (header_byte >> 2) & 0x03;
let (regen_size, compressed_size, num_streams, header_size) = match size_format {
0 => {
if data.len() < 3 {
return Err(DecompressError::CorruptLiterals);
}
let b0 = header_byte as usize;
let b1 = data[1] as usize;
let b2 = data[2] as usize;
let both_sizes = (b0 >> 4) | (b1 << 4) | (b2 << 12);
let regen = both_sizes & 0x3FF;
let comp = both_sizes >> 10;
(regen, comp, 1usize, 3)
}
1 => {
if data.len() < 3 {
return Err(DecompressError::CorruptLiterals);
}
let b0 = header_byte as usize;
let b1 = data[1] as usize;
let b2 = data[2] as usize;
let both_sizes = (b0 >> 4) | (b1 << 4) | (b2 << 12);
let regen = both_sizes & 0x3FF;
let comp = both_sizes >> 10;
(regen, comp, 4usize, 3)
}
2 => {
if data.len() < 4 {
return Err(DecompressError::CorruptLiterals);
}
let b0 = header_byte as usize;
let b1 = data[1] as usize;
let b2 = data[2] as usize;
let b3 = data[3] as usize;
let both_sizes = (b0 >> 4) | (b1 << 4) | (b2 << 12) | (b3 << 20);
let regen = both_sizes & 0x3FFF;
let comp = both_sizes >> 14;
(regen, comp, 4usize, 4)
}
3 => {
if data.len() < 5 {
return Err(DecompressError::CorruptLiterals);
}
let b0 = header_byte as usize;
let b1 = data[1] as usize;
let b2 = data[2] as usize;
let b3 = data[3] as usize;
let b4 = data[4] as usize;
let both_sizes = (b0 >> 4) | (b1 << 4) | (b2 << 12) | (b3 << 20) | (b4 << 28);
let regen = both_sizes & 0x3FFFF;
let comp = both_sizes >> 18;
(regen, comp, 4usize, 5)
}
_ => unreachable!(),
};
if data.len() < header_size + compressed_size {
return Err(DecompressError::CorruptLiterals);
}
let stream_data = &data[header_size..header_size + compressed_size];
let (huffman_table, huf_consumed) = if treeless {
let table = prev_huffman
.ok_or(DecompressError::CorruptLiterals)?
.clone();
(table, 0)
} else {
let (weights, consumed) = parse_huffman_weights(stream_data)?;
let table = build_huffman_decode_table(&weights)?;
(table, consumed)
};
let compressed_stream = &stream_data[huf_consumed..];
let literals_data = if num_streams == 1 {
decode_single_stream(
&huffman_table.0,
huffman_table.1,
compressed_stream,
regen_size,
)?
} else {
decode_4_streams(
&huffman_table.0,
huffman_table.1,
compressed_stream,
regen_size,
)?
};
Ok((
LiteralsSection {
data: literals_data,
huffman_table: Some(huffman_table),
},
header_size + compressed_size,
))
}
fn parse_raw_rle_header(data: &[u8]) -> Result<(usize, usize), DecompressError> {
let header_byte = data[0];
let size_format = (header_byte >> 2) & 0x03;
match size_format {
0 | 2 => Ok(((header_byte >> 3) as usize, 1)),
1 => {
if data.len() < 2 {
return Err(DecompressError::CorruptLiterals);
}
Ok((((header_byte >> 4) as usize) | ((data[1] as usize) << 4), 2))
}
3 => {
if data.len() < 3 {
return Err(DecompressError::CorruptLiterals);
}
Ok((
((header_byte >> 4) as usize)
| ((data[1] as usize) << 4)
| ((data[2] as usize) << 12),
3,
))
}
_ => unreachable!(),
}
}
fn parse_compressed_header(data: &[u8]) -> Result<(usize, usize, usize, usize), DecompressError> {
let header_byte = data[0];
let size_format = (header_byte >> 2) & 0x03;
match size_format {
0 => {
if data.len() < 3 {
return Err(DecompressError::CorruptLiterals);
}
let both = (header_byte as usize >> 4)
| ((data[1] as usize) << 4)
| ((data[2] as usize) << 12);
Ok((both & 0x3FF, both >> 10, 1, 3))
}
1 => {
if data.len() < 3 {
return Err(DecompressError::CorruptLiterals);
}
let both = (header_byte as usize >> 4)
| ((data[1] as usize) << 4)
| ((data[2] as usize) << 12);
Ok((both & 0x3FF, both >> 10, 4, 3))
}
2 => {
if data.len() < 4 {
return Err(DecompressError::CorruptLiterals);
}
let both = (header_byte as usize >> 4)
| ((data[1] as usize) << 4)
| ((data[2] as usize) << 12)
| ((data[3] as usize) << 20);
Ok((both & 0x3FFF, both >> 14, 4, 4))
}
3 => {
if data.len() < 5 {
return Err(DecompressError::CorruptLiterals);
}
let both = (header_byte as usize >> 4)
| ((data[1] as usize) << 4)
| ((data[2] as usize) << 12)
| ((data[3] as usize) << 20)
| ((data[4] as usize) << 28);
Ok((both & 0x3FFFF, both >> 18, 4, 5))
}
_ => unreachable!(),
}
}
pub(crate) fn decode_literals_ws(
data: &[u8],
ws: &mut BlockDecodeWorkspace,
) -> Result<usize, DecompressError> {
if data.is_empty() {
return Err(DecompressError::CorruptLiterals);
}
let block_type = match data[0] & 0x03 {
0 => LiteralsBlockType::Raw,
1 => LiteralsBlockType::Rle,
2 => LiteralsBlockType::Compressed,
3 => LiteralsBlockType::Treeless,
_ => unreachable!(),
};
let consumed = match block_type {
LiteralsBlockType::Raw => {
let (regen_size, header_size) = parse_raw_rle_header(data)?;
if regen_size > crate::frame::MAX_BLOCK_SIZE {
return Err(DecompressError::CorruptLiterals);
}
if data.len() < header_size + regen_size {
return Err(DecompressError::CorruptLiterals);
}
ws.literal_buf.clear();
ws.literal_buf
.extend_from_slice(&data[header_size..header_size + regen_size]);
header_size + regen_size
}
LiteralsBlockType::Rle => {
let (regen_size, header_size) = parse_raw_rle_header(data)?;
if regen_size > crate::frame::MAX_BLOCK_SIZE {
return Err(DecompressError::CorruptLiterals);
}
if data.len() < header_size + 1 {
return Err(DecompressError::CorruptLiterals);
}
let byte = data[header_size];
ws.literal_buf.clear();
ws.literal_buf.resize(regen_size, byte);
header_size + 1
}
LiteralsBlockType::Compressed | LiteralsBlockType::Treeless => {
let treeless = matches!(block_type, LiteralsBlockType::Treeless);
let (regen_size, compressed_size, num_streams, header_size) =
parse_compressed_header(data)?;
if regen_size > crate::frame::MAX_BLOCK_SIZE {
return Err(DecompressError::CorruptLiterals);
}
if data.len() < header_size + compressed_size {
return Err(DecompressError::CorruptLiterals);
}
let stream_data = &data[header_size..header_size + compressed_size];
let huf_consumed = if treeless {
if !ws.huf_valid {
return Err(DecompressError::CorruptLiterals);
}
0
} else {
let (weights, consumed) = parse_huffman_weights(stream_data)?;
ws.huf_table_log = build_huffman_decode_table_into(
&weights,
&mut ws.huf_table,
&mut ws.huf_all_weights,
&mut ws.huf_rank_count,
&mut ws.huf_rank_start,
)?;
ws.huf_valid = true;
consumed
};
let compressed_stream = &stream_data[huf_consumed..];
if num_streams == 1 {
decode_single_stream_vec(
&ws.huf_table,
ws.huf_table_log,
compressed_stream,
regen_size,
&mut ws.literal_buf,
)?;
} else {
decode_4_streams_into(
&ws.huf_table,
ws.huf_table_log,
compressed_stream,
regen_size,
&mut ws.literal_buf,
)?;
}
header_size + compressed_size
}
};
let real_len = ws.literal_buf.len();
ws.literal_buf.resize(real_len + 32, 0);
ws.literal_buf.truncate(real_len);
Ok(consumed)
}