use crate::types::PIPELINE_PRODUCE_SIZE;
use crate::types::STEP_LABEL_INLINE;
use std::rc::Rc;
use smallvec::{smallvec, SmallVec};
use smol_str::SmolStr;
use crate::engine::volcano::steps::traits::ExplainNode;
use crate::{
engine::{
context::GraphCtx,
traverser::Traverser,
volcano::steps::traits::{CoreStep, StepRef},
},
types::error::StoreError,
};
#[derive(Debug)]
pub struct AsStep {
upstream: Option<StepRef>,
labels: SmallVec<[SmolStr; STEP_LABEL_INLINE]>,
}
impl AsStep {
pub fn new(labels: SmallVec<[SmolStr; STEP_LABEL_INLINE]>) -> Self {
Self { upstream: None, labels }
}
}
impl CoreStep for AsStep {
fn add_upper(&mut self, upstream: StepRef) {
self.upstream = Some(upstream);
}
fn produce(
&mut self,
ctx: &mut dyn GraphCtx,
) -> Result<Option<SmallVec<[Rc<Traverser>; PIPELINE_PRODUCE_SIZE]>>, StoreError> {
let Some(upstream) = self.upstream.as_ref() else {
return Ok(None);
};
let Some(t) = upstream.next(ctx)? else {
return Ok(None);
};
let labeled =
Rc::new(Traverser { value: t.value.clone(), parent: t.parent.clone(), labels: Some(self.labels.clone()) });
Ok(Some(smallvec![labeled]))
}
fn reset(&mut self) {
if let Some(up) = &self.upstream {
up.reset();
}
}
fn upper(&self) -> Option<StepRef> {
self.upstream.clone()
}
fn explain(&self) -> ExplainNode {
let params = vec![("label", format!("{:?}", self.labels))];
ExplainNode::new("AsStep").with_params(params)
}
}
#[derive(Debug)]
pub struct SelectStep {
upstream: Option<StepRef>,
labels: SmallVec<[SmolStr; STEP_LABEL_INLINE]>,
}
impl SelectStep {
pub fn new(labels: SmallVec<[SmolStr; STEP_LABEL_INLINE]>) -> Self {
Self { upstream: None, labels }
}
}
impl CoreStep for SelectStep {
fn add_upper(&mut self, upstream: StepRef) {
self.upstream = Some(upstream);
}
fn produce(
&mut self,
ctx: &mut dyn GraphCtx,
) -> Result<Option<SmallVec<[Rc<Traverser>; PIPELINE_PRODUCE_SIZE]>>, StoreError> {
loop {
let Some(upstream) = self.upstream.as_ref() else {
return Ok(None);
};
let Some(t) = upstream.next(ctx)? else {
return Ok(None);
};
if let Some(found) = find_labeled(&t, &self.labels) {
return Ok(Some(smallvec![found]));
}
}
}
fn reset(&mut self) {
if let Some(up) = &self.upstream {
up.reset();
}
}
fn upper(&self) -> Option<StepRef> {
self.upstream.clone()
}
fn explain(&self) -> ExplainNode {
let params = vec![("label", format!("{:?}", self.labels))];
ExplainNode::new("SelectStep").with_params(params)
}
}
fn find_labeled(t: &Rc<Traverser>, targets: &SmallVec<[SmolStr; STEP_LABEL_INLINE]>) -> Option<Rc<Traverser>> {
let mut cur = Some(Rc::clone(t));
while let Some(node) = cur {
if let Some(ref labels) = node.labels {
for target in targets {
if labels.contains(target) {
return Some(Rc::clone(&node));
}
}
}
cur = node.parent.as_ref().map(Rc::clone);
}
None
}