j2k_jpeg/internal/
scratch.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Reusable scratch buffers for the decode path.
4//!
5//! A [`ScratchPool`] owns every `Vec` that the sequential scan decoder
6//! would otherwise allocate on each call: the three rolling MCU stripe
7//! buffers, the per-component DC predictor, the chroma upsample rows, and
8//! the RGB row buffers used by [`j2k_core::RowSink`] drivers.
9//!
10//! Use [`Decoder::decode_into_with_scratch`](crate::Decoder::decode_into_with_scratch)
11//! / [`decode_rows_with_scratch`](crate::Decoder::decode_rows_with_scratch)
12//! with a single long-lived pool to pay the allocation cost once across a
13//! tile batch. The pool grows monotonically; it never shrinks.
14
15use crate::entropy::sequential::{PreparedDecodePlan, StripeBuffer};
16use alloc::vec::Vec;
17use j2k_core::ScratchPool as CoreScratchPool;
18
19#[derive(Debug, Default)]
20pub(crate) struct YCbCr420Rows {
21    pub(crate) cb_top: Vec<u8>,
22    pub(crate) cb_bot: Vec<u8>,
23    pub(crate) cr_top: Vec<u8>,
24    pub(crate) cr_bot: Vec<u8>,
25}
26
27impl YCbCr420Rows {
28    fn resize_width(&mut self, width: usize) {
29        self.cb_top.resize(width, 0);
30        self.cb_bot.resize(width, 0);
31        self.cr_top.resize(width, 0);
32        self.cr_bot.resize(width, 0);
33    }
34}
35
36#[derive(Debug, Default)]
37pub(crate) struct YCbCrGenericRows {
38    pub(crate) cb_up: Vec<u8>,
39    pub(crate) cr_up: Vec<u8>,
40}
41
42impl YCbCrGenericRows {
43    fn resize_width(&mut self, width: usize) {
44        self.cb_up.resize(width, 0);
45        self.cr_up.resize(width, 0);
46    }
47}
48
49#[derive(Debug, Default)]
50pub(crate) struct RgbGenericRows {
51    pub(crate) r: Vec<u8>,
52    pub(crate) g: Vec<u8>,
53    pub(crate) b: Vec<u8>,
54    pub(crate) k: Vec<u8>,
55}
56
57impl RgbGenericRows {
58    fn resize_width(&mut self, width: usize) {
59        self.r.resize(width, 0);
60        self.g.resize(width, 0);
61        self.b.resize(width, 0);
62        self.k.resize(width, 0);
63    }
64}
65
66#[derive(Debug, Default)]
67pub(crate) struct SinkRows {
68    pub(crate) top_row: Vec<u8>,
69    pub(crate) bottom_row: Vec<u8>,
70}
71
72impl SinkRows {
73    fn resize_width(&mut self, width: usize) {
74        let rgb_len = width.saturating_mul(3);
75        self.top_row.resize(rgb_len, 0);
76        self.bottom_row.resize(rgb_len, 0);
77    }
78}
79
80/// Pool of decoder-internal scratch buffers, reusable across many
81/// [`Decoder::decode_into_with_scratch`](crate::Decoder::decode_into_with_scratch)
82/// / [`decode_rows_with_scratch`](crate::Decoder::decode_rows_with_scratch)
83/// calls.
84#[derive(Debug, Default)]
85pub struct ScratchPool {
86    pub(crate) prev_dc: Vec<i32>,
87    pub(crate) stripe_a: StripeBuffer,
88    pub(crate) stripe_b: StripeBuffer,
89    pub(crate) stripe_c: StripeBuffer,
90    pub(crate) ycbcr_420_rows: YCbCr420Rows,
91    pub(crate) ycbcr_generic_rows: YCbCrGenericRows,
92    pub(crate) rgb_generic_rows: RgbGenericRows,
93    pub(crate) lossless_prev_row: Vec<u8>,
94    pub(crate) lossless_curr_row: Vec<u8>,
95    sink_rows: SinkRows,
96}
97
98impl ScratchPool {
99    /// Create an empty pool. The first decode that uses it pays the full
100    /// allocation cost; subsequent decodes at the same-or-smaller shape
101    /// reuse the underlying `Vec`s with zero allocations.
102    #[must_use]
103    pub fn new() -> Self {
104        Self::default()
105    }
106
107    /// Grow every internal scratch buffer to the shape required by `plan`
108    /// and zero the predictor so each decode starts clean.
109    pub(crate) fn prepare_for(
110        &mut self,
111        plan: &PreparedDecodePlan,
112        mcus_per_row: u32,
113        block_size: u32,
114    ) {
115        let n = plan.sampling.len();
116        if self.prev_dc.len() < n {
117            self.prev_dc.resize(n, 0);
118        }
119        for dc in &mut self.prev_dc[..n] {
120            *dc = 0;
121        }
122        self.stripe_a.resize_for(plan, mcus_per_row, block_size);
123        self.stripe_b.resize_for(plan, mcus_per_row, block_size);
124        self.stripe_c.resize_for(plan, mcus_per_row, block_size);
125        let denominator = 8 / block_size.max(1);
126        let width = plan.dimensions.0.div_ceil(denominator) as usize;
127        self.ycbcr_420_rows.resize_width(width);
128        self.ycbcr_generic_rows.resize_width(width);
129        self.rgb_generic_rows.resize_width(width);
130        self.sink_rows.resize_width(width);
131    }
132
133    pub(crate) fn take_sink_rows(&mut self, width: usize) -> SinkRows {
134        let mut rows = core::mem::take(&mut self.sink_rows);
135        rows.resize_width(width);
136        rows
137    }
138
139    pub(crate) fn restore_sink_rows(&mut self, rows: SinkRows) {
140        self.sink_rows = rows;
141    }
142}
143
144impl CoreScratchPool for ScratchPool {
145    fn bytes_allocated(&self) -> usize {
146        fn vec_bytes<T>(vec: &Vec<T>) -> usize {
147            vec.capacity().saturating_mul(core::mem::size_of::<T>())
148        }
149
150        fn stripe_bytes(stripe: &StripeBuffer) -> usize {
151            let mut total = 0usize;
152            for plane in &stripe.planes {
153                total = total.saturating_add(vec_bytes(plane));
154            }
155            total = total
156                .saturating_add(vec_bytes(&stripe.plane_strides))
157                .saturating_add(vec_bytes(&stripe.plane_rows));
158            total
159        }
160
161        let mut total = vec_bytes(&self.prev_dc);
162        total = total
163            .saturating_add(stripe_bytes(&self.stripe_a))
164            .saturating_add(stripe_bytes(&self.stripe_b))
165            .saturating_add(stripe_bytes(&self.stripe_c))
166            .saturating_add(vec_bytes(&self.ycbcr_420_rows.cb_top))
167            .saturating_add(vec_bytes(&self.ycbcr_420_rows.cb_bot))
168            .saturating_add(vec_bytes(&self.ycbcr_420_rows.cr_top))
169            .saturating_add(vec_bytes(&self.ycbcr_420_rows.cr_bot))
170            .saturating_add(vec_bytes(&self.ycbcr_generic_rows.cb_up))
171            .saturating_add(vec_bytes(&self.ycbcr_generic_rows.cr_up))
172            .saturating_add(vec_bytes(&self.rgb_generic_rows.r))
173            .saturating_add(vec_bytes(&self.rgb_generic_rows.g))
174            .saturating_add(vec_bytes(&self.rgb_generic_rows.b))
175            .saturating_add(vec_bytes(&self.rgb_generic_rows.k))
176            .saturating_add(vec_bytes(&self.lossless_prev_row))
177            .saturating_add(vec_bytes(&self.lossless_curr_row))
178            .saturating_add(vec_bytes(&self.sink_rows.top_row))
179            .saturating_add(vec_bytes(&self.sink_rows.bottom_row));
180        total
181    }
182
183    fn reset(&mut self) {
184        fn clear_stripe(stripe: &mut StripeBuffer) {
185            for plane in &mut stripe.planes {
186                plane.clear();
187            }
188            stripe.plane_strides.clear();
189            stripe.plane_rows.clear();
190        }
191
192        self.prev_dc.clear();
193        clear_stripe(&mut self.stripe_a);
194        clear_stripe(&mut self.stripe_b);
195        clear_stripe(&mut self.stripe_c);
196        self.ycbcr_420_rows.cb_top.clear();
197        self.ycbcr_420_rows.cb_bot.clear();
198        self.ycbcr_420_rows.cr_top.clear();
199        self.ycbcr_420_rows.cr_bot.clear();
200        self.ycbcr_generic_rows.cb_up.clear();
201        self.ycbcr_generic_rows.cr_up.clear();
202        self.rgb_generic_rows.r.clear();
203        self.rgb_generic_rows.g.clear();
204        self.rgb_generic_rows.b.clear();
205        self.rgb_generic_rows.k.clear();
206        self.lossless_prev_row.clear();
207        self.lossless_curr_row.clear();
208        self.sink_rows.top_row.clear();
209        self.sink_rows.bottom_row.clear();
210    }
211}
j2k_jpeg/internal/scratch.rs

j2k_jpeg/internal/
scratch.rs
