1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106

#![no_std]

use num_traits::{
    int::PrimInt,
    ops::wrapping::WrappingNeg,
};

#[derive(Clone)]
pub struct Luby<T: PrimInt + WrappingNeg>(T, T);

impl<T: PrimInt + WrappingNeg> Luby<T> {
    /// Initialises a new Luby sequence.
    pub fn new() -> Luby<T> {
        // Initial values picked so that iterating yields 1, 1, 2, ...
        Luby(T::zero(), T::zero())
    }
}

impl<T: PrimInt + WrappingNeg> Iterator for Luby<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        let (u, v) = (&mut self.0, &mut self.1);

        if *u & u.wrapping_neg() == *v {
            // Reset the sequence: u += 1, v = 1
            *u = u.checked_add(&T::one())?;
            *v = T::one();
        } else {
            // Double v
            *v = v.checked_add(v)?;
        }

        Some(*v)
    }
}

#[cfg(test)]
mod tests {
    use crate::generic_test;
    use core::fmt::Debug;
    use num_traits::{
        int::PrimInt,
        ops::wrapping::WrappingNeg,
    };

    /// `Luby<T>` ends the sequence cleanly, rather than overflowing.
    fn exhaust<T: PrimInt + WrappingNeg>() {
        let luby = super::Luby::<u8>::new();
        for _ in luby {}
    }
    generic_test!(exhaust, u8, u16, i8, i16);

    /// `Luby<T>` produces a sequence that matches the known prefix
    fn prefix<T: PrimInt + WrappingNeg + Debug>() {
        let luby = super::Luby::<T>::new();

        for (i, j) in luby.zip(LUBY_PREFIX) {
            assert_eq!(i, T::from(*j).unwrap());
        }
    }
    generic_test!(prefix, u8, u16, u32, u64, i8, i16, i32, i64);

    /// Sequence prefix taken from https://oeis.org/A182105
    const LUBY_PREFIX: &[u8] = &[1, 1, 2,
                                 1, 1, 2, 4,
                                 1, 1, 2,
                                 1, 1, 2, 4, 8,
                                 1, 1, 2,
                                 1, 1, 2, 4,
                                 1, 1, 2,
                                 1, 1, 2, 4, 8, 16,
                                 1, 1, 2,
                                 1, 1, 2, 4,
                                 1, 1, 2,
                                 1, 1, 2, 4, 8,
                                 1, 1, 2,
                                 1, 1, 2, 4,
                                 1, 1, 2,
                                 1, 1, 2, 4, 8, 16, 32,
                                 1, 1, 2,
                                 1, 1, 2, 4,
                                 1, 1, 2,
                                 1, 1, 2, 4, 8,
                                 1, 1, 2,
                                 1, 1, 2, 4,
                                 1, 1, 2,
                                 1, 1, 2, 4, 8, 16,
                                 1, 1, 2,
                                 1, 1, 2, 4,
                                 1, 1, 2,
                                 1, 1, 2, 4, 8];
}

#[cfg(test)]
#[macro_export]
macro_rules! generic_test {
    ($f:ident, $($t:ty),+) => {
        paste::item! {
            $(#[test]
              fn [< $f _ $t >]() {
                  $f::<$t>()
              })+
        }
    };
}