Skip to main content

nekolib/utils/
bitop.rs

1#[derive(Clone, Copy)]
2pub struct PdepPextMaskU8([u8; 3], u8);
3#[derive(Clone, Copy)]
4pub struct PdepPextMaskU16([u16; 4], u16);
5#[derive(Clone, Copy)]
6pub struct PdepPextMaskU32([u32; 5], u32);
7#[derive(Clone, Copy)]
8pub struct PdepPextMaskU64([u64; 6], u64);
9#[derive(Clone, Copy)]
10pub struct PdepPextMaskU128([u128; 7], u128);
11
12pub trait Pext<R> {
13    fn pext(self, mask: R) -> Self;
14}
15
16macro_rules! impl_pext {
17    ( ($maskty:ident, $basety:ident, $lg:literal) ) => {
18        impl Pext<$basety> for $basety {
19            fn pext(self, mut m: $basety) -> $basety {
20                let mut x = self & m;
21                let mut mk = !m << 1;
22                for i in 0..$lg {
23                    let mut mp = mk ^ (mk << 1);
24                    for j in 1..$lg {
25                        mp ^= mp << (1 << j);
26                    }
27                    let mv = mp & m;
28                    m = m ^ mv | (mv >> (1 << i));
29                    let t = x & mv;
30                    x = (x ^ t) | (t >> (1 << i));
31                    mk &= !mp;
32                }
33                x
34            }
35        }
36        impl Pext<$maskty> for $basety {
37            fn pext(self, mask: $maskty) -> $basety {
38                let mut x = self & mask.1;
39                for i in 0..$lg {
40                    let mv = mask.0[i];
41                    let t = x & mv;
42                    x = (x ^ t) | (t >> (1 << i));
43                }
44                x
45            }
46        }
47    };
48    ( $( ( $($tt:tt)* ), )* ) => { $( impl_pext!( ( $($tt)* ) ); )* };
49}
50
51impl_pext! {
52    (PdepPextMaskU8, u8, 3),
53    (PdepPextMaskU16, u16, 4),
54    (PdepPextMaskU32, u32, 5),
55    (PdepPextMaskU64, u64, 6),
56    (PdepPextMaskU128, u128, 7),
57}
58
59pub trait Pdep<R> {
60    fn pdep(self, mask: R) -> Self;
61}
62
63macro_rules! impl_pdep {
64    ( ($maskty:ident, $basety:ident, $lg:literal) ) => {
65        impl Pdep<$basety> for $basety {
66            fn pdep(self, m: $basety) -> $basety {
67                self.pdep(<$maskty>::new(m))
68            }
69        }
70        impl Pdep<$maskty> for $basety {
71            fn pdep(self, mask: $maskty) -> $basety {
72                let mut x = self;
73                for i in (0..$lg).rev() {
74                    let mv = mask.0[i];
75                    let t = x << (1 << i);
76                    x = (x & !mv) | (t & mv);
77                }
78                x & mask.1
79            }
80        }
81    };
82    ( $( ( $($tt:tt)* ), )* ) => { $( impl_pdep!( ( $($tt)* ) ); )* };
83}
84
85impl_pdep! {
86    (PdepPextMaskU8, u8, 3),
87    (PdepPextMaskU16, u16, 4),
88    (PdepPextMaskU32, u32, 5),
89    (PdepPextMaskU64, u64, 6),
90    (PdepPextMaskU128, u128, 7),
91}
92
93macro_rules! pext_loop {
94    ( $mk:ident, $m:ident, $sh:expr ) => {{
95        let mp = Self::mp($mk);
96        let mv = mp & $m;
97        $m = $m ^ mv | (mv >> (1 << $sh));
98        $mk &= !mp;
99        mv
100    }};
101}
102
103macro_rules! impl_pdep_pext_mask {
104    ( ($maskty:ident, $basety:ident, [$($i:literal),*], $lg:literal) ) => {
105        impl $maskty {
106            pub const fn new(mut m: $basety) -> Self {
107                let m0 = m;
108                let mut res = [0; $lg];
109                let mut mk = !m << 1;
110                $( res[$i] = pext_loop!(mk, m, $i) );*;
111                res[$lg - 1] = Self::mp(mk) & m;
112                Self(res, m0)
113            }
114            const fn mp(mk: $basety) -> $basety {
115                let mut mp = mk ^ (mk << 1);
116                $( mp ^= mp << (1 << (1 + $i)) );*;
117                mp
118            }
119            pub const fn get(self) -> $basety { self.1 }
120        }
121    };
122    ( $( ( $($tt:tt)* ), )* ) => { $( impl_pdep_pext_mask!( ( $($tt)* ) ); )* };
123}
124
125impl_pdep_pext_mask! {
126    (PdepPextMaskU8, u8, [0, 1], 3),
127    (PdepPextMaskU16, u16, [0, 1, 2], 4),
128    (PdepPextMaskU32, u32, [0, 1, 2, 3], 5),
129    (PdepPextMaskU64, u64, [0, 1, 2, 3, 4], 6),
130    (PdepPextMaskU128, u128, [0, 1, 2, 3, 4, 5], 7),
131}
132
133#[test]
134fn test() {
135    let x = 0b_0101_0111_0000_1001_1110_1010_0000_0010_u32;
136    let m = 0b_0100_1001_1001_1010_0100_0100_0101_0100_u32;
137    //          1   0  1 0  0 1 0   1    0    0 0  0
138    let ext = 0b_1010_0101_0000;
139
140    assert_eq!(x.pext(m), ext);
141    assert_eq!(x.pext(PdepPextMaskU32::new(m)), ext);
142    assert_eq!(ext.pdep(m), x & m);
143}