use std::io::Read;
use crate::codec::{apply_transform, div_ceil, split_tiles, tile_sig, Atlas, Dims, TILE_SIZE};
use crate::error::{Error, Result};
use crate::mnf6::{
compress_stream, decode_record_walk, med_predict, modcolor_forward_iter, pack9, planar_pack,
proxy_cost, ycocg_forward, Coder, MnF6Stats, FLAG_PLANAR,
};
pub const HEADER8: usize = 32;
pub const FLAG_PRED_MED: u32 = 4; pub const FLAGS_RESERVED8: u32 = !0b111;
pub const BACKEND_ZLIB: u32 = 0;
pub const BACKEND_ZSTD19: u32 = 1;
pub const BACKEND_RESERVED2: u32 = 2; pub const BACKEND_RESERVED3: u32 = 3;
pub const ZSTD_WINDOW_LOG_MAX: u32 = 23;
use crate::codec::RANK_CANDIDATES;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Profile {
Balanced,
Maximum,
}
fn backend_candidates() -> Vec<(u32, Coder)> {
vec![
(BACKEND_ZSTD19, Coder::Zstd(19)),
(BACKEND_ZLIB, Coder::Zlib),
]
}
fn backend_compress(backend: u32, buf: &[u8]) -> Result<Vec<u8>> {
match backend {
BACKEND_ZLIB => Ok(compress_stream(&Coder::Zlib, buf)),
BACKEND_ZSTD19 => Ok(compress_stream(&Coder::Zstd(19), buf)),
other => Err(Error::fmt_msg(format!(
"backend {other} not a compressed-output coder"
))),
}
}
fn backend_decompress(backend: u32, buf: &[u8], limit: u64) -> Result<Vec<u8>> {
match backend {
BACKEND_ZLIB => {
let mut d = flate2::read::ZlibDecoder::new(buf).take(limit + 1);
let mut v = Vec::new();
d.read_to_end(&mut v)
.map_err(|_| Error::Decompress("corrupt zlib stream"))?;
if v.len() as u64 > limit {
return Err(Error::fmt_msg("bomb"));
}
if d.get_ref().total_in() as usize != buf.len() {
return Err(Error::fmt_msg("trailing bytes after zlib stream"));
}
Ok(v)
}
BACKEND_ZSTD19 => {
let cursor = std::io::Cursor::new(buf);
let mut d = zstd::stream::read::Decoder::new(cursor)
.map_err(|_| Error::Decompress("zstd init"))?;
d.window_log_max(ZSTD_WINDOW_LOG_MAX)
.map_err(|_| Error::Decompress("zstd window param"))?;
let mut v = Vec::new();
d.single_frame()
.take(limit + 1)
.read_to_end(&mut v)
.map_err(|_| Error::Decompress("corrupt zstd stream"))?;
if v.len() as u64 > limit {
return Err(Error::fmt_msg("bomb"));
}
Ok(v)
}
other => Err(Error::fmt_msg(format!(
"backend {other} not a decompressable coder"
))),
}
}
pub fn encode_mnf8(
pixels: &[u8],
dims: Dims,
use_geom: bool,
profile: Profile,
) -> (Vec<u8>, MnF6Stats, u32) {
let candidates: Vec<(u32, Coder)> = match profile {
Profile::Balanced => vec![(BACKEND_ZSTD19, Coder::Zstd(19))],
Profile::Maximum => backend_candidates(),
};
let mut best: Option<(Vec<u8>, u32, MnF6Stats, u32)> = None;
for (backend, coder) in candidates {
for flags in [
0u32,
FLAG_PLANAR,
FLAG_PRED_MED,
FLAG_PLANAR | FLAG_PRED_MED,
] {
let planar = flags & FLAG_PLANAR != 0;
let med = flags & FLAG_PRED_MED != 0;
let (file, st) = encode_stream_coder(
pixels, dims, true, use_geom, flags, planar, med, &coder, backend,
)
.expect("encoding an in-memory image cannot fail");
let size = file.len();
if best.as_ref().is_none_or(|b| size < b.0.len()) {
best = Some((file, flags, st, backend));
}
}
}
let (file, flags, mut st, backend) = best.expect("at least one candidate");
st.chosen_flags = flags;
(file, st, backend)
}
#[allow(clippy::too_many_arguments)]
fn encode_stream_coder(
pixels: &[u8],
dims: Dims,
use_atlas: bool,
use_geom: bool,
flags: u32,
planar: bool,
med: bool,
coder: &Coder,
backend: u32,
) -> Result<(Vec<u8>, MnF6Stats)> {
let med_pred = if med {
med_predict(pixels, dims.width as usize, dims.height as usize)
} else {
Vec::new()
};
let tiles = split_tiles(pixels, dims);
let mut raw: Vec<u8> = Vec::with_capacity(tiles.len() * 774);
let mut stats = MnF6Stats::default();
let mut atlas = Atlas::new();
for (idx, tile) in tiles.iter().enumerate() {
let n = tile.nbytes();
let pack = |rgb: &[u8]| -> Vec<u8> {
if planar {
planar_pack(rgb)
} else {
rgb.to_vec()
}
};
let lit_src: Vec<u8> = if med {
let (tw, th) = (tile.width as usize, tile.height as usize);
let cols_px = dims.width as usize; let tx = (idx % ((dims.width as usize).div_ceil(16))) * 16;
let ty = (idx / ((dims.width as usize).div_ceil(16))) * 16;
let mut src = Vec::with_capacity(n);
for y in 0..th {
let row_off = ((ty + y) * cols_px + tx) * 3;
src.extend_from_slice(&med_pred[row_off..row_off + tw * 3]);
}
tile.data
.iter()
.zip(&src)
.map(|(x, m)| x.wrapping_sub(*m))
.collect()
} else {
tile.data.clone()
};
let lit_body = pack(&lit_src);
let mut choice_body = lit_body.clone();
let mut choice_mode = 0u8;
let mut choice_ref: Option<(u32, u8)> = None;
let mut best_cost = 1 + proxy_cost(coder, &lit_body);
{
let (cy, cco, ccg) = ycocg_forward(&tile.data);
let mut m3 = Vec::with_capacity(n + 160);
m3.extend_from_slice(&cy);
m3.extend_from_slice(&pack9(&cco));
m3.extend_from_slice(&pack9(&ccg));
let m3_cost = 1 + proxy_cost(coder, &m3);
if m3_cost < best_cost {
best_cost = m3_cost;
choice_body = m3;
choice_mode = 3;
}
}
let tile_t = modcolor_forward_iter(&tile.data);
if use_atlas && tiles.len() > 1 {
let sig = tile_sig(tile.width, tile.height, &tile.data);
for prior in atlas.rank_by_sig(&sig, tile.width, tile.height, RANK_CANDIDATES) {
let (_, _, ref_rgb) = &atlas.patches[prior as usize];
let transforms: &[u8] = if !use_geom {
&[0]
} else if tile.width == tile.height {
&[0, 1, 2, 3, 4, 5, 6, 7]
} else {
&[0, 1, 2, 3]
};
for &t in transforms {
let base = if t == 0 {
ref_rgb.clone()
} else {
apply_transform(t, tile.width, tile.height, ref_rgb)
};
let mut r = vec![0u8; n];
for i in 0..n {
r[i] = tile.data[i].wrapping_sub(base[i]);
}
let body = pack(&r);
let meta: u64 = if t == 0 { 5 } else { 6 };
let cost = meta + proxy_cost(coder, &body);
if cost < best_cost {
best_cost = cost;
choice_body = body;
choice_mode = if t == 0 { 1 } else { 2 };
choice_ref = Some((prior, t));
}
let base_t = modcolor_forward_iter(&base);
let mut r_c = vec![0u8; n];
for i in 0..n {
r_c[i] = tile_t[i].wrapping_sub(base_t[i]);
}
let body_c = pack(&r_c);
let cost_c = meta + proxy_cost(coder, &body_c);
if cost_c < best_cost {
best_cost = cost_c;
choice_body = body_c;
choice_mode = if t == 0 { 4 } else { 5 };
choice_ref = Some((prior, t));
}
}
}
}
stats.modes.push(choice_mode);
stats.refs.push(choice_ref);
match choice_mode {
0 => stats.literal += 1,
1 => stats.ref_exact += 1,
2 => {
stats.ref_geom += 1;
if choice_ref.map(|(_, t)| t != 0).unwrap_or(false) {
stats.geom_nontrivial += 1;
}
}
3 => stats.ycocg_literal += 1,
4 => stats.ref_exact_color += 1,
5 => {
stats.ref_geom_color += 1;
if choice_ref.map(|(_, t)| t != 0).unwrap_or(false) {
stats.geom_nontrivial += 1;
}
}
_ => unreachable!("mode byte constrained above"),
}
raw.push(choice_mode);
if let Some((p, t)) = choice_ref {
raw.extend_from_slice(&p.to_le_bytes());
if t != 0 {
raw.push(t);
}
}
raw.extend_from_slice(&choice_body);
atlas.insert(tile);
}
let payload = backend_compress(backend, &raw)?;
let payload_len =
u32::try_from(payload.len()).map_err(|_| Error::fmt_msg("payload exceeds u32"))?; let mut file = Vec::with_capacity(HEADER8 + payload.len());
file.extend_from_slice(b"MNF8");
file.extend_from_slice(&dims.width.to_le_bytes());
file.extend_from_slice(&dims.height.to_le_bytes());
file.extend_from_slice(&TILE_SIZE.to_le_bytes());
let tile_count = u32::try_from(tiles.len()).map_err(|_| Error::fmt_msg("too many tiles"))?;
file.extend_from_slice(&tile_count.to_le_bytes());
file.extend_from_slice(&payload_len.to_le_bytes());
file.extend_from_slice(&flags.to_le_bytes());
file.extend_from_slice(&backend.to_le_bytes());
file.extend_from_slice(&payload);
Ok((file, stats))
}
pub fn decode_mnf8(buf: &[u8]) -> Result<(Dims, Vec<u8>, u32)> {
decode_mnf8_with_limits(buf, &crate::mnf6::DecodeLimits::default())
}
pub fn decode_mnf8_with_limits(
buf: &[u8],
lim: &crate::mnf6::DecodeLimits,
) -> Result<(Dims, Vec<u8>, u32)> {
if buf.len() < HEADER8 {
return Err(Error::fmt_msg("short"));
}
if buf[0..4] != *b"MNF8" {
return Err(Error::fmt_msg("magic"));
}
let u32le = |o: usize| u32::from_le_bytes(buf[o..o + 4].try_into().unwrap());
let width = u32le(4);
let height = u32le(8);
if width == 0 || height == 0 {
return Err(Error::fmt_msg("dims"));
}
if width as u64 * height as u64 > lim.max_pixels {
return Err(Error::Limits {
which: "max_pixels",
limit: lim.max_pixels,
actual: width as u64 * height as u64,
});
}
if u32le(12) != TILE_SIZE {
return Err(Error::fmt_msg("tilesize"));
}
let cols = div_ceil(width, TILE_SIZE);
let rows = div_ceil(height, TILE_SIZE);
let tile_count = u32le(16);
if tile_count != cols * rows {
return Err(Error::fmt_msg("tilecount"));
}
if tile_count > lim.max_tiles {
return Err(Error::Limits {
which: "max_tiles",
limit: u64::from(lim.max_tiles),
actual: u64::from(tile_count),
});
}
let payload_len = u32le(20);
if payload_len > (1 << 31) - 1 {
return Err(Error::fmt_msg("payloadlen"));
}
if buf.len() != HEADER8 + payload_len as usize {
return Err(Error::fmt_msg("filelen"));
}
let flags = u32le(24);
if flags & FLAGS_RESERVED8 != 0 {
return Err(Error::fmt_msg("reserved-flags"));
}
let backend = u32le(28);
if !matches!(
backend,
BACKEND_ZLIB | BACKEND_ZSTD19 | BACKEND_RESERVED2 | BACKEND_RESERVED3
) {
return Err(Error::fmt_msg(format!("unknown backend {backend}")));
}
if matches!(backend, BACKEND_RESERVED2 | BACKEND_RESERVED3) {
return Err(Error::fmt_msg(format!(
"reserved backend {backend} (not shipped)"
)));
}
let planar = flags & FLAG_PLANAR != 0;
let mode3_max: u64 =
1 + (TILE_SIZE * TILE_SIZE) as u64 + 2 * ((TILE_SIZE * TILE_SIZE) as u64 * 9 + 7) / 8 + 2;
let px_bytes_total: u64 = width as u64 * height as u64 * 3;
let limit = px_bytes_total
.checked_add(6 * tile_count as u64)
.and_then(|l| l.checked_add(2 * mode3_max + 2))
.ok_or_else(|| Error::fmt_msg("decompressed size overflow"))?;
let payload = &buf[HEADER8..];
let raw = match backend {
BACKEND_ZLIB | BACKEND_ZSTD19 => backend_decompress(backend, payload, limit)?,
other => return Err(Error::fmt_msg(format!("unsupported backend {other}"))),
};
let pixels = decode_record_walk(
&raw,
width,
height,
tile_count,
planar,
flags & FLAG_PRED_MED != 0,
)?;
Ok((Dims { width, height }, pixels, backend))
}
#[cfg(test)]
mod tests {
use super::*;
fn roundtrip(px: &[u8], dims: Dims, profile: Profile) -> u32 {
let (buf, st, backend) = encode_mnf8(px, dims, true, profile);
let (d, out, _) = decode_mnf8(&buf).expect("decode");
assert_eq!(d, dims);
assert_eq!(
out, px,
"profile {profile:?} backend {backend} flags {st:?}"
);
backend
}
#[test]
fn mnf8_balanced_roundtrip_all_classes() {
for (w, h, seed) in [(16u32, 16u32, 1), (48, 33, 2), (31, 17, 4), (1, 1, 5)] {
let mut s: u64 = (seed as u64).wrapping_mul(0x9E3779B97F4A7C15) ^ 0xDEADBEEF;
let mut rnd = move || {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
s
};
let px: Vec<u8> = (0..w as usize * h as usize * 3)
.map(|_| (rnd() & 0xFF) as u8)
.collect();
roundtrip(
&px,
Dims {
width: w,
height: h,
},
Profile::Balanced,
);
}
}
#[test]
fn mnf8_maximum_never_loses_to_mnf7() {
for name in [
"screenshot_640x400.png",
"solid_256x256.png",
"gradient_320x240.png",
] {
let path = format!("{}/testdata/{name}", env!("CARGO_MANIFEST_DIR"));
let img = image::open(&path)
.unwrap_or_else(|e| panic!("{name}: {e}"))
.to_rgb8();
let d = Dims {
width: img.width(),
height: img.height(),
};
let px = img.into_raw();
let (buf8, _, backend) = encode_mnf8(&px, d, true, Profile::Maximum);
let (buf7, _) = crate::mnf6::encode_mnf7(&px, d, true);
assert!(
buf8.len() <= buf7.len() + 4,
"{name}: MNF8({backend}) {} > MNF7(+4) {}",
buf8.len(),
buf7.len()
);
let (_, out, _) = decode_mnf8(&buf8).expect("decode");
assert_eq!(out, px, "{name}: exact");
}
}
#[test]
fn mnf8_rejects() {
let mut f = vec![0u8; 32];
f[0..4].copy_from_slice(b"MNF8");
f[4..8].copy_from_slice(&16u32.to_le_bytes());
f[8..12].copy_from_slice(&16u32.to_le_bytes());
f[12..16].copy_from_slice(&16u32.to_le_bytes());
f[16..20].copy_from_slice(&1u32.to_le_bytes());
f[28..32].copy_from_slice(&7u32.to_le_bytes()); assert!(decode_mnf8(&f).is_err(), "unknown backend rejected");
f[28..32].copy_from_slice(&3u32.to_le_bytes()); assert!(decode_mnf8(&f).is_err(), "reserved backend rejected");
f[28..32].copy_from_slice(&2u32.to_le_bytes()); assert!(decode_mnf8(&f).is_err(), "reserved brotli backend rejected");
f[28..32].copy_from_slice(&2u32.to_le_bytes());
let r = std::panic::catch_unwind(|| decode_mnf8(&f));
assert!(r.is_ok() && r.unwrap().is_err(), "must error, not panic");
f[28..32].copy_from_slice(&0u32.to_le_bytes());
f[24] = 0b100; assert!(decode_mnf8(&f).is_err(), "reserved flag rejected");
f[0] = b'X';
f[24] = 0;
assert!(decode_mnf8(&f).is_err(), "magic");
}
#[test]
fn mnf8_balanced_picks_zstd_not_hardcoded() {
let (w, h) = (64u32, 64u32);
let mut px = vec![0u8; (w * h * 3) as usize];
for y in 0..h {
for x in 0..w {
let i = ((y * w + x) * 3) as usize;
px[i] = (x * 4) as u8;
px[i + 1] = (y * 4) as u8;
px[i + 2] = 128;
}
}
let (buf, st, backend) = encode_mnf8(
&px,
Dims {
width: w,
height: h,
},
false,
Profile::Balanced,
);
assert_eq!(backend, BACKEND_ZSTD19);
assert!(
st.chosen_flags & !0b111 == 0,
"balanced wrote a legal flags id, got {:#x}",
st.chosen_flags
);
let (_, out, _) = decode_mnf8(&buf).expect("decode");
assert_eq!(out, px);
}
#[test]
fn mnf8_no_geom_no_flips() {
let (w, h) = (32u32, 32u32); let mut px = vec![0u8; (w * h * 3) as usize];
for y in 0..h {
for x in 0..w {
let i = ((y * w + x) * 3) as usize;
px[i] = (x * 7 % 256) as u8;
px[i + 1] = (y * 3 % 256) as u8;
px[i + 2] = 5;
}
}
let tile_row = 16usize;
let tile_col = 0usize;
let src_off = (tile_row * w as usize + tile_col) * 3;
let tile = px[src_off..src_off + 16 * 16 * 3].to_vec();
let flipped = apply_transform(1, 16, 16, &tile);
px[src_off..src_off + 16 * 16 * 3].copy_from_slice(&flipped);
let stacked = px.clone();
let d = Dims {
width: w,
height: h,
};
let (buf, st, _) = encode_mnf8(&stacked, d, false, Profile::Balanced);
assert_eq!(st.ref_geom, 0, "no flips allowed with use_geom=false");
assert_eq!(st.geom_nontrivial, 0);
let (_, out, _) = decode_mnf8(&buf).expect("decode");
assert_eq!(out, stacked);
}
#[test]
fn mnf8_no_atlas_arm() {
let d = Dims {
width: 32,
height: 32,
};
let px: Vec<u8> = (0..d.width as usize * d.height as usize * 3)
.map(|i| (i * 7 % 256) as u8)
.collect();
let (buf, st, _) = encode_mnf8_noclimb(&px, d);
assert_eq!(
st.ref_exact + st.ref_geom + st.ref_exact_color + st.ref_geom_color,
0,
"no-atlas arm must emit zero references"
);
let (_, out, _) = decode_mnf8(&buf).expect("decode");
assert_eq!(out, px);
}
fn encode_mnf8_noclimb(px: &[u8], d: Dims) -> (Vec<u8>, MnF6Stats, u32) {
let mut best: Option<(Vec<u8>, MnF6Stats, u32)> = None;
for flags in [0u32, FLAG_PLANAR] {
let planar = flags & FLAG_PLANAR != 0;
let (file, st) = encode_stream_coder(
px,
d,
false,
false,
flags,
planar,
false,
&Coder::Zstd(19),
BACKEND_ZSTD19,
)
.expect("encode");
if best.as_ref().is_none_or(|b| file.len() < b.0.len()) {
best = Some((file, st, BACKEND_ZSTD19));
}
}
let (f, s, b) = best.expect("variant");
(f, s, b)
}
#[test]
fn mnf8_med_flag_roundtrip_photo_like() {
let (w, h) = (64u32, 64u32);
let mut px = vec![0u8; (w * h * 3) as usize];
for y in 0..h {
for x in 0..w {
let i = ((y * w + x) * 3) as usize;
px[i] = (x * 4) as u8;
px[i + 1] = (y * 4) as u8;
px[i + 2] = ((x + y) / 2) as u8;
}
}
let d = Dims {
width: w,
height: h,
};
let (buf, st, _) = encode_mnf8(&px, d, false, Profile::Balanced);
assert!(
st.chosen_flags & !0b111 == 0,
"chosen flags must be a legal MNF8 id, got {:#x}",
st.chosen_flags
);
let (_, out, _) = decode_mnf8(&buf).expect("decode");
if out != px {
let first = out
.iter()
.zip(&px)
.position(|(a, b)| a != b)
.expect("differs");
let (b_y, rem) = (first / 3 / (w as usize), (first / 3) % (w as usize));
panic!(
"not exact: first diff at byte {first} (pixel ({b_y},{rem}), ch {}): out={:?} px={:?}",
first % 3,
&out[first.saturating_sub(4)..first + 4],
&px[first.saturating_sub(4)..first + 4]
);
}
assert_eq!(out, px, "exact reconstruction under the measured flags");
}
#[test]
fn mnf8_med_flag_picked_on_photo_like_smoothing() {
let (w, h) = (64u32, 64u32);
let mut px = vec![0u8; (w * h * 3) as usize];
for y in 0..h {
for x in 0..w {
let i = ((y * w + x) * 3) as usize;
px[i] = (x * 4 + y * 2) as u8;
px[i + 1] = (y * 6) as u8;
px[i + 2] = ((x + y) * 3) as u8;
}
}
let d = Dims {
width: w,
height: h,
};
let (_, st, _) = encode_mnf8(
&px,
d,
false, Profile::Balanced,
);
assert_eq!(
st.chosen_flags & FLAG_PRED_MED,
FLAG_PRED_MED,
"smooth no-repeat bodies must select the MED predictor (picked {:#x})",
st.chosen_flags
);
}
#[test]
fn mnf8_forced_mode4_mode5_records() {
let mut s: u64 = 42;
let mut rnd = move || {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
s
};
let a: Vec<u8> = (0..16 * 16 * 3).map(|_| (rnd() & 0xFF) as u8).collect();
let ta_t = crate::mnf6::modcolor_forward(&a);
let mut tb_t = ta_t.clone();
let ramp = [0u8, 1, 2, 255];
for i in (1..tb_t.len()).step_by(3) {
tb_t[i] = tb_t[i].wrapping_add(ramp[(i / 3) % 4]);
}
let b = crate::mnf6::modcolor_inverse(&tb_t);
let mut two = a.clone();
two.extend_from_slice(&b);
let d2 = Dims {
width: 16,
height: 32,
};
let (buf, st, _) = encode_mnf8(&two, d2, true, Profile::Balanced);
assert!(
st.ref_exact_color > 0,
"expected mode 4 (or a color-geom below) selected, stats: exact={} geom={} ec={} gc={}",
st.ref_exact,
st.ref_geom,
st.ref_exact_color,
st.ref_geom_color
);
let (_, out, _) = decode_mnf8(&buf).expect("decode");
assert_eq!(out, two, "mode-4 record reconstructs exactly through MNF8");
let flipped = apply_transform(1, 16, 16, &a);
let tb_t2 = crate::mnf6::modcolor_forward(&flipped);
let mut tc_t = tb_t2;
for i in (1..tc_t.len()).step_by(3) {
tc_t[i] = tc_t[i].wrapping_add(ramp[(i / 3) % 4]);
}
let c = crate::mnf6::modcolor_inverse(&tc_t);
let mut three = flipped.clone();
three.extend_from_slice(&c);
let d3 = Dims {
width: 16,
height: 32,
};
let (buf3, st3, _) = encode_mnf8(&three, d3, true, Profile::Balanced);
assert!(
st3.ref_exact_color > 0 || st3.ref_geom_color > 0,
"expected a color-domain reference (exact or geom), stats: ec={} gc={}",
st3.ref_exact_color,
st3.ref_geom_color
);
let (_, out3, _) = decode_mnf8(&buf3).expect("decode");
assert_eq!(out3, three);
}
}