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::fmt::Debug;
use num_traits::FromPrimitive;
use crate::{
error::{Error, Result},
frame::modular::{
ChannelInfo, ModularBufferInfo, ModularChannel, ModularGridKind, Predictor,
borrowed_buffers::with_buffers,
},
headers::{
self,
frame_header::FrameHeader,
modular::{TransformId, WeightedHeader},
},
image::Rect,
util::{AtomicRef, AtomicRefMut, tracing_wrappers::*},
};
use std::ops::Deref;
use std::ops::DerefMut;
use super::{RctOp, RctPermutation};
#[derive(Debug, Clone)]
pub enum TransformStep {
Rct {
buf_in: [usize; 3],
buf_out: [usize; 3],
op: RctOp,
perm: RctPermutation,
},
Palette {
buf_in: usize,
buf_pal: usize,
buf_out: Vec<usize>,
num_colors: usize,
num_deltas: usize,
predictor: Predictor,
wp_header: WeightedHeader,
},
HSqueeze {
buf_in: [usize; 2],
buf_out: usize,
},
VSqueeze {
buf_in: [usize; 2],
buf_out: usize,
},
}
#[derive(Debug)]
pub struct TransformStepChunk {
pub(super) step: TransformStep,
// Grid position this transform should produce.
// Note that this is a lie for Palette with AverageAll or Weighted, as the transform with
// position (0, y) will produce the entire row of blocks (*, y) (and there will be no
// transforms with position (x, y) with x > 0).
pub(super) grid_pos: (usize, usize),
// Number of inputs that are not yet available.
pub(super) incomplete_deps: usize,
}
impl TransformStepChunk {
// Marks that one dependency of this transform is ready, and potentially runs the transform,
// returning the new buffers that are now ready.
#[instrument(level = "trace", skip_all)]
pub fn dep_ready(
&mut self,
frame_header: &FrameHeader,
buffers: &mut [ModularBufferInfo],
) -> Result<Vec<(usize, usize)>> {
self.incomplete_deps = self.incomplete_deps.checked_sub(1).unwrap();
if self.incomplete_deps > 0 {
trace!(
"skipping transform chunk because incomplete_deps = {}",
self.incomplete_deps
);
return Ok(vec![]);
}
let buf_out: &[usize] = match &self.step {
TransformStep::Rct { buf_out, .. } => buf_out,
TransformStep::Palette { buf_out, .. } => buf_out,
TransformStep::HSqueeze { buf_out, .. } | TransformStep::VSqueeze { buf_out, .. } => {
&[*buf_out]
}
};
let out_grid_kind = buffers[buf_out[0]].grid_kind;
let out_grid = buffers[buf_out[0]].get_grid_idx(out_grid_kind, self.grid_pos);
let out_size = buffers[buf_out[0]].info.size;
for bo in buf_out {
assert_eq!(out_grid_kind, buffers[*bo].grid_kind);
assert_eq!(out_size, buffers[*bo].info.size);
}
match &self.step {
TransformStep::Rct {
buf_in,
buf_out,
op,
perm,
} => {
for i in 0..3 {
assert_eq!(out_grid_kind, buffers[buf_in[i]].grid_kind);
assert_eq!(out_size, buffers[buf_in[i]].info.size);
// Optimistically move the buffers to the output if possible.
// If not, creates buffers in the output that are a copy of the input buffers.
// This should be rare.
*buffers[buf_out[i]].buffer_grid[out_grid].data.borrow_mut() =
Some(buffers[buf_in[i]].buffer_grid[out_grid].get_buffer()?);
}
with_buffers(buffers, buf_out, out_grid, false, |mut bufs| {
super::rct::do_rct_step(&mut bufs, *op, *perm);
Ok(())
})?;
Ok(buf_out.iter().map(|x| (*x, out_grid)).collect())
}
TransformStep::Palette {
buf_in,
buf_pal,
buf_out,
..
} if buffers[*buf_in].info.size.0 == 0 => {
// Nothing to do, just bookkeeping.
buffers[*buf_in].buffer_grid[out_grid].mark_used();
buffers[*buf_pal].buffer_grid[0].mark_used();
with_buffers(buffers, buf_out, out_grid, false, |_| Ok(()))?;
Ok(buf_out.iter().map(|x| (*x, out_grid)).collect())
}
TransformStep::Palette {
buf_in,
buf_pal,
buf_out,
num_colors,
num_deltas,
predictor,
..
} if !predictor.requires_full_row() => {
assert_eq!(out_grid_kind, buffers[*buf_in].grid_kind);
assert_eq!(out_size, buffers[*buf_in].info.size);
{
let img_in =
AtomicRef::map(buffers[*buf_in].buffer_grid[out_grid].data.borrow(), |x| {
x.as_ref().unwrap()
});
let img_pal =
AtomicRef::map(buffers[*buf_pal].buffer_grid[0].data.borrow(), |x| {
x.as_ref().unwrap()
});
// Ensure that the output buffers are present.
// TODO(szabadka): Extend the callback to support many grid points.
with_buffers(buffers, buf_out, out_grid, false, |_| Ok(()))?;
let grid_shape = buffers[buf_out[0]].grid_shape;
let grid_x = out_grid % grid_shape.0;
let grid_y = out_grid / grid_shape.0;
let border = if *predictor == Predictor::Zero { 0 } else { 1 };
let grid_x0 = grid_x.saturating_sub(border);
let grid_y0 = grid_y.saturating_sub(border);
let grid_x1 = grid_x + 1;
let grid_y1 = grid_y + 1;
let mut out_bufs = vec![];
for i in buf_out {
for gy in grid_y0..grid_y1 {
for gx in grid_x0..grid_x1 {
let grid = gy * grid_shape.0 + gx;
let buf = &buffers[*i];
let b = &buf.buffer_grid[grid];
let data = b.data.borrow_mut();
out_bufs.push(AtomicRefMut::map(data, |x| x.as_mut().unwrap()));
}
}
}
let mut out_buf_refs: Vec<&mut ModularChannel> =
out_bufs.iter_mut().map(|x| x.deref_mut()).collect();
super::palette::do_palette_step_one_group(
&img_in,
&img_pal,
&mut out_buf_refs,
grid_x - grid_x0,
grid_y - grid_y0,
grid_x1 - grid_x0,
grid_y1 - grid_y0,
*num_colors,
*num_deltas,
*predictor,
);
}
buffers[*buf_in].buffer_grid[out_grid].mark_used();
buffers[*buf_pal].buffer_grid[0].mark_used();
Ok(buf_out.iter().map(|x| (*x, out_grid)).collect())
}
TransformStep::Palette {
buf_in,
buf_pal,
buf_out,
num_colors,
num_deltas,
predictor,
wp_header,
} => {
assert_eq!(out_grid_kind, buffers[*buf_in].grid_kind);
assert_eq!(out_size, buffers[*buf_in].info.size);
let mut generated_chunks = Vec::<(usize, usize)>::new();
let grid_shape = buffers[buf_out[0]].grid_shape;
{
assert_eq!(out_grid % grid_shape.0, 0);
let grid_y = out_grid / grid_shape.0;
let grid_y0 = grid_y.saturating_sub(1);
let grid_y1 = grid_y + 1;
let mut in_bufs = vec![];
for grid_x in 0..grid_shape.0 {
let grid = grid_y * grid_shape.0 + grid_x;
in_bufs.push(AtomicRef::map(
buffers[*buf_in].buffer_grid[grid].data.borrow(),
|x| x.as_ref().unwrap(),
));
// Ensure that the output buffers are present.
// TODO(szabadka): Extend the callback to support many grid points.
with_buffers(buffers, buf_out, out_grid + grid_x, false, |_| Ok(()))?;
}
let in_buf_refs: Vec<&ModularChannel> =
in_bufs.iter().map(|x| x.deref()).collect();
let img_pal =
AtomicRef::map(buffers[*buf_pal].buffer_grid[0].data.borrow(), |x| {
x.as_ref().unwrap()
});
let mut out_bufs = vec![];
for i in buf_out {
for grid_y in grid_y0..grid_y1 {
for grid_x in 0..grid_shape.0 {
let grid = grid_y * grid_shape.0 + grid_x;
let buf = &buffers[*i];
let b = &buf.buffer_grid[grid];
let data = b.data.borrow_mut();
out_bufs.push(AtomicRefMut::map(data, |x| x.as_mut().unwrap()));
}
}
}
let mut out_buf_refs: Vec<&mut ModularChannel> =
out_bufs.iter_mut().map(|x| x.deref_mut()).collect();
super::palette::do_palette_step_group_row(
&in_buf_refs,
&img_pal,
&mut out_buf_refs,
grid_y - grid_y0,
grid_shape.0,
*num_colors,
*num_deltas,
*predictor,
wp_header,
)?;
}
buffers[*buf_pal].buffer_grid[0].mark_used();
for grid_x in 0..grid_shape.0 {
buffers[*buf_in].buffer_grid[out_grid + grid_x].mark_used();
for buf in buf_out {
generated_chunks.push((*buf, out_grid + grid_x));
}
}
Ok(generated_chunks)
}
TransformStep::HSqueeze { buf_in, buf_out } => {
let buf_avg = &buffers[buf_in[0]];
let buf_res = &buffers[buf_in[1]];
let in_grid = buf_avg.get_grid_idx(out_grid_kind, self.grid_pos);
let res_grid = buf_res.get_grid_idx(out_grid_kind, self.grid_pos);
{
trace!(
"HSqueeze {:?} -> {:?}, grid {out_grid} grid pos {:?}",
buf_in, buf_out, self.grid_pos
);
let (gx, gy) = self.grid_pos;
let in_avg = AtomicRef::map(buf_avg.buffer_grid[in_grid].data.borrow(), |x| {
x.as_ref().unwrap()
});
let has_next = gx + 1 < buffers[*buf_out].grid_shape.0;
let gx_next = if has_next { gx + 1 } else { gx };
let next_avg_grid = buf_avg.get_grid_idx(out_grid_kind, (gx_next, gy));
let in_next_avg =
AtomicRef::map(buf_avg.buffer_grid[next_avg_grid].data.borrow(), |x| {
x.as_ref().unwrap()
});
let in_next_avg_rect = if has_next {
Some(in_next_avg.data.get_rect(buf_avg.get_grid_rect(
frame_header,
out_grid_kind,
(gx_next, gy),
)))
} else {
None
};
let in_res = AtomicRef::map(buf_res.buffer_grid[res_grid].data.borrow(), |x| {
x.as_ref().unwrap()
});
let out_prev = if gx == 0 {
None
} else {
let prev_out_grid =
buffers[*buf_out].get_grid_idx(out_grid_kind, (gx - 1, gy));
Some(AtomicRef::map(
buffers[*buf_out].buffer_grid[prev_out_grid].data.borrow(),
|x| x.as_ref().unwrap(),
))
};
with_buffers(buffers, &[*buf_out], out_grid, false, |mut bufs| {
super::squeeze::do_hsqueeze_step(
&in_avg.data.get_rect(buf_avg.get_grid_rect(
frame_header,
out_grid_kind,
(gx, gy),
)),
&in_res.data.get_rect(buf_res.get_grid_rect(
frame_header,
out_grid_kind,
(gx, gy),
)),
&in_next_avg_rect,
&out_prev,
&mut bufs,
);
Ok(())
})?;
}
buffers[buf_in[0]].buffer_grid[in_grid].mark_used();
buffers[buf_in[1]].buffer_grid[res_grid].mark_used();
Ok(vec![(*buf_out, out_grid)])
}
TransformStep::VSqueeze { buf_in, buf_out } => {
let buf_avg = &buffers[buf_in[0]];
let buf_res = &buffers[buf_in[1]];
let in_grid = buf_avg.get_grid_idx(out_grid_kind, self.grid_pos);
let res_grid = buf_res.get_grid_idx(out_grid_kind, self.grid_pos);
{
trace!(
"VSqueeze {:?} -> {:?} grid: {out_grid:?} grid pos: {:?}",
buf_in, buf_out, self.grid_pos
);
let (gx, gy) = self.grid_pos;
let in_avg = AtomicRef::map(buf_avg.buffer_grid[in_grid].data.borrow(), |x| {
x.as_ref().unwrap()
});
let has_next = gy + 1 < buffers[*buf_out].grid_shape.1;
let gy_next = if has_next { gy + 1 } else { gy };
let next_avg_grid = buf_avg.get_grid_idx(out_grid_kind, (gx, gy_next));
let in_next_avg =
AtomicRef::map(buf_avg.buffer_grid[next_avg_grid].data.borrow(), |x| {
x.as_ref().unwrap()
});
let in_next_avg_rect = if has_next {
Some(in_next_avg.data.get_rect(buf_avg.get_grid_rect(
frame_header,
out_grid_kind,
(gx, gy_next),
)))
} else {
None
};
let in_res = AtomicRef::map(buf_res.buffer_grid[res_grid].data.borrow(), |x| {
x.as_ref().unwrap()
});
let out_prev = if gy == 0 {
None
} else {
let prev_out_grid =
buffers[*buf_out].get_grid_idx(out_grid_kind, (gx, gy - 1));
Some(AtomicRef::map(
buffers[*buf_out].buffer_grid[prev_out_grid].data.borrow(),
|x| x.as_ref().unwrap(),
))
};
let avg_grid_rect =
buf_avg.get_grid_rect(frame_header, out_grid_kind, (gx, gy));
let res_grid_rect =
buf_res.get_grid_rect(frame_header, out_grid_kind, (gx, gy));
with_buffers(buffers, &[*buf_out], out_grid, false, |mut bufs| {
super::squeeze::do_vsqueeze_step(
&in_avg.data.get_rect(avg_grid_rect),
&in_res.data.get_rect(res_grid_rect),
&in_next_avg_rect,
&out_prev,
&mut bufs,
);
Ok(())
})?;
}
buffers[buf_in[0]].buffer_grid[in_grid].mark_used();
buffers[buf_in[1]].buffer_grid[res_grid].mark_used();
Ok(vec![(*buf_out, out_grid)])
}
}
}
}
#[instrument(level = "trace", err)]
fn check_equal_channels(
channels: &[(usize, ChannelInfo)],
first_chan: usize,
num: usize,
) -> Result<()> {
if first_chan + num > channels.len() {
return Err(Error::InvalidChannelRange(
first_chan,
first_chan + num,
channels.len(),
));
}
for inc in 1..num {
if !channels[first_chan]
.1
.is_equivalent(&channels[first_chan + inc].1)
{
return Err(Error::MixingDifferentChannels);
}
}
Ok(())
}
fn meta_apply_single_transform(
transform: &headers::modular::Transform,
header: &headers::modular::GroupHeader,
channels: &mut Vec<(usize, ChannelInfo)>,
transform_steps: &mut Vec<TransformStep>,
mut add_transform_buffer: impl FnMut(ChannelInfo, String) -> usize,
) -> Result<()> {
match transform.id {
TransformId::Rct => {
let begin_channel = transform.begin_channel as usize;
let op = RctOp::from_u32(transform.rct_type % 7).unwrap();
let perm = RctPermutation::from_u32(transform.rct_type / 7)
.expect("header decoding should ensure rct_type < 42");
check_equal_channels(channels, begin_channel, 3)?;
let mut buf_in = [0; 3];
let buf_out = [
channels[begin_channel].0,
channels[begin_channel + 1].0,
channels[begin_channel + 2].0,
];
for i in 0..3 {
let c = &mut channels[begin_channel + i];
let mut info = c.1;
info.output_channel_idx = -1;
c.0 = add_transform_buffer(
info,
format!(
"RCT (op {op:?} perm {perm:?}) starting at channel {begin_channel}, \
input {i}"
),
);
buf_in[i] = c.0;
}
transform_steps.push(TransformStep::Rct {
buf_out,
buf_in,
op,
perm,
});
trace!("applied RCT: {channels:?}");
}
TransformId::Squeeze => {
let steps = if transform.squeezes.is_empty() {
super::squeeze::default_squeeze(channels)
} else {
transform.squeezes.clone()
};
for step in steps {
super::squeeze::check_squeeze_params(channels, &step)?;
let in_place = step.in_place;
let horizontal = step.horizontal;
let begin_channel = step.begin_channel as usize;
let num_channels = step.num_channels as usize;
let end_channel = begin_channel + num_channels;
let new_chan_offset = if in_place {
end_channel
} else {
channels.len()
};
for ic in 0..num_channels {
let chan = &channels[begin_channel + ic].1;
let new_shift = if let Some(shift) = chan.shift {
if shift.0 > 30 || shift.1 > 30 {
return Err(Error::TooManySqueezes);
}
if horizontal {
Some((shift.0 + 1, shift.1))
} else {
Some((shift.0, shift.1 + 1))
}
} else {
None
};
let w = chan.size.0;
let h = chan.size.1;
let (new_size_0, new_size_1) = if horizontal {
((w.div_ceil(2), h), (w - w.div_ceil(2), h))
} else {
((w, h.div_ceil(2)), (w, h - h.div_ceil(2)))
};
let new_0 = ChannelInfo {
output_channel_idx: -1,
shift: new_shift,
size: new_size_0,
bit_depth: chan.bit_depth,
};
let buf_0 = add_transform_buffer(
new_0,
format!("Squeezed channel, original channel {}", begin_channel + ic),
);
let new_1 = ChannelInfo {
output_channel_idx: -1,
shift: new_shift,
size: new_size_1,
bit_depth: chan.bit_depth,
};
let buf_1 = add_transform_buffer(
new_1,
format!("Squeeze residual, original channel {}", begin_channel + ic),
);
if horizontal {
transform_steps.push(TransformStep::HSqueeze {
buf_in: [buf_0, buf_1],
buf_out: channels[begin_channel + ic].0,
});
} else {
transform_steps.push(TransformStep::VSqueeze {
buf_in: [buf_0, buf_1],
buf_out: channels[begin_channel + ic].0,
});
}
channels[begin_channel + ic] = (buf_0, new_0);
channels.insert(new_chan_offset + ic, (buf_1, new_1));
trace!("applied squeeze: {channels:?}");
}
}
}
TransformId::Palette => {
let begin_channel = transform.begin_channel as usize;
let num_channels = transform.num_channels as usize;
let num_colors = transform.num_colors as usize;
let num_deltas = transform.num_deltas as usize;
let pred = Predictor::from_u32(transform.predictor_id)
.expect("header decoding should ensure a valid predictor");
check_equal_channels(channels, begin_channel, num_channels)?;
// We already checked the bit_depth for all channels from `begin_channel` is
// equal in the line above.
let bit_depth = channels[begin_channel].1.bit_depth;
let pchan_info = ChannelInfo {
output_channel_idx: -1,
shift: None,
size: (num_colors + num_deltas, num_channels),
bit_depth,
};
let pchan = add_transform_buffer(
pchan_info,
format!(
"Palette for palette transform starting at channel {begin_channel} with \
{num_channels} channels"
),
);
let mut inchan_info = channels[begin_channel].1;
inchan_info.output_channel_idx = -1;
let inchan = add_transform_buffer(
inchan_info,
format!(
"Pixel data for palette transform starting at channel {begin_channel} with \
{num_channels} channels",
),
);
transform_steps.push(TransformStep::Palette {
buf_in: inchan,
buf_pal: pchan,
buf_out: channels[begin_channel..(begin_channel + num_channels)]
.iter()
.map(|x| x.0)
.collect(),
num_colors,
num_deltas,
predictor: pred,
wp_header: header.wp_header.clone(),
});
channels.drain(begin_channel + 1..begin_channel + num_channels);
channels[begin_channel].0 = inchan;
channels.insert(0, (pchan, pchan_info));
trace!("applied palette: {channels:?}");
}
TransformId::Invalid => {
unreachable!("header decoding for invalid transforms should fail");
}
}
Ok(())
}
#[instrument(level = "trace", ret)]
pub fn meta_apply_transforms(
channels: &[ChannelInfo],
header: &headers::modular::GroupHeader,
) -> Result<(Vec<ModularBufferInfo>, Vec<TransformStep>)> {
let mut buffer_info = vec![];
let mut transform_steps = vec![];
// (buffer id, channel info)
let mut channels: Vec<_> = channels.iter().cloned().enumerate().collect();
// First, add all the pre-transform channels to the buffer list.
for chan in channels.iter() {
buffer_info.push(ModularBufferInfo {
info: chan.1,
coded_channel_id: -1,
description: format!(
"Input channel {}, size {}x{}",
chan.0, chan.1.size.0, chan.1.size.1
),
// To be filled by make_grids.
grid_kind: ModularGridKind::None,
grid_shape: (0, 0),
buffer_grid: vec![],
});
}
let mut add_transform_buffer = |info, description| {
buffer_info.push(ModularBufferInfo {
info,
coded_channel_id: -1,
description,
// To be filled by make_grids.
grid_kind: ModularGridKind::None,
grid_shape: (0, 0),
buffer_grid: vec![],
});
buffer_info.len() - 1
};
// Apply transforms to the channel list.
for transform in &header.transforms {
meta_apply_single_transform(
transform,
header,
&mut channels,
&mut transform_steps,
&mut add_transform_buffer,
)?;
}
// All the channels left over at the end of applying transforms are the channels that are
// actually coded.
for (chid, chan) in channels.iter().enumerate() {
buffer_info[chan.0].coded_channel_id = chid as isize;
}
#[cfg(feature = "tracing")]
for (i, transform) in transform_steps.iter().enumerate() {
trace!("Transform step {i}: {transform:?}");
}
Ok((buffer_info, transform_steps))
}
#[derive(Debug)]
pub enum LocalTransformBuffer<'a> {
// This channel has been consumed by some transform.
Empty,
// This channel has not been written to yet.
Placeholder(ChannelInfo),
// Temporary, locally-allocated channel.
Owned(ModularChannel),
// Channel belonging to the global image.
Borrowed(&'a mut ModularChannel),
}
impl LocalTransformBuffer<'_> {
fn channel_info(&self) -> ChannelInfo {
match self {
LocalTransformBuffer::Empty => unreachable!("an empty buffer has no channel info"),
LocalTransformBuffer::Owned(m) => m.channel_info(),
LocalTransformBuffer::Placeholder(c) => *c,
LocalTransformBuffer::Borrowed(m) => m.channel_info(),
}
}
fn borrow_mut(&mut self) -> &mut ModularChannel {
match self {
LocalTransformBuffer::Owned(m) => m,
LocalTransformBuffer::Borrowed(m) => m,
LocalTransformBuffer::Empty => unreachable!("tried to borrow an empty channel"),
LocalTransformBuffer::Placeholder(_) => {
unreachable!("tried to borrow a placeholder channel")
}
}
}
fn take(&mut self) -> Self {
assert!(!matches!(self, LocalTransformBuffer::Empty));
let mut r = LocalTransformBuffer::Empty;
std::mem::swap(self, &mut r);
r
}
fn allocate_if_needed(&mut self) -> Result<()> {
if let LocalTransformBuffer::Placeholder(c) = self {
*self = LocalTransformBuffer::Owned(ModularChannel::new_with_shift(
c.size,
c.shift,
c.bit_depth,
)?);
}
Ok(())
}
}
#[instrument(level = "trace", ret)]
pub fn meta_apply_local_transforms<'a, 'b>(
channels_in: Vec<&'a mut ModularChannel>,
buffer_storage: &'b mut Vec<LocalTransformBuffer<'a>>,
header: &headers::modular::GroupHeader,
) -> Result<(Vec<&'b mut ModularChannel>, Vec<TransformStep>)> {
let mut transform_steps = vec![];
// (buffer id, channel info)
let mut channels: Vec<_> = channels_in
.iter()
.map(|x| x.channel_info())
.enumerate()
.collect();
debug!(?channels, "initial channels");
// First, add all the pre-transform channels to the buffer list.
buffer_storage.extend(channels_in.into_iter().map(LocalTransformBuffer::Borrowed));
#[allow(unused_variables)]
let mut add_transform_buffer = |info, description| {
trace!(description, ?info, "adding channel buffer");
buffer_storage.push(LocalTransformBuffer::Placeholder(info));
buffer_storage.len() - 1
};
// Apply transforms to the channel list.
for transform in &header.transforms {
meta_apply_single_transform(
transform,
header,
&mut channels,
&mut transform_steps,
&mut add_transform_buffer,
)?;
}
debug!(?channels, ?buffer_storage, "channels after transforms");
debug!(?transform_steps);
// Ensure that the buffer indices in `channels` appear in increasing order, by reordering them
// if necessary.
if !channels.iter().map(|x| x.0).is_sorted() {
let mut buf_new_position: Vec<_> = channels.iter().map(|x| x.0).collect();
buf_new_position.sort();
let buf_tmp: Vec<_> = channels
.iter()
.map(|x| {
let mut b = LocalTransformBuffer::Empty;
std::mem::swap(&mut b, &mut buffer_storage[x.0]);
b
})
.collect();
let mut buf_remap: Vec<_> = (0..buffer_storage.len()).collect();
for (new_pos, (ch_info, buf)) in buf_new_position
.iter()
.cloned()
.zip(channels.iter_mut().zip(buf_tmp.into_iter()))
{
assert!(matches!(
buffer_storage[new_pos],
LocalTransformBuffer::Empty
));
buf_remap[ch_info.0] = new_pos;
buffer_storage[new_pos] = buf;
ch_info.0 = new_pos;
}
for step in transform_steps.iter_mut() {
use std::iter::once;
match step {
TransformStep::Rct {
buf_in, buf_out, ..
} => {
for b in buf_in.iter_mut().chain(buf_out.iter_mut()) {
*b = buf_remap[*b];
}
}
TransformStep::Palette {
buf_in,
buf_pal,
buf_out,
..
} => {
for b in once(buf_in).chain(once(buf_pal)).chain(buf_out.iter_mut()) {
*b = buf_remap[*b];
}
}
TransformStep::HSqueeze { buf_in, buf_out }
| TransformStep::VSqueeze { buf_in, buf_out } => {
for b in once(buf_out).chain(buf_in.iter_mut()) {
*b = buf_remap[*b];
}
}
}
}
}
debug!(?channels, ?buffer_storage, "sorted channels");
debug!(?transform_steps);
// Since RCT steps will try to transfer buffers from the source channels to the destination
// channels, make sure we do the reverse transformation here (to have the caller-provided
// buffers be used for writing temporary data).
for ts in transform_steps.iter() {
if let TransformStep::Rct {
buf_in, buf_out, ..
} = ts
{
for c in 0..3 {
assert_eq!(
buffer_storage[buf_in[c]].channel_info(),
buffer_storage[buf_out[c]].channel_info()
);
assert!(matches!(
buffer_storage[buf_in[c]],
LocalTransformBuffer::Placeholder(_)
));
buffer_storage.swap(buf_in[c], buf_out[c]);
}
}
}
debug!(?channels, ?buffer_storage, "RCT-adjusted channels");
// Allocate all the coded channels if they aren't yet.
for (buf, _) in channels.iter() {
buffer_storage[*buf].allocate_if_needed()?;
}
debug!(?channels, ?buffer_storage, "allocated buffers");
// Extract references to to-be-decoded buffers.
let mut coded_buffers = Vec::with_capacity(channels.len());
let mut buffer_tail = &mut buffer_storage[..];
let mut last_buffer = None;
for (buf, _) in channels {
let offset = if let Some(lb) = last_buffer {
buf.checked_sub(lb).unwrap()
} else {
buf + 1
};
let cur_buf;
(cur_buf, buffer_tail) = buffer_tail.split_at_mut(offset);
coded_buffers.push(cur_buf.last_mut().unwrap().borrow_mut());
last_buffer = Some(buf);
}
Ok((coded_buffers, transform_steps))
}
impl TransformStep {
// Marks that one dependency of this transform is ready, and potentially runs the transform,
// returning the new buffers that are now ready.
pub fn local_apply(&self, buffers: &mut [LocalTransformBuffer]) -> Result<()> {
match self {
TransformStep::Rct {
buf_in,
buf_out,
op,
perm,
} => {
for i in 0..3 {
assert_eq!(
buffers[buf_in[i]].channel_info(),
buffers[buf_out[i]].channel_info()
);
}
let [mut a, mut b, mut c] = [
buffers[buf_in[0]].take(),
buffers[buf_in[1]].take(),
buffers[buf_in[2]].take(),
];
{
let mut bufs = [a.borrow_mut(), b.borrow_mut(), c.borrow_mut()];
super::rct::do_rct_step(&mut bufs, *op, *perm);
}
buffers[buf_out[0]] = a;
buffers[buf_out[1]] = b;
buffers[buf_out[2]] = c;
}
TransformStep::Palette {
buf_in,
buf_pal,
buf_out,
num_colors,
num_deltas,
predictor,
wp_header,
} => {
for b in buf_out.iter() {
assert_eq!(
buffers[*b].channel_info().size,
buffers[*buf_in].channel_info().size
);
buffers[*b].allocate_if_needed()?;
}
let mut img_in = buffers[*buf_in].take();
let mut img_pal = buffers[*buf_pal].take();
let mut out_bufs: Vec<_> = buf_out.iter().map(|x| buffers[*x].take()).collect();
{
let mut bufs: Vec<_> = out_bufs.iter_mut().map(|x| x.borrow_mut()).collect();
super::palette::do_palette_step_general(
img_in.borrow_mut(),
img_pal.borrow_mut(),
&mut bufs,
*num_colors,
*num_deltas,
*predictor,
wp_header,
);
}
for (pos, buf) in buf_out.iter().zip(out_bufs.into_iter()) {
buffers[*pos] = buf;
}
}
TransformStep::HSqueeze { buf_in, buf_out } => {
buffers[*buf_out].allocate_if_needed()?;
let mut out_buf = buffers[*buf_out].take();
let mut in_avg = buffers[buf_in[0]].take();
let mut in_res = buffers[buf_in[1]].take();
{
let mut bufs: Vec<_> = vec![out_buf.borrow_mut()];
let in_avg = &in_avg.borrow_mut().data;
let in_res = &in_res.borrow_mut().data;
super::squeeze::do_hsqueeze_step(
&in_avg.get_rect(Rect {
size: in_avg.size(),
origin: (0, 0),
}),
&in_res.get_rect(Rect {
size: in_res.size(),
origin: (0, 0),
}),
&None,
&None,
&mut bufs,
);
}
buffers[*buf_out] = out_buf;
}
TransformStep::VSqueeze { buf_in, buf_out } => {
buffers[*buf_out].allocate_if_needed()?;
let mut out_buf = buffers[*buf_out].take();
let mut in_avg = buffers[buf_in[0]].take();
let mut in_res = buffers[buf_in[1]].take();
{
let mut bufs: Vec<_> = vec![out_buf.borrow_mut()];
let in_avg = &in_avg.borrow_mut().data;
let in_res = &in_res.borrow_mut().data;
super::squeeze::do_vsqueeze_step(
&in_avg.get_rect(Rect {
size: in_avg.size(),
origin: (0, 0),
}),
&in_res.get_rect(Rect {
size: in_res.size(),
origin: (0, 0),
}),
&None,
&None,
&mut bufs,
);
}
buffers[*buf_out] = out_buf;
}
};
Ok(())
}
}