entropy decoder

src/bitreader.rs

@@ -1,4 +1,3 @@
-use crate::scan_header::EncodingOrder;
 use anyhow::Result;
 
 pub struct BitReader<'a> {

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<u8, NPtr> {
+    pub(crate) fn huffman_map(&self) -> HashMap<HuffmanMapKey, NPtr> {
         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!(),
         };
 

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<Vec<u8>> {
         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)
     }
 }

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!(),

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,

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();

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 {