use crate::api::{JxlColorType, JxlDataFormat};
use crate::error::{Error, Result};
use crate::headers::Orientation;
use crate::render::internal::ChannelInfo;
use crate::render::save::SaveStage;
use crate::util::{ShiftRightCeil, tracing_wrappers::*};
use super::internal::{RenderPipelineShared, Stage};
use super::stages::ExtendToImageDimensionsStage;
use super::{RenderPipeline, RenderPipelineInOutStage, RenderPipelineInPlaceStage};
pub(crate) struct RenderPipelineBuilder<Pipeline: RenderPipeline> {
shared: RenderPipelineShared<Pipeline::Buffer>,
}
impl<Pipeline: RenderPipeline> RenderPipelineBuilder<Pipeline> {
#[instrument(level = "debug")]
pub(super) fn new_with_chunk_size(
num_channels: usize,
size: (usize, usize),
downsampling_shift: usize,
mut log_group_size: usize,
num_passes: usize,
chunk_size: usize,
) -> Self {
info!("creating render pipeline");
assert!(chunk_size <= u16::MAX as usize);
assert_ne!(chunk_size, 0);
log_group_size += downsampling_shift;
Self {
shared: RenderPipelineShared {
channel_info: vec![vec![
ChannelInfo {
ty: None,
downsample: (0, 0)
};
num_channels
]],
input_size: size,
log_group_size,
group_count: (size.0.shrc(log_group_size), size.1.shrc(log_group_size)),
stages: vec![],
group_chan_ready_passes: vec![
vec![0; num_channels];
size.0.shrc(log_group_size)
* size.1.shrc(log_group_size)
],
num_passes,
chunk_size,
extend_stage_index: None,
},
}
}
pub(super) fn add_stage_internal(mut self, stage: Stage<Pipeline::Buffer>) -> Result<Self> {
let input_type = stage.input_type();
let output_type = stage.output_type();
let shift = stage.shift();
let border = stage.border();
let is_extend = matches!(stage, Stage::Extend(_));
let current_info = self.shared.channel_info.last().unwrap().clone();
debug!(
last_stage_channel_info = ?current_info,
extend_stage_index= ?self.shared.extend_stage_index,
"adding stage '{stage}'",
);
let mut after_info = vec![];
for (c, info) in current_info.iter().enumerate() {
if !stage.uses_channel(c) {
after_info.push(ChannelInfo {
ty: info.ty,
downsample: (0, 0),
});
} else {
if let Some(ty) = info.ty
&& ty != input_type
{
return Err(Error::PipelineChannelTypeMismatch(
stage.to_string(),
c,
input_type,
ty,
));
}
after_info.push(ChannelInfo {
ty: Some(output_type.unwrap_or(input_type)),
downsample: shift,
});
}
}
if self.shared.extend_stage_index.is_some()
&& (shift != (0, 0) || border != (0, 0) || is_extend)
{
return Err(Error::PipelineInvalidStageAfterExtend(stage.to_string()));
}
if is_extend {
self.shared.extend_stage_index = Some(self.shared.stages.len());
}
debug!(
new_channel_info = ?after_info,
extend_stage_index= ?self.shared.extend_stage_index,
"added stage '{stage}'",
);
self.shared.channel_info.push(after_info);
self.shared.stages.push(stage);
Ok(self)
}
pub fn new(
num_channels: usize,
size: (usize, usize),
downsampling_shift: usize,
log_group_size: usize,
num_passes: usize,
) -> Self {
Self::new_with_chunk_size(
num_channels,
size,
downsampling_shift,
log_group_size,
num_passes,
1 << (log_group_size + downsampling_shift),
)
}
#[instrument(skip_all, err)]
pub fn add_save_stage(
self,
channels: &[usize],
orientation: Orientation,
output_buffer_index: usize,
color_type: JxlColorType,
data_format: JxlDataFormat,
fill_opaque_alpha: bool,
) -> Result<Self> {
let stage = SaveStage::new(
channels,
orientation,
output_buffer_index,
color_type,
data_format,
fill_opaque_alpha,
);
self.add_stage_internal(Stage::Save(stage))
}
#[instrument(skip_all, err)]
pub fn add_extend_stage(self, extend: ExtendToImageDimensionsStage) -> Result<Self> {
self.add_stage_internal(Stage::Extend(extend))
}
#[instrument(skip_all, err)]
pub fn add_inplace_stage<S: RenderPipelineInPlaceStage>(self, stage: S) -> Result<Self> {
self.add_stage_internal(Stage::InPlace(Pipeline::box_inplace_stage(stage)))
}
#[instrument(skip_all, err)]
pub fn add_inout_stage<S: RenderPipelineInOutStage>(self, stage: S) -> Result<Self> {
self.add_stage_internal(Stage::InOut(Pipeline::box_inout_stage(stage)))
}
#[instrument(skip_all, err)]
pub fn build(mut self) -> Result<Box<Pipeline>> {
let channel_info = &mut self.shared.channel_info;
let num_channels = channel_info[0].len();
let mut cur_downsamples = vec![(0u8, 0u8); num_channels];
for (s, stage) in self.shared.stages.iter().enumerate().rev() {
let [current_info, next_info, ..] = &mut channel_info[s..] else {
unreachable!()
};
let mut save_downsample = None;
for chan in 0..num_channels {
let cur_chan = &mut current_info[chan];
let next_chan = &mut next_info[chan];
let uses_channel = stage.uses_channel(chan);
let input_type = stage.input_type();
if cur_chan.ty.is_none() {
cur_chan.ty = if uses_channel {
Some(input_type)
} else {
next_chan.ty
}
}
let cur_downsample = &mut cur_downsamples[chan];
if matches!(stage, Stage::Save(_))
&& save_downsample.is_some_and(|x| x != *cur_downsample)
{
save_downsample = Some(*cur_downsample);
return Err(Error::SaveDifferentDownsample(
save_downsample.unwrap(),
*cur_downsample,
));
}
let next_downsample = &mut next_chan.downsample;
let next_total_downsample = *cur_downsample;
cur_downsample.0 = cur_downsample
.0
.checked_add(next_downsample.0)
.ok_or(Error::ArithmeticOverflow)?;
cur_downsample.1 = cur_downsample
.1
.checked_add(next_downsample.1)
.ok_or(Error::ArithmeticOverflow)?;
*next_downsample = next_total_downsample;
}
}
for (chan, cur_downsample) in cur_downsamples.iter().enumerate() {
channel_info[0][chan].downsample = *cur_downsample;
}
#[cfg(feature = "tracing")]
{
for (s, (current_info, stage)) in channel_info
.iter()
.zip(self.shared.stages.iter())
.enumerate()
{
debug!("final channel info before stage {s} '{stage}': {current_info:?}");
}
debug!(
"final channel info after all stages {:?}",
channel_info.last().unwrap()
);
}
for (c, chinfo) in channel_info.iter().flat_map(|x| x.iter().enumerate()) {
if chinfo.ty.is_none() {
return Err(Error::PipelineChannelUnused(c));
}
}
Ok(Box::new(Pipeline::new_from_shared(self.shared)?))
}
}