Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
use std::{
ops::DerefMut,
sync::atomic::{AtomicUsize, Ordering},
};
use crate::{
error::Result,
frame::{
DataStatus,
modular::{ChannelInfo, IMAGE_OFFSET, IMAGE_PADDING},
},
headers::bit_depth::BitDepth,
image::Image,
util::{AtomicRefCell, AtomicRefMut},
};
use super::ModularBufferInfo;
// All the information on a specific buffer needed by Modular decoding.
#[derive(Debug)]
pub(super) struct ModularChannel {
// Actual pixel buffer.
pub(super) data: Image<i32>,
// Holds additional information such as the weighted predictor's error channel's last row for
// the transform chunk that produced this buffer.
pub(super) auxiliary_data: Option<Image<i32>>,
// Shift of the channel (None if this is a meta-channel).
pub(super) shift: Option<(usize, usize)>,
pub(super) bit_depth: BitDepth,
}
impl ModularChannel {
pub fn new(size: (usize, usize), bit_depth: BitDepth) -> Result<Self> {
Self::new_with_shift(size, Some((0, 0)), bit_depth)
}
pub fn new_with_shift(
size: (usize, usize),
shift: Option<(usize, usize)>,
bit_depth: BitDepth,
) -> Result<Self> {
Ok(ModularChannel {
data: Image::new_with_padding(size, IMAGE_OFFSET, IMAGE_PADDING)?,
auxiliary_data: None,
shift,
bit_depth,
})
}
fn try_clone(&self) -> Result<Self> {
Ok(ModularChannel {
data: self.data.try_clone()?,
auxiliary_data: self
.auxiliary_data
.as_ref()
.map(Image::try_clone)
.transpose()?,
shift: self.shift,
bit_depth: self.bit_depth,
})
}
pub fn channel_info(&self) -> ChannelInfo {
ChannelInfo {
output_channel_idx: None,
size: self.data.size(),
shift: self.shift,
bit_depth: self.bit_depth,
}
}
}
#[derive(Debug)]
pub(super) struct ModularBuffer {
pub(super) data: AtomicRefCell<Option<ModularChannel>>,
// Number of times this buffer will be used, *including* when it is used for output.
pub(super) remaining_uses: AtomicUsize,
// Transform steps that use the image data in this buffer for final renders.
pub(super) used_by_transforms_final: Vec<usize>,
// Transform steps that depend on this buffer for the current rendering pass.
pub(super) used_by_transforms_current: Vec<usize>,
// Transform step that will produce this channel (None if the channel is final).
pub(super) produced_by_step: Option<usize>,
pub(super) size: (usize, usize),
// Status of the data in this buffer. Note that the distinction between "Zero"
// and "partial" is only meaningful for section0 coded buffers.
pub(super) data_status: DataStatus,
}
const DISABLE_MODULAR_BUFFER_DEALLOCATION_FOR_DEBUG: bool = false;
impl ModularBuffer {
pub fn new(size: (usize, usize)) -> Self {
ModularBuffer {
data: AtomicRefCell::new(None),
remaining_uses: AtomicUsize::new(0),
used_by_transforms_final: vec![],
used_by_transforms_current: vec![],
size,
data_status: DataStatus::Zero,
produced_by_step: None,
}
}
pub fn has_buffer(&self) -> bool {
self.data.borrow().is_some()
}
pub fn make_buffer(&self, info: &ChannelInfo) -> Result<ModularChannel> {
Ok(ModularChannel {
data: Image::new_with_padding(self.size, IMAGE_OFFSET, IMAGE_PADDING)?,
auxiliary_data: None,
shift: info.shift,
bit_depth: info.bit_depth,
})
}
pub fn ensure_buffer(&self, info: &ChannelInfo) -> Result<()> {
if !self.has_buffer() {
let buf = self.make_buffer(info)?;
*self.data.borrow_mut() = Some(buf);
}
Ok(())
}
// Gives out a copy of the buffer + auxiliary buffer, marking the buffer as used.
// If this was the last usage of the buffer, does not actually copy the buffer.
pub fn get_buffer(&self, can_consume: bool) -> Result<ModularChannel> {
if !can_consume || DISABLE_MODULAR_BUFFER_DEALLOCATION_FOR_DEBUG {
return ModularChannel::try_clone(self.data.borrow().as_ref().unwrap());
}
let mut ret = None;
let _ = self.remaining_uses.fetch_update(
Ordering::Release,
Ordering::Acquire,
|remaining_pre| {
let remaining = remaining_pre.checked_sub(1).unwrap();
if ret.is_none() {
if remaining == 0 {
ret = Some(Ok(self.data.borrow_mut().take().unwrap()))
} else {
ret = self.data.borrow().as_ref().map(ModularChannel::try_clone);
}
} else if remaining == 0 {
*self.data.borrow_mut() = None;
}
Some(remaining)
},
);
Ok(ret.transpose()?.unwrap())
}
pub fn mark_used(&self, can_consume: bool) {
if !can_consume || DISABLE_MODULAR_BUFFER_DEALLOCATION_FOR_DEBUG {
return;
}
let _ = self.remaining_uses.fetch_update(
Ordering::Release,
Ordering::Acquire,
|remaining_pre: usize| {
let remaining = remaining_pre.checked_sub(1).unwrap();
if remaining == 0 {
*self.data.borrow_mut() = None;
}
Some(remaining)
},
);
}
}
pub fn with_buffers<T>(
buffers: &[ModularBufferInfo],
indices: &[usize],
grid: usize,
f: impl FnOnce(Vec<&mut ModularChannel>) -> Result<T>,
) -> Result<T> {
let mut bufs = vec![];
for i in indices {
// Allocate buffers if they are not present.
let buf = &buffers[*i];
let b = &buf.buffer_grid[grid];
b.ensure_buffer(&buf.info)?;
let data = b.data.borrow_mut();
// Skip zero-sized *tiles*.
//
// Note that some bitstreams can contain channels with one dimension being 0 (e.g. palette
// meta-channel with 0 colors has size (0, 3)). Those must still participate in channel
// numbering (but carry no entropy-coded pixels), so we only skip when both dimensions are 0.
// TODO(veluca): figure out if this is the best approach or we should instead pass through
// empty buffers.
if b.size.0 == 0 && b.size.1 == 0 {
continue;
}
bufs.push(AtomicRefMut::map(data, |x| x.as_mut().unwrap()));
}
f(bufs.iter_mut().map(|x| x.deref_mut()).collect())
}