reddb_server/storage/query/engine/

iterator.rs

//! Query Result Iterators
//!
//! Lazy evaluation of query results through binding streams.
//!
//! # Design
//!
//! - Pull-based iteration (demand-driven)
//! - Composable iterator wrappers
//! - Memory-efficient streaming
//! - Early termination support

use super::binding::{Binding, Value, Var};
use super::cancel::CancelToken;
use std::collections::HashSet;
use std::fmt::Debug;

/// Issue #808 / 750d — cooperative cancellation check shared by every
/// pull-based operator. Returns [`IterError::Cancelled`] the moment the
/// stream layer raises the token, so a long-running scan / filter / join
/// loop stops within one row of the cancel signal rather than running to
/// completion. A `None` token (the default) never cancels, so operators
/// built without a token behave exactly as before.
#[inline]
fn check_shared_cancel(token: &Option<CancelToken>) -> Result<(), IterError> {
    if let Some(token) = token {
        if token.is_cancelled() {
            return Err(IterError::Cancelled);
        }
    }
    Ok(())
}

/// Result from iterator operations
pub type IterResult = Result<Option<Binding>, IterError>;

/// Iterator errors
#[derive(Debug, Clone)]
pub enum IterError {
    /// Execution error
    Execution(String),
    /// Timeout
    Timeout,
    /// Cancelled
    Cancelled,
    /// Resource exhausted
    ResourceExhausted(String),
}

impl std::fmt::Display for IterError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            IterError::Execution(msg) => write!(f, "Execution error: {}", msg),
            IterError::Timeout => write!(f, "Query timeout"),
            IterError::Cancelled => write!(f, "Query cancelled"),
            IterError::ResourceExhausted(msg) => write!(f, "Resource exhausted: {}", msg),
        }
    }
}

impl std::error::Error for IterError {}

/// Core trait for binding iterators
pub trait BindingIterator: Debug + Send {
    /// Get next binding
    fn next_binding(&mut self) -> IterResult;

    /// Check if there are more results (optional hint)
    fn has_next(&self) -> bool {
        true // Default: unknown, caller should try next()
    }

    /// Get variables produced by this iterator
    fn vars(&self) -> Vec<Var>;

    /// Cancel iteration
    fn cancel(&mut self);

    /// Reset iterator to beginning (if supported)
    fn reset(&mut self) -> bool {
        false // Default: not supported
    }
}

/// Base query iterator wrapping a binding source
#[derive(Debug)]
pub struct QueryIterBase {
    bindings: Vec<Binding>,
    index: usize,
    vars: Vec<Var>,
    cancelled: bool,
    /// Issue #808 / 750d — shared cancel token observed on every pull so a
    /// scan over a large binding set stops promptly when the stream is
    /// cancelled. `None` (the default) preserves the legacy behaviour.
    cancel: Option<CancelToken>,
}

impl QueryIterBase {
    /// Create from binding list
    pub fn new(bindings: Vec<Binding>) -> Self {
        let vars = if let Some(first) = bindings.first() {
            first.all_vars().into_iter().cloned().collect()
        } else {
            Vec::new()
        };

        Self {
            bindings,
            index: 0,
            vars,
            cancelled: false,
            cancel: None,
        }
    }

    /// Create empty iterator
    pub fn empty() -> Self {
        Self {
            bindings: Vec::new(),
            index: 0,
            vars: Vec::new(),
            cancelled: false,
            cancel: None,
        }
    }

    /// Create single-result iterator
    pub fn single(binding: Binding) -> Self {
        let vars = binding.all_vars().into_iter().cloned().collect();
        Self {
            bindings: vec![binding],
            index: 0,
            vars,
            cancelled: false,
            cancel: None,
        }
    }

    /// Issue #808 / 750d — attach a shared cancel token. The scan checks it
    /// on every `next_binding`, so cancelling the token (from the stream
    /// layer on disconnect or explicit cancel) halts production within one
    /// row regardless of how many bindings remain.
    pub fn with_cancel(mut self, token: CancelToken) -> Self {
        self.cancel = Some(token);
        self
    }
}

