datum/stream/

scalar.rs

//! Planner-visible scalar integer operators and their safe chunk executor.
//!
//! Opaque `map`/`filter` closures deliberately do not enter this module.  Only
//! the built-in operators represented by [`ScalarArithmeticOp`] and
//! [`ScalarCompareOp`] are eligible for speculative chunk execution.

use super::flow::FlowTransform;
use super::{BoxStream, Flow, FlowHints, Source, StreamError, StreamResult};
use std::sync::Arc;

pub(crate) const SCALAR_CHUNK_SIZE: usize = 128;

/// A planner-visible checked integer map operation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ScalarArithmeticOp {
    Add,
    Subtract,
    Multiply,
}

/// A planner-visible integer comparison used by built-in filters.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ScalarCompareOp {
    Equal,
    NotEqual,
    LessThan,
    LessOrEqual,
    GreaterThan,
    GreaterOrEqual,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ScalarI64Op {
    Arithmetic(ScalarArithmeticOp),
    Compare(ScalarCompareOp),
}

impl ScalarI64Op {
    fn apply_checked(self, value: i64, operand: i64) -> StreamResult<Option<i64>> {
        match self {
            Self::Arithmetic(ScalarArithmeticOp::Add) => value
                .checked_add(operand)
                .map(Some)
                .ok_or_else(|| map_overflow("add", value, "+", operand)),
            Self::Arithmetic(ScalarArithmeticOp::Subtract) => value
                .checked_sub(operand)
                .map(Some)
                .ok_or_else(|| map_overflow("subtract", value, "-", operand)),
            Self::Arithmetic(ScalarArithmeticOp::Multiply) => value
                .checked_mul(operand)
                .map(Some)
                .ok_or_else(|| map_overflow("multiply", value, "*", operand)),
            Self::Compare(compare) => Ok(compare_i64(compare, value, operand).then_some(value)),
        }
    }
}

fn map_overflow(name: &str, lhs: i64, symbol: &str, rhs: i64) -> StreamError {
    StreamError::Failed(format!(
        "integer overflow in map_{name}: {lhs} {symbol} {rhs}"
    ))
}

fn compare_i64(op: ScalarCompareOp, lhs: i64, rhs: i64) -> bool {
    match op {
        ScalarCompareOp::Equal => lhs == rhs,
        ScalarCompareOp::NotEqual => lhs != rhs,
        ScalarCompareOp::LessThan => lhs < rhs,
        ScalarCompareOp::LessOrEqual => lhs <= rhs,
        ScalarCompareOp::GreaterThan => lhs > rhs,
        ScalarCompareOp::GreaterOrEqual => lhs >= rhs,
    }
}

/// Monomorphic step metadata stored by both linear streams and GraphDSL maps.
///
/// The function pointers let the generic graph planner downcast a
/// `ScalarChunkStep<In>` only when `In == i64`, while keeping the production
/// kernels themselves monomorphic and inspectable in generated code.
pub(crate) struct ScalarChunkStep<T> {
    op: ScalarI64Op,
    operand: T,
    clone_value: fn(&T) -> T,
    copy_chunk: fn(&[T], &mut Vec<T>),
    checked: fn(ScalarI64Op, T, T) -> StreamResult<Option<T>>,
    map_chunk: fn(ScalarArithmeticOp, &mut [T], T) -> bool,
    compare_chunk: fn(ScalarCompareOp, &[T], &mut [u8], T),
}

impl ScalarChunkStep<i64> {
    pub(crate) fn new(op: ScalarI64Op, operand: i64) -> Self {
        Self {
            op,
            operand,
            clone_value: |value| *value,
            copy_chunk: |input, output| output.extend_from_slice(input),
            checked: |op, value, operand| op.apply_checked(value, operand),
            map_chunk: map_i64_chunk,
            compare_chunk: compare_i64_chunk,
        }
    }
}

impl<T> Clone for ScalarChunkStep<T> {
    fn clone(&self) -> Self {
        Self {
            op: self.op,
            operand: (self.clone_value)(&self.operand),
            clone_value: self.clone_value,
            copy_chunk: self.copy_chunk,
            checked: self.checked,
            map_chunk: self.map_chunk,
            compare_chunk: self.compare_chunk,
        }
    }
}

impl<T> ScalarChunkStep<T> {
    pub(crate) fn apply_checked(&self, value: T) -> StreamResult<Option<T>> {
        (self.checked)(self.op, value, (self.clone_value)(&self.operand))
    }

