use crate::transpose::Transposition;
use crate::util::{DctSample, try_vec};
use crate::{PxdctError, PxdctExecutor};
use std::sync::Arc;
pub trait MultidimensionalDctExecutor<T> {
fn execute(&self, data: &mut [T]) -> Result<(), PxdctError>;
fn execute_with_scratch(&self, data: &mut [T], scratch: &mut [T]) -> Result<(), PxdctError>;
fn width(&self) -> usize;
fn height(&self) -> usize;
fn required_scratch_size(&self) -> usize;
}
pub(crate) struct TwoDimensionalDct<T> {
pub(crate) width_executor: Arc<dyn PxdctExecutor<T> + Send + Sync>,
pub(crate) height_executor: Arc<dyn PxdctExecutor<T> + Send + Sync>,
pub(crate) transpose_width_to_height: Arc<dyn Transposition<T> + Send + Sync>,
pub(crate) width: usize,
pub(crate) height: usize,
}
impl<T: DctSample> MultidimensionalDctExecutor<T> for TwoDimensionalDct<T> {
fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
let mut scratch = try_vec![T::default(); self.required_scratch_size()];
self.execute_with_scratch(data, &mut scratch)
}
fn execute_with_scratch(&self, data: &mut [T], scratch: &mut [T]) -> Result<(), PxdctError> {
if scratch.len() < self.required_scratch_size() {
return Err(PxdctError::ScratchBufferIsTooSmall(
scratch.len(),
self.required_scratch_size(),
));
}
let full_size = self.width * self.height;
if !data.len().is_multiple_of(full_size) {
return Err(PxdctError::InvalidSizeMultiplier(data.len(), full_size));
}
let (scratch, _) = scratch.split_at_mut(self.required_scratch_size());
for chunk in data.chunks_exact_mut(full_size) {
scratch.copy_from_slice(chunk);
scratch.chunks_exact_mut(self.width).for_each(|row| {
_ = self.width_executor.execute(row);
});
self.transpose_width_to_height.transpose(scratch, chunk);
chunk.chunks_exact_mut(self.height).for_each(|column| {
_ = self.height_executor.execute(column);
});
}
Ok(())
}
#[inline]
fn width(&self) -> usize {
self.width
}
#[inline]
fn height(&self) -> usize {
self.height
}
#[inline]
fn required_scratch_size(&self) -> usize {
self.width * self.height
}
}