impl BindingIterator for QueryIterBase {
    fn next_binding(&mut self) -> IterResult {
        if self.cancelled {
            return Err(IterError::Cancelled);
        }
        check_shared_cancel(&self.cancel)?;

        if self.index < self.bindings.len() {
            // Each binding is emitted exactly once, so move it out of the
            // backing store instead of cloning per row. `reset()` remains
            // structurally valid (index rewinds); it just re-emits the moved
            // slots as defaults — and it is never invoked on this iterator.
            let binding = std::mem::take(&mut self.bindings[self.index]);
            self.index += 1;
            Ok(Some(binding))
        } else {
            Ok(None)
        }
    }

    fn has_next(&self) -> bool {
        !self.cancelled && self.index < self.bindings.len()
    }

    fn vars(&self) -> Vec<Var> {
        self.vars.clone()
    }

    fn cancel(&mut self) {
        self.cancelled = true;
    }

    fn reset(&mut self) -> bool {
        self.index = 0;
        self.cancelled = false;
        true
    }
}

/// Filter iterator - applies predicate to upstream
pub struct QueryIterFilter {
    upstream: Box<dyn BindingIterator>,
    predicate: Box<dyn Fn(&Binding) -> bool + Send + Sync>,
    cancelled: bool,
    /// Issue #808 / 750d — shared cancel token. A highly selective filter
    /// can pull many upstream rows before yielding one; the token is
    /// re-checked on every upstream pull so the rejection loop cannot
    /// outrun the cancel signal.
    cancel: Option<CancelToken>,
}

impl std::fmt::Debug for QueryIterFilter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("QueryIterFilter")
            .field("upstream", &self.upstream)
            .field("predicate", &"<filter fn>")
            .field("cancelled", &self.cancelled)
            .finish()
    }
}

impl QueryIterFilter {
    /// Create filter iterator
    pub fn new<F>(upstream: Box<dyn BindingIterator>, predicate: F) -> Self
    where
        F: Fn(&Binding) -> bool + Send + Sync + 'static,
    {
        Self {
            upstream,
            predicate: Box::new(predicate),
            cancelled: false,
            cancel: None,
        }
    }

    /// Issue #808 / 750d — attach a shared cancel token (see [`QueryIterBase::with_cancel`]).
    pub fn with_cancel(mut self, token: CancelToken) -> Self {
        self.cancel = Some(token);
        self
    }
}

impl BindingIterator for QueryIterFilter {
    fn next_binding(&mut self) -> IterResult {
        if self.cancelled {
            return Err(IterError::Cancelled);
        }

        loop {
            check_shared_cancel(&self.cancel)?;
            match self.upstream.next_binding()? {
                Some(binding) => {
                    if (self.predicate)(&binding) {
                        return Ok(Some(binding));
                    }
                    // Continue to next
                }
                None => return Ok(None),
            }
        }
    }

    fn vars(&self) -> Vec<Var> {
        self.upstream.vars()
    }

    fn cancel(&mut self) {
        self.cancelled = true;
        self.upstream.cancel();
    }
}

/// Join iterator - nested loop join
pub struct QueryIterJoin {
    left: Box<dyn BindingIterator>,
    right_factory: Box<dyn Fn() -> Box<dyn BindingIterator> + Send + Sync>,
    current_left: Option<Binding>,
    current_right: Option<Box<dyn BindingIterator>>,
    vars: Vec<Var>,
    cancelled: bool,
    /// Issue #808 / 750d — shared cancel token. A nested-loop join is the
    /// worst case for unbounded work (|left| × |right|); the token is
    /// re-checked on every iteration of the merge loop so a cancel stops
    /// the join without draining the cross product.
    cancel: Option<CancelToken>,
}

impl std::fmt::Debug for QueryIterJoin {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("QueryIterJoin")
            .field("left", &self.left)
            .field("right_factory", &"<factory fn>")
            .field("current_left", &self.current_left)
            .field("current_right", &self.current_right)
            .field("vars", &self.vars)
            .field("cancelled", &self.cancelled)
            .finish()
    }
}

impl QueryIterJoin {
    /// Create join iterator
    pub fn new<F>(left: Box<dyn BindingIterator>, right_factory: F, right_vars: Vec<Var>) -> Self
    where
        F: Fn() -> Box<dyn BindingIterator> + Send + Sync + 'static,
    {
        let mut vars = left.vars();
        for v in right_vars {
            if !vars.contains(&v) {
                vars.push(v);
            }
        }

        Self {
            left,
            right_factory: Box::new(right_factory),
            current_left: None,
            current_right: None,
            vars,
            cancelled: false,
            cancel: None,
        }
    }

