use crate::error::Error;
pub const MIN_MATCH_PX: usize = 7;
pub const WINDOW_PX: usize = 1 << 24;
pub const TAG_LITERAL: u8 = 0;
pub const TAG_MATCH: u8 = 1;
pub const TAG_NEAR: u8 = 2;
const MAX_MATCH_BYTES: usize = 255;
pub const MAX_CANDIDATES: usize = 1024;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenStats {
pub literals: u64,
pub matches: u64,
pub near_matches: u64,
}
#[inline]
fn key3(b: &[u8]) -> u32 {
u32::from(b[0]) | (u32::from(b[1]) << 8) | (u32::from(b[2]) << 16)
}
pub fn encode_plane(res: &[u8], w: usize, h: usize) -> Result<(Vec<u8>, TokenStats), Error> {
encode_plane_cap(res, w, h, MAX_CANDIDATES)
}
pub fn encode_plane_cap(
res: &[u8],
w: usize,
h: usize,
cap: usize,
) -> Result<(Vec<u8>, TokenStats), Error> {
let npx = w
.checked_mul(h)
.ok_or_else(|| Error::fmt_msg("dims overflow"))?;
if res.len() != npx * 3 {
return Err(Error::fmt_msg("token plane size mismatch"));
}
let mut wire = Vec::with_capacity(res.len() / 2);
let mut table: std::collections::HashMap<u32, Vec<u32>> = std::collections::HashMap::new();
let mut stats = TokenStats::default();
let mut pos: usize = 0;
while pos < npx {
let key = key3(&res[pos * 3..pos * 3 + 3]);
let mut best_len = 0usize;
let mut best_dist = 0usize;
let mut ncand = 0usize;
if let Some(chain) = table.get(&key) {
for &q in chain.iter().rev() {
let q = q as usize;
let dist = pos - q;
if dist > WINDOW_PX || ncand >= cap {
break;
}
ncand += 1;
let mut l = 0usize;
while l < MAX_MATCH_BYTES
&& pos * 3 + l < res.len()
&& res[q * 3 + l] == res[pos * 3 + l]
{
l += 1;
}
let lpx = l / 3;
if lpx >= MIN_MATCH_PX {
let (dx, dy) = (dist % w, dist / w);
if dx <= 0xFFFF && dy <= 0xFFFF {
let near = dx < 256 && dy < 256;
let cur_near = best_len > 0 && {
let (cdx, cdy) = (best_dist % w, best_dist / w);
cdx < 256 && cdy < 256
};
if lpx > best_len || (lpx == best_len && near && !cur_near) {
best_len = lpx;
best_dist = dist;
}
}
}
}
}
if best_len >= MIN_MATCH_PX {
let (dx, dy) = (best_dist % w, best_dist / w);
if dx < 256 && dy < 256 {
wire.push(TAG_NEAR);
wire.push(dx as u8);
wire.push(dy as u8);
wire.push(best_len as u8);
stats.near_matches += 1;
} else {
wire.push(TAG_MATCH);
wire.extend_from_slice(&(dx as u16).to_le_bytes());
wire.extend_from_slice(&(dy as u16).to_le_bytes());
wire.push(best_len as u8);
}
for p in pos..pos + best_len {
let k = key3(&res[p * 3..p * 3 + 3]);
table.entry(k).or_default().push(p as u32);
}
pos += best_len;
stats.matches += 1;
} else {
wire.push(TAG_LITERAL);
wire.extend_from_slice(&res[pos * 3..pos * 3 + 3]);
table.entry(key).or_default().push(pos as u32);
pos += 1;
stats.literals += 1;
}
}
Ok((wire, stats))
}
pub fn decode_plane(wire: &[u8], w: usize, h: usize) -> Result<Vec<u8>, Error> {
let npx = w
.checked_mul(h)
.ok_or_else(|| Error::fmt_msg("dims overflow"))?;
let mut out: Vec<u8> = Vec::with_capacity(npx * 3);
let mut c = 0usize;
while c < wire.len() {
match wire[c] {
TAG_LITERAL => {
if c + 4 > wire.len() {
return Err(Error::fmt_msg("token wire: truncated literal"));
}
out.extend_from_slice(&wire[c + 1..c + 4]);
c += 4;
}
TAG_MATCH => {
if c + 6 > wire.len() {
return Err(Error::fmt_msg("token wire: truncated match"));
}
let dx = u16::from_le_bytes([wire[c + 1], wire[c + 2]]) as usize;
let dy = u16::from_le_bytes([wire[c + 3], wire[c + 4]]) as usize;
let len = wire[c + 5] as usize;
c += 6;
if len < MIN_MATCH_PX {
return Err(Error::fmt_msg("token wire: match below minimum"));
}
let p = out.len() / 3;
let dist = dy * w + dx;
if dist == 0 || dist > p {
return Err(Error::fmt_msg("token wire: match source out of range"));
}
if out.len() + len * 3 > npx * 3 {
return Err(Error::fmt_msg("token wire: match overruns plane"));
}
let q = (p - dist) * 3;
for k in 0..len * 3 {
let b = out[q + k];
out.push(b);
}
}
TAG_NEAR => {
if c + 4 > wire.len() {
return Err(Error::fmt_msg("token wire: truncated near match"));
}
let dx = wire[c + 1] as usize;
let dy = wire[c + 2] as usize;
let len = wire[c + 3] as usize;
c += 4;
if len < MIN_MATCH_PX {
return Err(Error::fmt_msg("token wire: near match below minimum"));
}
let p = out.len() / 3;
let dist = dy * w + dx;
if dist == 0 || dist > p {
return Err(Error::fmt_msg("token wire: near match source out of range"));
}
if out.len() + len * 3 > npx * 3 {
return Err(Error::fmt_msg("token wire: near match overruns plane"));
}
let q = (p - dist) * 3;
for k in 0..len * 3 {
let b = out[q + k];
out.push(b);
}
}
_ => return Err(Error::fmt_msg("token wire: unknown tag")),
}
}
if out.len() != npx * 3 {
return Err(Error::fmt_msg("token wire: plane size mismatch"));
}
Ok(out)
}
#[test]
fn near_matches_fire_and_roundtrip_exactly() {
let (w, h) = (300usize, 200usize);
let mut plane = vec![0u8; w * h * 3];
for i in 0..w * h {
plane[i * 3] = (i % 251) as u8;
plane[i * 3 + 1] = ((i * 7) % 251) as u8;
plane[i * 3 + 2] = (i * 11 % 251) as u8;
}
let block: Vec<u8> = (0..16 * 16 * 3).map(|k| (k * 13 % 251) as u8).collect();
for (bx, by) in [(10usize, 5usize), (20usize, 10usize), (250usize, 30usize)] {
for y in 0..16 {
for x in 0..16 {
let dst = ((by + y) * w + bx + x) * 3;
plane[dst..dst + 3].copy_from_slice(&block[(y * 16 + x) * 3..(y * 16 + x) * 3 + 3]);
}
}
}
let (wire, st) = encode_plane(&plane, w, h).expect("encode");
assert!(st.near_matches >= 2, "near repeats must fire tag 2: {st:?}");
let back = decode_plane(&wire, w, h).expect("decode");
assert_eq!(back, plane, "near+far mixed wire exact");
}
#[test]
fn near_preferred_on_equal_length() {
let (w, h) = (600usize, 60usize);
let mut plane = vec![0u8; w * h * 3];
for i in 0..w * h {
plane[i * 3] = (i % 251) as u8;
plane[i * 3 + 1] = ((i * 7) % 251) as u8;
plane[i * 3 + 2] = ((i * 3) % 251) as u8;
}
let block: Vec<u8> = (0..16 * 16 * 3).map(|k| (k * 13 % 251) as u8).collect();
for y in 0..16 {
for x in 0..16 {
let dst = (y * w + 500 + x) * 3;
plane[dst..dst + 3].copy_from_slice(&block[(y * 16 + x) * 3..(y * 16 + x) * 3 + 3]);
}
}
for y in 0..16 {
for x in 0..16 {
let dst = ((20 + y) * w + 30 + x) * 3;
plane[dst..dst + 3].copy_from_slice(&block[(y * 16 + x) * 3..(y * 16 + x) * 3 + 3]);
}
}
for y in 0..16 {
for x in 0..16 {
let dst = ((40 + y) * w + 60 + x) * 3;
plane[dst..dst + 3].copy_from_slice(&block[(y * 16 + x) * 3..(y * 16 + x) * 3 + 3]);
}
}
let (wire, st) = encode_plane(&plane, w, h).expect("encode");
assert!(
st.near_matches >= 1,
"near candidates must emit tag 2: {st:?}"
);
let back = decode_plane(&wire, w, h).expect("decode");
assert_eq!(back, plane);
}
#[test]
fn near_malformed_rejected() {
let (w, h) = (64usize, 8usize);
let plane: Vec<u8> = (0..w * h * 3).map(|k| (k % 251) as u8).collect();
let (wire, _) = encode_plane(&plane, w, h).expect("encode");
let mut bad = wire.clone();
bad.push(TAG_NEAR);
bad.push(5);
bad.push(1);
assert!(decode_plane(&bad, w, h).is_err());
bad.push(7);
assert!(decode_plane(&bad, w, h).is_err()); assert!(decode_plane(
&{
let mut v = wire.clone();
v.push(TAG_NEAR);
v.extend_from_slice(&[5, 1, 3]);
v
},
w,
h
)
.is_err());
}