use super::bitwriter::BoolEncoder;
use crate::prob::inv_remap_prob;
pub fn forward_remap_prob(new_p: u8, old_p: u8) -> u8 {
let mut best = 0u8;
let mut best_err = i32::MAX;
for d in 0u8..255 {
let p = inv_remap_prob(d as i32, old_p as i32);
if p == new_p {
return d;
}
let err = (p as i32 - new_p as i32).abs();
if err < best_err {
best_err = err;
best = d;
}
}
best }
pub fn encode_term_subexp(enc: &mut BoolEncoder, d: u32) {
debug_assert!(d <= 254, "term-subexp delta out of range: {d}");
if d < 16 {
enc.write_bool(0, 128);
enc.write_literal(d, 4);
} else if d < 32 {
enc.write_bool(1, 128);
enc.write_bool(0, 128);
enc.write_literal(d - 16, 4);
} else if d < 64 {
enc.write_bool(1, 128);
enc.write_bool(1, 128);
enc.write_bool(0, 128);
enc.write_literal(d - 32, 5);
} else {
enc.write_bool(1, 128);
enc.write_bool(1, 128);
enc.write_bool(1, 128);
let w = d - 64;
if w < 65 {
enc.write_literal(w, 7);
} else {
let val = d + 1; enc.write_literal(val >> 1, 7);
enc.write_bool(val & 1, 128);
}
}
}
pub fn diff_update_encode(enc: &mut BoolEncoder, old_p: u8, new_p: u8) -> u8 {
if new_p == old_p {
enc.write_bool(0, 252);
old_p
} else {
enc.write_bool(1, 252);
encode_term_subexp(enc, forward_remap_prob(new_p, old_p) as u32);
new_p
}
}
pub fn update_mv_prob_encode(enc: &mut BoolEncoder, old_p: u8, new_p: u8) -> u8 {
if new_p == old_p {
enc.write_bool(0, 252);
old_p
} else {
enc.write_bool(1, 252);
let coded = (new_p & !1) | 1;
enc.write_literal((coded >> 1) as u32, 7);
coded
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::bits::BoolDecoder;
fn xs(s: &mut u64) -> u64 {
*s ^= *s << 13;
*s ^= *s >> 7;
*s ^= *s << 17;
*s
}
#[test]
fn forward_remap_inverts_every_reachable_target() {
for old in 1u8..=255 {
let mut reachable = [false; 256];
for d in 0u8..255 {
reachable[inv_remap_prob(d as i32, old as i32) as usize] = true;
}
for new in 1u8..=255 {
let d = forward_remap_prob(new, old);
let got = inv_remap_prob(d as i32, old as i32);
if reachable[new as usize] {
assert_eq!(got, new, "reachable old={old} new={new} d={d}");
} else {
let best_err = (1u8..=255)
.filter(|&p| reachable[p as usize])
.map(|p| (p as i32 - new as i32).abs())
.min()
.unwrap();
assert_eq!(
(got as i32 - new as i32).abs(),
best_err,
"unreachable old={old} new={new}: got {got}"
);
}
}
}
}
#[test]
fn term_subexp_roundtrips_through_decoder() {
let mut enc = BoolEncoder::new();
for d in 0u32..=254 {
encode_term_subexp(&mut enc, d);
}
let bytes = enc.finish();
let mut dec = BoolDecoder::new(&bytes).unwrap();
for d in 0u32..=254 {
assert_eq!(crate::decode::decode_term_subexp(&mut dec), d, "delta {d}");
}
}
#[test]
fn diff_update_roundtrips_through_decoder() {
let mut s = 0xabcd_1234_5678_9999u64;
let cases: Vec<(u8, u8)> = (0..4000)
.map(|_| ((1 + xs(&mut s) % 255) as u8, (1 + xs(&mut s) % 255) as u8))
.collect();
let mut enc = BoolEncoder::new();
for &(old, new) in &cases {
diff_update_encode(&mut enc, old, new);
}
let bytes = enc.finish();
let mut dec = BoolDecoder::new(&bytes).unwrap();
for &(old, new) in &cases {
let mut p = old;
crate::decode::diff_update(&mut dec, &mut p);
assert_eq!(p, new, "old={old} new={new}");
}
}
#[test]
fn mv_prob_update_roundtrips() {
let mut enc = BoolEncoder::new();
let targets: Vec<u8> = (0..127).map(|i| ((i << 1) | 1) as u8).collect();
let mut coded = Vec::new();
for &t in &targets {
coded.push(update_mv_prob_encode(&mut enc, 128, t));
}
let bytes = enc.finish();
let mut dec = BoolDecoder::new(&bytes).unwrap();
for (&t, &c) in targets.iter().zip(&coded) {
assert_eq!(c, t, "odd target should code exactly");
let mut p = 128u8;
if dec.read_bool(252) == 1 {
p = ((dec.literal(7) << 1) | 1) as u8;
}
assert_eq!(p, t);
}
}
}