    /// Issue #808 / 750d — attach a shared cancel token (see [`QueryIterBase::with_cancel`]).
    pub fn with_cancel(mut self, token: CancelToken) -> Self {
        self.cancel = Some(token);
        self
    }
}

impl BindingIterator for QueryIterJoin {
    fn next_binding(&mut self) -> IterResult {
        if self.cancelled {
            return Err(IterError::Cancelled);
        }

        loop {
            check_shared_cancel(&self.cancel)?;
            // Try to get next from current right
            if let Some(ref mut right) = self.current_right {
                if let Some(right_binding) = right.next_binding()? {
                    // Merge with current left
                    if let Some(ref left_binding) = self.current_left {
                        if let Some(merged) = left_binding.merge(&right_binding) {
                            return Ok(Some(merged));
                        }
                        // Conflict, continue to next right
                        continue;
                    }
                }
            }

            // Need new left binding
            match self.left.next_binding()? {
                Some(left_binding) => {
                    self.current_left = Some(left_binding);
                    self.current_right = Some((self.right_factory)());
                }
                None => return Ok(None),
            }
        }
    }

    fn vars(&self) -> Vec<Var> {
        self.vars.clone()
    }

    fn cancel(&mut self) {
        self.cancelled = true;
        self.left.cancel();
        if let Some(ref mut right) = self.current_right {
            right.cancel();
        }
    }
}

/// Union iterator - concatenates multiple iterators
#[derive(Debug)]
pub struct QueryIterUnion {
    iterators: Vec<Box<dyn BindingIterator>>,
    current_index: usize,
    vars: Vec<Var>,
    cancelled: bool,
}

impl QueryIterUnion {
    /// Create union iterator
    pub fn new(iterators: Vec<Box<dyn BindingIterator>>) -> Self {
        let mut vars = Vec::new();
        for iter in &iterators {
            for v in iter.vars() {
                if !vars.contains(&v) {
                    vars.push(v);
                }
            }
        }

        Self {
            iterators,
            current_index: 0,
            vars,
            cancelled: false,
        }
    }
}

impl BindingIterator for QueryIterUnion {
    fn next_binding(&mut self) -> IterResult {
        if self.cancelled {
            return Err(IterError::Cancelled);
        }

        while self.current_index < self.iterators.len() {
            match self.iterators[self.current_index].next_binding()? {
                Some(binding) => return Ok(Some(binding)),
                None => {
                    self.current_index += 1;
                }
            }
        }

        Ok(None)
    }

    fn vars(&self) -> Vec<Var> {
        self.vars.clone()
    }

    fn cancel(&mut self) {
        self.cancelled = true;
        for iter in &mut self.iterators {
            iter.cancel();
        }
    }
}

/// Project iterator - selects subset of variables
#[derive(Debug)]
pub struct QueryIterProject {
    upstream: Box<dyn BindingIterator>,
    project_vars: Vec<Var>,
    cancelled: bool,
}

impl QueryIterProject {
    /// Create project iterator
    pub fn new(upstream: Box<dyn BindingIterator>, vars: Vec<Var>) -> Self {
        Self {
            upstream,
            project_vars: vars,
            cancelled: false,
        }
    }
}

impl BindingIterator for QueryIterProject {
    fn next_binding(&mut self) -> IterResult {
        if self.cancelled {
            return Err(IterError::Cancelled);
        }

        match self.upstream.next_binding()? {
            Some(binding) => Ok(Some(binding.project(&self.project_vars))),
            None => Ok(None),
        }
    }

    fn vars(&self) -> Vec<Var> {
        self.project_vars.clone()
    }

    fn cancel(&mut self) {
        self.cancelled = true;
        self.upstream.cancel();
    }
}

/// Slice iterator - limit and offset
#[derive(Debug)]
pub struct QueryIterSlice {
    upstream: Box<dyn BindingIterator>,
    offset: u64,
    limit: Option<u64>,
    skipped: u64,
    returned: u64,
    cancelled: bool,
}

impl QueryIterSlice {
    /// Create slice iterator
    pub fn new(upstream: Box<dyn BindingIterator>, offset: u64, limit: Option<u64>) -> Self {
        Self {
            upstream,
            offset,
            limit,
            skipped: 0,
            returned: 0,
            cancelled: false,
        }
    }

