polars_utils/
floor_divmod.rs

1use crate::float16::pf16;
2
3pub trait FloorDivMod: Sized {
4    // Returns the flooring division and associated modulo of lhs / rhs.
5    // This is the same division / modulo combination as Python.
6    //
7    // Returns (0, 0) if other == 0.
8    fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self);
9}
10
11macro_rules! impl_float_div_mod {
12    ($T:ty) => {
13        impl FloorDivMod for $T {
14            #[inline]
15            fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self) {
16                let div = (self / other).floor();
17                let mod_ = self - other * div;
18                (div, mod_)
19            }
20        }
21    };
22}
23
24macro_rules! impl_unsigned_div_mod {
25    ($T:ty) => {
26        impl FloorDivMod for $T {
27            #[inline]
28            fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self) {
29                if other == 0 {
30                    return (0, 0);
31                }
32
33                (self / other, self % other)
34            }
35        }
36    };
37}
38
39macro_rules! impl_signed_div_mod {
40    ($T:ty) => {
41        impl FloorDivMod for $T {
42            #[inline]
43            fn wrapping_floor_div_mod(self, other: Self) -> (Self, Self) {
44                if other == 0 {
45                    return (0, 0);
46                }
47
48                // Rust/C-style remainder is in the correct congruence
49                // class, but may not have the right sign. We want a
50                // remainder with the same sign as the RHS, which we
51                // can get by adding RHS to the remainder if the sign of
52                // the non-zero remainder differs from our RHS.
53                //
54                // Similarly, Rust/C-style division truncates instead of floors.
55                // If the remainder was non-zero and the signs were different
56                // (we'd have a negative result before truncating), we need to
57                // subtract 1 from the result.
58                let mut div = self.wrapping_div(other);
59                let mut mod_ = self.wrapping_rem(other);
60                if mod_ != 0 && (self < 0) != (other < 0) {
61                    div -= 1;
62                    mod_ += other;
63                }
64                (div, mod_)
65            }
66        }
67    };
68}
69
70impl_unsigned_div_mod!(u8);
71impl_unsigned_div_mod!(u16);
72impl_unsigned_div_mod!(u32);
73impl_unsigned_div_mod!(u64);
74impl_unsigned_div_mod!(u128);
75impl_unsigned_div_mod!(usize);
76impl_signed_div_mod!(i8);
77impl_signed_div_mod!(i16);
78impl_signed_div_mod!(i32);
79impl_signed_div_mod!(i64);
80impl_signed_div_mod!(i128);
81impl_signed_div_mod!(isize);
82impl_float_div_mod!(pf16);
83impl_float_div_mod!(f32);
84impl_float_div_mod!(f64);
85
86#[cfg(test)]
87mod test {
88    use super::*;
89
90    #[test]
91    fn test_signed_wrapping_div_mod() {
92        // Test for all i8, should transfer to other values.
93        for lhs in i8::MIN..=i8::MAX {
94            for rhs in i8::MIN..=i8::MAX {
95                let ans = if rhs != 0 {
96                    let fdiv = (lhs as f64 / rhs as f64).floor();
97                    let fmod = lhs as f64 - rhs as f64 * fdiv;
98
99                    // float -> int conversion saturates, we want wrapping, double convert.
100                    ((fdiv as i32) as i8, (fmod as i32) as i8)
101                } else {
102                    (0, 0)
103                };
104
105                assert_eq!(lhs.wrapping_floor_div_mod(rhs), ans);
106            }
107        }
108    }
109}
polars_utils/floor_divmod.rs

polars_utils/
floor_divmod.rs
