j2k_jpeg/
context.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Shared decode context for tile-oriented workloads.
4
5use crate::allocation::{checked_add_allocation_bytes, try_reserve_for_len_with_live_budget};
6use crate::entropy::huffman::HuffmanTable;
7use crate::entropy::sequential::PreparedDecodePlan;
8use crate::error::JpegError;
9use crate::parse::tables::RawHuffmanTable;
10use alloc::vec::Vec;
11use core::mem::size_of;
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 MAX_DECODE_PLAN_CACHE_BYTES: usize = 16 * 1024 * 1024;
18const TABLE_CACHE_ALLOCATION_RESERVE_BYTES: usize = 1024 * 1024;
19
20/// Conservative heap reservation used by decode workspace planning. Decode
21/// plan keys/entries are hard-capped at 16 MiB; the remaining MiB covers the
22/// fallibly reserved inline Huffman-cache arena and allocator bookkeeping.
23pub(crate) const MAX_DECODER_CONTEXT_ALLOCATION_BYTES: usize =
24    MAX_DECODE_PLAN_CACHE_BYTES + TABLE_CACHE_ALLOCATION_RESERVE_BYTES;
25
26#[derive(Debug, Clone)]
27struct CachedQuantTable {
28    digest: u64,
29    table: [u16; 64],
30}
31
32#[derive(Debug)]
33struct CachedHuffmanTable {
34    digest: u64,
35    raw: RawHuffmanTable,
36    table: HuffmanTable,
37}
38
39#[derive(Debug)]
40struct CachedDecodePlan {
41    digest: u64,
42    header_prefix: Vec<u8>,
43    plan: PreparedDecodePlan,
44    allocation_bytes: usize,
45}
46
47/// Shared decode context for WSI tile batches.
48///
49/// Reuse one context across many related JPEG tiles to amortize Huffman-table
50/// construction and quant-table cloning when the stream family repeats the same
51/// DHT/DQT definitions across tiles.
52#[derive(Debug, Default)]
53pub struct DecoderContext {
54    quant_tables: [Option<CachedQuantTable>; QUANT_CACHE_SLOTS],
55    huffman_tables: Vec<Option<CachedHuffmanTable>>,
56    decode_plans: [Option<CachedDecodePlan>; PLAN_CACHE_SLOTS],
57    decode_plan_cache_bytes: usize,
58    cache_hits: u64,
59    cache_misses: u64,
60    cache_evictions: u64,
61}
62
63impl DecoderContext {
64    /// Create an empty decode context.
65    #[must_use]
66    pub fn new() -> Self {
67        Self {
68            quant_tables: core::array::from_fn(|_| None),
69            huffman_tables: Vec::new(),
70            decode_plans: core::array::from_fn(|_| None),
71            decode_plan_cache_bytes: 0,
72            cache_hits: 0,
73            cache_misses: 0,
74            cache_evictions: 0,
75        }
76    }
77
78    pub(crate) fn resolve_quant_table(&mut self, table: [u16; 64]) -> [u16; 64] {
79        let digest = digest_quant_table(&table);
80        self.resolve_quant_table_with_digest(table, digest)
81    }
82
83    #[expect(
84        clippy::cast_possible_truncation,
85        reason = "the digest cast only selects a cache shard and intentionally uses the native word"
86    )]
87    fn resolve_quant_table_with_digest(&mut self, table: [u16; 64], digest: u64) -> [u16; 64] {
88        let start = (digest as usize) % self.quant_tables.len();
89        for probe in 0..self.quant_tables.len() {
90            let slot = (start + probe) % self.quant_tables.len();
91            match &self.quant_tables[slot] {
92                Some(cached) if cached.digest == digest && cached.table == table => {
93                    self.cache_hits = self.cache_hits.saturating_add(1);
94                    return cached.table;
95                }
96                None => {
97                    self.quant_tables[slot] = Some(CachedQuantTable { digest, table });
98                    self.cache_misses = self.cache_misses.saturating_add(1);
99                    return table;
100                }
101                Some(_) => {}
102            }
103        }
104
105        let slot = start;
106        self.quant_tables[slot] = Some(CachedQuantTable { digest, table });
107        self.cache_misses = self.cache_misses.saturating_add(1);
108        self.cache_evictions = self.cache_evictions.saturating_add(1);
109        table
110    }
111
112    pub(crate) fn resolve_huffman_table_with_live_budget(
113        &mut self,
114        raw: &RawHuffmanTable,
115        live_bytes: &mut usize,
116        cap: usize,
117    ) -> Result<HuffmanTable, JpegError> {
118        let digest = digest_huffman_table(raw);
119        self.resolve_huffman_table_with_digest_and_live_budget(raw, digest, live_bytes, cap)
120    }
121
122    #[expect(
123        clippy::cast_possible_truncation,
124        reason = "the digest cast only selects a cache shard and intentionally uses the native word"
125    )]
126    fn resolve_huffman_table_with_digest_and_live_budget(
127        &mut self,
128        raw: &RawHuffmanTable,
129        digest: u64,
130        live_bytes: &mut usize,
131        cap: usize,
132    ) -> Result<HuffmanTable, JpegError> {
133        self.ensure_huffman_cache_slots(live_bytes, cap)?;
134        let start = (digest as usize) % self.huffman_tables.len();
135        for probe in 0..self.huffman_tables.len() {
136            let slot = (start + probe) % self.huffman_tables.len();
137            match &self.huffman_tables[slot] {
138                Some(cached) if cached.digest == digest && &cached.raw == raw => {
139                    self.cache_hits = self.cache_hits.saturating_add(1);
140                    return Ok(cached.table.clone());
141                }
142                None => {
143                    let table = HuffmanTable::from_raw(raw)?;
144                    self.huffman_tables[slot] = Some(CachedHuffmanTable {
145                        digest,
146                        raw: raw.clone(),
147                        table: table.clone(),
148                    });
149                    self.cache_misses = self.cache_misses.saturating_add(1);
150                    return Ok(table);
151                }
152                Some(_) => {}
153            }
154        }
155
156        let slot = start;
157        let table = HuffmanTable::from_raw(raw)?;
158        self.huffman_tables[slot] = Some(CachedHuffmanTable {
159            digest,
160            raw: raw.clone(),
161            table: table.clone(),
162        });
163        self.cache_misses = self.cache_misses.saturating_add(1);
164        self.cache_evictions = self.cache_evictions.saturating_add(1);
165        Ok(table)
166    }
167
168    fn ensure_huffman_cache_slots(
169        &mut self,
170        live_bytes: &mut usize,
171        cap: usize,
172    ) -> Result<(), JpegError> {
173        if self.huffman_tables.len() == HUFFMAN_CACHE_SLOTS {
174            return Ok(());
175        }
176        try_reserve_for_len_with_live_budget(
177            &mut self.huffman_tables,
178            HUFFMAN_CACHE_SLOTS,
179            live_bytes,
180            cap,
181        )?;
182        self.huffman_tables
183            .resize_with(HUFFMAN_CACHE_SLOTS, || None);
184        Ok(())
185    }
186
187    pub(crate) fn resolve_decode_plan<F>(
188        &mut self,
189        header_prefix: &[u8],
190        retained_external_bytes: usize,
191        build: F,
192    ) -> Result<PreparedDecodePlan, JpegError>
193    where
194        F: FnOnce(&mut Self) -> Result<PreparedDecodePlan, JpegError>,
195    {
196        let digest = digest_bytes(header_prefix);
197        self.resolve_decode_plan_with_digest(header_prefix, digest, retained_external_bytes, build)
198    }
199
200    #[expect(
201        clippy::cast_possible_truncation,
202        reason = "the digest cast only selects a cache shard and intentionally uses the native word"
203    )]
204    fn resolve_decode_plan_with_digest<F>(
205        &mut self,
206        header_prefix: &[u8],
207        digest: u64,
208        retained_external_bytes: usize,
209        build: F,
210    ) -> Result<PreparedDecodePlan, JpegError>
211    where
212        F: FnOnce(&mut Self) -> Result<PreparedDecodePlan, JpegError>,
213    {
214        let start = (digest as usize) % self.decode_plans.len();
215        let retained_context_bytes = self.retained_allocation_bytes();
216        let initial_live_bytes =
217            checked_add_allocation_bytes(retained_external_bytes, retained_context_bytes)?;
218        for probe in 0..self.decode_plans.len() {
219            let slot = (start + probe) % self.decode_plans.len();
220            match &self.decode_plans[slot] {
221                Some(cached)
222                    if cached.digest == digest
223                        && cached.header_prefix.as_slice() == header_prefix =>
224                {
225                    self.cache_hits = self.cache_hits.saturating_add(1);
226                    let mut live_bytes = initial_live_bytes;
227                    return try_clone_decode_plan(&cached.plan, &mut live_bytes, None)?.ok_or(
228                        JpegError::InternalInvariant {
229                            reason: "cached decode plan unexpectedly bypassed cloning",
230                        },
231                    );
232                }
233                Some(_) | None => {}
234            }
235        }
236
237        let built = build(self)?;
238        self.cache_misses = self.cache_misses.saturating_add(1);
239        let predicted_bytes = decode_plan_entry_bytes(header_prefix.len(), &built)?;
240        if predicted_bytes > MAX_DECODE_PLAN_CACHE_BYTES {
241            // The cache is an optimization. A key that cannot fit its entire
242            // exact entry under the aggregate cache budget is decoded but not
243            // retained.
244            return Ok(built);
245        }
246
247        self.evict_decode_plans_until_fits(start, predicted_bytes);
248        let slot = self.first_empty_decode_plan_slot(start).unwrap_or_else(|| {
249            self.evict_decode_plan_slot(start);
250            start
251        });
252        let mut live_bytes = checked_add_allocation_bytes(
253            retained_external_bytes,
254            self.retained_allocation_bytes(),
255        )?;
256        live_bytes = checked_add_allocation_bytes(live_bytes, built.retained_allocation_bytes()?)?;
257        let mut owned_prefix = Vec::new();
258        try_reserve_for_len_with_live_budget(
259            &mut owned_prefix,
260            header_prefix.len(),
261            &mut live_bytes,
262            j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES,
263        )?;
264        owned_prefix.extend_from_slice(header_prefix);
265        let prefix_bytes = owned_prefix.capacity();
266        let Some(cached_plan) = try_clone_decode_plan(&built, &mut live_bytes, Some(prefix_bytes))?
267        else {
268            return Ok(built);
269        };
270        let allocation_bytes = owned_prefix
271            .capacity()
272            .checked_add(cached_plan.retained_allocation_bytes()?)
273            .ok_or_else(context_cap_error)?;
274        if allocation_bytes > MAX_DECODE_PLAN_CACHE_BYTES {
275            return Ok(built);
276        }
277        self.evict_decode_plans_until_fits(start, allocation_bytes);
278        if self.decode_plans[slot].is_some() {
279            return Err(JpegError::InternalInvariant {
280                reason: "decode-plan cache selected an occupied insertion slot",
281            });
282        }
283        let new_cache_bytes = self
284            .decode_plan_cache_bytes
285            .checked_add(allocation_bytes)
286            .ok_or(JpegError::InternalInvariant {
287                reason: "decode-plan cache byte accounting overflowed",
288            })?;
289        let huffman_bytes = self
290            .huffman_tables
291            .capacity()
292            .checked_mul(size_of::<Option<CachedHuffmanTable>>())
293            .ok_or(JpegError::InternalInvariant {
294                reason: "Huffman cache byte accounting overflowed",
295            })?;
296        let retained_after_insert =
297            new_cache_bytes
298                .checked_add(huffman_bytes)
299                .ok_or(JpegError::InternalInvariant {
300                    reason: "decoder context byte accounting overflowed",
301                })?;
302        if retained_after_insert > MAX_DECODER_CONTEXT_ALLOCATION_BYTES {
303            // The cache is optional. Allocator capacity rounding can make an
304            // otherwise valid entry exceed the context-only retention cap;
305            // decode with `built` and release the attempted cache copy.
306            return Ok(built);
307        }
308        self.decode_plans[slot] = Some(CachedDecodePlan {
309            digest,
310            header_prefix: owned_prefix,
311            plan: cached_plan,
312            allocation_bytes,
313        });
314        self.decode_plan_cache_bytes = new_cache_bytes;
315        Ok(built)
316    }
317
318    fn first_empty_decode_plan_slot(&self, start: usize) -> Option<usize> {
319        (0..self.decode_plans.len())
320            .map(|probe| (start + probe) % self.decode_plans.len())
321            .find(|&slot| self.decode_plans[slot].is_none())
322    }
323
324    fn evict_decode_plans_until_fits(&mut self, start: usize, incoming_bytes: usize) {
325        for probe in 0..self.decode_plans.len() {
326            if self.decode_plan_cache_bytes.saturating_add(incoming_bytes)
327                <= MAX_DECODE_PLAN_CACHE_BYTES
328            {
329                break;
330            }
331            let slot = (start + probe) % self.decode_plans.len();
332            self.evict_decode_plan_slot(slot);
333        }
334    }
335
336    fn evict_decode_plan_slot(&mut self, slot: usize) {
337        if let Some(cached) = self.decode_plans[slot].take() {
338            self.decode_plan_cache_bytes = self
339                .decode_plan_cache_bytes
340                .saturating_sub(cached.allocation_bytes);
341            self.cache_evictions = self.cache_evictions.saturating_add(1);
342        }
343    }
344
345    pub(crate) fn retained_allocation_bytes(&self) -> usize {
346        let huffman_bytes = self
347            .huffman_tables
348            .capacity()
349            .saturating_mul(size_of::<Option<CachedHuffmanTable>>());
350        self.decode_plan_cache_bytes.saturating_add(huffman_bytes)
351    }
352
353    fn occupied_cache_slots(&self) -> u64 {
354        let occupied = self
355            .quant_tables
356            .iter()
357            .filter(|slot| slot.is_some())
358            .count()
359            + self
360                .huffman_tables
361                .iter()
362                .filter(|slot| slot.is_some())
363                .count()
364            + self
365                .decode_plans
366                .iter()
367                .filter(|slot| slot.is_some())
368                .count();
369        occupied as u64
370    }
371}
372
373#[doc(hidden)]
374impl CodecContext for DecoderContext {
375    fn clear(&mut self) {
376        *self = Self::new();
377    }
378
379    fn cache_stats(&self) -> CacheStats {
380        CacheStats::with_slots(
381            self.cache_hits,
382            self.cache_misses,
383            self.occupied_cache_slots(),
384            self.cache_evictions,
385        )
386    }
387}
388
389fn digest_bytes(bytes: &[u8]) -> u64 {
390    j2k_core::__j2k_fnv1a64_bytes!(bytes)
391}
392
393fn digest_quant_table(table: &[u16; 64]) -> u64 {
394    let mut hash = j2k_core::__j2k_fnv1a64_init!();
395    for &entry in table {
396        for byte in entry.to_le_bytes() {
397            j2k_core::__j2k_fnv1a64_update!(hash, byte);
398        }
399    }
400    hash
401}
402
403fn digest_huffman_table(raw: &RawHuffmanTable) -> u64 {
404    let mut hash = digest_bytes(&raw.bits);
405    for &byte in raw.values.as_slice() {
406        j2k_core::__j2k_fnv1a64_update!(hash, byte);
407    }
408    hash
409}
410
411fn decode_plan_entry_bytes(
412    header_prefix_len: usize,
413    plan: &PreparedDecodePlan,
414) -> Result<usize, JpegError> {
415    header_prefix_len
416        .checked_add(plan.retained_allocation_bytes()?)
417        .ok_or_else(context_cap_error)
418}
419
420fn context_cap_error() -> JpegError {
421    JpegError::MemoryCapExceeded {
422        requested: usize::MAX,
423        cap: j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES,
424    }
425}
426
427fn try_clone_decode_plan(
428    plan: &PreparedDecodePlan,
429    live_bytes: &mut usize,
430    cache_prefix_bytes: Option<usize>,
431) -> Result<Option<PreparedDecodePlan>, JpegError> {
432    let mut components = Vec::new();
433    try_reserve_for_len_with_live_budget(
434        &mut components,
435        plan.components.len(),
436        live_bytes,
437        j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES,
438    )?;
439    components.extend(plan.components.iter().cloned());
440    if let Some(prefix_bytes) = cache_prefix_bytes {
441        let projected = prefix_bytes
442            .checked_add(PreparedDecodePlan::allocation_bytes_for_counts(
443                components.capacity(),
444                plan.huffman_tables.len(),
445            )?)
446            .ok_or_else(context_cap_error)?;
447        if projected > MAX_DECODE_PLAN_CACHE_BYTES {
448            return Ok(None);
449        }
450    }
451    let huffman_tables = plan
452        .huffman_tables
453        .try_clone_with_live_budget(live_bytes, j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES)?;
454    Ok(Some(PreparedDecodePlan {
455        components,
456        huffman_tables,
457        sampling: plan.sampling,
458        color_space: plan.color_space,
459        restart_interval: plan.restart_interval,
460        dimensions: plan.dimensions,
461        scan_offset: plan.scan_offset,
462        scratch_bytes: plan.scratch_bytes,
463    }))
464}
465
466#[cfg(test)]
467mod tests {
468    use super::*;
469    use crate::entropy::sequential::PreparedComponentPlan;
470    use crate::info::{ColorSpace, SamplingFactors};
471    use alloc::vec;
472
473    fn empty_plan(scan_offset: usize) -> PreparedDecodePlan {
474        PreparedDecodePlan {
475            components: vec![],
476            huffman_tables: crate::entropy::huffman::PreparedHuffmanTables::try_with_capacity(0)
477                .expect("empty arena"),
478            sampling: SamplingFactors::from_validated_components(&[(1, 1)]),
479            color_space: ColorSpace::Grayscale,
480            restart_interval: None,
481            dimensions: (16, 16),
482            scan_offset,
483            scratch_bytes: 0,
484        }
485    }
486
487    fn resolve_huffman_table(
488        ctx: &mut DecoderContext,
489        raw: &RawHuffmanTable,
490    ) -> Result<HuffmanTable, JpegError> {
491        let mut live_bytes = ctx.retained_allocation_bytes();
492        ctx.resolve_huffman_table_with_live_budget(
493            raw,
494            &mut live_bytes,
495            MAX_DECODER_CONTEXT_ALLOCATION_BYTES,
496        )
497    }
498
499    fn resolve_huffman_table_with_digest(
500        ctx: &mut DecoderContext,
501        raw: &RawHuffmanTable,
502        digest: u64,
503    ) -> Result<HuffmanTable, JpegError> {
504        let mut live_bytes = ctx.retained_allocation_bytes();
505        ctx.resolve_huffman_table_with_digest_and_live_budget(
506            raw,
507            digest,
508            &mut live_bytes,
509            MAX_DECODER_CONTEXT_ALLOCATION_BYTES,
510        )
511    }
512
513    #[test]
514    fn quant_table_cache_hits_return_same_value() {
515        let mut ctx = DecoderContext::new();
516        let first = ctx.resolve_quant_table([7; 64]);
517        let second = ctx.resolve_quant_table([7; 64]);
518        assert_eq!(first, second);
519
520        let stats = ctx.cache_stats();
521        assert_eq!(stats.hits, 1);
522        assert_eq!(stats.misses, 1);
523        assert_eq!(stats.occupied_slots, 1);
524        assert_eq!(stats.evictions, 0);
525    }
526
527    #[test]
528    fn huffman_table_cache_hits_return_same_value() {
529        let raw = RawHuffmanTable {
530            bits: [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
531            values: crate::parse::tables::HuffmanValues::from_slice(&[0]),
532        };
533        let mut ctx = DecoderContext::new();
534        let first = resolve_huffman_table(&mut ctx, &raw).unwrap();
535        let second = resolve_huffman_table(&mut ctx, &raw).unwrap();
536        assert_eq!(first, second);
537    }
538
539    #[test]
540    fn quant_table_digest_collision_compares_full_table_contents() {
541        let mut ctx = DecoderContext::new();
542        let first = ctx.resolve_quant_table_with_digest([7; 64], 0);
543        let second = ctx.resolve_quant_table_with_digest([8; 64], 0);
544
545        assert_ne!(first, second);
546        assert_eq!(first, [7; 64]);
547        assert_eq!(second, [8; 64]);
548        assert_eq!(ctx.cache_stats().misses, 2);
549    }
550
551    #[test]
552    fn huffman_table_digest_collision_compares_full_raw_table_contents() {
553        let first_raw = RawHuffmanTable {
554            bits: [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
555            values: crate::parse::tables::HuffmanValues::from_slice(&[0]),
556        };
557        let second_raw = RawHuffmanTable {
558            bits: [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
559            values: crate::parse::tables::HuffmanValues::from_slice(&[1]),
560        };
561        let mut ctx = DecoderContext::new();
562
563        let first = resolve_huffman_table_with_digest(&mut ctx, &first_raw, 0).unwrap();
564        let second = resolve_huffman_table_with_digest(&mut ctx, &second_raw, 0).unwrap();
565
566        assert_ne!(first, second);
567        assert_eq!(ctx.cache_stats().misses, 2);
568    }
569
570    #[test]
571    fn prepared_plan_cache_hits_skip_rebuild() {
572        let mut ctx = DecoderContext::new();
573        let prefix = [0xFF, 0xD8, 0xFF, 0xDA];
574        let mut builds = 0usize;
575
576        let first = ctx
577            .resolve_decode_plan(&prefix, 0, |_| {
578                builds += 1;
579                Ok(empty_plan(42))
580            })
581            .unwrap();
582
583        let second = ctx
584            .resolve_decode_plan(&prefix, 0, |_| {
585                builds += 1;
586                unreachable!("cache hit should bypass rebuild")
587            })
588            .unwrap();
589
590        assert_eq!(builds, 1);
591        assert_eq!(first.scan_offset, second.scan_offset);
592    }
593
594    #[test]
595    fn cache_hit_clone_shares_one_exact_external_live_budget() {
596        let raw = RawHuffmanTable {
597            bits: [0; 16],
598            values: crate::parse::tables::HuffmanValues::default(),
599        };
600        let mut huffman_tables =
601            crate::entropy::huffman::PreparedHuffmanTables::try_with_capacity(1)
602                .expect("bounded arena");
603        let table = huffman_tables
604            .push(HuffmanTable::from_raw(&raw).expect("empty table"))
605            .expect("reserved arena");
606        let mut plan = empty_plan(7);
607        plan.huffman_tables = huffman_tables;
608        plan.components.push(PreparedComponentPlan {
609            h: 1,
610            v: 1,
611            output_index: 0,
612            quant: [1; 64],
613            dc_table: Some(table),
614            ac_table: Some(table),
615        });
616
617        let mut ctx = DecoderContext::new();
618        let prefix = [0xFF, 0xD8, 0xFF, 0xDA];
619        ctx.resolve_decode_plan(&prefix, 0, |_| Ok(plan))
620            .expect("initial cache insertion");
621        let cached_plan_bytes = ctx
622            .decode_plans
623            .iter()
624            .flatten()
625            .next()
626            .expect("cached plan")
627            .plan
628            .retained_allocation_bytes()
629            .expect("cached plan bytes");
630        let exact_external = j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES
631            .checked_sub(ctx.retained_allocation_bytes())
632            .and_then(|remaining| remaining.checked_sub(cached_plan_bytes))
633            .expect("fixture leaves an external budget");
634
635        ctx.resolve_decode_plan(&prefix, exact_external, |_| {
636            unreachable!("cache hit must bypass rebuild")
637        })
638        .expect("exact live boundary");
639        assert!(matches!(
640            ctx.resolve_decode_plan(&prefix, exact_external + 1, |_| {
641                unreachable!("cache hit must bypass rebuild")
642            }),
643            Err(JpegError::MemoryCapExceeded { .. })
644        ));
645    }
646
647    #[test]
648    fn prepared_plan_digest_collision_compares_full_header_prefix() {
649        let mut ctx = DecoderContext::new();
650        let first = ctx
651            .resolve_decode_plan_with_digest(b"first", 0, 0, |_| Ok(empty_plan(1)))
652            .unwrap();
653        let second = ctx
654            .resolve_decode_plan_with_digest(b"second", 0, 0, |_| Ok(empty_plan(2)))
655            .unwrap();
656        let first_hit = ctx
657            .resolve_decode_plan_with_digest(b"first", 0, 0, |_| {
658                unreachable!("full-key cache hit must bypass rebuild")
659            })
660            .unwrap();
661
662        assert_eq!(first.scan_offset, 1);
663        assert_eq!(second.scan_offset, 2);
664        assert_eq!(first_hit.scan_offset, 1);
665        assert_eq!(ctx.cache_stats().hits, 1);
666    }
667
668    #[test]
669    fn prepared_plan_cache_full_eviction_is_deterministic() {
670        let mut ctx = DecoderContext::new();
671        let cache_slots = u8::try_from(PLAN_CACHE_SLOTS).expect("plan cache slot count fits u8");
672        for key in 0..cache_slots {
673            ctx.resolve_decode_plan_with_digest(&[key], 0, 0, |_| Ok(empty_plan(usize::from(key))))
674                .unwrap();
675        }
676        ctx.resolve_decode_plan_with_digest(&[cache_slots], 0, 0, |_| {
677            Ok(empty_plan(PLAN_CACHE_SLOTS))
678        })
679        .unwrap();
680
681        assert_eq!(ctx.cache_stats().evictions, 1);
682        let mut rebuilt = false;
683        let first = ctx
684            .resolve_decode_plan_with_digest(&[0], 0, 0, |_| {
685                rebuilt = true;
686                Ok(empty_plan(99))
687            })
688            .unwrap();
689        assert!(rebuilt, "the start slot must be the deterministic victim");
690        assert_eq!(first.scan_offset, 99);
691    }
692
693    #[test]
694    fn decode_plan_cache_entry_boundary_bypasses_oversized_keys() {
695        let plan = empty_plan(0);
696        assert_eq!(
697            decode_plan_entry_bytes(MAX_DECODE_PLAN_CACHE_BYTES, &plan).unwrap(),
698            MAX_DECODE_PLAN_CACHE_BYTES
699        );
700        assert!(
701            decode_plan_entry_bytes(MAX_DECODE_PLAN_CACHE_BYTES + 1, &plan).unwrap()
702                > MAX_DECODE_PLAN_CACHE_BYTES
703        );
704    }
705
706    #[test]
707    fn decode_plan_cache_entry_counts_tables_retained_after_table_cache_eviction() {
708        let raw = RawHuffmanTable {
709            bits: [0; 16],
710            values: crate::parse::tables::HuffmanValues::default(),
711        };
712        let mut huffman_tables =
713            crate::entropy::huffman::PreparedHuffmanTables::try_with_capacity(1)
714                .expect("bounded arena");
715        let table = huffman_tables
716            .push(HuffmanTable::from_raw(&raw).expect("empty table"))
717            .expect("reserved arena");
718        let mut plan = empty_plan(0);
719        plan.huffman_tables = huffman_tables;
720        plan.components.push(PreparedComponentPlan {
721            h: 1,
722            v: 1,
723            output_index: 0,
724            quant: [1; 64],
725            dc_table: Some(table),
726            ac_table: Some(table),
727        });
728
729        let entry_bytes = decode_plan_entry_bytes(0, &plan).expect("bounded plan");
730        assert_eq!(entry_bytes, plan.retained_allocation_bytes().unwrap());
731        let logical_bytes = PreparedDecodePlan::allocation_bytes_for_counts(
732            plan.components.len(),
733            plan.huffman_tables.len(),
734        )
735        .expect("logical plan bytes");
736        // `push` from an empty vector commonly retains spare component slots.
737        // Cache entry accounting must use that allocator-returned capacity,
738        // not the one logical component requested by this fixture.
739        let component_spare_bytes = (plan.components.capacity() - plan.components.len())
740            * size_of::<PreparedComponentPlan>();
741        assert_eq!(entry_bytes - logical_bytes, component_spare_bytes);
742        assert!(entry_bytes > size_of::<PreparedComponentPlan>());
743    }
744
745    #[test]
746    fn oversized_decode_plan_key_is_not_retained() {
747        let prefix = vec![0u8; MAX_DECODE_PLAN_CACHE_BYTES + 1];
748        let mut ctx = DecoderContext::new();
749        let mut builds = 0usize;
750        for _ in 0..2 {
751            ctx.resolve_decode_plan(&prefix, 0, |_| {
752                builds += 1;
753                Ok(empty_plan(builds))
754            })
755            .unwrap();
756        }
757
758        assert_eq!(builds, 2, "oversized keys must bypass the cache");
759        assert_eq!(ctx.decode_plan_cache_bytes, 0);
760        assert!(ctx.decode_plans.iter().all(Option::is_none));
761    }
762
763    #[test]
764    fn context_reserve_covers_all_fixed_table_cache_allocations() {
765        let maximum_table_bytes =
766            HUFFMAN_CACHE_SLOTS.saturating_mul(size_of::<Option<CachedHuffmanTable>>());
767        assert!(maximum_table_bytes <= TABLE_CACHE_ALLOCATION_RESERVE_BYTES);
768
769        let ctx = DecoderContext::new();
770        assert_eq!(ctx.retained_allocation_bytes(), 0);
771        assert!(ctx.decode_plan_cache_bytes <= MAX_DECODE_PLAN_CACHE_BYTES);
772    }
773}
j2k_jpeg/context.rs

j2k_jpeg/
context.rs