    /// Create limit-only iterator
    pub fn limit(upstream: Box<dyn BindingIterator>, limit: u64) -> Self {
        Self::new(upstream, 0, Some(limit))
    }

    /// Create offset-only iterator
    pub fn offset(upstream: Box<dyn BindingIterator>, offset: u64) -> Self {
        Self::new(upstream, offset, None)
    }
}

impl BindingIterator for QueryIterSlice {
    fn next_binding(&mut self) -> IterResult {
        if self.cancelled {
            return Err(IterError::Cancelled);
        }

        // Check limit
        if let Some(limit) = self.limit {
            if self.returned >= limit {
                return Ok(None);
            }
        }

        // Skip offset
        while self.skipped < self.offset {
            match self.upstream.next_binding()? {
                Some(_) => {
                    self.skipped += 1;
                }
                None => return Ok(None),
            }
        }

        // Return result
        match self.upstream.next_binding()? {
            Some(binding) => {
                self.returned += 1;
                Ok(Some(binding))
            }
            None => Ok(None),
        }
    }

    fn vars(&self) -> Vec<Var> {
        self.upstream.vars()
    }

    fn cancel(&mut self) {
        self.cancelled = true;
        self.upstream.cancel();
    }
}

/// Sort iterator - orders results
#[derive(Debug)]
pub struct QueryIterSort {
    upstream: Box<dyn BindingIterator>,
    comparators: Vec<SortKey>,
    sorted: Option<Vec<Binding>>,
    index: usize,
    cancelled: bool,
}

/// Sort key specification
#[derive(Debug, Clone)]
pub struct SortKey {
    /// Variable to sort by
    pub var: Var,
    /// Ascending order
    pub ascending: bool,
}

impl SortKey {
    /// Create ascending sort key
    pub fn asc(var: Var) -> Self {
        Self {
            var,
            ascending: true,
        }
    }

    /// Create descending sort key
    pub fn desc(var: Var) -> Self {
        Self {
            var,
            ascending: false,
        }
    }
}

impl QueryIterSort {
    /// Create sort iterator
    pub fn new(upstream: Box<dyn BindingIterator>, comparators: Vec<SortKey>) -> Self {
        Self {
            upstream,
            comparators,
            sorted: None,
            index: 0,
            cancelled: false,
        }
    }

    /// Materialize and sort all results
    fn materialize(&mut self) -> Result<(), IterError> {
        if self.sorted.is_some() {
            return Ok(());
        }

        let mut bindings = Vec::new();
        while let Some(b) = self.upstream.next_binding()? {
            bindings.push(b);
        }

        // Sort by comparators
        let comparators = self.comparators.clone();
        bindings.sort_by(|a, b| {
            for key in &comparators {
                let a_val = a.get(&key.var);
                let b_val = b.get(&key.var);

                let ordering = compare_values(a_val, b_val);
                if ordering != std::cmp::Ordering::Equal {
                    return if key.ascending {
                        ordering
                    } else {
                        ordering.reverse()
                    };
                }
            }
            std::cmp::Ordering::Equal
        });

        self.sorted = Some(bindings);
        Ok(())
    }
}

impl BindingIterator for QueryIterSort {
    fn next_binding(&mut self) -> IterResult {
        if self.cancelled {
            return Err(IterError::Cancelled);
        }

        self.materialize()?;

        if let Some(ref sorted) = self.sorted {
            if self.index < sorted.len() {
                let binding = sorted[self.index].clone();
                self.index += 1;
                return Ok(Some(binding));
            }
        }

        Ok(None)
    }

    fn vars(&self) -> Vec<Var> {
        self.upstream.vars()
    }

    fn cancel(&mut self) {
        self.cancelled = true;
        self.upstream.cancel();
    }
}

/// Compare two optional values
fn compare_values(a: Option<&Value>, b: Option<&Value>) -> std::cmp::Ordering {
    match (a, b) {
        (None, None) => std::cmp::Ordering::Equal,
        (None, Some(_)) => std::cmp::Ordering::Less,
        (Some(_), None) => std::cmp::Ordering::Greater,
        (Some(a_val), Some(b_val)) => compare_value(a_val, b_val),
    }
}

