Skip to main content

phasm_core/codec/jpeg/
tables.rs

1// Copyright (c) 2026 Christoph Gaffga
2// SPDX-License-Identifier: GPL-3.0-only
3// https://github.com/cgaffga/phasmcore
4
5//! Quantization and Huffman table parsing/serialization.
6//!
7//! Handles DQT (Define Quantization Table) and DHT (Define Huffman Table)
8//! marker segments. Supports both 8-bit and 16-bit quantization precision
9//! and multiple tables per marker segment.
10
11use super::dct::QuantTable;
12use super::error::{JpegError, Result};
13use super::zigzag::ZIGZAG_TO_NATURAL;
14
15/// Parse a DQT marker segment body (after the 2-byte length).
16///
17/// Returns a list of (table_id, QuantTable) pairs. A single DQT segment
18/// can contain multiple tables.
19pub fn parse_dqt(data: &[u8]) -> Result<Vec<(u8, QuantTable)>> {
20    let mut tables = Vec::new();
21    let mut pos = 0;
22
23    while pos < data.len() {
24        if pos >= data.len() {
25            break;
26        }
27        let pq_tq = data[pos];
28        pos += 1;
29        let precision = pq_tq >> 4;
30        let table_id = pq_tq & 0x0F;
31
32        if table_id > 3 {
33            return Err(JpegError::InvalidQuantTableId(table_id));
34        }
35
36        let mut values = [0u16; 64];
37        if precision == 0 {
38            // 8-bit values
39            if pos + 64 > data.len() {
40                return Err(JpegError::UnexpectedEof);
41            }
42            for zi in 0..64 {
43                let ni = ZIGZAG_TO_NATURAL[zi];
44                values[ni] = data[pos + zi] as u16;
45            }
46            pos += 64;
47        } else if precision == 1 {
48            // 16-bit values
49            if pos + 128 > data.len() {
50                return Err(JpegError::UnexpectedEof);
51            }
52            for zi in 0..64 {
53                let ni = ZIGZAG_TO_NATURAL[zi];
54                values[ni] = u16::from_be_bytes([data[pos + zi * 2], data[pos + zi * 2 + 1]]);
55            }
56            pos += 128;
57        } else {
58            return Err(JpegError::InvalidMarkerData("invalid DQT precision"));
59        }
60
61        tables.push((table_id, QuantTable::new(values)));
62    }
63
64    Ok(tables)
65}
66
67/// Write a DQT marker segment (including 0xFFDB marker and length).
68pub fn write_dqt(table_id: u8, qt: &QuantTable) -> Vec<u8> {
69    let mut out = Vec::new();
70    out.push(0xFF);
71    out.push(0xDB);
72
73    // Check if all values fit in 8 bits
74    let precision = if qt.values.iter().all(|&v| v <= 255) { 0u8 } else { 1u8 };
75    let data_len = if precision == 0 { 64 } else { 128 };
76    let length = 2 + 1 + data_len; // length field + pq_tq + values
77    out.push((length >> 8) as u8);
78    out.push(length as u8);
79    out.push((precision << 4) | (table_id & 0x0F));
80
81    for zi in 0..64 {
82        let ni = ZIGZAG_TO_NATURAL[zi];
83        if precision == 0 {
84            out.push(qt.values[ni] as u8);
85        } else {
86            out.extend_from_slice(&qt.values[ni].to_be_bytes());
87        }
88    }
89
90    out
91}
92
93/// Parsed Huffman table specification.
94#[derive(Debug, Clone)]
95pub struct HuffmanSpec {
96    /// Table class: 0 = DC, 1 = AC.
97    pub class: u8,
98    /// Table ID (0–3).
99    pub id: u8,
100    /// Number of codes of each length (1–16).
101    pub bits: [u8; 16],
102    /// Symbol values in order of increasing code length.
103    pub huffval: Vec<u8>,
104}
105
106/// Parse a DHT marker segment body (after the 2-byte length).
107///
108/// Returns a list of HuffmanSpec. A single DHT segment can contain multiple tables.
109pub fn parse_dht(data: &[u8]) -> Result<Vec<HuffmanSpec>> {
110    let mut specs = Vec::new();
111    let mut pos = 0;
112
113    while pos < data.len() {
114        if pos >= data.len() {
115            break;
116        }
117        let tc_th = data[pos];
118        pos += 1;
119        let class = tc_th >> 4;
120        let id = tc_th & 0x0F;
121
122        if class > 1 || id > 3 {
123            return Err(JpegError::InvalidHuffmanTableId(tc_th));
124        }
125
126        if pos + 16 > data.len() {
127            return Err(JpegError::UnexpectedEof);
128        }
129        let mut bits = [0u8; 16];
130        bits.copy_from_slice(&data[pos..pos + 16]);
131        pos += 16;
132
133        let total: usize = bits.iter().map(|&b| b as usize).sum();
134        if pos + total > data.len() {
135            return Err(JpegError::UnexpectedEof);
136        }
137        let huffval = data[pos..pos + total].to_vec();
138        pos += total;
139
140        specs.push(HuffmanSpec {
141            class,
142            id,
143            bits,
144            huffval,
145        });
146    }
147
148    Ok(specs)
149}
150
151/// Write a DHT marker segment (including 0xFFC4 marker and length).
152pub fn write_dht(spec: &HuffmanSpec) -> Vec<u8> {
153    let mut out = Vec::new();
154    out.push(0xFF);
155    out.push(0xC4);
156
157    let total: usize = spec.bits.iter().map(|&b| b as usize).sum();
158    let length = 2 + 1 + 16 + total;
159    out.push((length >> 8) as u8);
160    out.push(length as u8);
161    out.push((spec.class << 4) | (spec.id & 0x0F));
162    out.extend_from_slice(&spec.bits);
163    out.extend_from_slice(&spec.huffval);
164
165    out
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    #[test]
173    fn parse_8bit_dqt() {
174        // Build a DQT body: precision=0, id=0, 64 values 1..64 in zigzag order
175        let mut body = vec![0x00u8]; // pq=0, tq=0
176        for i in 0..64u8 {
177            body.push(i + 1);
178        }
179        let tables = parse_dqt(&body).unwrap();
180        assert_eq!(tables.len(), 1);
181        let (id, qt) = &tables[0];
182        assert_eq!(*id, 0);
183        // Zigzag index 0 maps to natural index 0, value should be 1
184        assert_eq!(qt.values[0], 1);
185        // Zigzag index 1 maps to natural index 1, value should be 2
186        assert_eq!(qt.values[1], 2);
187        // Zigzag index 2 maps to natural index 8, value should be 3
188        assert_eq!(qt.values[8], 3);
189    }
190
191    #[test]
192    fn dqt_roundtrip() {
193        let mut values = [0u16; 64];
194        for i in 0..64 {
195            values[i] = (i + 1) as u16;
196        }
197        let qt = QuantTable::new(values);
198        let written = write_dqt(0, &qt);
199        // Skip marker (2 bytes) and length (2 bytes)
200        let body = &written[4..];
201        let tables = parse_dqt(body).unwrap();
202        assert_eq!(tables.len(), 1);
203        assert_eq!(tables[0].1.values, values);
204    }
205
206    #[test]
207    fn parse_dht_basic() {
208        // Build DHT body: class=0, id=0, standard DC luminance
209        let mut body = vec![0x00u8]; // tc=0, th=0
210        let bits = [0u8, 1, 5, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0];
211        body.extend_from_slice(&bits);
212        let vals: Vec<u8> = (0..12).collect();
213        body.extend_from_slice(&vals);
214
215        let specs = parse_dht(&body).unwrap();
216        assert_eq!(specs.len(), 1);
217        assert_eq!(specs[0].class, 0);
218        assert_eq!(specs[0].id, 0);
219        assert_eq!(specs[0].bits, bits);
220        assert_eq!(specs[0].huffval, vals);
221    }
222
223    #[test]
224    fn dht_roundtrip() {
225        let spec = HuffmanSpec {
226            class: 1,
227            id: 0,
228            bits: [0, 2, 1, 3, 3, 2, 4, 3, 5, 5, 4, 4, 0, 0, 1, 125],
229            huffval: (0..162).collect(),
230        };
231        let written = write_dht(&spec);
232        let body = &written[4..];
233        let specs = parse_dht(body).unwrap();
234        assert_eq!(specs.len(), 1);
235        assert_eq!(specs[0].class, spec.class);
236        assert_eq!(specs[0].id, spec.id);
237        assert_eq!(specs[0].bits, spec.bits);
238        assert_eq!(specs[0].huffval, spec.huffval);
239    }
240}