    fn is_multiply(&self) -> bool {
        self.op == ScalarI64Op::Arithmetic(ScalarArithmeticOp::Multiply)
    }

    pub(crate) fn is_add_or_subtract(&self) -> bool {
        matches!(
            self.op,
            ScalarI64Op::Arithmetic(ScalarArithmeticOp::Add | ScalarArithmeticOp::Subtract)
        )
    }
}

/// Process one private chunk through a pure built-in segment.
///
/// `original` retains the unmodified input.  `values`, `compact`, and `mask`
/// are allocated once per materialization and reused.  A flagged add/sub
/// chunk is discarded and replayed element-major through `apply_checked`,
/// preserving the first failing element/stage and the valid output prefix.
pub(crate) struct ScalarChunkProcessor<T> {
    original: Vec<T>,
    values: Vec<T>,
    compact: Vec<T>,
    mask: Vec<u8>,
}

pub(crate) struct ScalarChunkResult {
    pub(crate) failure: Option<StreamError>,
}

impl<T> ScalarChunkProcessor<T> {
    pub(crate) fn new() -> Self {
        Self {
            original: Vec::with_capacity(SCALAR_CHUNK_SIZE),
            values: Vec::with_capacity(SCALAR_CHUNK_SIZE),
            compact: Vec::with_capacity(SCALAR_CHUNK_SIZE),
            mask: Vec::with_capacity(SCALAR_CHUNK_SIZE),
        }
    }

    pub(crate) fn load_from(&mut self, input: &mut dyn Iterator<Item = T>) -> usize {
        self.original.clear();
        self.original.extend(input.take(SCALAR_CHUNK_SIZE));
        self.original.len()
    }

    pub(crate) fn output(&self) -> &[T] {
        &self.values
    }

    pub(crate) fn drain_output(&mut self) -> std::vec::Drain<'_, T> {
        self.values.drain(..)
    }

    pub(crate) fn process(&mut self, steps: &[ScalarChunkStep<T>]) -> ScalarChunkResult {
        assert!(
            !steps.is_empty(),
            "scalar chunk plan must contain an operator"
        );
        if steps.iter().any(ScalarChunkStep::is_multiply) {
            return self.scalar_replay(steps);
        }

        self.values.clear();
        (steps[0].copy_chunk)(&self.original, &mut self.values);
        let mut overflow = false;

        for step in steps {
            match step.op {
                ScalarI64Op::Arithmetic(op) => {
                    overflow |=
                        (step.map_chunk)(op, &mut self.values, (step.clone_value)(&step.operand));
                }
                ScalarI64Op::Compare(op) => {
                    self.mask.resize(self.values.len(), 0);
                    (step.compare_chunk)(
                        op,
                        &self.values,
                        &mut self.mask,
                        (step.clone_value)(&step.operand),
                    );
                    self.compact.clear();
                    for (value, &keep) in self.values.iter().zip(&self.mask) {
                        if keep != 0 {
                            self.compact.push((step.clone_value)(value));
                        }
                    }
                    std::mem::swap(&mut self.values, &mut self.compact);
                }
            }
        }

        if overflow {
            self.scalar_replay(steps)
        } else {
            ScalarChunkResult { failure: None }
        }
    }

    fn scalar_replay(&mut self, steps: &[ScalarChunkStep<T>]) -> ScalarChunkResult {
        self.values.clear();
        for input in self
            .original
            .iter()
            .map(|value| (steps[0].clone_value)(value))
        {
            let mut value = Some(input);
            for step in steps {
                let Some(current) = value else {
                    break;
                };
                match step.apply_checked(current) {
                    Ok(next) => value = next,
                    Err(error) => {
                        return ScalarChunkResult {
                            failure: Some(error),
                        };
                    }
                }
            }
            if let Some(value) = value {
                self.values.push(value);
            }
        }
        ScalarChunkResult { failure: None }
    }
}

#[inline(never)]
fn add_i64_chunk(values: &mut [i64], operand: i64) -> bool {
    let mut overflow_bits = 0_u64;
    for value in values {
        let lhs = *value;
        let result = lhs.wrapping_add(operand);
        *value = result;
        overflow_bits |= ((lhs ^ result) & (operand ^ result)) as u64;
    }
    overflow_bits & (1_u64 << 63) != 0
}

