j2k_jpeg/
context.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Shared decode context for tile-oriented workloads.
4
5use crate::entropy::huffman::HuffmanTable;
6use crate::entropy::sequential::PreparedDecodePlan;
7use crate::error::JpegError;
8use crate::error::Warning;
9use crate::info::Info;
10use crate::parse::tables::RawHuffmanTable;
11use alloc::sync::Arc;
12use j2k_core::{CacheStats, CodecContext};
13
14const QUANT_CACHE_SLOTS: usize = 8;
15const HUFFMAN_CACHE_SLOTS: usize = 8;
16const PLAN_CACHE_SLOTS: usize = 8;
17const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
18const FNV_PRIME: u64 = 0x0000_0100_0000_01B3;
19
20#[derive(Debug, Clone)]
21struct CachedQuantTable {
22    digest: u64,
23    table: Arc<[u16; 64]>,
24}
25
26#[derive(Debug, Clone)]
27struct CachedHuffmanTable {
28    digest: u64,
29    raw: RawHuffmanTable,
30    table: Arc<HuffmanTable>,
31}
32
33#[derive(Debug, Clone)]
34struct CachedDecodePlan {
35    digest: u64,
36    header_prefix: Arc<[u8]>,
37    info: Info,
38    warnings: Arc<[Warning]>,
39    plan: PreparedDecodePlan,
40}
41
42/// Shared decode context for WSI tile batches.
43///
44/// Reuse one context across many related JPEG tiles to amortize Huffman-table
45/// construction and quant-table cloning when the stream family repeats the same
46/// DHT/DQT definitions across tiles.
47#[derive(Debug, Default)]
48pub struct DecoderContext {
49    quant_tables: [Option<CachedQuantTable>; QUANT_CACHE_SLOTS],
50    huffman_tables: [Option<CachedHuffmanTable>; HUFFMAN_CACHE_SLOTS],
51    decode_plans: [Option<CachedDecodePlan>; PLAN_CACHE_SLOTS],
52    cache_hits: u64,
53    cache_misses: u64,
54    cache_evictions: u64,
55}
56
57impl DecoderContext {
58    /// Create an empty decode context.
59    #[must_use]
60    pub fn new() -> Self {
61        Self {
62            quant_tables: core::array::from_fn(|_| None),
63            huffman_tables: core::array::from_fn(|_| None),
64            decode_plans: core::array::from_fn(|_| None),
65            cache_hits: 0,
66            cache_misses: 0,
67            cache_evictions: 0,
68        }
69    }
70
71    pub(crate) fn resolve_quant_table(&mut self, table: [u16; 64]) -> Arc<[u16; 64]> {
72        let digest = digest_quant_table(&table);
73        self.resolve_quant_table_with_digest(table, digest)
74    }
75
76    fn resolve_quant_table_with_digest(&mut self, table: [u16; 64], digest: u64) -> Arc<[u16; 64]> {
77        let start = (digest as usize) % self.quant_tables.len();
78        for probe in 0..self.quant_tables.len() {
79            let slot = (start + probe) % self.quant_tables.len();
80            match &self.quant_tables[slot] {
81                Some(cached) if cached.digest == digest && cached.table.as_ref() == &table => {
82                    self.cache_hits = self.cache_hits.saturating_add(1);
83                    return Arc::clone(&cached.table);
84                }
85                None => {
86                    let table = Arc::new(table);
87                    self.quant_tables[slot] = Some(CachedQuantTable {
88                        digest,
89                        table: Arc::clone(&table),
90                    });
91                    self.cache_misses = self.cache_misses.saturating_add(1);
92                    return table;
93                }
94                Some(_) => {}
95            }
96        }
97
98        let slot = start;
99        let table = Arc::new(table);
100        self.quant_tables[slot] = Some(CachedQuantTable {
101            digest,
102            table: Arc::clone(&table),
103        });
104        self.cache_misses = self.cache_misses.saturating_add(1);
105        self.cache_evictions = self.cache_evictions.saturating_add(1);
106        table
107    }
108
109    pub(crate) fn resolve_huffman_table(
110        &mut self,
111        raw: &RawHuffmanTable,
112    ) -> Result<Arc<HuffmanTable>, JpegError> {
113        let digest = digest_huffman_table(raw);
114        self.resolve_huffman_table_with_digest(raw, digest)
115    }
116
117    fn resolve_huffman_table_with_digest(
118        &mut self,
119        raw: &RawHuffmanTable,
120        digest: u64,
121    ) -> Result<Arc<HuffmanTable>, JpegError> {
122        let start = (digest as usize) % self.huffman_tables.len();
123        for probe in 0..self.huffman_tables.len() {
124            let slot = (start + probe) % self.huffman_tables.len();
125            match &self.huffman_tables[slot] {
126                Some(cached) if cached.digest == digest && &cached.raw == raw => {
127                    self.cache_hits = self.cache_hits.saturating_add(1);
128                    return Ok(Arc::clone(&cached.table));
129                }
130                None => {
131                    let table = Arc::new(HuffmanTable::from_raw(raw)?);
132                    self.huffman_tables[slot] = Some(CachedHuffmanTable {
133                        digest,
134                        raw: raw.clone(),
135                        table: Arc::clone(&table),
136                    });
137                    self.cache_misses = self.cache_misses.saturating_add(1);
138                    return Ok(table);
139                }
140                Some(_) => {}
141            }
142        }
143
144        let slot = start;
145        let table = Arc::new(HuffmanTable::from_raw(raw)?);
146        self.huffman_tables[slot] = Some(CachedHuffmanTable {
147            digest,
148            raw: raw.clone(),
149            table: Arc::clone(&table),
150        });
151        self.cache_misses = self.cache_misses.saturating_add(1);
152        self.cache_evictions = self.cache_evictions.saturating_add(1);
153        Ok(table)
154    }
155
156    pub(crate) fn resolve_decode_plan<F>(
157        &mut self,
158        header_prefix: &[u8],
159        build: F,
160    ) -> Result<(Info, Arc<[Warning]>, PreparedDecodePlan), JpegError>
161    where
162        F: FnOnce(&mut Self) -> Result<(Info, Arc<[Warning]>, PreparedDecodePlan), JpegError>,
163    {
164        let digest = digest_bytes(header_prefix);
165        let start = (digest as usize) % self.decode_plans.len();
166        let mut empty_slot = None;
167        for probe in 0..self.decode_plans.len() {
168            let slot = (start + probe) % self.decode_plans.len();
169            match &self.decode_plans[slot] {
170                Some(cached)
171                    if cached.digest == digest
172                        && cached.header_prefix.as_ref() == header_prefix =>
173                {
174                    self.cache_hits = self.cache_hits.saturating_add(1);
175                    return Ok((
176                        cached.info.clone(),
177                        Arc::clone(&cached.warnings),
178                        cached.plan.clone(),
179                    ));
180                }
181                None => {
182                    empty_slot = Some(slot);
183                    break;
184                }
185                Some(_) => {}
186            }
187        }
188
189        let built = build(self)?;
190        let slot = empty_slot.unwrap_or(start);
191        self.decode_plans[slot] = Some(CachedDecodePlan {
192            digest,
193            header_prefix: Arc::<[u8]>::from(header_prefix),
194            info: built.0.clone(),
195            warnings: Arc::clone(&built.1),
196            plan: built.2.clone(),
197        });
198        self.cache_misses = self.cache_misses.saturating_add(1);
199        if empty_slot.is_none() {
200            self.cache_evictions = self.cache_evictions.saturating_add(1);
201        }
202        Ok(built)
203    }
204
205    fn occupied_cache_slots(&self) -> u64 {
206        let occupied = self
207            .quant_tables
208            .iter()
209            .filter(|slot| slot.is_some())
210            .count()
211            + self
212                .huffman_tables
213                .iter()
214                .filter(|slot| slot.is_some())
215                .count()
216            + self
217                .decode_plans
218                .iter()
219                .filter(|slot| slot.is_some())
220                .count();
221        occupied as u64
222    }
223}
224
225impl CodecContext for DecoderContext {
226    fn clear(&mut self) {
227        *self = Self::new();
228    }
229
230    fn cache_stats(&self) -> CacheStats {
231        CacheStats::with_slots(
232            self.cache_hits,
233            self.cache_misses,
234            self.occupied_cache_slots(),
235            self.cache_evictions,
236        )
237    }
238}
239
240fn digest_bytes(bytes: &[u8]) -> u64 {
241    let mut hash = FNV_OFFSET;
242    for &byte in bytes {
243        hash ^= u64::from(byte);
244        hash = hash.wrapping_mul(FNV_PRIME);
245    }
246    hash
247}
248
249fn digest_quant_table(table: &[u16; 64]) -> u64 {
250    let mut hash = FNV_OFFSET;
251    for &entry in table {
252        for byte in entry.to_le_bytes() {
253            hash ^= u64::from(byte);
254            hash = hash.wrapping_mul(FNV_PRIME);
255        }
256    }
257    hash
258}
259
260fn digest_huffman_table(raw: &RawHuffmanTable) -> u64 {
261    let mut hash = digest_bytes(&raw.bits);
262    for &byte in raw.values.as_slice() {
263        hash ^= u64::from(byte);
264        hash = hash.wrapping_mul(FNV_PRIME);
265    }
266    hash
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272    use crate::info::{ColorSpace, SamplingFactors, SofKind};
273    use alloc::vec;
274
275    #[test]
276    fn quant_table_cache_hits_return_same_arc() {
277        let mut ctx = DecoderContext::new();
278        let first = ctx.resolve_quant_table([7; 64]);
279        let second = ctx.resolve_quant_table([7; 64]);
280        assert!(Arc::ptr_eq(&first, &second));
281
282        let stats = ctx.cache_stats();
283        assert_eq!(stats.hits, 1);
284        assert_eq!(stats.misses, 1);
285        assert_eq!(stats.occupied_slots, 1);
286        assert_eq!(stats.evictions, 0);
287    }
288
289    #[test]
290    fn huffman_table_cache_hits_return_same_arc() {
291        let raw = RawHuffmanTable {
292            bits: [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
293            values: crate::parse::tables::HuffmanValues::from_slice(&[0]),
294        };
295        let mut ctx = DecoderContext::new();
296        let first = ctx.resolve_huffman_table(&raw).unwrap();
297        let second = ctx.resolve_huffman_table(&raw).unwrap();
298        assert!(Arc::ptr_eq(&first, &second));
299    }
300
301    #[test]
302    fn quant_table_digest_collision_compares_full_table_contents() {
303        let mut ctx = DecoderContext::new();
304        let first = ctx.resolve_quant_table_with_digest([7; 64], 0);
305        let second = ctx.resolve_quant_table_with_digest([8; 64], 0);
306
307        assert!(!Arc::ptr_eq(&first, &second));
308        assert_eq!(*first, [7; 64]);
309        assert_eq!(*second, [8; 64]);
310        assert_eq!(ctx.cache_stats().misses, 2);
311    }
312
313    #[test]
314    fn huffman_table_digest_collision_compares_full_raw_table_contents() {
315        let first_raw = RawHuffmanTable {
316            bits: [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
317            values: crate::parse::tables::HuffmanValues::from_slice(&[0]),
318        };
319        let second_raw = RawHuffmanTable {
320            bits: [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
321            values: crate::parse::tables::HuffmanValues::from_slice(&[1]),
322        };
323        let mut ctx = DecoderContext::new();
324
325        let first = ctx
326            .resolve_huffman_table_with_digest(&first_raw, 0)
327            .unwrap();
328        let second = ctx
329            .resolve_huffman_table_with_digest(&second_raw, 0)
330            .unwrap();
331
332        assert!(!Arc::ptr_eq(&first, &second));
333        assert_eq!(ctx.cache_stats().misses, 2);
334    }
335
336    #[test]
337    fn prepared_plan_cache_hits_skip_rebuild() {
338        let mut ctx = DecoderContext::new();
339        let prefix = [0xFF, 0xD8, 0xFF, 0xDA];
340        let warnings = Arc::<[Warning]>::from([]);
341        let mut builds = 0usize;
342
343        let first = ctx
344            .resolve_decode_plan(&prefix, |_| {
345                builds += 1;
346                Ok((
347                    Info {
348                        dimensions: (16, 16),
349                        color_space: ColorSpace::YCbCr,
350                        sampling: SamplingFactors::from_validated_components(&[
351                            (2, 2),
352                            (1, 1),
353                            (1, 1),
354                        ]),
355                        sof_kind: SofKind::Baseline8,
356                        bit_depth: 8,
357                        restart_interval: None,
358                        mcu_geometry: crate::info::McuGeometry {
359                            width: 16,
360                            height: 16,
361                            columns: 1,
362                            rows: 1,
363                            count: 1,
364                        },
365                        scan_count: 1,
366                    },
367                    Arc::clone(&warnings),
368                    PreparedDecodePlan {
369                        components: vec![],
370                        sampling: SamplingFactors::from_validated_components(&[
371                            (2, 2),
372                            (1, 1),
373                            (1, 1),
374                        ]),
375                        color_space: ColorSpace::YCbCr,
376                        restart_interval: None,
377                        dimensions: (16, 16),
378                        scan_offset: 42,
379                        scratch_bytes: 0,
380                    },
381                ))
382            })
383            .unwrap();
384
385        let second = ctx
386            .resolve_decode_plan(&prefix, |_| {
387                builds += 1;
388                unreachable!("cache hit should bypass rebuild")
389            })
390            .unwrap();
391
392        assert_eq!(builds, 1);
393        assert_eq!(first.0, second.0);
394        assert!(Arc::ptr_eq(&first.1, &second.1));
395        assert_eq!(first.2.scan_offset, second.2.scan_offset);
396    }
397}
j2k_jpeg/context.rs

j2k_jpeg/
context.rs
