arcium-primitives 0.4.4

pub mod heap_array;
pub mod identifiers;

use std::{error::Error, fmt::Debug, ops::Add};

/// Marker trait for scalar element types used in broadcast arithmetic on [`HeapArray`].
///
/// This trait is intentionally **not** implemented for `HeapArray` itself, which allows
/// both element-wise (`HeapArray * HeapArray`) and broadcast (`HeapArray * scalar`)
/// operator impls to coexist without coherence conflicts.
///
/// Blanket-implemented for all `Copy` types. Since `HeapArray` contains a `Box<[T]>`,
/// it can never be `Copy` and therefore can never satisfy `Element`.
pub trait Element {}
impl<T: Copy> Element for T {}

pub use heap_array::{HeapArray, HeapMatrix, RowMajorHeapMatrix};
pub use identifiers::{FaultyPeer, PeerId, PeerIndex, PeerNumber, ProtocolInfo, SessionId};
use subtle::Choice;
use typenum::{NonZero, Unsigned, B1};

use crate::utils::IntoExactSizeIterator;

// ---------- Typenum ---------- //

/// A trait to represent positive nonzero unsigned integer typenum constants.
pub trait Positive: Unsigned + NonZero + Debug + Eq + Send + Clone {
    const SIZE: usize;
}
impl<T: Unsigned + NonZero + Debug + Eq + Send + Clone> Positive for T {
    const SIZE: usize = <T as Unsigned>::USIZE;
}

/// A trait to represent nonnegative (zero or positive) unsigned integer typenum constants.
pub trait NonNegative: Unsigned + Debug + Eq + Send + Clone {}
impl<T: Unsigned + Debug + Eq + Send + Clone> NonNegative for T {}

/// A trait to represent positive nonzero unsigned integer typenum constants which accept "+1"
/// operation.
pub trait PositivePlusOne: Positive + Add<B1, Output: Positive> {}
impl<T: Positive + Add<B1, Output: Positive>> PositivePlusOne for T {}

// ---------- Traits ---------- //

/// A trait to define batching types that contain a compile-time known size of elements.
pub trait Batched: Sized + IntoExactSizeIterator<Item = <Self as Batched>::Item> {
    /// The type of elements contained in the batch.
    type Item;

    /// The size of the batch.
    type Size: Positive;

    /// Returns the size of the batch as a usize.
    fn batch_size() -> usize {
        Self::Size::SIZE
    }
}

/// A type which can be conditionally selected with a loose promise of constant time.
/// Compared to `subtle::ConditionallySelectable`, this trait does not require
/// Copy, but the internal elements should be Copy for the promise to loosely hold.
pub trait ConditionallySelectable: Sized {
    /// Select `a` or `b` according to `choice`.
    ///
    /// # Returns
    ///
    /// * `a` if `choice == Choice(0)`;
    /// * `b` if `choice == Choice(1)`.
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self;
}

impl<T: subtle::ConditionallySelectable> ConditionallySelectable for T {
    #[inline]
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        <T as subtle::ConditionallySelectable>::conditional_select(a, b, choice)
    }
}

/// Collect an iterator of Results<T, E> into a container of T and a vector of E,
/// then pack all the errors in a single error of the same type. Requires that E can be
/// constructed from a Vec<E>.
pub trait CollectAll<T, E: Error + From<Vec<E>>>: Iterator<Item = Result<T, E>> + Sized {
    /// Collect the results into a container of successes, or pack all errors.
    /// Contrary to `collect<Result<Vec<T>, E>>`, this method collects all errors
    /// instead of stopping at the first one.
    fn collect_all<TC: FromIterator<T> + Extend<T> + Default>(self) -> Result<TC, E> {
        let (values, errors): (TC, Vec<E>) =
            itertools::Itertools::partition_map(self, |v| match v {
                Ok(v) => itertools::Either::Left(v),
                Err(e) => itertools::Either::Right(e),
            });
        match errors.len() {
            0 => Ok(values),
            1 => Err(errors.into_iter().next().unwrap()),
            _ => Err(E::from(errors)),
        }
    }

    /// Collect the results into a container of successes, or pack all errors.
    /// Contrary to `collect<Result<Vec<T>, E>>`, this method collects all errors
    /// instead of stopping at the first one.
    fn collect_all_vec(self) -> Result<Vec<T>, E> {
        self.collect_all::<Vec<T>>()
    }

    /// Collect the results, returning Ok(()) if all succeeded, or packing all errors
    fn collect_errors(self) -> Result<(), E> {
        let errors: Vec<E> = self.filter_map(Result::err).collect();
        match errors.len() {
            0 => Ok(()),
            1 => Err(errors.into_iter().next().unwrap()),
            _ => Err(E::from(errors)),
        }
    }
}

impl<T, E: Error + From<Vec<E>>, I: Iterator<Item = Result<T, E>>> CollectAll<T, E> for I {}

