Skip to main content

yui_core/conc/misc/
u256.rs

1//! [`U256`]: a 256-bit word, used as a wide [`BitStorage`] for [`BitMap`](super::bitmap::BitMap).
2
3use std::ops::{BitAnd, BitOr, BitOrAssign, Shl, Sub};
4
5use super::bitmap::BitStorage;
6
7/// 256-bit storage: two little-endian `u128` limbs `[low, high]`.
8#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Default)]
9pub struct U256([u128; 2]);
10
11impl BitAnd for U256 {
12    type Output = Self;
13    fn bitand(self, rhs: Self) -> Self {
14        Self([self.0[0] & rhs.0[0], self.0[1] & rhs.0[1]])
15    }
16}
17
18impl BitOr for U256 {
19    type Output = Self;
20    fn bitor(self, rhs: Self) -> Self {
21        Self([self.0[0] | rhs.0[0], self.0[1] | rhs.0[1]])
22    }
23}
24
25impl BitOrAssign for U256 {
26    fn bitor_assign(&mut self, rhs: Self) {
27        self.0[0] |= rhs.0[0];
28        self.0[1] |= rhs.0[1];
29    }
30}
31
32impl Shl<u32> for U256 {
33    type Output = Self;
34    fn shl(self, rhs: u32) -> Self {
35        let [lo, hi] = self.0;
36        if rhs == 0 {
37            self
38        } else if rhs < 128 {
39            Self([lo << rhs, (hi << rhs) | (lo >> (128 - rhs))])
40        } else {
41            Self([0, lo << (rhs - 128)])
42        }
43    }
44}
45
46impl Sub for U256 {
47    type Output = Self;
48    fn sub(self, rhs: Self) -> Self {
49        let (lo, borrow) = self.0[0].overflowing_sub(rhs.0[0]);
50        let hi = self.0[1].wrapping_sub(rhs.0[1]).wrapping_sub(borrow as u128);
51        Self([lo, hi])
52    }
53}
54
55impl BitStorage for U256 {
56    const WIDTH: u32 = 256;
57
58    fn one() -> Self {
59        Self([1, 0])
60    }
61
62    fn is_zero(self) -> bool {
63        self.0 == [0, 0]
64    }
65
66    fn count_ones(self) -> u32 {
67        self.0[0].count_ones() + self.0[1].count_ones()
68    }
69
70    fn trailing_zeros(self) -> u32 {
71        if self.0[0] != 0 {
72            self.0[0].trailing_zeros()
73        } else if self.0[1] != 0 {
74            128 + self.0[1].trailing_zeros()
75        } else {
76            256
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use super::super::bitmap::BitMap;
85
86    #[test]
87    fn shl_across_limbs() {
88        let one = U256::one();
89        let x = one << 200;
90        assert_eq!(x.trailing_zeros(), 200);
91        assert_eq!(x.count_ones(), 1);
92
93        let y = one << 127;
94        assert_eq!((y << 1).trailing_zeros(), 128);
95    }
96
97    #[test]
98    fn sub_borrows() {
99        let x = U256::one() << 128;   // lowest bit of the high limb
100        let y = x - U256::one();      // all 128 low bits set
101        assert_eq!(y.count_ones(), 128);
102        assert_eq!(y.trailing_zeros(), 0);
103    }
104
105    #[test]
106    fn bitmap_over_128() {
107        let mut m: BitMap<u16, U256> = BitMap::new();
108        m.insert(3);
109        m.insert(130);
110        m.insert(255);
111        assert_eq!(m.len(), 3);
112        assert!(m.contains(130));
113        assert!(!m.contains(129));
114        assert_eq!(m.iter().collect::<Vec<u16>>(), vec![3, 130, 255]);
115    }
116}