Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-utils/src/floor_divmod.rs
6939 views
1
pub 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
9
macro_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
22
macro_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
33
macro_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
64
impl_unsigned_div_mod!(u8);
65
impl_unsigned_div_mod!(u16);
66
impl_unsigned_div_mod!(u32);
67
impl_unsigned_div_mod!(u64);
68
impl_unsigned_div_mod!(u128);
69
impl_unsigned_div_mod!(usize);
70
impl_signed_div_mod!(i8);
71
impl_signed_div_mod!(i16);
72
impl_signed_div_mod!(i32);
73
impl_signed_div_mod!(i64);
74
impl_signed_div_mod!(i128);
75
impl_signed_div_mod!(isize);
76
impl_float_div_mod!(f32);
77
impl_float_div_mod!(f64);
78
79
#[cfg(test)]
80
mod 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
}
103
104