use std::convert::Infallible;
use std::io::{self, Bytes, Read};
use crate::maps::{black, mode, white, Mode, EDFB_HALF, EOL};
use crate::{BitReader, ByteReader, Color, Transitions};
fn with_markup<D, R>(decoder: D, reader: &mut R) -> Option<u16>
where
D: Fn(&mut R) -> Option<u16>,
{
let mut sum: u16 = 0;
while let Some(n) = decoder(reader) {
sum = sum.checked_add(n)?;
if n < 64 {
return Some(sum);
}
}
None
}
fn colored(current: Color, reader: &mut impl BitReader) -> Option<u16> {
match current {
Color::Black => with_markup(black::decode, reader),
Color::White => with_markup(white::decode, reader),
}
}
pub fn pels(line: &[u16], width: u16) -> impl Iterator<Item = Color> + '_ {
use std::iter::repeat;
let mut color = Color::White;
let mut last = 0;
let pad_color = if line.len() & 1 == 1 { !color } else { color };
line.iter()
.flat_map(move |&p| {
let c = color;
color = !color;
let n = p.saturating_sub(last);
last = p;
repeat(c).take(n as usize)
})
.chain(repeat(pad_color))
.take(width as usize)
}
pub fn decode_g3(input: impl Iterator<Item = u8>, mut line_cb: impl FnMut(&[u16])) -> Option<()> {
let reader = input.map(Result::<u8, Infallible>::Ok);
let mut decoder = Group3Decoder::new(reader).ok()?;
while let Ok(status) = decoder.advance() {
line_cb(decoder.transitions());
if status == DecodeStatus::End {
return Some(());
}
}
None
}
#[derive(PartialEq, Eq, Debug, Copy, Clone)]
pub enum DecodeStatus {
Incomplete,
End,
}
pub struct Group3Decoder<R> {
reader: ByteReader<R>,
current: Vec<u16>,
}
impl<E: std::fmt::Debug, R: Iterator<Item = Result<u8, E>>> Group3Decoder<R> {
pub fn new(reader: R) -> Result<Self, DecodeError<E>> {
let mut reader = ByteReader::new(reader).map_err(DecodeError::Reader)?;
skip_to_eol(&mut reader).map_err(|_| DecodeError::Invalid)?;
Ok(Group3Decoder {
reader,
current: vec![],
})
}
pub fn advance(&mut self) -> Result<DecodeStatus, DecodeError<E>> {
self.current.clear();
let mut a0: u16 = 0;
let mut color = Color::White;
loop {
if is_eol_ahead(&self.reader) {
break;
}
match colored(color, &mut self.reader) {
Some(p) => {
a0 = a0.checked_add(p).ok_or(DecodeError::Invalid)?;
self.current.push(a0);
color = !color;
}
None => break,
}
}
skip_to_eol(&mut self.reader).map_err(|_| DecodeError::Invalid)?;
for _ in 0..5 {
if is_eol_ahead(&self.reader) {
skip_to_eol(&mut self.reader).map_err(|_| DecodeError::Invalid)?;
} else {
return Ok(DecodeStatus::Incomplete);
}
}
Ok(DecodeStatus::End)
}
pub fn transitions(&self) -> &[u16] {
&self.current
}
}
fn is_eol_ahead<E, R: Iterator<Item = Result<u8, E>>>(reader: &ByteReader<R>) -> bool {
reader.peek(9) == Some(0)
}
fn skip_to_eol<E: std::fmt::Debug, R: Iterator<Item = Result<u8, E>>>(
reader: &mut ByteReader<R>,
) -> Result<(), DecodeError<E>> {
while reader.peek(1) == Some(0) {
reader.consume(1).map_err(DecodeError::Reader)?;
}
if reader.peek(1) == Some(1) {
reader.consume(1).map_err(DecodeError::Reader)?;
Ok(())
} else {
Err(DecodeError::Invalid)
}
}
pub fn decode_g4(
input: impl Iterator<Item = u8>,
width: u16,
height: Option<u16>,
mut line_cb: impl FnMut(&[u16]),
) -> Option<()> {
let reader = input.map(Result::<u8, Infallible>::Ok);
let mut decoder = Group4Decoder::new(reader, width).ok()?;
let max_lines = height.unwrap_or(u16::MAX);
let mut lines_emitted: u16 = 0;
while lines_emitted < max_lines {
let status = decoder.advance().ok()?;
if status == DecodeStatus::End {
break;
}
line_cb(decoder.transition());
lines_emitted += 1;
}
if let Some(h) = height {
while lines_emitted < h {
line_cb(&[]);
lines_emitted += 1;
}
}
Some(())
}
#[derive(Debug)]
pub enum DecodeError<E> {
Reader(E),
Invalid,
Unsupported,
}
impl<E> std::fmt::Display for DecodeError<E> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Decode Error")
}
}
impl<E: std::error::Error> std::error::Error for DecodeError<E> {}
pub struct Group4Decoder<R> {
reader: ByteReader<R>,
reference: Vec<u16>,
current: Vec<u16>,
width: u16,
}
impl<E, R: Iterator<Item = Result<u8, E>>> Group4Decoder<R> {
pub fn new(reader: R, width: u16) -> Result<Self, E> {
Ok(Group4Decoder {
reader: ByteReader::new(reader)?,
reference: Vec::new(),
current: Vec::new(),
width,
})
}
pub fn advance(&mut self) -> Result<DecodeStatus, DecodeError<E>> {
let mut transitions = Transitions::new(&self.reference);
let mut a0 = 0;
let mut color = Color::White;
let mut start_of_row = true;
loop {
let mode = match mode::decode(&mut self.reader) {
Some(mode) => mode,
None => return Err(DecodeError::Invalid),
};
match mode {
Mode::Pass => {
if start_of_row && color == Color::White {
transitions.pos += 1;
} else {
transitions
.next_color(a0, !color, false)
.ok_or(DecodeError::Invalid)?;
}
if let Some(b2) = transitions.next() {
a0 = b2;
}
}
Mode::Vertical(delta) => {
let b1 = transitions
.next_color(a0, !color, start_of_row)
.unwrap_or(self.width);
let a1_i32 = b1 as i32 + delta as i32;
if a1_i32 < 0 || a1_i32 > self.width as i32 {
break;
}
let a1 = a1_i32 as u16;
if a1 < self.width {
self.current.push(a1);
}
color = !color;
a0 = a1;
if delta < 0 {
transitions.seek_back(a0);
}
}
Mode::Horizontal => {
let a0a1 = colored(color, &mut self.reader).ok_or(DecodeError::Invalid)?;
let a1a2 = colored(!color, &mut self.reader).ok_or(DecodeError::Invalid)?;
let a1 = a0.checked_add(a0a1).ok_or(DecodeError::Invalid)?;
let a2 = a1.checked_add(a1a2).ok_or(DecodeError::Invalid)?;
if a1 < self.width {
self.current.push(a1);
}
if a2 >= self.width {
break;
}
self.current.push(a2);
a0 = a2;
}
Mode::Extension => {
let _ext = self.reader.peek(3).ok_or(DecodeError::Invalid)?;
let _ = self.reader.consume(3);
return Err(DecodeError::Unsupported);
}
Mode::EOF => return Ok(DecodeStatus::End),
}
start_of_row = false;
if a0 >= self.width {
break;
}
}
std::mem::swap(&mut self.reference, &mut self.current);
self.current.clear();
Ok(DecodeStatus::Incomplete)
}
pub fn transition(&self) -> &[u16] {
&self.reference
}
pub fn line(&self) -> Line {
Line {
transitions: &self.reference,
width: self.width,
}
}
}
pub struct Line<'a> {
pub transitions: &'a [u16],
pub width: u16,
}
impl<'a> Line<'a> {
pub fn pels(&self) -> impl Iterator<Item = Color> + 'a {
pels(&self.transitions, self.width)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn g4_fuzz_crash_horizontal_overflow() {
let data: Vec<u8> = vec![0xe8, 0x05, 0x00, 0x00, 0x00];
let mut lines = 0u32;
let result = decode_g4(data.into_iter(), 100, Some(10), |_| {
lines += 1;
});
assert!(
result.is_some(),
"decoder should recover from caught overflow"
);
assert!(lines <= 10, "should not exceed requested height");
}
#[test]
fn g3_fuzz_crash_run_length_overflow() {
let mut data = vec![
0x10, 0x10, 0x00, 0x04, 0x00, 0x10, 0x00, 0xb3, 0x00, 0x00, 0x10, 0x00, 0xb3, 0x00,
0x10, 0x10,
];
data.extend_from_slice(&[0xce; 103]);
let result = decode_g3(data.into_iter(), |_| {});
assert_eq!(result, None, "corrupt G3 data should return None");
}
#[test]
fn g4_large_width_no_overflow() {
let data: Vec<u8> = vec![0x00; 512];
let result = decode_g4(data.into_iter(), 40000, Some(1), |_| {});
let _ = result; }
#[test]
fn g4_zero_width_no_panic() {
let data: Vec<u8> = vec![0x00; 64];
let result = decode_g4(data.into_iter(), 0, Some(1), |_| {});
let _ = result; }
#[test]
fn g3_random_bytes_no_panic() {
let data: Vec<u8> = (0..512).map(|i| (i * 37 + 13) as u8).collect();
let result = decode_g3(data.into_iter(), |_| {});
let _ = result; }
#[test]
fn g4_roundtrip_width_boundary_transition() {
let transitions = vec![3u16, 4];
let width = 4u16;
let input_pels: Vec<_> = super::pels(&transitions, width).collect();
let writer = crate::VecWriter::new();
let mut encoder = crate::encoder::Encoder::new(writer);
let _ = encoder.encode_line(input_pels.iter().copied(), width);
let encoded = encoder.finish().unwrap().finish();
let mut decoded = Vec::new();
let _ = decode_g4(encoded.into_iter(), width, Some(1), |line| {
decoded.push(line.to_vec());
});
let decoded_line = decoded.first().expect("decoded one line");
let output_pels: Vec<_> = super::pels(decoded_line, width).collect();
assert_eq!(
input_pels, output_pels,
"pels must roundtrip (decoded transitions: {:?})",
decoded_line
);
}
#[test]
fn g4_roundtrip_canonical_form() {
for &(width, ref transitions) in &[
(10u16, vec![5]),
(2000, vec![10]),
(2000, vec![3, 51]),
(4, vec![3]),
(100, vec![50]),
(100, vec![1]),
(100, vec![99]),
] {
let input_pels: Vec<_> = super::pels(transitions, width).collect();
let writer = crate::VecWriter::new();
let mut encoder = crate::encoder::Encoder::new(writer);
let _ = encoder.encode_line(input_pels.iter().copied(), width);
let encoded = encoder.finish().unwrap().finish();
let mut decoded = Vec::new();
let _ = decode_g4(encoded.into_iter(), width, Some(1), |line| {
decoded.push(line.to_vec());
});
let decoded_line = decoded.first().expect("decoded one line");
let output_pels: Vec<_> = super::pels(decoded_line, width).collect();
assert_eq!(
input_pels, output_pels,
"pels must roundtrip for width={width} transitions={transitions:?}, \
got decoded transitions {decoded_line:?}"
);
assert!(
decoded_line.iter().all(|&t| t < width),
"decoder produced non-canonical transition list {decoded_line:?} \
(contains width={width}); transitions should all be < width"
);
}
}
}