From c5cff90c632b157df10b9fb11c6626d702f93566 Mon Sep 17 00:00:00 2001 From: Matthew Kim <38759997+friendlymatthew@users.noreply.github.com> Date: Tue, 11 Jun 2024 16:22:43 -0400 Subject: [PATCH] entropy decoder --- src/bitreader.rs | 1 - src/coding.rs | 15 ++++---- src/entropy_decoder.rs | 78 +++++++++++++++++++++++------------------- src/frame_header.rs | 2 -- src/huffman_tree.rs | 2 +- src/parser.rs | 2 +- src/scan_header.rs | 4 +-- 7 files changed, 54 insertions(+), 50 deletions(-) diff --git a/src/bitreader.rs b/src/bitreader.rs index ce53655..844a319 100644 --- a/src/bitreader.rs +++ b/src/bitreader.rs @@ -1,4 +1,3 @@ -use crate::scan_header::EncodingOrder; use anyhow::Result; pub struct BitReader<'a> { diff --git a/src/coding.rs b/src/coding.rs index 3e3fa15..12a3552 100644 --- a/src/coding.rs +++ b/src/coding.rs @@ -1,4 +1,4 @@ -use crate::huffman_tree::{HuffmanTree, NPtr}; +use crate::huffman_tree::{HuffmanClass, HuffmanTree, NPtr}; use std::collections::HashMap; pub(crate) enum Operation { @@ -11,15 +11,16 @@ pub(crate) enum EntropyCoding { Arithmetic(Vec<()>), } +/// (table_class, destination_id) +type HuffmanMapKey = (HuffmanClass, u8); + impl EntropyCoding { - pub(crate) fn huffman_map(&self) -> HashMap { + pub(crate) fn huffman_map(&self) -> HashMap { let mut map = HashMap::new(); match self { - EntropyCoding::Huffman(hts) => { - hts.iter().for_each(|ht| { - map.insert(ht.destination_id, ht.root); - }) - }, + EntropyCoding::Huffman(hts) => hts.iter().for_each(|ht| { + map.insert((ht.class, ht.destination_id), ht.root); + }), _ => panic!(), }; diff --git a/src/entropy_decoder.rs b/src/entropy_decoder.rs index be0f079..a864ebe 100644 --- a/src/entropy_decoder.rs +++ b/src/entropy_decoder.rs @@ -1,6 +1,6 @@ use crate::coding::EntropyCoding; use crate::frame_header::FrameHeader; -use crate::huffman_tree::HuffmanNode; +use crate::huffman_tree::{HuffmanClass, }; use crate::scan_header::ScanHeader; use anyhow::{anyhow, Result}; @@ -39,60 +39,66 @@ impl<'a> EntropyDecoder<'a> { fn decode_huffman(&mut self) -> Result> { let mut image_data = vec![]; - let huffman_map = self.entropy_coding.huffman_map(); - println!("huffman map keys: {:?}", huffman_map.keys()); - let component_ord_ids: Vec<_> = self + let ac_dc_destination_ids: Vec<_> = self .scan_header .scan_component_selectors .iter() - .map(|s| s.component_id) + .map(|s| (s.dc_destination_id, s.ac_destination_id)) .collect(); - println!("number of components {:?}", component_ord_ids); - let mut component_ptr = 0; + let mut num_coeffs = 0; + + let mut node_cursor = *huffman_map + .get(&(HuffmanClass::DC, ac_dc_destination_ids[component_ptr].0)) + .ok_or(anyhow!(format!( + "failed to find a component with id: {component_ptr}" + )))?; - let mut current_root = *huffman_map - .get(&component_ord_ids[component_ptr]) - .ok_or(anyhow!(format!("failed to find a component with id: {component_ptr}")))?; while self.cursor < self.data.len() { - match self.data[self.cursor] { - 0 => { - if let Some(ptr) = current_root { - current_root = unsafe { (*ptr.as_ptr()).left }; - } - } - 1 => { - if let Some(node) = current_root { - current_root = unsafe { (*node.as_ptr()).right }; - } - } - _ => return Err(anyhow!("data input should only be 1's and 0's")), - } + if let Some(node) = node_cursor { + unsafe { + if (*node.as_ptr()).code != u8::MAX { + image_data.push((*node.as_ptr()).code); - if HuffmanNode::is_leaf(current_root) { - let decompressed_value = if let Some(node) = current_root { - unsafe { (*node.as_ptr()).code } - } else { - return Err(anyhow!("unexpected None pointer after checking valid leaf")); - }; + component_ptr += 1; + if component_ptr == ac_dc_destination_ids.len() { + component_ptr = 0; + num_coeffs += 1; + } - image_data.push(decompressed_value); - component_ptr += 1; + let (next_class, next_destination_id) = if num_coeffs % 64 == 0 { + (HuffmanClass::DC, ac_dc_destination_ids[component_ptr].0) + } else { + (HuffmanClass::AC, ac_dc_destination_ids[component_ptr].1) + }; - if component_ptr >= component_ord_ids.len() { - component_ptr = 0; + node_cursor = + *huffman_map + .get(&(next_class, next_destination_id)) + .ok_or(anyhow!(format!( + "failed to find a component with id: {component_ptr}" + )))?; + } else { + match self.data[self.cursor] { + 0 => { + node_cursor = (*node.as_ptr()).left; + } + 1 => { + node_cursor = (*node.as_ptr()).right; + } + _ => unreachable!(), + }; + } } - current_root = *huffman_map - .get(&component_ord_ids[component_ptr]) - .ok_or(anyhow!(format!("failed to find a component with id {}", component_ptr)))?; } self.cursor += 1; } + println!("image data: {:?}", image_data.len()); Ok(image_data) } } diff --git a/src/frame_header.rs b/src/frame_header.rs index c95b5da..e9bc9fa 100644 --- a/src/frame_header.rs +++ b/src/frame_header.rs @@ -33,8 +33,6 @@ impl ComponentType { pub(crate) fn from(b: u8) -> (Self, EncodingOrder) { match b { 1 => (ComponentType::Grayscale, EncodingOrder::NonInterleaved), - /// When `ComponentType` is 2, that means both component types are provided. Defaults - /// to Color 2 => (ComponentType::Color, EncodingOrder::Interleaved), 3 => (ComponentType::Color, EncodingOrder::Interleaved), 4 => todo!(), diff --git a/src/huffman_tree.rs b/src/huffman_tree.rs index a770162..c92a310 100644 --- a/src/huffman_tree.rs +++ b/src/huffman_tree.rs @@ -4,7 +4,7 @@ use std::marker::PhantomData; use std::ptr::NonNull; /// https://www.youtube.com/watch?v=wLoWd2KyUro -#[derive(Debug, PartialEq, Copy, Clone)] +#[derive(Debug, PartialEq, Copy, Clone, Hash, Eq)] pub enum HuffmanClass { AC = 1, DC = 0, diff --git a/src/parser.rs b/src/parser.rs index aabc846..3067f37 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -45,7 +45,7 @@ impl Parser { let ht_numbers = ht_informations & ht_number_mask; // extract ht type (bit 4) - let ht_type_mask= Simd::splat(0b11110000); + let ht_type_mask = Simd::splat(0b11110000); let ht_types = (ht_informations & ht_type_mask) >> 4; let ht_numbers = ht_numbers.to_array(); diff --git a/src/scan_header.rs b/src/scan_header.rs index 5066f57..3f56501 100644 --- a/src/scan_header.rs +++ b/src/scan_header.rs @@ -61,11 +61,11 @@ pub struct ScanComponentSelector { /// Tdj: Specifies one of four possible DC entropy coding table destinations from which the entropy /// table needed for decoding of the DC coefficients of component selector j is retrieved. - dc_destination_id: u8, + pub(crate) dc_destination_id: u8, /// Taj: Specifies one of four possible AC entropy coding table destinations from which the entropy /// table needed for decoding of the AC coefficients of component selector j is retrieved. - ac_destination_id: u8, + pub(crate) ac_destination_id: u8, } impl ScanComponentSelector {