qdrant-edge 0.7.0

#[cfg(any(test, feature = "testing"))]
use std::fmt::Debug;
use std::sync::atomic::AtomicBool;

use check_stopped::CheckStopped;
use on_final_count::OnFinalCount;

use crate::common::iterator_ext::stoppable_iter::StoppableIter;

pub(super) mod on_final_count;

mod check_stopped;
mod fallible;
pub mod ordering_iterator;
pub mod stoppable_iter;

pub use fallible::{FallibleIteratorExt, TransposeResultIter};

pub trait IteratorExt: Iterator {
    /// Periodically check if the iteration should be stopped.
    /// The closure `f` is called every `every` iterations, and should return `true` if the iteration should be stopped.
    fn check_stop_every<F>(self, every: usize, f: F) -> CheckStopped<Self, F>
    where
        F: FnMut() -> bool,
        Self: Sized,
    {
        CheckStopped::new(self, every, f)
    }

    /// Stops the iterator if `is_stopped` is set to true
    #[inline]
    fn stop_if<'a>(self, is_stopped: &'a AtomicBool) -> StoppableIter<'a, Self>
    where
        Self: Sized,
    {
        StoppableIter::new(self, is_stopped)
    }

    /// Will execute the callback when the iterator is dropped.
    ///
    /// The callback receives the total number of times `.next()` was called on the iterator,
    /// including the final one where it usually returns `None`.
    ///
    /// Consider subtracting 1 if the final `None` is not needed.
    fn on_final_count<F>(self, f: F) -> OnFinalCount<Self, F>
    where
        F: FnMut(usize),
        Self: Sized,
    {
        OnFinalCount::new(self, f)
    }

    /// Consume the iterator and call `black_box` on each item, for benchmarking purposes.
    fn black_box(self)
    where
        Self: Sized,
    {
        self.for_each(|p| {
            std::hint::black_box(p);
        });
    }

    /// [`Iterator::any()`] but for fallible predicates.
    fn try_any<F, E>(&mut self, mut f: F) -> Result<bool, E>
    where
        F: FnMut(Self::Item) -> Result<bool, E>,
        Self: Sized,
    {
        self.find_map(|item| match f(item) {
            Ok(true) => Some(Ok(true)),
            Ok(false) => None,
            Err(e) => Some(Err(e)),
        })
        .unwrap_or(Ok(false))
    }
}

impl<I: Iterator> IteratorExt for I {}

/// Checks that [`Iterator::fold()`] yields same values as [`Iterator::next()`].
/// Panics if it is not.
#[cfg(any(test, feature = "testing"))]
pub fn check_iterator_fold<I: Iterator, F: Fn() -> I>(mk_iter: F)
where
    I::Item: PartialEq + Debug,
{
    const EXTRA_COUNT: usize = 3;

    // Treat values returned by `next()` as reference.
    let mut reference_values = Vec::new();
    let mut iter = mk_iter();
    #[expect(
        clippy::while_let_on_iterator,
        reason = "Reference implementation: call bare-bones `next()` explicitly"
    )]
    while let Some(value) = iter.next() {
        reference_values.push(value);
    }

    // Check that `next()` after exhaustion returns None.
    for _ in 0..EXTRA_COUNT {
        assert!(
            iter.next().is_none(),
            "Iterator returns values after it's exhausted",
        );
    }
    drop(iter);

    // Check `fold()` yields same values as `next()`.
    let mut values_for_fold = Vec::new();
    for split_at in 0..reference_values.len() + EXTRA_COUNT {
        let mut iter = mk_iter();
        values_for_fold.clear();

        for _ in 0..split_at.min(reference_values.len()) {
            values_for_fold.push(iter.next().expect("not enough values"));
        }
        // Call `next()` a few times to check that these extra calls won't break
        // `fold()`.
        for _ in 0..split_at.saturating_sub(reference_values.len()) {
            assert!(iter.next().is_none());
        }

        let acc = iter.fold(values_for_fold.len(), |acc, value| {
            assert_eq!(acc, values_for_fold.len());
            values_for_fold.push(value);
            acc + 1
        });
        assert_eq!(reference_values, values_for_fold);
        assert_eq!(acc, values_for_fold.len());
    }
}

/// Checks that [`ExactSizeIterator::len()`] returns correct length.
/// Panics if it is not.
#[cfg(any(test, feature = "testing"))]
pub fn check_exact_size_iterator_len<I: ExactSizeIterator>(mut iter: I) {
    for expected_len in (0..iter.len()).rev() {
        iter.next();
        assert_eq!(iter.len(), expected_len);
    }
    assert!(iter.next().is_none());
    assert_eq!(iter.len(), 0);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn try_any() {
        let is_even = |n: i32| match n {
            n if n < 0 => Err("negative"),
            n => Ok(n % 2 == 0),
        };

        assert_eq!([1, 3, 4, 5].into_iter().try_any(is_even), Ok(true));
        assert_eq!([1, 3, 5, 7].into_iter().try_any(is_even), Ok(false));
        assert_eq!(std::iter::empty().try_any(is_even), Ok(false));
        assert_eq!([1, 3, -1, 4].into_iter().try_any(is_even), Err("negative"));

        // Short-circuits on first `Ok(true)` without evaluating the rest.
        let mut iter = [1, 2, 3, -1].into_iter();
        assert_eq!(iter.try_any(is_even), Ok(true));
        assert_eq!(iter.next(), Some(3));
    }
}