polars_utils/
floor_divmod.rs

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

polars_utils/
floor_divmod.rs