#[inline(never)]
fn subtract_i64_chunk(values: &mut [i64], operand: i64) -> bool {
    let mut overflow_bits = 0_u64;
    for value in values {
        let lhs = *value;
        let result = lhs.wrapping_sub(operand);
        *value = result;
        overflow_bits |= ((lhs ^ operand) & (lhs ^ result)) as u64;
    }
    overflow_bits & (1_u64 << 63) != 0
}

fn map_i64_chunk(op: ScalarArithmeticOp, values: &mut [i64], operand: i64) -> bool {
    match op {
        ScalarArithmeticOp::Add => add_i64_chunk(values, operand),
        ScalarArithmeticOp::Subtract => subtract_i64_chunk(values, operand),
        // Multiplication is deliberately kept on scalar replay: baseline x86-64
        // synthesizes packed i64 multiply and loses on the real cost model.
        ScalarArithmeticOp::Multiply => true,
    }
}

#[inline(never)]
fn compare_i64_chunk(op: ScalarCompareOp, values: &[i64], mask: &mut [u8], operand: i64) {
    assert_eq!(values.len(), mask.len());
    for (value, keep) in values.iter().zip(mask) {
        *keep = u8::from(compare_i64(op, *value, operand));
    }
}

/// Production add-kernel probe used by the checked-in S2 codegen evidence.
#[doc(hidden)]
#[inline(never)]
pub fn scalar_add_codegen_probe(values: &mut [i64], operand: i64) -> bool {
    add_i64_chunk(values, operand)
}

/// Production subtract-kernel probe used by the checked-in S2 codegen evidence.
#[doc(hidden)]
#[inline(never)]
pub fn scalar_subtract_codegen_probe(values: &mut [i64], operand: i64) -> bool {
    subtract_i64_chunk(values, operand)
}

/// Production comparison-kernel probe used by the checked-in S2 codegen evidence.
#[doc(hidden)]
#[inline(never)]
pub fn scalar_compare_codegen_probe(values: &[i64], mask: &mut [u8], operand: i64) {
    compare_i64_chunk(ScalarCompareOp::GreaterThan, values, mask, operand);
}

struct ScalarI64Stream {
    input: BoxStream<i64>,
    steps: Arc<[ScalarChunkStep<i64>]>,
    processor: Option<ScalarChunkProcessor<i64>>,
    output_index: usize,
    pending_upstream_error: Option<StreamError>,
    pending_failure: Option<StreamError>,
    upstream_done: bool,
    chunked: bool,
}

impl ScalarI64Stream {
    fn new(input: BoxStream<i64>, steps: Arc<[ScalarChunkStep<i64>]>) -> Self {
        let chunked = steps.iter().any(ScalarChunkStep::is_add_or_subtract);
        Self {
            input,
            steps,
            processor: chunked.then(ScalarChunkProcessor::new),
            output_index: 0,
            pending_upstream_error: None,
            pending_failure: None,
            upstream_done: false,
            chunked,
        }
    }

    fn new_scalar(input: BoxStream<i64>, steps: Arc<[ScalarChunkStep<i64>]>) -> Self {
        Self {
            input,
            steps,
            processor: None,
            output_index: 0,
            pending_upstream_error: None,
            pending_failure: None,
            upstream_done: false,
            chunked: false,
        }
    }

    fn next_scalar(&mut self) -> Option<StreamResult<i64>> {
        loop {
            let item = self.input.next()?;
            let mut value = match item {
                Ok(value) => Some(value),
                Err(error) => return Some(Err(error)),
            };
            for step in self.steps.iter() {
                let Some(current) = value else {
                    break;
                };
                match step.op.apply_checked(current, step.operand) {
                    Ok(next) => value = next,
                    Err(error) => return Some(Err(error)),
                }
            }
            if let Some(value) = value {
                return Some(Ok(value));
            }
        }
    }

    fn fill_chunk(&mut self) {
        let mut input = self.input.by_ref().map_while(|item| match item {
            Ok(value) => Some(value),
            Err(error) => {
                self.pending_upstream_error = Some(error);
                self.upstream_done = true;
                None
            }
        });
        let processor = self
            .processor
            .as_mut()
            .expect("chunk processor exists for add/subtract segments");
        let loaded = processor.load_from(&mut input);
        if loaded < SCALAR_CHUNK_SIZE && self.pending_upstream_error.is_none() {
            self.upstream_done = true;
        }
        let result = processor.process(&self.steps);
        self.output_index = 0;
        self.pending_failure = result.failure;
    }
}

