use num_traits::Float;
use crate::Iterstats;
pub trait FillNan<A = Self>: Sized {
type Output;
fn fill_nan<I: Iterator<Item = A>>(
iter: I,
repl: Self::Output,
) -> FillNanIter<impl Iterator<Item = Self::Output>, Self::Output>;
}
#[derive(Debug, Clone, Copy)]
pub struct FillNanIter<I, A> {
iter: I,
repl: A,
}
impl<I, A> Iterator for FillNanIter<I, A>
where
I: Iterator<Item = A>,
A: Float,
{
type Item = A;
fn next(&mut self) -> Option<Self::Item> {
self.iter
.next()
.map(|item| if item.is_nan() { self.repl } else { item })
}
}
macro_rules! fillnan_impl {
($typ:ty) => {
impl FillNan for $typ {
type Output = $typ;
fn fill_nan<I: Iterator<Item = Self>>(
iter: I,
repl: Self::Output,
) -> FillNanIter<impl Iterator<Item = Self::Output>, Self::Output> {
FillNanIter { iter, repl }
}
}
impl FillNan for &$typ {
type Output = $typ;
fn fill_nan<I: Iterator<Item = Self>>(
iter: I,
repl: Self::Output,
) -> FillNanIter<impl Iterator<Item = Self::Output>, Self::Output> {
iter.map(|i| *i).fill_nan(repl)
}
}
};
}
fillnan_impl!(f64);
fillnan_impl!(f32);
#[cfg(test)]
mod tests {
use super::*;
use paste::paste;
macro_rules! test_fillnan {
($name:ident: $typ:ty as $iter:expr; into_iter => $repl:literal $expected:expr) => {
paste! {
#[test]
fn [<$name _into_iter>]() {
let filled = <$typ>::fill_nan($iter.into_iter(), $repl).collect::<Vec<_>>();
assert_eq!(filled, $expected);
}
}
};
($name:ident: $typ:ty as $iter:expr; iter => $repl:literal $expected:expr) => {
paste! {
#[test]
fn [<$name _iter>]() {
let filled = <&$typ>::fill_nan($iter.iter(), $repl).collect::<Vec<_>>();
assert_eq!(filled, $expected);
}
}
};
($name:ident: $typ:ty as $iter:expr => $repl:literal $expected:expr) => {
test_fillnan!($name: $typ as $iter; iter => $repl $expected);
test_fillnan!($name: $typ as $iter; into_iter => $repl $expected);
};
}
test_fillnan!(f64: f64 as [0.1, -2.3, f64::NAN, f64::INFINITY, 2.1] => 0. vec![0.1, -2.3, 0.0, f64::INFINITY, 2.1]);
test_fillnan!(f32: f32 as [0.1, -2.3, f32::NAN, f32::INFINITY, 2.1] => 0. vec![0.1, -2.3, 0.0, f32::INFINITY, 2.1]);
}