use constriction::stream::{
model::ContiguousCategoricalEntropyModel, queue::DefaultRangeDecoder,
queue::DefaultRangeEncoder, Decode, Encode,
};
const SCALE: u64 = 1 << 24;
pub const MIX_FAST_CAP: u64 = 4096;
pub const MIX_SLOW_CAP: u64 = 65536;
const EWA_LR: f64 = 0.15;
const COUNT_INC: u64 = 32;
pub const MIX5_FAST_CAP: u64 = 32768;
pub const MIX5_SLOW_CAP: u64 = 1048576;
pub const MIX5_NCTX: usize = 16;
#[derive(Clone, Copy)]
pub struct MixParams {
pub nctx: usize,
pub cap_fast: u64,
pub cap_slow: u64,
}
impl MixParams {
pub const BACKEND4: Self = Self {
nctx: 8,
cap_fast: MIX_FAST_CAP,
cap_slow: MIX_SLOW_CAP,
};
pub const BACKEND5: Self = Self {
nctx: MIX5_NCTX,
cap_fast: MIX5_FAST_CAP,
cap_slow: MIX5_SLOW_CAP,
};
}
#[inline]
pub fn mix_ctx_bucket(prev: u8) -> u8 {
match prev {
0 => 0,
1..=2 => 1,
3..=5 => 2,
6..=11 => 3,
12..=24 => 4,
25..=48 => 5,
49..=96 => 6,
_ => 7,
}
}
#[derive(Clone)]
struct MixCtx {
fast: Vec<u64>,
slow: Vec<u64>,
log_w: f64,
cap_fast: u64,
cap_slow: u64,
}
impl MixCtx {
fn new(cap_fast: u64, cap_slow: u64) -> Self {
Self {
fast: vec![1; 256],
slow: vec![1; 256],
log_w: 0.0,
cap_fast,
cap_slow,
}
}
fn probs(&self) -> Vec<u32> {
let tf = self.fast.iter().sum::<u64>() as f64;
let ts = self.slow.iter().sum::<u64>() as f64;
let wf = 1.0 / (1.0 + (-self.log_w).exp());
let mut p: Vec<u64> = self
.fast
.iter()
.zip(&self.slow)
.map(|(&cf, &cs)| {
((wf * (cf as f64 / tf) + (1.0 - wf) * (cs as f64 / ts)) * SCALE as f64) as u64
})
.collect();
let sum: u64 = p.iter().sum();
p[0] += SCALE - sum;
let mut dense = 0usize;
for (i, &v) in p.iter().enumerate() {
if v > p[dense] {
dense = i;
}
}
let zeros = p.iter().filter(|&&v| v == 0).count();
p[dense] = p[dense].max(1 + zeros as u64);
for v in p.iter_mut() {
if *v == 0 {
*v = 1;
}
}
p.into_iter().map(|v| v as u32).collect()
}
fn model(&self) -> Option<ContiguousCategoricalEntropyModel<u32, Vec<u32>, 24>> {
ContiguousCategoricalEntropyModel::<u32, Vec<u32>, 24>::from_nonzero_fixed_point_probabilities(
self.probs(),
false,
)
.ok()
}
fn step(&mut self, s: u8) {
let si = usize::from(s);
let tf = self.fast.iter().sum::<u64>() as f64;
let ts = self.slow.iter().sum::<u64>() as f64;
let pf = (self.fast[si] as f64 / tf).max(1e-9);
let ps = (self.slow[si] as f64 / ts).max(1e-9);
self.log_w = (self.log_w + EWA_LR * (pf.ln() - ps.ln())).clamp(-20.0, 20.0);
self.fast[si] += COUNT_INC;
self.slow[si] += COUNT_INC;
if self.fast.iter().sum::<u64>() > self.cap_fast {
for c in &mut self.fast {
*c = (*c + 1) >> 1;
}
}
if self.slow.iter().sum::<u64>() > self.cap_slow {
for c in &mut self.slow {
*c = (*c + 1) >> 1;
}
}
}
}
#[inline]
fn mix_ctx_key(params: &MixParams, prev: u8, prev2: u8) -> usize {
let pb = usize::from(mix_ctx_bucket(prev));
if params.nctx == 8 {
pb
} else {
let band = usize::from(mix_ctx_bucket(prev2) >= 3);
pb | (band << 3)
}
}
pub fn mix_encode_with(params: &MixParams, data: &[u8]) -> Vec<u8> {
let mut ctxs: Vec<MixCtx> = (0..params.nctx)
.map(|_| MixCtx::new(params.cap_fast, params.cap_slow))
.collect();
let mut enc = DefaultRangeEncoder::new();
let mut prev = 0u8;
let mut prev2 = 0u8;
for &b in data {
let c = mix_ctx_key(params, prev, prev2);
let m = ctxs[c]
.model()
.expect("mix model with all-nonzero probabilities is always buildable");
enc.encode_symbols(std::iter::once((usize::from(b), &m)))
.expect("mix encode symbol");
ctxs[c].step(b);
prev2 = prev;
prev = b;
}
let compressed = enc.into_compressed().expect("mix seal");
let mut out = Vec::with_capacity(compressed.len() * 4 + 4);
out.extend_from_slice(
&u32::try_from(data.len())
.expect("stream len fits u32")
.to_le_bytes(),
);
for w in compressed {
out.extend_from_slice(&w.to_le_bytes());
}
out
}
pub fn mix_encode(data: &[u8]) -> Vec<u8> {
mix_encode_with(&MixParams::BACKEND4, data)
}
pub fn mix_decode_with(
params: &MixParams,
data: &[u8],
limit: usize,
) -> Result<Vec<u8>, crate::error::Error> {
if data.len() < 4 || !(data.len() - 4).is_multiple_of(4) {
return Err(crate::error::Error::fmt_msg("mix payload not word-aligned"));
}
let stream_len = usize::try_from(u32::from_le_bytes(data[..4].try_into().expect("4 bytes")))
.map_err(|_| crate::error::Error::fmt_msg("mix stream length"))?;
if stream_len > limit {
return Err(crate::error::Error::fmt_msg("bomb"));
}
let mut words = Vec::with_capacity((data.len() - 4) / 4);
for chunk in data[4..].chunks_exact(4) {
words.push(u32::from_le_bytes(chunk.try_into().expect("4 bytes")));
}
let mut dec = DefaultRangeDecoder::from_compressed(words)
.map_err(|e| crate::error::Error::fmt_msg(format!("mix decoder open: {e:?}")))?;
let mut ctxs: Vec<MixCtx> = (0..params.nctx)
.map(|_| MixCtx::new(params.cap_fast, params.cap_slow))
.collect();
let mut out = Vec::with_capacity(stream_len);
let mut prev = 0u8;
let mut prev2 = 0u8;
for _i in 0..stream_len {
let c = mix_ctx_key(params, prev, prev2);
let m = ctxs[c]
.model()
.ok_or_else(|| crate::error::Error::fmt_msg("mix model build failed"))?;
let s = dec
.decode_symbols(std::iter::once(&m))
.next()
.ok_or_else(|| crate::error::Error::fmt_msg("mix stream exhausted"))?
.map_err(|e| crate::error::Error::fmt_msg(format!("mix decode symbol: {e:?}")))?;
let s = u8::try_from(s)
.map_err(|_| crate::error::Error::fmt_msg("mix symbol out of byte range"))?;
ctxs[c].step(s);
out.push(s);
prev2 = prev;
prev = s;
}
Ok(out)
}
pub fn mix_decode(data: &[u8], limit: usize) -> Result<Vec<u8>, crate::error::Error> {
mix_decode_with(&MixParams::BACKEND4, data, limit)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn mix_backend4_roundtrip_runs() {
let data: Vec<u8> = std::iter::repeat_n(0u8, 4096)
.chain(std::iter::repeat_n(7u8, 512))
.collect();
let wire = mix_encode(&data);
assert_eq!(mix_decode(&wire, data.len()).unwrap(), data);
}
#[test]
fn mix_backend5_roundtrip_and_distinct() {
let data: Vec<u8> = (0..8192u32).map(|i| (i * 7 + (i >> 5)) as u8).collect();
let w4 = mix_encode(&data);
let w5 = mix_encode_with(&MixParams::BACKEND5, &data);
assert_eq!(mix_decode(&w4, data.len()).unwrap(), data);
assert_eq!(
mix_decode_with(&MixParams::BACKEND5, &w5, data.len()).unwrap(),
data
);
assert_ne!(w4, w5, "backend 5's evolution must differ from backend 4's");
}
#[test]
fn mix_backend5_rejects_malformed() {
let data: Vec<u8> = (0..4096u32).map(|i| ((i * 31) >> 3) as u8).collect();
let w5 = mix_encode_with(&MixParams::BACKEND5, &data);
assert!(mix_decode_with(&MixParams::BACKEND5, &w5[..17], 4096).is_err());
let mut bomb = w5.clone();
bomb[0..4].copy_from_slice(&u32::MAX.to_le_bytes());
assert!(mix_decode_with(&MixParams::BACKEND5, &bomb, 4096).is_err());
let r = mix_decode_with(&MixParams::BACKEND5, &w5[..12], 4096);
assert!(r.is_ok() || r.is_err(), "no panic reachable"); let _ = r.map(|v| assert!(v.len() <= 4096));
assert_eq!(
mix_decode_with(&MixParams::BACKEND5, &w5, data.len()).unwrap(),
data
);
}
}