impl Iterator for ScalarI64Stream {
    type Item = StreamResult<i64>;

    fn next(&mut self) -> Option<Self::Item> {
        if !self.chunked {
            return self.next_scalar();
        }
        loop {
            if let Some(&value) = self
                .processor
                .as_ref()
                .expect("chunk processor exists for add/subtract segments")
                .output()
                .get(self.output_index)
            {
                self.output_index += 1;
                return Some(Ok(value));
            }
            if let Some(error) = self.pending_failure.take() {
                self.upstream_done = true;
                return Some(Err(error));
            }
            if let Some(error) = self.pending_upstream_error.take() {
                return Some(Err(error));
            }
            if self.upstream_done {
                return None;
            }
            self.fill_chunk();
        }
    }
}

fn filter_i64_stream<F>(input: BoxStream<i64>, predicate: F) -> BoxStream<i64>
where
    F: Fn(i64) -> bool + Send + Sync + 'static,
{
    Box::new(input.filter_map(move |item| match item {
        Ok(value) if predicate(value) => Some(Ok(value)),
        Ok(_) => None,
        Err(error) => Some(Err(error)),
    }))
}

fn scalar_i64_stream(input: BoxStream<i64>, steps: Arc<[ScalarChunkStep<i64>]>) -> BoxStream<i64> {
    if let [step] = steps.as_ref() {
        let operand = step.operand;
        return match step.op {
            ScalarI64Op::Compare(ScalarCompareOp::Equal) => {
                filter_i64_stream(input, move |value| value == operand)
            }
            ScalarI64Op::Compare(ScalarCompareOp::NotEqual) => {
                filter_i64_stream(input, move |value| value != operand)
            }
            ScalarI64Op::Compare(ScalarCompareOp::LessThan) => {
                filter_i64_stream(input, move |value| value < operand)
            }
            ScalarI64Op::Compare(ScalarCompareOp::LessOrEqual) => {
                filter_i64_stream(input, move |value| value <= operand)
            }
            ScalarI64Op::Compare(ScalarCompareOp::GreaterThan) => {
                filter_i64_stream(input, move |value| value > operand)
            }
            ScalarI64Op::Compare(ScalarCompareOp::GreaterOrEqual) => {
                filter_i64_stream(input, move |value| value >= operand)
            }
            ScalarI64Op::Arithmetic(ScalarArithmeticOp::Multiply) => {
                Box::new(input.map(move |item| {
                    item.and_then(|value| {
                        ScalarI64Op::Arithmetic(ScalarArithmeticOp::Multiply)
                            .apply_checked(value, operand)
                            .map(|value| value.expect("multiplication retains its input"))
                    })
                }))
            }
            ScalarI64Op::Arithmetic(ScalarArithmeticOp::Add | ScalarArithmeticOp::Subtract) => {
                Box::new(ScalarI64Stream::new(input, steps))
            }
        };
    }
    Box::new(ScalarI64Stream::new(input, steps))
}

pub(crate) struct ScalarI64SourceFactory<Mat> {
    source: Arc<dyn super::SourceFactory<i64, Mat>>,
    steps: Arc<[ScalarChunkStep<i64>]>,
}

impl<Mat> ScalarI64SourceFactory<Mat> {
    fn append(&self, next: &[ScalarChunkStep<i64>]) -> Arc<dyn super::SourceFactory<i64, Mat>>
    where
        Mat: Send + 'static,
    {
        let mut steps = Vec::with_capacity(self.steps.len() + next.len());
        steps.extend_from_slice(&self.steps);
        steps.extend_from_slice(next);
        Arc::new(Self {
            source: Arc::clone(&self.source),
            steps: steps.into(),
        })
    }
}

impl<Mat> super::SourceFactory<i64, Mat> for ScalarI64SourceFactory<Mat>
where
    Mat: Send + 'static,
{
    fn create(
        self: Arc<Self>,
        materializer: &super::Materializer,
    ) -> StreamResult<(BoxStream<i64>, Mat)> {
        let (input, mat) = Arc::clone(&self.source).create(materializer)?;
        Ok((scalar_i64_stream(input, Arc::clone(&self.steps)), mat))
    }

    fn append_scalar_i64(
        self: Arc<Self>,
        steps: &[ScalarChunkStep<i64>],
    ) -> Option<Arc<dyn super::SourceFactory<i64, Mat>>> {
        Some(self.append(steps))
    }
}

