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 crate::util::SmallVec;
/// Multi-row channel accessor for immutable access.
///
/// Provides 2D indexing where `channels[ch]` returns `&[&[T]]` (all rows for a channel),
/// and `channels[ch][row]` returns `&[T]` (pixels for a specific row).
///
/// This eliminates nested Vec collections while maintaining the same indexing syntax.
pub struct Channels<'a, T> {
// The number of input rows should be maximized by the EPF0 stage, which has 21.
pub(crate) row_data: SmallVec<&'a [T], 32>,
num_channels: usize,
pub(crate) rows_per_channel: usize,
}
impl<'a, T> Channels<'a, T> {
/// Create a new Channels accessor.
///
/// # Arguments
/// * `row_data` - Flat vector of all rows for all channels (length = num_channels * rows_per_channel)
/// * `num_channels` - Number of channels
/// * `rows_per_channel` - Number of rows per channel (typically 2*BORDER+1)
pub fn new(
row_data: SmallVec<&'a [T], 32>,
num_channels: usize,
rows_per_channel: usize,
) -> Self {
debug_assert_eq!(
row_data.len(),
num_channels * rows_per_channel,
"row_data length must equal num_channels * rows_per_channel"
);
Self {
row_data,
num_channels,
rows_per_channel,
}
}
/// Returns the number of channels.
pub fn len(&self) -> usize {
self.num_channels
}
/// Returns true if there are no channels.
pub fn is_empty(&self) -> bool {
self.num_channels == 0
}
/// Returns an iterator over channel slices.
pub fn iter(&self) -> impl Iterator<Item = &[&'a [T]]> {
(0..self.num_channels).map(move |ch| &self[ch])
}
}
/// Implement indexing: channels[ch] returns &[&[T]]
impl<'a, T> std::ops::Index<usize> for Channels<'a, T> {
type Output = [&'a [T]];
fn index(&self, ch: usize) -> &[&'a [T]] {
let start = ch * self.rows_per_channel;
&self.row_data[start..start + self.rows_per_channel]
}
}
/// Multi-row channel accessor for mutable access.
///
/// Provides 2D indexing where `channels[ch]` returns `&[&mut [T]]` or `&mut [&mut [T]]`,
/// and `channels[ch][row]` returns `&mut [T]` (pixels for a specific row).
pub struct ChannelsMut<'a, T> {
// The number of output rows should be maximized by the Upsample8 stage, which has 8.
pub(crate) row_data: SmallVec<&'a mut [T], 8>,
num_channels: usize,
pub(crate) rows_per_channel: usize,
}
impl<'a, T> ChannelsMut<'a, T> {
/// Create a new ChannelsMut accessor.
///
/// # Arguments
/// * `row_data` - Flat vector of all mutable rows for all channels
/// * `num_channels` - Number of channels
/// * `rows_per_channel` - Number of rows per channel (typically 1 << SHIFT)
pub fn new(
row_data: SmallVec<&'a mut [T], 8>,
num_channels: usize,
rows_per_channel: usize,
) -> Self {
debug_assert_eq!(
row_data.len(),
num_channels * rows_per_channel,
"row_data length must equal num_channels * rows_per_channel"
);
Self {
row_data,
num_channels,
rows_per_channel,
}
}
/// Returns the number of channels.
pub fn len(&self) -> usize {
self.num_channels
}
/// Returns true if there are no channels.
pub fn is_empty(&self) -> bool {
self.num_channels == 0
}
/// Splits the first 3 channels into separate mutable slices.
/// Returns a tuple containing mutable references to each channel's rows.
#[allow(clippy::type_complexity)]
pub fn split_first_3_mut(
&mut self,
) -> (&mut [&'a mut [T]], &mut [&'a mut [T]], &mut [&'a mut [T]]) {
assert!(
3 <= self.num_channels,
"requested 3 channels but only have {}",
self.num_channels
);
let rpc = self.rows_per_channel;
let (first, rest) = self.row_data.split_at_mut(rpc);
let (second, rest) = rest.split_at_mut(rpc);
let (third, _) = rest.split_at_mut(rpc);
(first, second, third)
}
/// Returns a mutable iterator over all channels.
/// Each item is a mutable slice of rows for that channel.
pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut [&'a mut [T]]> {
let rpc = self.rows_per_channel;
self.row_data.chunks_mut(rpc)
}
}
/// Implement immutable indexing: channels[ch] returns &[&mut [T]]
impl<'a, T> std::ops::Index<usize> for ChannelsMut<'a, T> {
type Output = [&'a mut [T]];
fn index(&self, ch: usize) -> &[&'a mut [T]] {
let start = ch * self.rows_per_channel;
&self.row_data[start..start + self.rows_per_channel]
}
}
/// Implement mutable indexing: &mut channels[ch] returns &mut [&mut [T]]
impl<'a, T> std::ops::IndexMut<usize> for ChannelsMut<'a, T> {
fn index_mut(&mut self, ch: usize) -> &mut [&'a mut [T]] {
let start = ch * self.rows_per_channel;
&mut self.row_data[start..start + self.rows_per_channel]
}
}