use alloc::vec::Vec;
use super::super::data::Stream;
use super::super::error::ExecError;
use super::super::params::LocalParams;
use super::super::registry::{GraphId, RuntimeParams};
use super::Exec;
pub(crate) struct SegmenterCx<'s, 'e, 'a> {
pub(crate) ex: &'s mut Exec<'e, 'a>,
pub(crate) params: &'s LocalParams,
pub(crate) custom_graphs: &'s [GraphId],
inputs: Vec<Stream<'a>>,
consumed: Vec<usize>,
chunks: usize,
}
impl<'s, 'e, 'a> SegmenterCx<'s, 'e, 'a> {
pub(crate) fn new(
ex: &'s mut Exec<'e, 'a>,
params: &'s LocalParams,
custom_graphs: &'s [GraphId],
n_inputs: usize,
) -> Result<Self, ExecError> {
let inputs = (0..n_inputs as u32)
.map(|n| ex.rt.stream(n).cloned().ok_or(ExecError::GraphInvalid))
.collect::<Result<Vec<_>, _>>()?;
ex.rt.reset();
Ok(Self {
ex,
params,
custom_graphs,
consumed: alloc::vec![0; n_inputs],
inputs,
chunks: 0,
})
}
pub(crate) fn n_inputs(&self) -> usize {
self.inputs.len()
}
pub(crate) const fn format_version(&self) -> u32 {
self.ex.p.version
}
pub(crate) fn input(&self, index: usize) -> Option<Stream<'a>> {
let input = self.inputs.get(index)?;
let done = *self.consumed.get(index)?;
input.slice(done, input.n_elts() - done)
}
pub(crate) fn remaining(&self) -> Vec<usize> {
self.inputs
.iter()
.zip(&self.consumed)
.map(|(input, done)| input.n_elts() - done)
.collect()
}
pub(crate) fn process_chunk(
&mut self,
n_elts: &[usize],
start: GraphId,
params: Option<&RuntimeParams>,
) -> Result<(), ExecError> {
if n_elts.len() != self.inputs.len() {
return Err(ExecError::SuccessorInputs);
}
let mut chunk = Vec::with_capacity(n_elts.len());
for (index, &count) in n_elts.iter().enumerate() {
let slice = self.inputs[index]
.slice(self.consumed[index], count)
.ok_or(ExecError::Parameter("chunk larger than the input"))?;
chunk.push(slice);
}
for (done, &count) in self.consumed.iter_mut().zip(n_elts) {
*done += count;
}
self.ex.rt.reset();
let rtsids: Vec<u32> = chunk
.iter()
.map(|stream| self.ex.rt.ref_input(stream.clone()))
.collect();
let depth = self.ex.segmenter_depth + 1;
self.ex.run_successor(start, params, &rtsids, depth)?;
self.ex.flush_chunk(&chunk)?;
self.chunks += 1;
Ok(())
}
pub(crate) fn check_done(&self) -> Result<(), ExecError> {
if self.chunks == 0 {
return Err(ExecError::SegmenterNoChunks);
}
if self
.inputs
.iter()
.zip(&self.consumed)
.any(|(input, &done)| done < input.n_elts())
{
return Err(ExecError::SegmenterNotConsumed);
}
Ok(())
}
}