fn steps_with(next: ScalarChunkStep<i64>) -> Arc<[ScalarChunkStep<i64>]> {
    Arc::from([next])
}

impl<Mat> Source<i64, Mat>
where
    Mat: Send + 'static,
{
    fn scalar_op(self, op: ScalarI64Op, operand: i64) -> Self {
        let steps = steps_with(ScalarChunkStep::new(op, operand));
        if let Some(factory) = Arc::clone(&self.factory).append_scalar_i64(&steps) {
            return Source {
                factory,
                terminal_factory: None,
                hints: self.hints.without_inline_micro(),
                attributes: self.attributes,
                split_hook: None,
            };
        }
        if self.hints.inline_micro.is_none() {
            return match op {
                ScalarI64Op::Arithmetic(_) => self.try_map(move |value| {
                    op.apply_checked(value, operand)
                        .map(|value| value.expect("arithmetic maps retain their input"))
                }),
                ScalarI64Op::Compare(compare) => match compare {
                    ScalarCompareOp::Equal => self.filter(move |value| *value == operand),
                    ScalarCompareOp::NotEqual => self.filter(move |value| *value != operand),
                    ScalarCompareOp::LessThan => self.filter(move |value| *value < operand),
                    ScalarCompareOp::LessOrEqual => self.filter(move |value| *value <= operand),
                    ScalarCompareOp::GreaterThan => self.filter(move |value| *value > operand),
                    ScalarCompareOp::GreaterOrEqual => self.filter(move |value| *value >= operand),
                },
            };
        }
        let Source {
            factory,
            terminal_factory: _,
            hints,
            attributes,
            split_hook: _,
        } = self;
        let next_factory = Arc::new(ScalarI64SourceFactory {
            source: factory,
            steps,
        });
        Source {
            factory: next_factory,
            terminal_factory: None,
            hints: hints.without_inline_micro(),
            attributes,
            split_hook: None,
        }
    }

    /// Apply a checked built-in integer operation.
    #[must_use]
    pub fn map_op(self, op: ScalarArithmeticOp, operand: i64) -> Self {
        self.scalar_op(ScalarI64Op::Arithmetic(op), operand)
    }

    #[must_use]
    pub fn map_add(self, operand: i64) -> Self {
        self.map_op(ScalarArithmeticOp::Add, operand)
    }

    #[must_use]
    pub fn map_subtract(self, operand: i64) -> Self {
        self.map_op(ScalarArithmeticOp::Subtract, operand)
    }

    #[must_use]
    pub fn map_multiply(self, operand: i64) -> Self {
        self.map_op(ScalarArithmeticOp::Multiply, operand)
    }

    /// Retain values matching a built-in integer comparison.
    #[must_use]
    pub fn filter_op(self, op: ScalarCompareOp, operand: i64) -> Self {
        let step = ScalarChunkStep::new(ScalarI64Op::Compare(op), operand);
        let steps = steps_with(step);
        if let Some(factory) = Arc::clone(&self.factory).append_scalar_i64(&steps) {
            return Source {
                factory,
                terminal_factory: None,
                hints: self.hints.without_inline_micro(),
                attributes: self.attributes,
                split_hook: None,
            };
        }
        if self.hints.inline_micro.is_some() {
            return Source {
                factory: Arc::new(ScalarI64SourceFactory {
                    source: self.factory,
                    steps,
                }),
                terminal_factory: None,
                hints: self.hints.without_inline_micro(),
                attributes: self.attributes,
                split_hook: None,
            };
        }
        match op {
            ScalarCompareOp::Equal => self.filter(move |value| *value == operand),
            ScalarCompareOp::NotEqual => self.filter(move |value| *value != operand),
            ScalarCompareOp::LessThan => self.filter(move |value| *value < operand),
            ScalarCompareOp::LessOrEqual => self.filter(move |value| *value <= operand),
            ScalarCompareOp::GreaterThan => self.filter(move |value| *value > operand),
            ScalarCompareOp::GreaterOrEqual => self.filter(move |value| *value >= operand),
        }
    }
}

pub(super) struct ScalarI64FlowPlan<In> {
    pub(super) prefix: super::PureTransform<In, i64>,
    pub(super) steps: Arc<[ScalarChunkStep<i64>]>,
}

impl<In> Clone for ScalarI64FlowPlan<In> {
    fn clone(&self) -> Self {
        Self {
            prefix: Arc::clone(&self.prefix),
            steps: Arc::clone(&self.steps),
        }
    }
}

