use alloc::boxed::Box;
use alloc::vec::Vec;
use fearless_simd::Level;
use super::fast::FastEncoder;
use super::greedy::encoder::GreedyEncoder;
use super::greedy::params::GreedyParams;
use super::hq::encoder::HqEncoder;
use super::hq::params::HqParams;
use super::rfc9841::context::SharedContextInner;
use super::stream::{Destination, finish};
use crate::compressor::shared::SharedBrotliError;
use crate::compressor::{BrotliCompressError, BrotliResult, CompressParams, QualityLevel};
pub(crate) enum Encoder {
Fast(FastEncoder),
Greedy(Box<GreedyEncoder>),
Hq(Box<HqEncoder>),
}
impl Encoder {
pub(crate) fn new(
level: Level,
params: &CompressParams,
size_hint: usize,
) -> BrotliResult<Self> {
match FastEncoder::new(level, params) {
Ok(encoder) => return Ok(Self::Fast(encoder)),
Err(BrotliCompressError::UnsupportedQuality(_)) => {}
Err(error) => return Err(error),
}
let size_hint = params.size_hint().unwrap_or(size_hint);
match GreedyEncoder::new(level, params, size_hint) {
Ok(encoder) => return Ok(Self::Greedy(Box::new(encoder))),
Err(BrotliCompressError::UnsupportedQuality(_)) => {}
Err(error) => return Err(error),
}
Ok(Self::Hq(Box::new(HqEncoder::new(level, params)?)))
}
pub(crate) const fn block_size_limit(&self) -> usize {
match self {
Self::Fast(encoder) => encoder.block_size_limit(),
Self::Greedy(encoder) => encoder.block_size_limit(),
Self::Hq(encoder) => encoder.block_size_limit(),
}
}
pub(crate) fn reset_for(&mut self, params: &CompressParams, size_hint: usize) -> bool {
let size_hint = params.size_hint().unwrap_or(size_hint);
match self {
Self::Fast(encoder) => {
if !encoder.matches(params) {
return false;
}
encoder.reset();
true
}
Self::Greedy(encoder) => {
let Ok(fresh) = GreedyParams::new(params, size_hint) else {
return false;
};
if !encoder.retarget(fresh) {
return false;
}
encoder.reset();
true
}
Self::Hq(encoder) => {
let Ok(fresh) = HqParams::new(params) else {
return false;
};
if fresh != *encoder.params() {
return false;
}
encoder.reset();
true
}
}
}
pub(crate) fn encode_block_with(
&mut self,
input: &[u8],
is_last: bool,
attached: Option<&SharedContextInner>,
) -> BrotliResult<&[u8]> {
match self {
Self::Fast(encoder) => encoder.encode_block(input, is_last),
Self::Greedy(encoder) => encoder.encode_block_with(input, is_last, attached),
Self::Hq(encoder) => encoder.encode_block_with(input, is_last, attached),
}
}
pub(crate) fn flush_block(
&mut self,
input: &[u8],
attached: Option<&SharedContextInner>,
) -> BrotliResult<&[u8]> {
match self {
Self::Fast(encoder) => encoder.flush_block(input),
Self::Greedy(encoder) => encoder.flush_block(input, attached),
Self::Hq(encoder) => encoder.flush_block(input, attached),
}
}
pub(crate) fn retained_bytes(&self) -> usize {
match self {
Self::Fast(encoder) => encoder.retained_bytes(),
Self::Greedy(encoder) => size_of::<GreedyEncoder>() + encoder.retained_bytes(),
Self::Hq(encoder) => size_of::<HqEncoder>() + encoder.retained_bytes(),
}
}
}
#[derive(Default)]
pub(crate) struct EncoderCache {
encoder: Option<Encoder>,
level: Option<core::mem::Discriminant<Level>>,
}
impl core::fmt::Debug for EncoderCache {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("EncoderCache")
.field("retained", &self.encoder.is_some())
.finish_non_exhaustive()
}
}
impl EncoderCache {
#[cfg_attr(all(feature = "hotpath", not(feature = "no_std")), hotpath::measure)]
pub(crate) fn acquire(
&mut self,
level: Level,
params: &CompressParams,
size_hint: usize,
) -> BrotliResult<&mut Encoder> {
let level_key = core::mem::discriminant(&level);
let reusable = self.level == Some(level_key)
&& self
.encoder
.as_mut()
.is_some_and(|encoder| encoder.reset_for(params, size_hint));
if !reusable {
self.encoder = Some(Encoder::new(level, params, size_hint)?);
self.level = Some(level_key);
}
match self.encoder.as_mut() {
Some(encoder) => Ok(encoder),
None => Err(BrotliCompressError::BufferOverflow),
}
}
pub(crate) fn retained_bytes(&self) -> usize {
self.encoder.as_ref().map_or(0, Encoder::retained_bytes)
}
pub(crate) const fn encoder(&mut self) -> Option<&mut Encoder> {
self.encoder.as_mut()
}
pub(crate) fn invalidate(&mut self) {
self.encoder = None;
self.level = None;
}
}
const fn check_large_window(params: &CompressParams) -> BrotliResult<()> {
if !params.lgwin().is_large() {
return Ok(());
}
match params.quality() {
QualityLevel::Q0 => Err(BrotliCompressError::Shared(
SharedBrotliError::UnsupportedLargeWindow { quality: 0 },
)),
QualityLevel::Q1 => Err(BrotliCompressError::Shared(
SharedBrotliError::UnsupportedLargeWindow { quality: 1 },
)),
QualityLevel::Q2 => Err(BrotliCompressError::Shared(
SharedBrotliError::UnsupportedLargeWindow { quality: 2 },
)),
_ => Ok(()),
}
}
pub(crate) const fn quality_reads_a_prefix(quality: QualityLevel) -> bool {
matches!(
quality,
QualityLevel::Q5
| QualityLevel::Q6
| QualityLevel::Q7
| QualityLevel::Q8
| QualityLevel::Q9
| QualityLevel::Q10
| QualityLevel::Q11
)
}
const fn empty_stream(params: &CompressParams) -> ([u8; 2], usize) {
let window = super::rfc9841::window::ResolvedWindow::new(params);
let fast = matches!(params.quality(), QualityLevel::Q0 | QualityLevel::Q1);
#[cfg(feature = "experimental")]
if !fast && params.stream_offset != 0 {
return ([3, 0], 1);
}
let window = if fast {
window.at_least(super::fast::constants::WINDOW_BITS_FAST)
} else {
window
};
let (header, bits) = window.header();
let complete = header | (3 << bits);
(complete.to_le_bytes(), ((bits + 9) >> 3) as usize)
}
pub(crate) fn compress_to_vec_attached(
cache: &mut EncoderCache,
level: Level,
params: &CompressParams,
attached: Option<&SharedContextInner>,
src: &[u8],
out: &mut Vec<u8>,
) -> BrotliResult<()> {
check_large_window(params)?;
if src.is_empty() {
let (bytes, length) = empty_stream(params);
out.extend_from_slice(&bytes[..length]);
return Ok(());
}
let start = out.len();
let encoder = match cache.acquire(level, params, src.len()) {
Ok(encoder) => encoder,
Err(error) => {
cache.invalidate();
return Err(error);
}
};
match finish(encoder, attached, src, Destination::Append(out)) {
Ok(_) => Ok(()),
Err(error) => {
out.truncate(start);
cache.invalidate();
Err(error)
}
}
}
pub(crate) fn compress_to_slice_attached(
cache: &mut EncoderCache,
level: Level,
params: &CompressParams,
attached: Option<&SharedContextInner>,
src: &[u8],
dst: &mut [u8],
) -> BrotliResult<usize> {
check_large_window(params)?;
if src.is_empty() {
let (bytes, length) = empty_stream(params);
let Some(destination) = dst.get_mut(..length) else {
cache.invalidate();
return Err(BrotliCompressError::OutputTooSmall);
};
destination.copy_from_slice(&bytes[..length]);
return Ok(length);
}
let encoder = match cache.acquire(level, params, src.len()) {
Ok(encoder) => encoder,
Err(error) => {
cache.invalidate();
return Err(error);
}
};
match finish(encoder, attached, src, Destination::Slice(dst)) {
Ok(written) => Ok(written),
Err(error) => {
cache.invalidate();
Err(error)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::compressor::{QualityLevel, WindowBits};
#[test]
#[cfg(feature = "experimental")]
fn empty_continuation_headers_match_scheduler_finalization() {
for quality in [QualityLevel::Q0, QualityLevel::Q2, QualityLevel::Q11] {
let mut params = CompressParams::new(quality, WindowBits::DEFAULT);
params.stream_offset = 123;
let mut encoder = Encoder::new(Level::fallback(), ¶ms, 0).expect("encoder");
let mut expected = Vec::new();
finish(&mut encoder, None, &[], Destination::Append(&mut expected)).expect("finish");
let (bytes, length) = empty_stream(¶ms);
assert_eq!(&bytes[..length], expected);
}
}
#[test]
fn empty_one_shot_headers_match_the_scheduler_without_allocating_an_encoder() {
for quality in [
QualityLevel::Q0,
QualityLevel::Q1,
QualityLevel::Q2,
QualityLevel::Q3,
QualityLevel::Q4,
QualityLevel::Q5,
QualityLevel::Q6,
QualityLevel::Q7,
QualityLevel::Q8,
QualityLevel::Q9,
QualityLevel::Q10,
QualityLevel::Q11,
] {
for window in [10, 16, 17, 22] {
let params =
CompressParams::new(quality, WindowBits::standard(window).expect("window"));
let mut encoder = Encoder::new(Level::fallback(), ¶ms, 0).expect("encoder");
let mut expected = Vec::new();
finish(&mut encoder, None, &[], Destination::Append(&mut expected))
.expect("finish");
let mut cache = EncoderCache::default();
let mut actual = b"prefix".to_vec();
compress_to_vec_attached(
&mut cache,
Level::fallback(),
¶ms,
None,
&[],
&mut actual,
)
.expect("append");
assert_eq!(&actual[..6], b"prefix");
assert_eq!(&actual[6..], expected);
assert_eq!(cache.retained_bytes(), 0);
let mut destination = [0; 2];
let written = compress_to_slice_attached(
&mut cache,
Level::fallback(),
¶ms,
None,
&[],
&mut destination,
)
.expect("slice");
assert_eq!(&destination[..written], expected);
assert!(
compress_to_slice_attached(
&mut cache,
Level::fallback(),
¶ms,
None,
&[],
&mut destination[..written - 1]
)
.is_err()
);
}
}
}
fn compress_to_vec(
level: Level,
params: &CompressParams,
src: &[u8],
out: &mut Vec<u8>,
) -> BrotliResult<()> {
compress_to_vec_attached(&mut EncoderCache::default(), level, params, None, src, out)
}
fn compress_to_slice(
level: Level,
params: &CompressParams,
src: &[u8],
dst: &mut [u8],
) -> BrotliResult<usize> {
compress_to_slice_attached(&mut EncoderCache::default(), level, params, None, src, dst)
}
fn params(quality: QualityLevel, lgwin: u8) -> CompressParams {
let lgwin = WindowBits::standard(lgwin).unwrap_or(WindowBits::DEFAULT);
CompressParams::new(quality, lgwin)
}
const IMPLEMENTED: [QualityLevel; 12] = [
QualityLevel::Q0,
QualityLevel::Q1,
QualityLevel::Q2,
QualityLevel::Q3,
QualityLevel::Q4,
QualityLevel::Q5,
QualityLevel::Q6,
QualityLevel::Q7,
QualityLevel::Q8,
QualityLevel::Q9,
QualityLevel::Q10,
QualityLevel::Q11,
];
#[test]
fn quality_routing_reaches_both_encoders() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
for quality in IMPLEMENTED {
let encoder = Encoder::new(level, ¶ms(quality, 22), 0).expect("routed");
match (quality, encoder) {
(QualityLevel::Q0 | QualityLevel::Q1, Encoder::Fast(_)) => {}
(
QualityLevel::Q2
| QualityLevel::Q3
| QualityLevel::Q4
| QualityLevel::Q5
| QualityLevel::Q6
| QualityLevel::Q7
| QualityLevel::Q8
| QualityLevel::Q9,
Encoder::Greedy(_),
) => {}
(QualityLevel::Q10 | QualityLevel::Q11, Encoder::Hq(_)) => {}
(quality, _) => panic!("quality {quality:?} routed to the wrong encoder"),
}
}
}
#[test]
fn every_quality_the_format_defines_now_compresses() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
for quality in IMPLEMENTED {
let mut out = Vec::new();
let outcome = compress_to_vec(level, ¶ms(quality, 22), b"data data data", &mut out);
assert!(outcome.is_ok(), "quality {quality:?} failed: {outcome:?}");
assert!(!out.is_empty(), "quality {quality:?} produced nothing");
}
}
#[test]
fn slice_output_reports_a_too_small_buffer() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
let mut dst = [0u8; 1];
for quality in IMPLEMENTED {
assert!(matches!(
compress_to_slice(
level,
¶ms(quality, 22),
b"hello world hello world",
&mut dst
),
Err(BrotliCompressError::OutputTooSmall)
));
}
}
#[test]
fn vector_and_slice_outputs_agree() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
for quality in IMPLEMENTED {
let params = params(quality, 22);
let input: Vec<u8> = (0..100_000u32).map(|i| (i % 251) as u8).collect();
let mut expected = Vec::new();
compress_to_vec(level, ¶ms, &input, &mut expected).expect("vector output");
let mut actual = vec![0u8; expected.len()];
let written =
compress_to_slice(level, ¶ms, &input, &mut actual).expect("slice output");
assert_eq!(written, expected.len(), "quality {quality:?}");
assert_eq!(actual, expected, "quality {quality:?}");
}
}
#[test]
fn an_empty_input_preserves_the_configured_stream_header() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
for quality in IMPLEMENTED {
let mut out = Vec::new();
compress_to_vec(level, ¶ms(quality, 22), b"", &mut out).expect("empty input");
assert_eq!(out, vec![59]);
let mut dst = [0u8; 4];
assert_eq!(
compress_to_slice(level, ¶ms(quality, 22), b"", &mut dst).ok(),
Some(1)
);
assert_eq!(dst[0], 59);
}
}
fn retained(cache: &mut EncoderCache) -> usize {
match cache.encoder.as_mut() {
Some(encoder) => core::ptr::from_mut(encoder) as usize,
None => 0,
}
}
#[test]
fn the_cache_reuses_an_encoder_of_the_same_shape() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
for quality in IMPLEMENTED {
let params = params(quality, 22).with_size_hint(Some(1 << 20));
let mut cache = EncoderCache::default();
cache.acquire(level, ¶ms, 100).expect("first acquire");
let first = retained(&mut cache);
assert_ne!(first, 0, "quality {quality:?}: nothing was retained");
cache.acquire(level, ¶ms, 5000).expect("second acquire");
assert_eq!(
retained(&mut cache),
first,
"quality {quality:?}: an identically shaped call rebuilt the encoder"
);
}
}
#[test]
fn the_cache_rebuilds_when_the_shape_changes() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
let mut cache = EncoderCache::default();
cache
.acquire(level, ¶ms(QualityLevel::Q0, 22), 1000)
.expect("first acquire");
assert!(matches!(cache.encoder, Some(Encoder::Fast(_))));
cache
.acquire(level, ¶ms(QualityLevel::Q11, 22), 1000)
.expect("second acquire");
assert!(
matches!(cache.encoder, Some(Encoder::Hq(_))),
"the cache reused a fast encoder for quality 11"
);
}
#[test]
fn the_cache_rebuilds_when_the_size_hint_moves_the_matcher() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
let small = params(QualityLevel::Q5, 22).with_size_hint(Some(1024));
let large = params(QualityLevel::Q5, 22).with_size_hint(Some(8 << 20));
let mut cache = EncoderCache::default();
cache.acquire(level, &small, 1024).expect("first acquire");
assert!(cache.acquire(level, &large, 1024).is_ok());
assert!(
!cache
.encoder
.as_mut()
.expect("retained")
.reset_for(&small, 1024),
"the encoder built for a large hint accepted a small one"
);
}
#[test]
fn invalidating_drops_the_retained_encoder() {
let level = Level::try_detect().unwrap_or_else(Level::baseline);
let mut cache = EncoderCache::default();
cache
.acquire(level, ¶ms(QualityLevel::Q5, 22), 1000)
.expect("acquire");
assert_ne!(retained(&mut cache), 0);
cache.invalidate();
assert_eq!(retained(&mut cache), 0, "invalidate kept the encoder");
assert!(cache.level.is_none(), "invalidate kept the level");
}
}