`as` free binary search for `RangeBounds<{Integer}>`

bisect.rs

//! `as` free binary search for `RangeBounds<{Integer}>`.
//!
//! This code is licensed under [CC0-1.0](https://creativecommons.org/publicdomain/zero/1.0).
//!
//! ```cargo
//! [package]
//! name = "bisect"
//! version = "0.0.0"
//! authors = ["Ryo Yamashita <[email protected]>"]
//! edition = "2018"
//! license = "CC0-1.0"
//! publish = false
//!
//! [dependencies]
//!
//! [dev-dependencies]
//! itertools = "0.9.0"
//! pretty_assertions = "0.6.1"
//! rand = "0.7.3"
//! ```

use std::{
    convert::Infallible,
    ops::{Add, Bound, Div, RangeBounds, Sub},
};

fn main() {}

pub trait RangeBoundsExt<T> {
    fn bisect_min<F: FnMut(T) -> bool>(&self, f: F) -> Option<T>;
    fn bisect_max<F: FnMut(T) -> bool>(&self, f: F) -> Option<T>;
}

impl<R: RangeBounds<T>, T: PrimitiveInteger> RangeBoundsExt<T> for R {
    fn bisect_min<F: FnMut(T) -> bool>(&self, mut f: F) -> Option<T> {
        return bisect::<_, BelowMax<_>, _>(self.end_bound(), self.start_bound(), &mut f);

        struct BelowMax<T>(Infallible, fn() -> T);

        impl<T: PrimitiveInteger> BisectBounds for BelowMax<T> {
            type Item = T;

            const OK_SIDE_BOUND: T = T::MAX_VALUE;
            const NG_SIDE_BOUND: T = T::MIN_VALUE;

            fn step_ok(ok: T) -> T {
                ok - T::ONE
            }

            fn minmax(ok: T, ng: T) -> (T, T) {
                (ng, ok)
            }
        }
    }

    fn bisect_max<F: FnMut(T) -> bool>(&self, mut f: F) -> Option<T> {
        return bisect::<_, MinAbove<_>, _>(self.start_bound(), self.end_bound(), &mut f);

        struct MinAbove<T>(Infallible, fn() -> T);

        impl<T: PrimitiveInteger> BisectBounds for MinAbove<T> {
            type Item = T;

            const OK_SIDE_BOUND: T = T::MIN_VALUE;
            const NG_SIDE_BOUND: T = T::MAX_VALUE;

            fn step_ok(ok: T) -> T {
                ok + T::ONE
            }

            fn minmax(ok: T, ng: T) -> (T, T) {
                (ok, ng)
            }
        }
    }
}

fn bisect<T: PrimitiveInteger, B: BisectBounds<Item = T>, F: FnMut(T) -> bool>(
    ok: Bound<&T>,
    ng: Bound<&T>,
    f: &mut F,
) -> Option<T> {
    let mut ok = match ok {
        Bound::Included(&x) if f(x) => x,
        Bound::Excluded(&x) if x != B::NG_SIDE_BOUND && f(B::step_ok(x)) => B::step_ok(x),
        Bound::Unbounded if f(B::OK_SIDE_BOUND) => B::OK_SIDE_BOUND,
        Bound::Included(_) | Bound::Excluded(_) | Bound::Unbounded => return None,
    };

    let mut ng = match ng {
        Bound::Included(&x) if f(x) => return Some(x),
        Bound::Unbounded if f(B::NG_SIDE_BOUND) => return Some(B::NG_SIDE_BOUND),
        Bound::Included(&x) | Bound::Excluded(&x) => x,
        Bound::Unbounded => B::NG_SIDE_BOUND,
    };

    if (ok < T::ZERO) == (ng < T::ZERO) {
        while B::max(ok, ng) - B::min(ok, ng) > T::ONE {
            let mid = B::min(ok, ng) + (B::max(ok, ng) - B::min(ok, ng)) / (T::ONE + T::ONE);
            *if f(mid) { &mut ok } else { &mut ng } = mid;
        }
        Some(ok)
    } else {
        let mut bisect = |ok, ng| bisect::<_, B, _>(Bound::Included(&ok), Bound::Excluded(&ng), f);
        bisect(T::ZERO, ng).or_else(|| bisect(ok, T::ZERO))
    }
}

trait BisectBounds {
    type Item: PrimitiveInteger;

    const OK_SIDE_BOUND: Self::Item;
    const NG_SIDE_BOUND: Self::Item;

    fn step_ok(ok: Self::Item) -> Self::Item;
    fn minmax(ok: Self::Item, ng: Self::Item) -> (Self::Item, Self::Item);

    fn min(ok: Self::Item, ng: Self::Item) -> Self::Item {
        let (min, _) = Self::minmax(ok, ng);
        min
    }

    fn max(ok: Self::Item, ng: Self::Item) -> Self::Item {
        let (_, max) = Self::minmax(ok, ng);
        max
    }
}

pub trait PrimitiveInteger:
    Copy + Ord + Add<Output = Self> + Sub<Output = Self> + Div<Output = Self>
{
    const ZERO: Self;
    const ONE: Self;
    const MIN_VALUE: Self;
    const MAX_VALUE: Self;
}

macro_rules! impl_primitive_integer(($($ty:ty),*) => {
    $(
        impl PrimitiveInteger for $ty {
            const ZERO: Self = 0;
            const ONE: Self = 1;
            const MIN_VALUE: Self = <$ty>::min_value();
            const MAX_VALUE: Self = <$ty>::max_value();
        }
    )*
});

impl_primitive_integer!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);

#[cfg(test)]
mod tests {
    use super::RangeBoundsExt as _;

    use itertools::Itertools as _;
    use pretty_assertions::assert_eq;
    use rand::Rng as _;

    #[test]
    fn lower_bound() {
        let mut rng = rand::thread_rng();

        let values = (0..1024)
            .map(|_| rng.gen_range(i64::min_value() + 1, i64::max_value()))
            .sorted()
            .collect::<Vec<_>>();

        let test = |target| {
            assert_eq!(
                values.iter().position(|&x| x >= target),
                (0..1024).bisect_min(|i| values[i] >= target),
            )
        };

        test(i64::min_value());
        test(i64::max_value());

        for &x in &values {
            test(x);
        }

        for _ in 0..128 {
            test(rng.gen::<i64>());
        }
    }

    #[test]
    fn i128_full() {
        assert_eq!(Some(42i128), (..).bisect_min(|i| i >= 42));
        assert_eq!(Some(57i128), (..).bisect_max(|i| i <= 57));
    }
}