fn scalar_flow_transform<In: Send + 'static>(
    prefix: super::PureTransform<In, i64>,
    steps: Arc<[ScalarChunkStep<i64>]>,
) -> super::PureTransform<In, i64> {
    Arc::new(move |input| {
        let input = prefix(input);
        scalar_i64_stream(input, Arc::clone(&steps))
    })
}

fn scalar_flow_fallback_transform<In: Send + 'static>(
    prefix: super::PureTransform<In, i64>,
    steps: Arc<[ScalarChunkStep<i64>]>,
) -> super::PureTransform<In, i64> {
    Arc::new(move |input| {
        Box::new(ScalarI64Stream::new_scalar(
            prefix(input),
            Arc::clone(&steps),
        ))
    })
}

impl<In, Mat> Flow<In, i64, Mat>
where
    In: Send + 'static,
    Mat: Send + 'static,
{
    fn scalar_op(self, op: ScalarI64Op, operand: i64) -> Self {
        if self.scalar_i64.is_none() && !self.hints.scalar_chunk_prefix {
            return match op {
                ScalarI64Op::Arithmetic(_) => self.try_map(move |value| {
                    op.apply_checked(value, operand)
                        .map(|value| value.expect("arithmetic maps retain their input"))
                }),
                ScalarI64Op::Compare(compare) => match compare {
                    ScalarCompareOp::Equal => self.filter(move |value| *value == operand),
                    ScalarCompareOp::NotEqual => self.filter(move |value| *value != operand),
                    ScalarCompareOp::LessThan => self.filter(move |value| *value < operand),
                    ScalarCompareOp::LessOrEqual => self.filter(move |value| *value <= operand),
                    ScalarCompareOp::GreaterThan => self.filter(move |value| *value > operand),
                    ScalarCompareOp::GreaterOrEqual => self.filter(move |value| *value >= operand),
                },
            };
        }
        let Flow {
            transform,
            materialize,
            hints,
            attributes,
            scalar_i64,
            scalar_i64_fallback,
        } = self;
        let next = ScalarChunkStep::new(op, operand);
        match (transform, scalar_i64) {
            (FlowTransform::Pure(_), Some(plan)) => {
                let plan = match Arc::downcast::<ScalarI64FlowPlan<In>>(plan) {
                    Ok(plan) => plan,
                    Err(_) => unreachable!("scalar i64 flow plan input type is preserved"),
                };
                let mut steps = Vec::with_capacity(plan.steps.len() + 1);
                steps.extend_from_slice(&plan.steps);
                steps.push(next.clone());
                let steps: Arc<[ScalarChunkStep<i64>]> = steps.into();
                let transform = scalar_flow_transform(Arc::clone(&plan.prefix), Arc::clone(&steps));
                let fallback_prefix = scalar_i64_fallback.unwrap_or_else(|| {
                    scalar_flow_fallback_transform(
                        Arc::clone(&plan.prefix),
                        Arc::clone(&plan.steps),
                    )
                });
                let fallback = scalar_flow_fallback_transform(fallback_prefix, steps_with(next));
                Flow {
                    transform: FlowTransform::Pure(transform),
                    materialize,
                    hints,
                    attributes,
                    scalar_i64: Some(Arc::new(ScalarI64FlowPlan {
                        prefix: Arc::clone(&plan.prefix),
                        steps,
                    })),
                    scalar_i64_fallback: Some(fallback),
                }
            }
            (FlowTransform::Pure(prefix), None) => {
                let steps = steps_with(next);
                let transform = scalar_flow_transform(Arc::clone(&prefix), Arc::clone(&steps));
                let fallback_prefix = scalar_i64_fallback.unwrap_or_else(|| Arc::clone(&prefix));
                let fallback = scalar_flow_fallback_transform(fallback_prefix, Arc::clone(&steps));
                Flow {
                    transform: FlowTransform::Pure(transform),
                    materialize,
                    hints,
                    attributes,
                    scalar_i64: Some(Arc::new(ScalarI64FlowPlan { prefix, steps })),
                    scalar_i64_fallback: Some(fallback),
                }
            }
            (FlowTransform::Runtime(runtime), _) => {
                let steps = steps_with(next);
                let scalar =
                    move |input: BoxStream<i64>| scalar_i64_stream(input, Arc::clone(&steps));
                Flow {
                    transform: FlowTransform::Runtime(Arc::new(move |input, materializer| {
                        runtime(input, materializer).map(&scalar)
                    })),
                    materialize,
                    hints: FlowHints::default(),
                    attributes,
                    scalar_i64: None,
                    scalar_i64_fallback: None,
                }
            }
        }
    }

    #[must_use]
    pub fn map_op(self, op: ScalarArithmeticOp, operand: i64) -> Self {
        self.scalar_op(ScalarI64Op::Arithmetic(op), operand)
    }

    #[must_use]
    pub fn map_add(self, operand: i64) -> Self {
        self.map_op(ScalarArithmeticOp::Add, operand)
    }

    #[must_use]
    pub fn map_subtract(self, operand: i64) -> Self {
        self.map_op(ScalarArithmeticOp::Subtract, operand)
    }

    #[must_use]
    pub fn map_multiply(self, operand: i64) -> Self {
        self.map_op(ScalarArithmeticOp::Multiply, operand)
    }

    #[must_use]
    pub fn filter_op(self, op: ScalarCompareOp, operand: i64) -> Self {
        if self.scalar_i64.is_some() || self.hints.scalar_chunk_prefix {
            return self.scalar_op(ScalarI64Op::Compare(op), operand);
        }
        match op {
            ScalarCompareOp::Equal => self.filter(move |value| *value == operand),
            ScalarCompareOp::NotEqual => self.filter(move |value| *value != operand),
            ScalarCompareOp::LessThan => self.filter(move |value| *value < operand),
            ScalarCompareOp::LessOrEqual => self.filter(move |value| *value <= operand),
            ScalarCompareOp::GreaterThan => self.filter(move |value| *value > operand),
            ScalarCompareOp::GreaterOrEqual => self.filter(move |value| *value >= operand),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Flow, Sink};
    use std::sync::{Arc, Mutex};

    #[test]
    fn first_failure_stage_and_valid_prefix_are_exact() {
        let emitted = Arc::new(Mutex::new(Vec::new()));
        let observed = Arc::clone(&emitted);
        let completion = Source::from_iter([0, i64::MIN + 5, i64::MAX])
            .map_add(1)
            .map_subtract(10)
            .run_with(Sink::foreach(move |value| {
                observed
                    .lock()
                    .expect("emitted values poisoned")
                    .push(value);
            }))
            .unwrap();

        assert_eq!(
            completion.wait(),
            Err(StreamError::Failed(format!(
                "integer overflow in map_subtract: {} - 10",
                i64::MIN + 6
            )))
        );
        assert_eq!(*emitted.lock().unwrap(), vec![-9]);
    }

    #[test]
    fn overflow_positions_around_chunk_boundary_match_scalar_oracle() {
        for len in [0, 1, 2, 3, 127, 128, 129, 255, 256, 257] {
            let mut input = vec![0_i64; len];
            if let Some(last) = input.last_mut() {
                *last = i64::MAX;
            }
            let chunked = Source::from_iter(input.clone())
                .map_add(1)
                .run_with(Sink::collect())
                .unwrap()
                .wait();
            let scalar = Source::from_iter(input)
                .try_map(|value| {
                    value.checked_add(1).ok_or_else(|| {
                        StreamError::Failed(format!("integer overflow in map_add: {value} + 1"))
                    })
                })
                .run_with(Sink::collect())
                .unwrap()
                .wait();
            assert_eq!(chunked, scalar, "length {len}");
        }
    }

    #[test]
    fn filters_compact_stably_before_and_after_fallible_maps() {
        let result = Source::from_iter(-300_i64..300)
            .filter_op(ScalarCompareOp::GreaterOrEqual, -250)
            .map_add(7)
            .filter_op(ScalarCompareOp::LessThan, 10)
            .run_with(Sink::collect())
            .unwrap()
            .wait()
            .unwrap();
        let expected: Vec<_> = (-300_i64..300)
            .filter(|value| *value >= -250)
            .map(|value| value + 7)
            .filter(|value| *value < 10)
            .collect();
        assert_eq!(result, expected);
    }

    #[test]
    fn flow_builtins_fuse_but_opaque_closures_remain_per_element() {
        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let observed = Arc::clone(&calls);
        let flow = Flow::identity()
            .map(move |value: i64| {
                observed.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                value
            })
            .map_add(2)
            .map_subtract(1);
        let output = Source::from_iter(0_i64..600)
            .via(flow)
            .run_with(Sink::collect())
            .unwrap()
            .wait()
            .unwrap();
        assert_eq!(calls.load(std::sync::atomic::Ordering::Relaxed), 600);
        assert_eq!(output[0], 1);
        assert_eq!(output[599], 600);

        let direct_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let observed = Arc::clone(&direct_calls);
        let head = Source::from_iter(0_i64..600)
            .map(move |value| {
                observed.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                value
            })
            .map_add(1)
            .run_with(Sink::head())
            .unwrap()
            .wait()
            .unwrap();
        assert_eq!(head, 1);
        assert_eq!(
            direct_calls.load(std::sync::atomic::Ordering::Relaxed),
            1,
            "an opaque prefix must not be pulled ahead by the chunk executor"
        );

        let flow_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let observed = Arc::clone(&flow_calls);
        let flow = Flow::identity()
            .map(move |value: i64| {
                observed.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                value
            })
            .map_add(1);
        let head = Source::from_iter(0_i64..600)
            .via(flow)
            .run_with(Sink::head())
            .unwrap()
            .wait()
            .unwrap();
        assert_eq!(head, 1);
        assert_eq!(
            flow_calls.load(std::sync::atomic::Ordering::Relaxed),
            1,
            "an opaque Flow prefix must remain one-element-at-a-time"
        );

        let dynamic_pulls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let observed = Arc::clone(&dynamic_pulls);
        let flow = Flow::identity()
            .map_add(1)
            .via(Flow::identity())
            .map_subtract(1)
            .map(|value| value);
        let head = Source::unfold(0_i64, move |value| {
            observed.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            Some((value + 1, value))
        })
        .via(flow)
        .run_with(Sink::head())
        .unwrap()
        .wait()
        .unwrap();
        assert_eq!(head, 0);
        assert_eq!(
            dynamic_pulls.load(std::sync::atomic::Ordering::Relaxed),
            1,
            "a dynamic source must select the scalar fallback for a chunked Flow"
        );

        let runtime_pulls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let observed = Arc::clone(&runtime_pulls);
        let flow = Flow::identity()
            .map_add(1)
            .via(Flow::from_runtime_transform(|input, _materializer| {
                Ok(input)
            }));
        let head = Source::unfold(0_i64, move |value| {
            observed.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            Some((value + 1, value))
        })
        .via(flow)
        .run_with(Sink::head())
        .unwrap()
        .wait()
        .unwrap();
        assert_eq!(head, 1);
        assert_eq!(
            runtime_pulls.load(std::sync::atomic::Ordering::Relaxed),
            1,
            "runtime Flow composition must scalarize an earlier chunked segment"
        );
    }

    #[test]
    fn multiplication_remains_checked_scalar() {
        let result = Source::from_iter([2_i64, i64::MAX])
            .map_multiply(2)
            .run_with(Sink::collect())
            .unwrap()
            .wait();
        assert_eq!(
            result,
            Err(StreamError::Failed(format!(
                "integer overflow in map_multiply: {} * 2",
                i64::MAX
            )))
        );
    }

    #[test]
    fn deterministic_random_add_sub_segments_match_scalar_oracle() {
        let mut state = 0x9e37_79b9_7f4a_7c15_u64;
        let mut random = || {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            state
        };
        for case in 0..200 {
            let len = (random() as usize) % 514;
            let add = random() as i64;
            let subtract = random() as i64;
            let mut input = Vec::with_capacity(len);
            for _ in 0..len {
                input.push(random() as i64);
            }
            if case % 4 == 0 && !input.is_empty() {
                input[(random() as usize) % len] = i64::MAX;
            }

            let chunked = Source::from_iter(input.clone())
                .map_add(add)
                .map_subtract(subtract)
                .run_with(Sink::collect())
                .unwrap()
                .wait();
            let scalar = Source::from_iter(input)
                .try_map(move |value| {
                    value.checked_add(add).ok_or_else(|| {
                        StreamError::Failed(format!("integer overflow in map_add: {value} + {add}"))
                    })
                })
                .try_map(move |value| {
                    value.checked_sub(subtract).ok_or_else(|| {
                        StreamError::Failed(format!(
                            "integer overflow in map_subtract: {value} - {subtract}"
                        ))
                    })
                })
                .run_with(Sink::collect())
                .unwrap()
                .wait();
            assert_eq!(chunked, scalar, "random case {case}, length {len}");
        }
    }
}