/// Compare two values
fn compare_value(a: &Value, b: &Value) -> std::cmp::Ordering {
    match (a, b) {
        (Value::Integer(a), Value::Integer(b)) => a.cmp(b),
        (Value::Float(a), Value::Float(b)) => a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal),
        (Value::String(a), Value::String(b)) => a.cmp(b),
        (Value::Boolean(a), Value::Boolean(b)) => a.cmp(b),
        (Value::Node(a), Value::Node(b)) => a.cmp(b),
        (Value::Edge(a), Value::Edge(b)) => a.cmp(b),
        (Value::Uri(a), Value::Uri(b)) => a.cmp(b),
        (Value::Null, Value::Null) => std::cmp::Ordering::Equal,
        // Cross-type comparison: types differ, so we return a consistent ordering
        _ => {
            let type_order = |v: &Value| -> u8 {
                match v {
                    Value::Null => 0,
                    Value::Boolean(_) => 1,
                    Value::Integer(_) => 2,
                    Value::Float(_) => 3,
                    Value::String(_) => 4,
                    Value::Node(_) => 5,
                    Value::Edge(_) => 6,
                    Value::Uri(_) => 7,
                }
            };
            type_order(a).cmp(&type_order(b))
        }
    }
}

/// Distinct iterator - removes duplicates
#[derive(Debug)]
pub struct QueryIterDistinct {
    upstream: Box<dyn BindingIterator>,
    seen: HashSet<u64>,
    cancelled: bool,
}

impl QueryIterDistinct {
    /// Create distinct iterator
    pub fn new(upstream: Box<dyn BindingIterator>) -> Self {
        Self {
            upstream,
            seen: HashSet::new(),
            cancelled: false,
        }
    }

    /// Hash a binding for deduplication
    fn hash_binding(binding: &Binding) -> u64 {
        use std::hash::{Hash, Hasher};
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        binding.hash(&mut hasher);
        hasher.finish()
    }
}

impl BindingIterator for QueryIterDistinct {
    fn next_binding(&mut self) -> IterResult {
        if self.cancelled {
            return Err(IterError::Cancelled);
        }

        loop {
            match self.upstream.next_binding()? {
                Some(binding) => {
                    let hash = Self::hash_binding(&binding);
                    if self.seen.insert(hash) {
                        return Ok(Some(binding));
                    }
                    // Already seen, continue
                }
                None => return Ok(None),
            }
        }
    }

    fn vars(&self) -> Vec<Var> {
        self.upstream.vars()
    }

    fn cancel(&mut self) {
        self.cancelled = true;
        self.upstream.cancel();
    }
}

/// Wrapper for boxed iterator to implement Iterator trait
pub struct QueryIter {
    inner: Box<dyn BindingIterator>,
}

impl QueryIter {
    /// Create from binding iterator
    pub fn new(inner: Box<dyn BindingIterator>) -> Self {
        Self { inner }
    }

    /// Get variables
    pub fn vars(&self) -> Vec<Var> {
        self.inner.vars()
    }

    /// Cancel iteration
    pub fn cancel(&mut self) {
        self.inner.cancel();
    }
}