pub trait TryFoldAll<T>: Iterator<Item = T> + Sized {
    /// Fold the results, returning the accumulated value if all succeeded,
    /// or packing all errors if any failed.
    fn try_fold_all<Acc, E: Error + From<Vec<E>>, F>(self, init: Acc, mut f: F) -> Result<Acc, E>
    where
        F: FnMut(Acc, T) -> (Acc, Option<E>),
    {
        let (acc, errors) = self.fold((init, Vec::new()), |(acc, mut errors), element| {
            let (new_acc, opt_err) = f(acc, element);
            if let Some(e) = opt_err {
                errors.push(e);
            }
            (new_acc, errors)
        });
        match errors.len() {
            0 => Ok(acc),
            1 => Err(errors.into_iter().next().unwrap()),
            _ => Err(E::from(errors)),
        }
    }
}

impl<T, I: Iterator<Item = T>> TryFoldAll<T> for I {}

#[cfg(test)]
mod tests {
    use std::fmt;

    use super::*;

    #[derive(Debug, PartialEq)]
    enum TestError {
        Single(u32),
        Multiple(Vec<TestError>),
    }

    impl fmt::Display for TestError {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            match self {
                TestError::Single(n) => write!(f, "error {n}"),
                TestError::Multiple(errs) => write!(f, "multiple errors: {}", errs.len()),
            }
        }
    }

    impl Error for TestError {}

    impl From<Vec<TestError>> for TestError {
        fn from(errors: Vec<TestError>) -> Self {
            TestError::Multiple(errors)
        }
    }

    // --- collect_all ---

    #[test]
    fn collect_all_all_ok() {
        let result: Result<Vec<_>, TestError> = vec![Ok(1), Ok(2), Ok(3)].into_iter().collect_all();
        assert_eq!(result, Ok(vec![1, 2, 3]));
    }

    #[test]
    fn collect_all_empty_iterator() {
        let result: Result<Vec<i32>, TestError> = std::iter::empty().collect_all();
        assert_eq!(result, Ok(vec![]));
    }

    #[test]
    fn collect_all_single_error() {
        let result: Result<Vec<_>, TestError> = vec![Ok(1), Err(TestError::Single(42)), Ok(3)]
            .into_iter()
            .collect_all();
        assert_eq!(result, Err(TestError::Single(42)));
    }

    #[test]
    fn collect_all_multiple_errors() {
        let result: Result<Vec<i32>, TestError> =
            vec![Err(TestError::Single(1)), Ok(2), Err(TestError::Single(3))]
                .into_iter()
                .collect_all();
        assert!(matches!(result, Err(TestError::Multiple(errs)) if errs.len() == 2));
    }

    // --- collect_errors ---

    #[test]
    fn collect_errors_all_ok() {
        let result: Result<(), TestError> = vec![Ok::<i32, TestError>(1), Ok(2), Ok(3)]
            .into_iter()
            .collect_errors();
        assert_eq!(result, Ok(()));
    }

    #[test]
    fn collect_errors_empty_iterator() {
        let result: Result<(), TestError> =
            std::iter::empty::<Result<i32, TestError>>().collect_errors();
        assert_eq!(result, Ok(()));
    }

    #[test]
    fn collect_errors_single_error() {
        let result: Result<(), TestError> = vec![Ok(1), Err(TestError::Single(7)), Ok(3)]
            .into_iter()
            .collect_errors();
        assert_eq!(result, Err(TestError::Single(7)));
    }

    #[test]
    fn collect_errors_multiple_errors() {
        let result: Result<(), TestError> =
            vec![Err(TestError::Single(1)), Ok(2), Err(TestError::Single(3))]
                .into_iter()
                .collect_errors();
        assert!(matches!(result, Err(TestError::Multiple(errs)) if errs.len() == 2));
    }

    // --- try_fold_all ---

    #[test]
    fn try_fold_all_all_ok() {
        let result: Result<i32, TestError> = vec![1, 2, 3]
            .into_iter()
            .try_fold_all(0, |acc, x| (acc + x, None));
        assert_eq!(result, Ok(6));
    }

    #[test]
    fn try_fold_all_empty_iterator() {
        let result: Result<i32, TestError> =
            std::iter::empty::<i32>().try_fold_all(42, |acc, x| (acc + x, None));
        assert_eq!(result, Ok(42));
    }

    #[test]
    fn try_fold_all_single_fn_error() {
        // f returns an error for value 2
        let result: Result<i32, TestError> = vec![1, 2, 3].into_iter().try_fold_all(0, |acc, x| {
            if x == 2 {
                (acc, Some(TestError::Single(x as u32)))
            } else {
                (acc + x, None)
            }
        });
        assert_eq!(result, Err(TestError::Single(2)));
    }

    #[test]
    fn try_fold_all_multiple_fn_errors() {
        // f returns errors for values 1 and 3
        let result: Result<i32, TestError> = vec![1, 2, 3].into_iter().try_fold_all(0, |acc, x| {
            if x == 1 || x == 3 {
                (acc, Some(TestError::Single(x as u32)))
            } else {
                (acc + x, None)
            }
        });
        assert!(matches!(result, Err(TestError::Multiple(errs)) if errs.len() == 2));
    }
}