use alloc::vec::Vec;
use super::allocation::{checked_include_elements, code_block_grid};
use super::{Decomposition, Precinct, SubBand, SubBandType};
use crate::error::{DecodingError, Result, ValidationError};
use crate::j2c::codestream::{Header, ProgressionOrder};
use crate::j2c::decode::{
DecodeAllocationBudget, DecompositionStorage, OutputRegion, TileDecompositions,
};
use crate::j2c::progression::{progression_iterator, ProgressionData};
use crate::j2c::rect::IntRect;
use crate::j2c::roi::RoiPlan;
use crate::j2c::tag_tree::TagNode;
use crate::j2c::tile::Tile;
use crate::DEFAULT_MAX_DECODE_BYTES;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub(crate) enum PrecinctUse {
Unused,
Parsed,
Required,
}
pub(crate) struct PrecinctSelection {
uses: Vec<PrecinctUse>,
runs: Vec<usize>,
component_runs: Vec<usize>,
records: Vec<u32>,
record_runs: Vec<usize>,
skip_by_length: bool,
parsed_in_scratch: bool,
required_packets: usize,
owned_bytes: usize,
scanned_packets: usize,
released_storage_bytes: usize,
}
impl PrecinctSelection {
pub(crate) fn run(&self, run: usize) -> Option<&[PrecinctUse]> {
let start = *self.runs.get(run)?;
let end = self.runs.get(run + 1).copied().unwrap_or(self.uses.len());
self.uses.get(start..end)
}
pub(crate) fn packet_use(&self, packet: &ProgressionData) -> Option<PrecinctUse> {
self.index(packet).map(|index| self.uses[index])
}
pub(crate) const fn skip_by_length(&self) -> bool {
self.skip_by_length
}
pub(crate) const fn parsed_in_scratch(&self) -> bool {
self.parsed_in_scratch
}
pub(crate) const fn required_packets(&self) -> usize {
self.required_packets
}
pub(crate) const fn owned_bytes(&self) -> usize {
self.owned_bytes
}
pub(crate) const fn scanned_packets(&self) -> usize {
self.scanned_packets
}
pub(crate) const fn released_storage_bytes(&self) -> usize {
self.released_storage_bytes
}
pub(crate) fn record_offset(&self, packet: &ProgressionData) -> Option<usize> {
let precinct = usize::try_from(packet.precinct).ok()?;
if self.record_runs.is_empty() {
return Some(precinct);
}
let run = self
.component_runs
.get(usize::from(packet.component))?
.checked_add(usize::from(packet.resolution))?;
let start = *self.record_runs.get(run)?;
let end = self
.record_runs
.get(run + 1)
.copied()
.unwrap_or(self.records.len());
if start == end {
return self
.packet_use(packet)
.filter(|usage| *usage != PrecinctUse::Unused)
.map(|_| precinct);
}
let number = u32::try_from(packet.precinct).ok()?;
self.records.get(start..end)?.binary_search(&number).ok()
}
fn index(&self, packet: &ProgressionData) -> Option<usize> {
let run = self
.component_runs
.get(usize::from(packet.component))?
.checked_add(usize::from(packet.resolution))?;
let start = *self.runs.get(run)?;
let end = self.runs.get(run + 1).copied().unwrap_or(self.uses.len());
let index = start.checked_add(usize::try_from(packet.precinct).ok()?)?;
(index < end).then_some(index)
}
}
struct SelectionCounts {
components: usize,
decompositions: usize,
sub_bands: usize,
runs: usize,
run_precincts: usize,
sub_band_precincts: usize,
}
pub(crate) fn select_region_precincts(
tile: &Tile<'_>,
header: &Header<'_>,
storage: &mut DecompositionStorage<'_>,
output_region: OutputRegion,
retained_baseline_bytes: usize,
) -> Result<Option<PrecinctSelection>> {
let released_storage_bytes = storage.retained_capacity_bytes()?;
let exact_integer_decode = storage.exact_integer_decode;
storage.release_all_allocations();
storage.exact_integer_decode = exact_integer_decode;
let counts = count(tile)?;
let owned_bytes = selection_owned_bytes(&counts)?;
preflight_unavoidable_minimum(&counts, retained_baseline_bytes, owned_bytes)?;
let mut budget = DecodeAllocationBudget::from_live_bytes(retained_baseline_bytes)?;
build_skeleton(tile, &counts, storage, &mut budget)?;
storage.roi_plan = RoiPlan::build(tile, header, storage, output_region, &mut budget)?;
let selection = match storage.roi_plan.as_ref() {
Some(plan) => Some(select(
tile,
storage,
plan,
&counts,
&mut budget,
released_storage_bytes,
)?),
None => None,
};
storage.sub_bands.clear();
storage.decompositions.clear();
storage.tile_decompositions.clear();
Ok(selection)
}
fn count(tile: &Tile<'_>) -> Result<SelectionCounts> {
let mut counts = SelectionCounts {
components: 0,
decompositions: 0,
sub_bands: 0,
runs: 0,
run_precincts: 0,
sub_band_precincts: 0,
};
let add = |total: &mut usize, count: usize| -> Result<()> {
*total = total
.checked_add(count)
.ok_or(ValidationError::ImageTooLarge)?;
Ok(())
};
for component_tile in tile.component_tiles() {
add(&mut counts.components, 1)?;
for resolution_tile in component_tile.resolution_tiles() {
let precincts = usize::try_from(resolution_tile.num_precincts())
.map_err(|_| ValidationError::ImageTooLarge)?;
let sub_bands = if resolution_tile.resolution == 0 {
1
} else {
add(&mut counts.decompositions, 1)?;
3
};
add(&mut counts.sub_bands, sub_bands)?;
add(&mut counts.runs, 1)?;
add(&mut counts.run_precincts, precincts)?;
add(
&mut counts.sub_band_precincts,
precincts
.checked_mul(sub_bands)
.ok_or(ValidationError::ImageTooLarge)?,
)?;
}
}
Ok(counts)
}
fn selection_owned_bytes(counts: &SelectionCounts) -> Result<usize> {
let mut bytes = 0usize;
checked_include_elements::<PrecinctUse>(&mut bytes, counts.run_precincts)?;
checked_include_elements::<usize>(&mut bytes, counts.runs)?;
checked_include_elements::<usize>(&mut bytes, counts.components)?;
Ok(bytes)
}
fn preflight_unavoidable_minimum(
counts: &SelectionCounts,
retained_baseline_bytes: usize,
owned_bytes: usize,
) -> Result<()> {
if retained_baseline_bytes > DEFAULT_MAX_DECODE_BYTES {
return Err(ValidationError::ImageTooLarge.into());
}
let mut bytes = retained_baseline_bytes;
let mut include_bytes = |count: usize| -> Result<()> {
bytes = bytes
.checked_add(count)
.filter(|bytes| *bytes <= DEFAULT_MAX_DECODE_BYTES)
.ok_or(ValidationError::ImageTooLarge)?;
Ok(())
};
include_bytes(owned_bytes)?;
let mut structure = 0usize;
checked_include_elements::<TileDecompositions>(&mut structure, counts.components)?;
checked_include_elements::<Decomposition>(&mut structure, counts.decompositions)?;
checked_include_elements::<SubBand>(&mut structure, counts.sub_bands)?;
checked_include_elements::<Option<IntRect>>(
&mut structure,
counts
.sub_bands
.checked_add(counts.decompositions)
.and_then(|count| count.checked_add(counts.components))
.ok_or(ValidationError::ImageTooLarge)?,
)?;
checked_include_elements::<Precinct>(&mut structure, counts.sub_band_precincts)?;
checked_include_elements::<TagNode>(
&mut structure,
counts
.sub_band_precincts
.checked_mul(2)
.ok_or(ValidationError::ImageTooLarge)?,
)?;
include_bytes(structure)
}
fn build_skeleton(
tile: &Tile<'_>,
counts: &SelectionCounts,
storage: &mut DecompositionStorage<'_>,
budget: &mut DecodeAllocationBudget,
) -> Result<()> {
budget.reserve_new(&mut storage.sub_bands, counts.sub_bands)?;
budget.reserve_new(&mut storage.decompositions, counts.decompositions)?;
budget.reserve_new(&mut storage.tile_decompositions, counts.components)?;
for component_tile in tile.component_tiles() {
let start = storage.decompositions.len();
let mut resolution_tiles = component_tile.resolution_tiles();
let ll = resolution_tiles
.next()
.ok_or(DecodingError::InvalidPrecinct)?;
let first_ll_sub_band =
push_skeleton_sub_band(storage, ll.sub_band_rect(SubBandType::LowLow))?;
for resolution_tile in resolution_tiles {
let mut sub_bands = [0; 3];
for (slot, sub_band_type) in sub_bands.iter_mut().zip([
SubBandType::HighLow,
SubBandType::LowHigh,
SubBandType::HighHigh,
]) {
*slot =
push_skeleton_sub_band(storage, resolution_tile.sub_band_rect(sub_band_type))?;
}
push_within_capacity(
&mut storage.decompositions,
Decomposition {
sub_bands,
rect: resolution_tile.rect,
},
)?;
}
push_within_capacity(
&mut storage.tile_decompositions,
TileDecompositions {
decompositions: start..storage.decompositions.len(),
first_ll_sub_band,
},
)?;
}
Ok(())
}
fn push_skeleton_sub_band(storage: &mut DecompositionStorage<'_>, rect: IntRect) -> Result<usize> {
let idx = storage.sub_bands.len();
push_within_capacity(
&mut storage.sub_bands,
SubBand {
sub_band_type: SubBandType::LowLow,
rect,
precincts: 0..0,
coefficients: 0..0,
storage_rect: rect,
},
)?;
Ok(idx)
}
fn push_within_capacity<T>(values: &mut Vec<T>, value: T) -> Result<()> {
if values.len() == values.capacity() {
return Err(DecodingError::InvalidPrecinct.into());
}
values.push(value);
Ok(())
}
fn select(
tile: &Tile<'_>,
storage: &DecompositionStorage<'_>,
plan: &RoiPlan,
counts: &SelectionCounts,
budget: &mut DecodeAllocationBudget,
released_storage_bytes: usize,
) -> Result<PrecinctSelection> {
let mut uses = Vec::new();
budget.reserve_new(&mut uses, counts.run_precincts)?;
let mut runs = Vec::new();
budget.reserve_new(&mut runs, counts.runs)?;
let mut component_runs = Vec::new();
budget.reserve_new(&mut component_runs, counts.components)?;
let layers = usize::from(tile.num_layers);
let mut required_packets = 0usize;
for (component_tile, tile_decompositions) in
tile.component_tiles().zip(&storage.tile_decompositions)
{
push_within_capacity(&mut component_runs, runs.len())?;
for (resolution, resolution_tile) in component_tile.resolution_tiles().enumerate() {
push_within_capacity(&mut runs, uses.len())?;
let (sub_band_types, sub_band_indices): (&[SubBandType], [usize; 3]) =
if resolution == 0 {
(
&[SubBandType::LowLow],
[tile_decompositions.first_ll_sub_band; 3],
)
} else {
let decomposition = storage
.decompositions
.get(tile_decompositions.decompositions.start + resolution - 1)
.ok_or(DecodingError::InvalidPrecinct)?;
(
&[
SubBandType::HighLow,
SubBandType::LowHigh,
SubBandType::HighHigh,
],
decomposition.sub_bands,
)
};
for precinct in resolution_tile
.precincts()
.ok_or(DecodingError::InvalidPrecinct)?
{
let mut required = false;
for (sub_band_type, sub_band_idx) in sub_band_types.iter().zip(sub_band_indices) {
let Some(window) = plan.sub_band_window(sub_band_idx) else {
continue;
};
let rect = resolution_tile.sub_band_rect(*sub_band_type);
let grid = code_block_grid(&resolution_tile, rect, precinct.rect)?;
if grid.columns > 0
&& grid.rows > 0
&& window.intersects(grid.area.intersect(rect))
{
required = true;
break;
}
}
if required {
required_packets = required_packets
.checked_add(layers)
.ok_or(ValidationError::ImageTooLarge)?;
}
push_within_capacity(
&mut uses,
if required {
PrecinctUse::Required
} else {
PrecinctUse::Unused
},
)?;
}
}
}
let skip_by_length = packets_skippable_by_length(tile, counts, layers)?;
let mut selection = PrecinctSelection {
uses,
runs,
component_runs,
records: Vec::new(),
record_runs: Vec::new(),
skip_by_length,
parsed_in_scratch: !skip_by_length && reads_precinct_layers_in_a_row(tile),
required_packets,
owned_bytes: 0,
scanned_packets: 0,
released_storage_bytes,
};
if !selection.skip_by_length {
mark_precincts_parsed_before_last_required(tile, &mut selection, budget.live_bytes())?;
}
index_records(&mut selection, budget)?;
Ok(selection)
}
fn index_records(
selection: &mut PrecinctSelection,
budget: &mut DecodeAllocationBudget,
) -> Result<()> {
let mut count = 0;
let mut sparse = false;
for run in 0..selection.runs.len() {
let uses = selection.run(run).ok_or(DecodingError::InvalidPrecinct)?;
if uses.contains(&PrecinctUse::Unused) {
sparse = true;
count += uses
.iter()
.filter(|usage| **usage != PrecinctUse::Unused)
.count();
}
}
if sparse {
budget.reserve_new(&mut selection.records, count)?;
budget.reserve_new(&mut selection.record_runs, selection.runs.len())?;
for (run, &start) in selection.runs.iter().enumerate() {
push_within_capacity(&mut selection.record_runs, selection.records.len())?;
let end = selection
.runs
.get(run + 1)
.copied()
.unwrap_or(selection.uses.len());
let uses = selection
.uses
.get(start..end)
.ok_or(DecodingError::InvalidPrecinct)?;
if !uses.contains(&PrecinctUse::Unused) {
continue;
}
for (number, precinct_use) in uses.iter().enumerate() {
if *precinct_use != PrecinctUse::Unused {
let number =
u32::try_from(number).map_err(|_| ValidationError::ImageTooLarge)?;
push_within_capacity(&mut selection.records, number)?;
}
}
}
}
let mut bytes = 0;
checked_include_elements::<PrecinctUse>(&mut bytes, selection.uses.capacity())?;
checked_include_elements::<usize>(&mut bytes, selection.runs.capacity())?;
checked_include_elements::<usize>(&mut bytes, selection.component_runs.capacity())?;
checked_include_elements::<u32>(&mut bytes, selection.records.capacity())?;
checked_include_elements::<usize>(&mut bytes, selection.record_runs.capacity())?;
selection.owned_bytes = bytes;
Ok(())
}
fn packets_skippable_by_length(
tile: &Tile<'_>,
counts: &SelectionCounts,
layers: usize,
) -> Result<bool> {
if !tile.progression_changes.is_empty() {
return Ok(false);
}
let mut lengths = 0usize;
for tile_part in &tile.tile_parts {
let Some(count) = tile_part.skippable_packet_count() else {
return Ok(false);
};
lengths = lengths
.checked_add(count)
.ok_or(ValidationError::ImageTooLarge)?;
}
let packets = counts
.run_precincts
.checked_mul(layers)
.ok_or(ValidationError::ImageTooLarge)?;
Ok(lengths == packets)
}
fn reads_precinct_layers_in_a_row(tile: &Tile<'_>) -> bool {
tile.progression_changes.is_empty()
&& matches!(
tile.progression_order,
ProgressionOrder::ResolutionPositionComponentLayer
| ProgressionOrder::PositionComponentResolutionLayer
| ProgressionOrder::ComponentPositionResolutionLayer
)
}
fn mark_precincts_parsed_before_last_required(
tile: &Tile<'_>,
selection: &mut PrecinctSelection,
retained_live_bytes: usize,
) -> Result<()> {
let mut remaining = selection.required_packets;
if remaining == 0 {
return Ok(());
}
let parseable_packets = tile
.tile_parts
.iter()
.try_fold(0usize, |total, tile_part| {
total.checked_add(tile_part.packet_header_bytes())
})
.ok_or(ValidationError::ImageTooLarge)?;
for packet in progression_iterator(tile, retained_live_bytes)?.take(parseable_packets) {
selection.scanned_packets += 1;
let Some(index) = selection.index(&packet) else {
continue;
};
match selection.uses[index] {
PrecinctUse::Required => {
remaining = remaining.saturating_sub(1);
if remaining == 0 {
break;
}
}
PrecinctUse::Unused => selection.uses[index] = PrecinctUse::Parsed,
PrecinctUse::Parsed => {}
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{encode, DecodeSettings, EncodeOptions, EncodeProgressionOrder, Image};
#[test]
fn mixed_runs_index_only_precincts_in_sparse_runs() {
use PrecinctUse::{Parsed, Required, Unused};
let mut selection = PrecinctSelection {
uses: vec![Required, Parsed, Required, Unused, Required, Unused, Unused],
runs: vec![0, 3, 5],
component_runs: vec![0, 2],
records: Vec::new(),
record_runs: Vec::new(),
skip_by_length: false,
parsed_in_scratch: false,
required_packets: 3,
owned_bytes: 0,
scanned_packets: 0,
released_storage_bytes: 0,
};
let cap = core::mem::size_of::<u32>() + 3 * core::mem::size_of::<usize>();
let mut budget = DecodeAllocationBudget::from_live_bytes_with_cap(0, cap).unwrap();
index_records(&mut selection, &mut budget).expect("only the sparse run is indexed");
for (component, resolution, expected) in [
(0, 0, vec![Some(0), Some(1), Some(2)]),
(0, 1, vec![None, Some(0)]),
(1, 0, vec![None, None]),
] {
for (precinct, offset) in expected.into_iter().enumerate() {
let packet = ProgressionData {
component,
resolution,
precinct: precinct as u64,
layer_num: 0,
};
assert_eq!(selection.record_offset(&packet), offset, "{packet:?}");
}
}
}
#[test]
fn dense_region_selection_keeps_the_existing_budget_boundary() {
const WIDTH: u32 = 120;
const HEIGHT: u32 = 64;
for use_ht_block_coding in [false, true] {
for (order, write_plt, roi) in [
(EncodeProgressionOrder::Lrcp, false, (0, 0, WIDTH, HEIGHT)),
(EncodeProgressionOrder::Lrcp, true, (0, 0, WIDTH, HEIGHT)),
(
EncodeProgressionOrder::Lrcp,
false,
(WIDTH - 1, HEIGHT - 1, 1, 1),
),
(
EncodeProgressionOrder::Rpcl,
false,
(WIDTH - 1, HEIGHT - 1, 1, 1),
),
] {
let options = EncodeOptions {
reversible: true,
num_decomposition_levels: 0,
code_block_width_exp: 0,
code_block_height_exp: 0,
precinct_exponents: vec![(2, 2)],
progression_order: order,
write_plt,
use_ht_block_coding,
..EncodeOptions::default()
};
let bytes = encode(
&vec![128; (WIDTH * HEIGHT) as usize],
WIDTH,
HEIGHT,
1,
8,
false,
&options,
)
.expect("encode fixture");
let image = Image::new(&bytes, &DecodeSettings::default()).expect("parse fixture");
let tiles = crate::j2c::tile::parse(
&mut crate::reader::BitReader::new(image.codestream),
&image.header,
0,
)
.expect("parse tile");
let tile = &tiles[0];
let counts = count(tile).unwrap();
let mut storage = DecompositionStorage::default();
let mut budget = DecodeAllocationBudget::from_live_bytes(0).unwrap();
build_skeleton(tile, &counts, &mut storage, &mut budget).unwrap();
let plan = RoiPlan::build(
tile,
&image.header,
&storage,
OutputRegion::from_tuple(roi),
&mut budget,
)
.unwrap()
.expect("region plan");
let live = budget.live_bytes();
let cap = live + selection_owned_bytes(&counts).unwrap();
let mut exact =
DecodeAllocationBudget::from_live_bytes_with_cap(live, cap).unwrap();
let selection = select(tile, &storage, &plan, &counts, &mut exact, 0)
.expect("dense selection fits the existing cap");
assert_eq!(exact.live_bytes(), cap);
assert_eq!(selection.owned_bytes(), cap - live);
for packet in progression_iterator(tile, 0).unwrap() {
assert_eq!(
selection.record_offset(&packet),
usize::try_from(packet.precinct).ok()
);
}
let mut short =
DecodeAllocationBudget::from_live_bytes_with_cap(live, cap - 1).unwrap();
assert!(matches!(
select(tile, &storage, &plan, &counts, &mut short, 0),
Err(crate::DecodeError::Validation(
ValidationError::ImageTooLarge
))
));
assert_eq!(
image.decode_native_region(roi).unwrap().data,
vec![128; (roi.2 * roi.3) as usize]
);
}
}
}
}