impl Iterator for QueryIter {
    type Item = Result<Binding, IterError>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.inner.next_binding() {
            Ok(Some(binding)) => Some(Ok(binding)),
            Ok(None) => None,
            Err(e) => Some(Err(e)),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::storage::query::engine::binding::BindingBuilder;

    fn make_binding(x: i64) -> Binding {
        BindingBuilder::new()
            .add_named("x", Value::Integer(x))
            .build()
    }

    #[test]
    fn test_base_iterator() {
        let bindings = vec![make_binding(1), make_binding(2), make_binding(3)];
        let mut iter = QueryIterBase::new(bindings);

        assert!(iter.has_next());
        assert!(iter.next_binding().unwrap().is_some());
        assert!(iter.next_binding().unwrap().is_some());
        assert!(iter.next_binding().unwrap().is_some());
        assert!(iter.next_binding().unwrap().is_none());
    }

    #[test]
    fn test_filter_iterator() {
        let bindings = vec![make_binding(1), make_binding(2), make_binding(3)];
        let base = Box::new(QueryIterBase::new(bindings));

        let mut iter = QueryIterFilter::new(base, |b| {
            b.get(&Var::new("x"))
                .and_then(|v| v.as_integer())
                .map(|i| i > 1)
                .unwrap_or(false)
        });

        // Should skip 1, return 2 and 3
        let b1 = iter.next_binding().unwrap().unwrap();
        assert_eq!(b1.get(&Var::new("x")), Some(&Value::Integer(2)));

        let b2 = iter.next_binding().unwrap().unwrap();
        assert_eq!(b2.get(&Var::new("x")), Some(&Value::Integer(3)));

        assert!(iter.next_binding().unwrap().is_none());
    }

    #[test]
    fn test_slice_iterator() {
        let bindings: Vec<_> = (1..=10).map(make_binding).collect();
        let base = Box::new(QueryIterBase::new(bindings));

        // Offset 2, limit 3
        let mut iter = QueryIterSlice::new(base, 2, Some(3));

        let b1 = iter.next_binding().unwrap().unwrap();
        assert_eq!(b1.get(&Var::new("x")), Some(&Value::Integer(3)));

        let b2 = iter.next_binding().unwrap().unwrap();
        assert_eq!(b2.get(&Var::new("x")), Some(&Value::Integer(4)));

        let b3 = iter.next_binding().unwrap().unwrap();
        assert_eq!(b3.get(&Var::new("x")), Some(&Value::Integer(5)));

        assert!(iter.next_binding().unwrap().is_none());
    }

    #[test]
    fn test_project_iterator() {
        let binding = BindingBuilder::new()
            .add_named("x", Value::Integer(1))
            .add_named("y", Value::Integer(2))
            .add_named("z", Value::Integer(3))
            .build();

        let base = Box::new(QueryIterBase::single(binding));
        let mut iter = QueryIterProject::new(base, vec![Var::new("x"), Var::new("z")]);

        let result = iter.next_binding().unwrap().unwrap();
        assert!(result.contains(&Var::new("x")));
        assert!(!result.contains(&Var::new("y")));
        assert!(result.contains(&Var::new("z")));
    }

    #[test]
    fn test_union_iterator() {
        let iter1 = Box::new(QueryIterBase::new(vec![make_binding(1), make_binding(2)]));
        let iter2 = Box::new(QueryIterBase::new(vec![make_binding(3), make_binding(4)]));

        let mut union = QueryIterUnion::new(vec![iter1, iter2]);

        let mut results = Vec::new();
        while let Some(b) = union.next_binding().unwrap() {
            results.push(b.get(&Var::new("x")).unwrap().as_integer().unwrap());
        }

        assert_eq!(results, vec![1, 2, 3, 4]);
    }

    #[test]
    fn test_sort_iterator() {
        let bindings = vec![make_binding(3), make_binding(1), make_binding(2)];
        let base = Box::new(QueryIterBase::new(bindings));

        let mut iter = QueryIterSort::new(base, vec![SortKey::asc(Var::new("x"))]);

        let mut results = Vec::new();
        while let Some(b) = iter.next_binding().unwrap() {
            results.push(b.get(&Var::new("x")).unwrap().as_integer().unwrap());
        }

        assert_eq!(results, vec![1, 2, 3]);
    }

    #[test]
    fn test_distinct_iterator() {
        let bindings = vec![
            make_binding(1),
            make_binding(2),
            make_binding(1),
            make_binding(3),
            make_binding(2),
        ];
        let base = Box::new(QueryIterBase::new(bindings));

        let mut iter = QueryIterDistinct::new(base);

        let mut results = Vec::new();
        while let Some(b) = iter.next_binding().unwrap() {
            results.push(b.get(&Var::new("x")).unwrap().as_integer().unwrap());
        }

        assert_eq!(results, vec![1, 2, 3]);
    }

    #[test]
    fn test_cancel_iterator() {
        let bindings: Vec<_> = (1..=100).map(make_binding).collect();
        let mut iter = QueryIterBase::new(bindings);

        // Read a few
        iter.next_binding().unwrap();
        iter.next_binding().unwrap();

        // Cancel
        iter.cancel();

        // Should return cancelled error
        assert!(matches!(iter.next_binding(), Err(IterError::Cancelled)));
    }

    #[test]
    fn test_query_iter_wrapper() {
        let bindings = vec![make_binding(1), make_binding(2)];
        let base = Box::new(QueryIterBase::new(bindings));

        let iter = QueryIter::new(base);
        let results: Vec<_> = iter.collect();

        assert_eq!(results.len(), 2);
        assert!(results.iter().all(|r| r.is_ok()));
    }

    // ---- Issue #808 / 750d — cooperative shared-token cancellation. ----

    /// An upstream that yields rows forever, so a downstream operator's
    /// loop only terminates when it observes the shared cancel token. Used
    /// to pin the latency budget for the cancel signal without depending on
    /// a finite (and slow-to-build) materialised input.
    #[derive(Debug)]
    struct InfiniteScan {
        n: i64,
    }

    impl BindingIterator for InfiniteScan {
        fn next_binding(&mut self) -> IterResult {
            self.n += 1;
            Ok(Some(make_binding(self.n)))
        }

        fn vars(&self) -> Vec<Var> {
            vec![Var::new("x")]
        }

        fn cancel(&mut self) {}
    }

    #[test]
    fn shared_token_cancels_scan() {
        let token = CancelToken::new();
        let bindings: Vec<_> = (1..=1000).map(make_binding).collect();
        let mut scan = QueryIterBase::new(bindings).with_cancel(token.clone());

        assert!(scan.next_binding().unwrap().is_some());
        token.cancel();
        // The scan still has hundreds of rows buffered, but the shared
        // token short-circuits the next pull.
        assert!(matches!(scan.next_binding(), Err(IterError::Cancelled)));
    }

    #[test]
    fn shared_token_cancels_filter_within_budget() {
        use std::time::{Duration, Instant};

        let token = CancelToken::new();
        // Always-false predicate: the filter pulls from an infinite scan
        // forever, never yielding. Only the shared token can stop it.
        let upstream = Box::new(InfiniteScan { n: 0 });
        let mut filter = QueryIterFilter::new(upstream, |_| false).with_cancel(token.clone());

        let canceller = token.clone();
        std::thread::spawn(move || {
            std::thread::sleep(Duration::from_millis(20));
            canceller.cancel();
        });

        let start = Instant::now();
        let result = filter.next_binding();
        let elapsed = start.elapsed();

        assert!(matches!(result, Err(IterError::Cancelled)));
        // Soft latency budget: the cancel must take effect well within a
        // second even though the filter would otherwise loop forever.
        assert!(
            elapsed < Duration::from_secs(1),
            "cancel took too long to take effect: {elapsed:?}"
        );
    }

    #[test]
    fn shared_token_cancels_join_within_budget() {
        use std::time::{Duration, Instant};

        let token = CancelToken::new();
        // Left side never yields a mergeable row (always-false filter over
        // an infinite scan), so the join's merge loop spins until the
        // shared token stops it.
        let left = Box::new(
            QueryIterFilter::new(Box::new(InfiniteScan { n: 0 }), |_| false)
                .with_cancel(token.clone()),
        );
        let mut join = QueryIterJoin::new(
            left,
            || Box::new(QueryIterBase::empty()),
            vec![Var::new("y")],
        )
        .with_cancel(token.clone());

        let canceller = token.clone();
        std::thread::spawn(move || {
            std::thread::sleep(Duration::from_millis(20));
            canceller.cancel();
        });

        let start = Instant::now();
        let result = join.next_binding();
        let elapsed = start.elapsed();

        assert!(matches!(result, Err(IterError::Cancelled)));
        assert!(
            elapsed < Duration::from_secs(1),
            "join cancel took too long: {elapsed:?}"
        );
    }

    #[test]
    fn cancellation_propagates_through_aggregate_consumer() {
        // An aggregate (e.g. COUNT) drains its upstream to completion. When
        // the upstream is cancel-aware, the aggregate's own pull surfaces
        // the Cancelled error rather than looping forever — i.e. the cancel
        // signal propagates *through* the aggregate without the aggregate
        // needing its own token.
        let token = CancelToken::new();
        let upstream: Box<dyn BindingIterator> = Box::new(
            QueryIterFilter::new(Box::new(InfiniteScan { n: 0 }), |_| false)
                .with_cancel(token.clone()),
        );

        // Simulate an aggregate's drain loop.
        fn count_all(mut it: Box<dyn BindingIterator>) -> Result<u64, IterError> {
            let mut n = 0;
            while it.next_binding()?.is_some() {
                n += 1;
            }
            Ok(n)
        }

        token.cancel();
        assert!(matches!(count_all(upstream), Err(IterError::Cancelled)));
    }
}