use crate::engine::volcano::steps::traits::ExplainNode;
use crate::types::PIPELINE_PRODUCE_SIZE;
use crate::{
engine::{
context::GraphCtx,
traverser::Traverser,
volcano::{
builder::PhysicalPlan,
steps::traits::{CoreStep, StepRef},
},
},
types::error::StoreError,
};
use smallvec::{smallvec, SmallVec};
use std::rc::Rc;
#[derive(Debug)]
pub struct ChooseStep {
upstream: Option<StepRef>,
predicate: PhysicalPlan,
true_choice: PhysicalPlan,
false_choice: Option<PhysicalPlan>,
active_plan: Option<PhysicalPlan>, }
impl ChooseStep {
pub fn new(predicate: PhysicalPlan, true_choice: PhysicalPlan, false_choice: Option<PhysicalPlan>) -> Self {
Self { upstream: None, predicate, true_choice, false_choice, active_plan: None }
}
}
impl CoreStep for ChooseStep {
fn add_upper(&mut self, upstream: StepRef) {
self.upstream = Some(upstream);
}
fn reset(&mut self) {
if let Some(u) = &self.upstream {
u.reset();
}
self.predicate.reset();
self.true_choice.reset();
if let Some(ref fc) = self.false_choice {
fc.reset();
}
}
fn upper(&self) -> Option<StepRef> {
self.upstream.clone()
}
fn produce(
&mut self,
ctx: &mut dyn GraphCtx,
) -> Result<Option<SmallVec<[Rc<Traverser>; PIPELINE_PRODUCE_SIZE]>>, StoreError> {
loop {
if let Some(ref plan) = self.active_plan {
match plan.next(ctx)? {
Some(t) => return Ok(Some(smallvec![t])),
None => self.active_plan = None,
}
}
let Some(upstream) = self.upstream.as_ref() else { return Ok(None) };
let Some(t) = upstream.next(ctx)? else { return Ok(None) };
self.predicate.reset();
self.predicate.inject(smallvec![Rc::clone(&t)]);
let matches = self.predicate.next(ctx)?.is_some();
let branch = if matches {
&self.true_choice
} else if let Some(ref fc) = self.false_choice {
fc
} else {
return Ok(Some(smallvec![t]));
};
branch.reset();
branch.inject(smallvec![Rc::clone(&t)]);
self.active_plan = Some(branch.clone());
}
}
fn explain(&self) -> ExplainNode {
let params = vec![("pred", format!("{:?}", self.predicate))];
let mut children = vec![("true".to_string(), self.true_choice.explain())];
if let Some(ref fc) = self.false_choice {
children.push(("false".to_string(), fc.explain()));
}
ExplainNode::new("ChooseStep").with_params(params).with_children(children)
}
}