use alloc::vec::Vec;
use core::convert::TryInto;
#[cfg(feature = "hash")]
use twox_hash::XxHash64;
#[cfg(feature = "hash")]
use core::hash::Hasher;
use super::{
CompressionLevel, Matcher, block_header::BlockHeader, frame_header::FrameHeader, levels::*,
match_generator::MatchGeneratorDriver,
};
use crate::common::MAX_BLOCK_SIZE;
use crate::fse::fse_encoder::{FSETable, default_ll_table, default_ml_table, default_of_table};
use crate::io::{Read, Write};
#[derive(Clone)]
pub struct EncoderDictionary {
pub(crate) inner: crate::decoding::Dictionary,
}
impl EncoderDictionary {
pub fn from_dictionary(dictionary: crate::decoding::Dictionary) -> Self {
Self { inner: dictionary }
}
pub fn from_bytes(
raw_dictionary: &[u8],
) -> Result<Self, crate::decoding::errors::DictionaryDecodeError> {
Ok(Self {
inner: crate::decoding::Dictionary::decode_dict_for_encoding(raw_dictionary)?,
})
}
pub fn id(&self) -> u32 {
self.inner.id
}
}
pub struct FrameCompressor<
R: Read = &'static [u8],
W: Write = Vec<u8>,
M: Matcher = MatchGeneratorDriver,
> {
uncompressed_data: Option<R>,
compressed_data: Option<W>,
compression_level: CompressionLevel,
dictionary: Option<EncoderDictionary>,
dictionary_entropy_cache: Option<CachedDictionaryEntropy>,
source_size_hint: Option<u64>,
state: CompressState<M>,
magicless: bool,
content_checksum: bool,
content_size_flag: bool,
dict_id_flag: bool,
target_block_size: Option<u32>,
#[cfg(feature = "hash")]
hasher: XxHash64,
#[cfg(feature = "lsm")]
frame_emit_info: Option<crate::encoding::frame_emit_info::FrameEmitInfo>,
#[cfg(all(feature = "lsm", feature = "hash"))]
per_block_checksums_enabled: bool,
#[cfg(all(feature = "lsm", feature = "hash"))]
block_checksums: Option<alloc::vec::Vec<u32>>,
#[cfg(feature = "lsm")]
block_decompressed_sizes: alloc::vec::Vec<u32>,
strategy_override: Option<crate::encoding::strategy::StrategyTag>,
}
#[derive(Clone, Default)]
pub(crate) struct CachedDictionaryEntropy {
pub(crate) huff: Option<crate::huff0::huff0_encoder::HuffmanTable>,
pub(crate) ll_previous: Option<PreviousFseTable>,
pub(crate) ml_previous: Option<PreviousFseTable>,
pub(crate) of_previous: Option<PreviousFseTable>,
}
impl CachedDictionaryEntropy {
pub(crate) fn heap_size(&self) -> usize {
let mut total = self.huff.as_ref().map_or(0, |h| h.heap_size());
for prev in [&self.ll_previous, &self.ml_previous, &self.of_previous] {
if let Some(PreviousFseTable::Custom(table)) = prev {
total +=
core::mem::size_of::<crate::fse::fse_encoder::FSETable>() + table.heap_size();
}
}
total
}
pub(crate) fn from_dictionary(dictionary: &crate::decoding::Dictionary) -> Self {
Self {
huff: dictionary.huf.table.to_encoder_table(),
ll_previous: dictionary
.fse
.literal_lengths
.to_encoder_table()
.map(|table| PreviousFseTable::Custom(SharedFseTable::new(table))),
ml_previous: dictionary
.fse
.match_lengths
.to_encoder_table()
.map(|table| PreviousFseTable::Custom(SharedFseTable::new(table))),
of_previous: dictionary
.fse
.offsets
.to_encoder_table()
.map(|table| PreviousFseTable::Custom(SharedFseTable::new(table))),
}
}
}
#[cfg(target_has_atomic = "ptr")]
pub(crate) type SharedFseTable = alloc::sync::Arc<FSETable>;
#[cfg(not(target_has_atomic = "ptr"))]
pub(crate) type SharedFseTable = alloc::rc::Rc<FSETable>;
#[derive(Clone)]
pub(crate) enum PreviousFseTable {
Default,
Custom(SharedFseTable),
Rle(u8),
}
impl PreviousFseTable {
pub(crate) fn as_table<'a>(&'a self, default: &'a FSETable) -> Option<&'a FSETable> {
match self {
Self::Default => Some(default),
Self::Custom(table) => Some(table),
Self::Rle(_) => None,
}
}
}
pub(crate) struct FseTables {
pub(crate) ll_default: crate::fse::fse_encoder::FseDefaultTable,
pub(crate) ll_previous: Option<PreviousFseTable>,
pub(crate) ml_default: crate::fse::fse_encoder::FseDefaultTable,
pub(crate) ml_previous: Option<PreviousFseTable>,
pub(crate) of_default: crate::fse::fse_encoder::FseDefaultTable,
pub(crate) of_previous: Option<PreviousFseTable>,
}
impl FseTables {
pub fn new() -> Self {
Self {
ll_default: default_ll_table(),
ll_previous: None,
ml_default: default_ml_table(),
ml_previous: None,
of_default: default_of_table(),
of_previous: None,
}
}
#[inline]
#[allow(clippy::borrow_deref_ref)]
pub(crate) fn ll_default_ref(&self) -> &FSETable {
&*self.ll_default
}
#[inline]
#[allow(clippy::borrow_deref_ref)]
pub(crate) fn ml_default_ref(&self) -> &FSETable {
&*self.ml_default
}
#[inline]
#[allow(clippy::borrow_deref_ref)]
pub(crate) fn of_default_ref(&self) -> &FSETable {
&*self.of_default
}
}
const PRESPLIT_BLOCK_MIN: usize = 3500;
const PRESPLIT_THRESHOLD_PENALTY_RATE: u64 = 16;
const PRESPLIT_THRESHOLD_BASE: u64 = PRESPLIT_THRESHOLD_PENALTY_RATE - 2;
const PRESPLIT_THRESHOLD_PENALTY: i32 = 3;
const PRESPLIT_CHUNK_SIZE: usize = 8 << 10;
const PRESPLIT_HASH_LOG_MAX: usize = 10;
const PRESPLIT_HASH_TABLE_SIZE: usize = 1 << PRESPLIT_HASH_LOG_MAX;
const PRESPLIT_KNUTH: u32 = 0x9E37_79B9;
const PRESPLIT_BORDERS_SEGMENT: usize = 512;
#[derive(Clone)]
struct PreSplitFingerprint {
events: [u32; PRESPLIT_HASH_TABLE_SIZE],
nb_events: usize,
}
impl Default for PreSplitFingerprint {
fn default() -> Self {
Self {
events: [0; PRESPLIT_HASH_TABLE_SIZE],
nb_events: 0,
}
}
}
fn reserve_for_next_block(
out: &mut Vec<u8>,
blocks_start: usize,
consumed: u64,
remaining: usize,
block_capacity: usize,
) {
let block_bound = remaining.min(block_capacity) + 3 + 16;
if out.capacity() - out.len() >= block_bound {
return;
}
let produced = (out.len() - blocks_start) as u64;
let estimate = if consumed == 0 {
block_bound
} else {
let scaled = ((remaining as u128 * produced as u128) / consumed as u128) as u64;
let headers = (remaining as u64 / block_capacity.max(1) as u64 + 1) * 3;
usize::try_from(scaled + scaled / 16 + headers + 64).unwrap_or(usize::MAX)
};
out.reserve_exact(estimate.max(block_bound + produced as usize));
}
fn presplit_hash2(bytes: &[u8], hash_log: usize) -> usize {
debug_assert!(hash_log >= 8);
if hash_log == 8 {
return bytes[0] as usize;
}
debug_assert!(hash_log <= PRESPLIT_HASH_LOG_MAX);
let value = u16::from_le_bytes([bytes[0], bytes[1]]) as u32;
(value.wrapping_mul(PRESPLIT_KNUTH) >> (32 - hash_log)) as usize
}
fn presplit_record_fingerprint(
fp: &mut PreSplitFingerprint,
src: &[u8],
sampling_rate: usize,
hash_log: usize,
) {
fp.events.fill(0);
fp.nb_events = 0;
if src.len() < 2 {
return;
}
let limit = src.len() - 1;
let mut n = 0usize;
while n < limit {
fp.events[presplit_hash2(&src[n..], hash_log)] += 1;
n += sampling_rate;
}
fp.nb_events += limit / sampling_rate;
}
fn presplit_record_byte_histogram(fp: &mut PreSplitFingerprint, src: &[u8]) {
fp.events.fill(0);
for &b in src {
fp.events[b as usize] += 1;
}
fp.nb_events = src.len();
}
fn presplit_distance(lhs: &PreSplitFingerprint, rhs: &PreSplitFingerprint, hash_log: usize) -> u64 {
let slots = 1usize << hash_log;
let mut distance = 0u64;
for idx in 0..slots {
let left = lhs.events[idx] as i128 * rhs.nb_events as i128;
let right = rhs.events[idx] as i128 * lhs.nb_events as i128;
distance += left.abs_diff(right) as u64;
}
distance
}
fn presplit_fingerprints_differ(
reference: &PreSplitFingerprint,
new_fp: &PreSplitFingerprint,
penalty: i32,
hash_log: usize,
) -> bool {
debug_assert!(reference.nb_events > 0);
debug_assert!(new_fp.nb_events > 0);
let p50 = reference.nb_events as u64 * new_fp.nb_events as u64;
let deviation = presplit_distance(reference, new_fp, hash_log);
let threshold =
p50 * (PRESPLIT_THRESHOLD_BASE + penalty as u64) / PRESPLIT_THRESHOLD_PENALTY_RATE;
deviation >= threshold
}
fn presplit_merge_events(acc: &mut PreSplitFingerprint, new_fp: &PreSplitFingerprint) {
for idx in 0..PRESPLIT_HASH_TABLE_SIZE {
acc.events[idx] += new_fp.events[idx];
}
acc.nb_events += new_fp.nb_events;
}
fn split_block_by_chunks(block: &[u8], level: usize) -> usize {
debug_assert_eq!(block.len(), MAX_BLOCK_SIZE as usize);
debug_assert!((1..=4).contains(&level));
let (sampling_rate, hash_log) = match level - 1 {
0 => (43, 8),
1 => (11, 9),
2 => (5, 10),
_ => (1, 10),
};
let mut past = PreSplitFingerprint::default();
let mut new_events = PreSplitFingerprint::default();
let mut penalty = PRESPLIT_THRESHOLD_PENALTY;
presplit_record_fingerprint(
&mut past,
&block[..PRESPLIT_CHUNK_SIZE],
sampling_rate,
hash_log,
);
let mut pos = PRESPLIT_CHUNK_SIZE;
while pos <= block.len() - PRESPLIT_CHUNK_SIZE {
presplit_record_fingerprint(
&mut new_events,
&block[pos..pos + PRESPLIT_CHUNK_SIZE],
sampling_rate,
hash_log,
);
if presplit_fingerprints_differ(&past, &new_events, penalty, hash_log) {
return pos;
}
presplit_merge_events(&mut past, &new_events);
if penalty > 0 {
penalty -= 1;
}
pos += PRESPLIT_CHUNK_SIZE;
}
block.len()
}
fn split_block_from_borders(block: &[u8]) -> usize {
debug_assert_eq!(block.len(), MAX_BLOCK_SIZE as usize);
let block_size = block.len();
let mut past = PreSplitFingerprint::default();
let mut new_fp = PreSplitFingerprint::default();
presplit_record_byte_histogram(&mut past, &block[..PRESPLIT_BORDERS_SEGMENT]);
presplit_record_byte_histogram(&mut new_fp, &block[block_size - PRESPLIT_BORDERS_SEGMENT..]);
if !presplit_fingerprints_differ(&past, &new_fp, 0, 8) {
return block_size;
}
let mut middle = PreSplitFingerprint::default();
let mid_start = block_size / 2 - PRESPLIT_BORDERS_SEGMENT / 2;
presplit_record_byte_histogram(
&mut middle,
&block[mid_start..mid_start + PRESPLIT_BORDERS_SEGMENT],
);
let dist_from_begin = presplit_distance(&past, &middle, 8);
let dist_from_end = presplit_distance(&new_fp, &middle, 8);
let min_distance = (PRESPLIT_BORDERS_SEGMENT as u64) * (PRESPLIT_BORDERS_SEGMENT as u64) / 3;
if dist_from_begin.abs_diff(dist_from_end) < min_distance {
return 64 * 1024;
}
if dist_from_begin > dist_from_end {
32 * 1024
} else {
96 * 1024
}
}
#[cfg(all(feature = "lsm", feature = "hash"))]
#[inline]
pub(crate) fn xxh64_block_low32(data: &[u8]) -> u32 {
let mut h = XxHash64::with_seed(0);
h.write(data);
h.finish() as u32
}
#[cfg(feature = "bench_internals")]
pub(crate) fn block_splitter_decision_for_bench(block: &[u8], split_level: usize) -> usize {
assert_eq!(
block.len(),
MAX_BLOCK_SIZE as usize,
"block_splitter_decision_for_bench expects exactly MAX_BLOCK_SIZE bytes"
);
assert!(
split_level <= 4,
"block_splitter_decision_for_bench: split_level must be in 0..=4, got {split_level}"
);
if split_level == 0 {
split_block_from_borders(block)
} else {
split_block_by_chunks(block, split_level)
}
}
#[inline]
fn warm_presplit_window(window: &[u8]) {
let mut acc = 0u8;
let mut i = 0usize;
while i < window.len() {
acc ^= window[i];
i += 64;
}
core::hint::black_box(acc);
}
pub(crate) fn optimal_block_size(
level: CompressionLevel,
block: &[u8],
remaining_src_size: usize,
block_size_max: usize,
savings: i64,
) -> usize {
let Some(split_level) = crate::encoding::levels::config::level_pre_split(level) else {
return remaining_src_size.min(block_size_max);
};
if remaining_src_size < MAX_BLOCK_SIZE as usize || block_size_max < MAX_BLOCK_SIZE as usize {
return remaining_src_size.min(block_size_max);
}
if savings < 3 {
return MAX_BLOCK_SIZE as usize;
}
if block.len() < MAX_BLOCK_SIZE as usize {
return remaining_src_size.min(block_size_max);
}
let raw_split = if split_level == 0 {
split_block_from_borders(&block[..MAX_BLOCK_SIZE as usize])
} else {
split_block_by_chunks(&block[..MAX_BLOCK_SIZE as usize], split_level)
};
raw_split
.max(PRESPLIT_BLOCK_MIN)
.min(MAX_BLOCK_SIZE as usize)
}
pub(crate) struct CompressState<M: Matcher> {
pub(crate) matcher: M,
pub(crate) last_huff_table: Option<crate::huff0::huff0_encoder::HuffmanTable>,
pub(crate) huff_table_spare: Option<crate::huff0::huff0_encoder::HuffmanTable>,
pub(crate) fse_tables: FseTables,
pub(crate) block_scratch: crate::encoding::blocks::CompressedBlockScratch,
pub(crate) offset_hist: [u32; 3],
pub(crate) strategy_tag: crate::encoding::strategy::StrategyTag,
pub(crate) huf_optimal_search: bool,
pub(crate) literal_compression_disabled: bool,
}
pub(crate) fn huf_search_enabled(
strategy: crate::encoding::strategy::StrategyTag,
_source_size: Option<u64>,
) -> bool {
use crate::encoding::strategy::StrategyTag;
matches!(strategy, StrategyTag::BtUltra | StrategyTag::BtUltra2)
}
impl<M: Matcher> CompressState<M> {
#[inline]
pub(crate) fn clear_huff_table(&mut self) {
if let Some(table) = self.last_huff_table.take() {
self.huff_table_spare = Some(table);
}
}
#[inline]
pub(crate) fn replace_huff_table(&mut self, table: crate::huff0::huff0_encoder::HuffmanTable) {
if let Some(old) = self.last_huff_table.replace(table) {
self.huff_table_spare = Some(old);
}
}
}
struct FramePrep {
window_size: u64,
use_dictionary_state: bool,
source_size_hint_known: bool,
initial_size_hint: Option<u64>,
}
fn initial_all_blocks_cap(initial_size_hint: Option<u64>, block_capacity: usize) -> usize {
const TINY_THRESHOLD: u64 = 4 * 1024;
const SMALL_THRESHOLD: u64 = 64 * 1024;
const TINY_CAP: usize = 4 * 1024;
const SMALL_CAP: usize = 16 * 1024;
const DEFAULT_CAP: usize = 130 * 1024;
let first_block_cap = block_capacity + 3 + 16;
match initial_size_hint {
Some(h) if h <= TINY_THRESHOLD => TINY_CAP.min(first_block_cap),
Some(h) if h <= SMALL_THRESHOLD => SMALL_CAP.min(first_block_cap),
_ => DEFAULT_CAP.min(first_block_cap),
}
}
pub(crate) trait OwnedBlockSource {
fn fill_block(
&mut self,
buf: &mut Vec<u8>,
block_capacity: usize,
size_hint_remaining: Option<u64>,
) -> (usize, bool);
}
impl OwnedBlockSource for &[u8] {
fn fill_block(
&mut self,
buf: &mut Vec<u8>,
block_capacity: usize,
_size_hint_remaining: Option<u64>,
) -> (usize, bool) {
let want = block_capacity - buf.len();
let take = want.min(self.len());
buf.extend_from_slice(&self[..take]);
*self = &self[take..];
(take, take < want || self.is_empty())
}
}
pub(crate) struct ReaderBlockSource<Rd> {
pub(crate) reader: Rd,
peeked: Option<u8>,
}
impl<Rd> ReaderBlockSource<Rd> {
pub(crate) fn new(reader: Rd) -> Self {
Self {
reader,
peeked: None,
}
}
}
impl<Rd: Read> OwnedBlockSource for ReaderBlockSource<Rd> {
fn fill_block(
&mut self,
buf: &mut Vec<u8>,
block_capacity: usize,
size_hint_remaining: Option<u64>,
) -> (usize, bool) {
let start = buf.len();
let mut filled = start;
let mut reached_eof = false;
if let Some(b) = self.peeked.take() {
buf.push(b);
filled += 1;
}
let initial_target = match size_hint_remaining {
Some(remaining) => {
let remaining = remaining.min(block_capacity as u64) as usize;
filled + remaining.min(block_capacity - filled)
}
None => block_capacity,
};
if buf.len() < initial_target {
buf.resize(initial_target, 0);
}
loop {
if reached_eof || filled == block_capacity {
break;
}
if filled == buf.len() {
let grow_to = (buf.len() * 2).clamp(filled + 1, block_capacity);
buf.resize(grow_to, 0);
}
let read_end = buf.len();
let new_bytes = self.reader.read(&mut buf[filled..read_end]).unwrap();
if new_bytes == 0 {
reached_eof = true;
break;
}
filled += new_bytes;
}
if !reached_eof && filled == block_capacity {
let mut probe = [0u8; 1];
if self.reader.read(&mut probe).unwrap() == 0 {
reached_eof = true;
} else {
self.peeked = Some(probe[0]);
}
}
buf.truncate(filled);
(filled - start, reached_eof)
}
}
impl<R: Read, W: Write> FrameCompressor<R, W, MatchGeneratorDriver> {
pub fn new(compression_level: CompressionLevel) -> Self {
Self {
uncompressed_data: None,
compressed_data: None,
compression_level,
dictionary: None,
dictionary_entropy_cache: None,
source_size_hint: None,
state: CompressState {
matcher: MatchGeneratorDriver::new(1024 * 128, 1),
last_huff_table: None,
huff_table_spare: None,
fse_tables: FseTables::new(),
block_scratch: crate::encoding::blocks::CompressedBlockScratch::new(),
offset_hist: [1, 4, 8],
strategy_tag: crate::encoding::strategy::StrategyTag::for_compression_level(
compression_level,
),
huf_optimal_search: true,
literal_compression_disabled: matches!(
compression_level,
crate::encoding::CompressionLevel::Level(n) if n < 0
),
},
magicless: false,
content_checksum: false,
content_size_flag: true,
dict_id_flag: true,
target_block_size: None,
#[cfg(feature = "hash")]
hasher: XxHash64::with_seed(0),
#[cfg(feature = "lsm")]
frame_emit_info: None,
#[cfg(all(feature = "lsm", feature = "hash"))]
per_block_checksums_enabled: false,
#[cfg(all(feature = "lsm", feature = "hash"))]
block_checksums: None,
#[cfg(feature = "lsm")]
block_decompressed_sizes: alloc::vec::Vec::new(),
strategy_override: None,
}
}
pub fn set_parameters(&mut self, params: &crate::encoding::CompressionParameters) {
self.compression_level = params.level();
let overrides = params.overrides();
self.strategy_override = overrides.strategy.map(|s| s.tag());
self.state.strategy_tag = self.strategy_override.unwrap_or_else(|| {
crate::encoding::levels::config::resolve_level_params(
self.compression_level,
self.source_size_hint,
)
.strategy_tag
});
self.state.huf_optimal_search =
huf_search_enabled(self.state.strategy_tag, self.source_size_hint);
self.state.literal_compression_disabled = self.state.strategy_tag
== crate::encoding::strategy::StrategyTag::Fast
&& overrides.target_length.map_or_else(
|| {
matches!(
self.compression_level,
crate::encoding::CompressionLevel::Level(n) if n < 0
)
},
|tl| tl > 0,
);
self.state.matcher.set_param_overrides(Some(overrides));
}
fn borrowed_eligible(&self, input_len: usize, prep: &FramePrep) -> bool {
if matches!(self.compression_level, CompressionLevel::Uncompressed)
|| input_len > u32::MAX as usize
{
return false;
}
if prep.use_dictionary_state {
let fits_u32 = self
.dictionary
.as_ref()
.and_then(|dict| dict.inner.dict_content.len().checked_add(input_len))
.is_some_and(|virtual_len| virtual_len <= u32::MAX as usize);
if !fits_u32 {
return false;
}
return self.state.matcher.borrowed_dict_supported();
}
self.state.matcher.borrowed_supported()
}
fn run_one_frame(&mut self, input: &[u8], prep: &FramePrep, out: &mut Vec<u8>) -> u64 {
if self.borrowed_eligible(input.len(), prep) {
self.run_borrowed_block_loop(input, out)
} else {
let mut cursor: &[u8] = input;
self.run_owned_block_loop(&mut cursor, prep.initial_size_hint, true, out)
}
}
pub fn compress_independent_frame_into(&mut self, input: &[u8], out: &mut Vec<u8>) {
self.source_size_hint = Some(input.len() as u64);
let prep = self.prepare_frame();
let total_uncompressed = input.len() as u64;
let emit_checksum = cfg!(feature = "hash") && self.content_checksum;
let checksum_len = if emit_checksum { 4 } else { 0 };
out.clear();
let first_block_bound = input.len().min(self.block_capacity()) + 3;
out.reserve(18 + first_block_bound + checksum_len);
self.append_frame_header(total_uncompressed, &prep, out);
let header_len = out.len();
let _ = self.run_one_frame(input, &prep, out);
#[cfg(feature = "hash")]
if self.content_checksum {
out.extend_from_slice(&(self.hasher.finish() as u32).to_le_bytes());
}
#[cfg(feature = "lsm")]
{
let blocks_end = out.len() - checksum_len;
self.populate_frame_emit_info(header_len, &out[header_len..blocks_end], emit_checksum);
}
#[cfg(not(feature = "lsm"))]
let _ = header_len;
}
pub fn compress_independent_frame(&mut self, input: &[u8]) -> Vec<u8> {
let mut out = Vec::new();
self.compress_independent_frame_into(input, &mut out);
out
}
fn run_borrowed_block_loop(&mut self, input: &[u8], out: &mut Vec<u8>) -> u64 {
let blocks_start = out.len();
let total_uncompressed = input.len() as u64;
if input.is_empty() {
let header = BlockHeader {
last_block: true,
block_type: crate::blocks::block::BlockType::Raw,
block_size: 0,
};
header.serialize(out);
#[cfg(feature = "lsm")]
self.block_decompressed_sizes.push(0);
#[cfg(all(feature = "lsm", feature = "hash"))]
if let Some(checksums) = self.block_checksums.as_mut() {
checksums.push(xxh64_block_low32(&[]));
}
return total_uncompressed;
}
unsafe {
self.state.matcher.set_borrowed_window(input);
}
struct ClearBorrowedOnDrop(*mut MatchGeneratorDriver);
impl Drop for ClearBorrowedOnDrop {
fn drop(&mut self) {
unsafe { (*self.0).clear_borrowed_window() };
}
}
let _clear_guard = ClearBorrowedOnDrop(core::ptr::addr_of_mut!(self.state.matcher));
let block_capacity = self.block_capacity();
let mut start = 0usize;
while start < input.len() {
reserve_for_next_block(
out,
blocks_start,
start as u64,
input.len() - start,
block_capacity,
);
let savings = start as i64 - (out.len() - blocks_start) as i64;
if savings >= 3
&& input.len() - start >= MAX_BLOCK_SIZE as usize
&& block_capacity >= MAX_BLOCK_SIZE as usize
&& crate::encoding::levels::config::level_pre_split(self.compression_level)
.is_some()
{
warm_presplit_window(&input[start..start + MAX_BLOCK_SIZE as usize]);
}
let block_len = optimal_block_size(
self.compression_level,
&input[start..],
input.len() - start,
block_capacity,
savings,
);
let end = (start + block_len).min(input.len());
let block = &input[start..end];
let last_block = end == input.len();
#[cfg(feature = "hash")]
if self.content_checksum {
self.hasher.write(block);
}
let dict_active =
self.dictionary.is_some() && self.state.matcher.supports_dictionary_priming();
crate::encoding::levels::compress_block_encoded_borrowed(
&mut self.state,
self.compression_level,
last_block,
block,
start,
end,
dict_active,
out,
#[cfg(feature = "lsm")]
Some(&mut self.block_decompressed_sizes),
#[cfg(all(feature = "lsm", feature = "hash"))]
self.block_checksums.as_mut(),
);
start = end;
}
total_uncompressed
}
}
impl<R: Read, W: Write, M: Matcher> FrameCompressor<R, W, M> {
pub fn new_with_matcher(matcher: M, compression_level: CompressionLevel) -> Self {
Self {
uncompressed_data: None,
compressed_data: None,
dictionary: None,
dictionary_entropy_cache: None,
source_size_hint: None,
state: CompressState {
matcher,
last_huff_table: None,
huff_table_spare: None,
fse_tables: FseTables::new(),
block_scratch: crate::encoding::blocks::CompressedBlockScratch::new(),
offset_hist: [1, 4, 8],
strategy_tag: crate::encoding::strategy::StrategyTag::for_compression_level(
compression_level,
),
huf_optimal_search: true,
literal_compression_disabled: matches!(
compression_level,
crate::encoding::CompressionLevel::Level(n) if n < 0
),
},
compression_level,
magicless: false,
content_checksum: false,
content_size_flag: true,
dict_id_flag: true,
target_block_size: None,
#[cfg(feature = "hash")]
hasher: XxHash64::with_seed(0),
#[cfg(feature = "lsm")]
frame_emit_info: None,
#[cfg(all(feature = "lsm", feature = "hash"))]
per_block_checksums_enabled: false,
#[cfg(all(feature = "lsm", feature = "hash"))]
block_checksums: None,
#[cfg(feature = "lsm")]
block_decompressed_sizes: alloc::vec::Vec::new(),
strategy_override: None,
}
}
pub fn set_magicless(&mut self, magicless: bool) {
self.magicless = magicless;
}
pub fn set_content_checksum(&mut self, emit: bool) {
self.content_checksum = emit;
}
pub fn set_content_size_flag(&mut self, emit: bool) {
self.content_size_flag = emit;
}
pub fn set_dictionary_id_flag(&mut self, emit: bool) {
self.dict_id_flag = emit;
}
pub fn set_target_block_size(&mut self, target: Option<u32>) {
self.target_block_size = target.map(|t| {
t.clamp(
crate::common::MIN_TARGET_BLOCK_SIZE,
crate::common::MAX_BLOCK_SIZE,
)
});
}
fn block_capacity(&self) -> usize {
self.target_block_size
.map_or(crate::common::MAX_BLOCK_SIZE as usize, |t| t as usize)
}
pub fn set_source(&mut self, uncompressed_data: R) -> Option<R> {
self.uncompressed_data.replace(uncompressed_data)
}
pub fn set_drain(&mut self, compressed_data: W) -> Option<W> {
self.compressed_data.replace(compressed_data)
}
pub fn set_source_size_hint(&mut self, size: u64) {
self.source_size_hint = Some(size);
}
pub fn heap_size(&self) -> usize {
let mut total = self.state.matcher.heap_size();
total += self
.state
.last_huff_table
.as_ref()
.map_or(0, |table| table.heap_size());
total += self
.state
.huff_table_spare
.as_ref()
.map_or(0, |table| table.heap_size());
total += self
.dictionary
.as_ref()
.map_or(0, |d| d.inner.dict_content.capacity());
total += self
.dictionary_entropy_cache
.as_ref()
.map_or(0, CachedDictionaryEntropy::heap_size);
#[cfg(all(feature = "lsm", feature = "hash"))]
{
total += self
.block_checksums
.as_ref()
.map_or(0, |v| v.capacity() * core::mem::size_of::<u32>());
}
#[cfg(feature = "lsm")]
{
total += self.block_decompressed_sizes.capacity() * core::mem::size_of::<u32>();
}
total
}
pub fn compress(&mut self) {
let prep = self.prepare_frame();
let mut source = self
.uncompressed_data
.take()
.expect("source must be set via set_source before compress()");
let mut all_blocks: Vec<u8> = Vec::with_capacity(initial_all_blocks_cap(
prep.initial_size_hint,
self.block_capacity(),
));
let mut block_source = ReaderBlockSource::new(&mut source);
let total_uncompressed = self.run_owned_block_loop(
&mut block_source,
prep.initial_size_hint,
false,
&mut all_blocks,
);
self.uncompressed_data = Some(source);
self.finish_frame(all_blocks, total_uncompressed, &prep);
}
fn prepare_frame(&mut self) -> FramePrep {
#[cfg(feature = "lsm")]
{
self.frame_emit_info = None;
self.block_decompressed_sizes.clear();
}
#[cfg(all(feature = "lsm", feature = "hash"))]
{
if self.per_block_checksums_enabled {
self.block_checksums = Some(alloc::vec::Vec::new());
} else {
self.block_checksums = None;
}
}
let initial_size_hint = self.source_size_hint;
let source_size_hint_known = initial_size_hint.is_some();
let use_dictionary_state =
!matches!(self.compression_level, CompressionLevel::Uncompressed)
&& self.state.matcher.supports_dictionary_priming()
&& self.dictionary.is_some();
if let Some(size_hint) = self.source_size_hint.take() {
self.state.matcher.set_source_size_hint(size_hint);
}
if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
self.state
.matcher
.set_dictionary_size_hint(dict.inner.dict_content.len());
}
self.state.matcher.reset(self.compression_level);
self.state.offset_hist = [1, 4, 8];
self.state.strategy_tag = self.strategy_override.unwrap_or_else(|| {
crate::encoding::levels::config::resolve_level_params(
self.compression_level,
initial_size_hint,
)
.strategy_tag
});
self.state.huf_optimal_search =
huf_search_enabled(self.state.strategy_tag, initial_size_hint);
let cached_entropy = if use_dictionary_state {
self.dictionary_entropy_cache.as_ref()
} else {
None
};
if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() {
self.state.offset_hist = dict.inner.offset_hist;
let cutoff_log = match self.state.strategy_tag {
crate::encoding::strategy::StrategyTag::Fast => {
crate::encoding::levels::config::FAST_ATTACH_DICT_CUTOFF_LOG
}
crate::encoding::strategy::StrategyTag::BtUltra
| crate::encoding::strategy::StrategyTag::BtUltra2 => 13,
crate::encoding::strategy::StrategyTag::Dfast => 14,
crate::encoding::strategy::StrategyTag::Greedy
| crate::encoding::strategy::StrategyTag::Lazy
| crate::encoding::strategy::StrategyTag::Btlazy2
| crate::encoding::strategy::StrategyTag::BtOpt => 15,
};
if self.state.matcher.dictionary_is_resident() {
self.state
.matcher
.reapply_resident_dictionary(dict.inner.offset_hist);
} else {
let prefer_copy_snapshot = initial_size_hint.is_some_and(|s| {
crate::encoding::levels::config::source_size_ceil_log(s) > cutoff_log
});
let restored = prefer_copy_snapshot
&& self
.state
.matcher
.restore_primed_dictionary(self.compression_level);
if !restored {
self.state.matcher.prime_with_dictionary(
dict.inner.dict_content.as_slice(),
dict.inner.offset_hist,
);
if prefer_copy_snapshot {
self.state
.matcher
.capture_primed_dictionary(self.compression_level);
}
}
}
}
if let Some(cache) = cached_entropy {
match &cache.huff {
Some(src) => {
if self.state.last_huff_table.is_none() {
self.state.last_huff_table = self.state.huff_table_spare.take();
}
match &mut self.state.last_huff_table {
Some(dst) => dst.clone_from(src),
slot => *slot = Some(src.clone()),
}
}
None => self.state.clear_huff_table(),
}
} else {
self.state.clear_huff_table();
}
if let Some(cache) = cached_entropy {
self.state
.fse_tables
.ll_previous
.clone_from(&cache.ll_previous);
self.state
.fse_tables
.ml_previous
.clone_from(&cache.ml_previous);
self.state
.fse_tables
.of_previous
.clone_from(&cache.of_previous);
} else {
self.state.fse_tables.ll_previous = None;
self.state.fse_tables.ml_previous = None;
self.state.fse_tables.of_previous = None;
}
let ll_entropy = cached_entropy.and_then(|cache| match cache.ll_previous.as_ref() {
Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
_ => None,
});
let ml_entropy = cached_entropy.and_then(|cache| match cache.ml_previous.as_ref() {
Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
_ => None,
});
let of_entropy = cached_entropy.and_then(|cache| match cache.of_previous.as_ref() {
Some(PreviousFseTable::Custom(table)) => Some(table.as_ref()),
_ => None,
});
self.state.matcher.seed_dictionary_entropy(
self.state.last_huff_table.as_ref(),
ll_entropy,
ml_entropy,
of_entropy,
);
#[cfg(feature = "hash")]
{
self.hasher = XxHash64::with_seed(0);
}
let window_size = self.state.matcher.window_size();
assert!(
window_size != 0,
"matcher reported window_size == 0, which is invalid"
);
FramePrep {
window_size,
use_dictionary_state,
source_size_hint_known,
initial_size_hint,
}
}
fn run_owned_block_loop<S: OwnedBlockSource>(
&mut self,
source: &mut S,
initial_size_hint: Option<u64>,
hint_is_exact: bool,
out: &mut Vec<u8>,
) -> u64 {
let blocks_start = out.len();
let mut total_uncompressed: u64 = 0;
let mut pending_input: Vec<u8> = Vec::new();
let mut reached_eof = false;
let mut savings = 0i64;
loop {
let block_capacity = self.block_capacity();
let mut uncompressed_data = self.state.matcher.get_next_space();
uncompressed_data.clear();
uncompressed_data.extend_from_slice(&pending_input);
pending_input.clear();
if !reached_eof {
let size_hint_remaining = match initial_size_hint {
Some(hint) if hint > total_uncompressed => Some(hint - total_uncompressed),
_ => None,
};
let (appended, eof) =
source.fill_block(&mut uncompressed_data, block_capacity, size_hint_remaining);
total_uncompressed += appended as u64;
reached_eof = eof;
}
let mut last_block = reached_eof;
let remaining_for_split = if reached_eof {
uncompressed_data.len()
} else {
block_capacity
};
if !matches!(self.compression_level, CompressionLevel::Uncompressed)
&& uncompressed_data.len() == block_capacity
{
let block_len = optimal_block_size(
self.compression_level,
&uncompressed_data,
remaining_for_split,
block_capacity,
savings,
);
if block_len < uncompressed_data.len() {
pending_input.clear();
pending_input.extend_from_slice(&uncompressed_data[block_len..]);
uncompressed_data.truncate(block_len);
last_block = false;
}
}
#[cfg(feature = "hash")]
if self.content_checksum {
self.hasher.write(&uncompressed_data);
}
let emitted =
total_uncompressed - uncompressed_data.len() as u64 - pending_input.len() as u64;
match initial_size_hint {
Some(hint) if hint >= total_uncompressed => {
let hint_remaining = hint - emitted;
let remaining = if hint_is_exact {
hint_remaining
} else {
let buffered = total_uncompressed - emitted;
const HINT_LOOKAHEAD: u64 = 64 * 1024;
hint_remaining.min(buffered + HINT_LOOKAHEAD)
};
reserve_for_next_block(
out,
blocks_start,
emitted,
remaining as usize,
self.block_capacity(),
);
}
_ => {
out.reserve(uncompressed_data.len() + 3 + 16);
}
}
if uncompressed_data.is_empty() {
let header = BlockHeader {
last_block: true,
block_type: crate::blocks::block::BlockType::Raw,
block_size: 0,
};
header.serialize(out);
#[cfg(feature = "lsm")]
self.block_decompressed_sizes.push(0);
#[cfg(all(feature = "lsm", feature = "hash"))]
if let Some(checksums) = self.block_checksums.as_mut() {
checksums.push(xxh64_block_low32(&[]));
}
break;
}
match self.compression_level {
CompressionLevel::Uncompressed => {
let header = BlockHeader {
last_block,
block_type: crate::blocks::block::BlockType::Raw,
block_size: uncompressed_data.len().try_into().unwrap(),
};
header.serialize(out);
#[cfg(feature = "lsm")]
self.block_decompressed_sizes
.push(uncompressed_data.len() as u32);
#[cfg(all(feature = "lsm", feature = "hash"))]
if let Some(checksums) = self.block_checksums.as_mut() {
checksums.push(xxh64_block_low32(&uncompressed_data));
}
out.extend_from_slice(&uncompressed_data);
savings +=
uncompressed_data.len() as i64 - (3 + uncompressed_data.len()) as i64;
}
CompressionLevel::Fastest
| CompressionLevel::Default
| CompressionLevel::Better
| CompressionLevel::Best
| CompressionLevel::Level(_) => {
let before_len = out.len();
let block_len = uncompressed_data.len();
let dict_active = self.dictionary.is_some()
&& self.state.matcher.supports_dictionary_priming();
compress_block_encoded(
&mut self.state,
self.compression_level,
last_block,
uncompressed_data,
out,
dict_active,
#[cfg(feature = "lsm")]
Some(&mut self.block_decompressed_sizes),
#[cfg(all(feature = "lsm", feature = "hash"))]
self.block_checksums.as_mut(),
);
savings += block_len as i64 - (out.len() - before_len) as i64;
}
}
if last_block && pending_input.is_empty() {
break;
}
}
total_uncompressed
}
fn append_frame_header(&self, total_uncompressed: u64, prep: &FramePrep, out: &mut Vec<u8>) {
let single_segment = self.content_size_flag
&& prep.source_size_hint_known
&& total_uncompressed <= prep.window_size;
let header = FrameHeader {
frame_content_size: self.content_size_flag.then_some(total_uncompressed),
single_segment,
content_checksum: cfg!(feature = "hash") && self.content_checksum,
dictionary_id: if prep.use_dictionary_state && self.dict_id_flag {
self.dictionary.as_ref().map(|dict| dict.inner.id as u64)
} else {
None
},
window_size: if single_segment {
None
} else {
Some(prep.window_size)
},
magicless: self.magicless,
};
header.serialize(out);
}
fn finish_frame(&mut self, all_blocks: Vec<u8>, total_uncompressed: u64, prep: &FramePrep) {
let mut header_buf: Vec<u8> = Vec::with_capacity(18);
self.append_frame_header(total_uncompressed, prep, &mut header_buf);
#[cfg(feature = "hash")]
let checksum_bytes = self
.content_checksum
.then(|| (self.hasher.finish() as u32).to_le_bytes());
let drain = self.compressed_data.as_mut().unwrap();
drain.write_all(&header_buf).unwrap();
drain.write_all(&all_blocks).unwrap();
#[cfg(feature = "hash")]
if let Some(checksum_bytes) = checksum_bytes {
drain.write_all(&checksum_bytes).unwrap();
}
#[cfg(feature = "lsm")]
{
let emit_checksum = cfg!(feature = "hash") && self.content_checksum;
self.populate_frame_emit_info(header_buf.len(), &all_blocks, emit_checksum);
}
}
#[cfg(feature = "lsm")]
fn populate_frame_emit_info(
&mut self,
header_len: usize,
all_blocks: &[u8],
emit_checksum: bool,
) {
use crate::blocks::block::BlockType as BT;
use crate::encoding::frame_emit_info::{FrameBlock, FrameEmitInfo};
let frame_header_len: u32 = match u32::try_from(header_len) {
Ok(v) => v,
Err(_) => return,
};
let all_blocks_len_u32: u32 = match u32::try_from(all_blocks.len()) {
Ok(v) => v,
Err(_) => return,
};
let mut blocks: Vec<FrameBlock> = Vec::new();
let mut cursor: usize = 0;
while cursor + 3 <= all_blocks.len() {
let mut header_u32 = [0u8; 4];
header_u32[..3].copy_from_slice(&all_blocks[cursor..cursor + 3]);
let raw = u32::from_le_bytes(header_u32);
let last_block = (raw & 1) != 0;
let block_type = match (raw >> 1) & 0b11 {
0 => BT::Raw,
1 => BT::RLE,
2 => BT::Compressed,
_ => BT::Reserved,
};
let block_size_field = raw >> 3;
let physical_body: u32 = match block_type {
BT::RLE => 1,
_ => block_size_field,
};
let cursor_u32: u32 = match u32::try_from(cursor) {
Ok(v) => v,
Err(_) => return,
};
let offset_in_frame = match frame_header_len.checked_add(cursor_u32) {
Some(v) => v,
None => return,
};
let decompressed_size = match self.block_decompressed_sizes.get(blocks.len()).copied() {
Some(size) => size,
None if matches!(block_type, BT::Raw | BT::RLE) => block_size_field,
None => {
debug_assert!(
false,
"missing decompressed-size sidecar entry for compressed block {}",
blocks.len()
);
return;
}
};
blocks.push(FrameBlock {
offset_in_frame,
header_size: 3,
body_size: physical_body,
block_size_field,
block_type,
last_block,
decompressed_size,
});
cursor += 3 + physical_body as usize;
if last_block {
break;
}
}
if cursor != all_blocks.len() || !blocks.last().is_some_and(|b| b.last_block) {
debug_assert!(
false,
"incomplete block scan in populate_frame_emit_info: cursor={} len={} last_block={:?}",
cursor,
all_blocks.len(),
blocks.last().map(|b| b.last_block)
);
return;
}
let checksum_range = if emit_checksum {
let cs_start = match frame_header_len.checked_add(all_blocks_len_u32) {
Some(v) => v,
None => return,
};
let cs_end = match cs_start.checked_add(4) {
Some(v) => v,
None => return,
};
Some(cs_start..cs_end)
} else {
None
};
let body_total = match frame_header_len.checked_add(all_blocks_len_u32) {
Some(v) => v,
None => return,
};
let total_size = if checksum_range.is_some() {
match body_total.checked_add(4) {
Some(v) => v,
None => return,
}
} else {
body_total
};
self.frame_emit_info = Some(FrameEmitInfo {
frame_header_range: 0..frame_header_len,
blocks,
checksum_range,
total_size,
});
}
#[cfg(feature = "lsm")]
pub fn last_frame_emit_info(&self) -> Option<&crate::encoding::frame_emit_info::FrameEmitInfo> {
self.frame_emit_info.as_ref()
}
#[cfg(all(feature = "lsm", feature = "hash"))]
pub fn enable_per_block_checksums(&mut self) {
self.per_block_checksums_enabled = true;
}
#[cfg(all(feature = "lsm", feature = "hash"))]
pub fn last_frame_block_checksums(&self) -> Option<&[u32]> {
self.block_checksums.as_deref()
}
pub fn source_mut(&mut self) -> Option<&mut R> {
self.uncompressed_data.as_mut()
}
pub fn drain_mut(&mut self) -> Option<&mut W> {
self.compressed_data.as_mut()
}
pub fn source(&self) -> Option<&R> {
self.uncompressed_data.as_ref()
}
pub fn drain(&self) -> Option<&W> {
self.compressed_data.as_ref()
}
pub fn take_source(&mut self) -> Option<R> {
self.uncompressed_data.take()
}
pub fn take_drain(&mut self) -> Option<W> {
self.compressed_data.take()
}
pub fn replace_matcher(&mut self, mut match_generator: M) -> M {
core::mem::swap(&mut match_generator, &mut self.state.matcher);
match_generator
}
pub fn set_compression_level(
&mut self,
compression_level: CompressionLevel,
) -> CompressionLevel {
let old = self.compression_level;
self.compression_level = compression_level;
self.state.literal_compression_disabled = matches!(
compression_level,
CompressionLevel::Level(n) if n < 0
);
self.strategy_override = None;
self.state.matcher.clear_param_overrides();
old
}
pub fn compression_level(&self) -> CompressionLevel {
self.compression_level
}
pub fn set_dictionary(
&mut self,
dictionary: crate::decoding::Dictionary,
) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
self.attach_dictionary(EncoderDictionary::from_dictionary(dictionary))
}
pub fn set_dictionary_from_bytes(
&mut self,
raw_dictionary: &[u8],
) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
self.attach_dictionary(EncoderDictionary::from_bytes(raw_dictionary)?)
}
pub fn set_encoder_dictionary(
&mut self,
dictionary: EncoderDictionary,
) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
self.attach_dictionary(dictionary)
}
pub fn clear_dictionary(&mut self) -> Option<EncoderDictionary> {
self.dictionary_entropy_cache = None;
self.state.matcher.invalidate_primed_dictionary();
self.dictionary.take()
}
fn attach_dictionary(
&mut self,
enc: EncoderDictionary,
) -> Result<Option<EncoderDictionary>, crate::decoding::errors::DictionaryDecodeError> {
let dictionary = &enc.inner;
if dictionary.id == 0 {
return Err(crate::decoding::errors::DictionaryDecodeError::ZeroDictionaryId);
}
if let Some(index) = dictionary.offset_hist.iter().position(|&rep| rep == 0) {
return Err(
crate::decoding::errors::DictionaryDecodeError::ZeroRepeatOffsetInDictionary {
index: index as u8,
},
);
}
self.dictionary_entropy_cache = Some(CachedDictionaryEntropy::from_dictionary(dictionary));
self.state.matcher.invalidate_primed_dictionary();
Ok(self.dictionary.replace(enc))
}
}
#[cfg(test)]
mod tests;