Rollup merge of #149097 - okaneco:gather_scatter_bits, r=Mark-Simulacrum

num: Implement `uint_gather_scatter_bits` feature for unsigned integers

Feature gate: `#![feature(uint_gather_scatter_bits)]`
Tracking issue: #149069
Accepted ACP: rust-lang/libs-team#695 (comment)

Implement `gather_bits`, `scatter_bits` functions on unsigned integers
Add tests to coretests

This implementation is a small improvement over the plain naive form (see the [solution sketch](rust-lang/libs-team#695)).
We only check the set bits in the mask instead of iterating over every bit.

Author: Zalathar

library/core/src/num/uint_macros.rs

@@ -479,6 +479,76 @@ macro_rules! uint_impl {
             intrinsics::bswap(self as $ActualT) as Self
         }
 
+        /// Returns an integer with the bit locations specified by `mask` packed
+        /// contiguously into the least significant bits of the result.
+        /// ```
+        /// #![feature(uint_gather_scatter_bits)]
+        #[doc = concat!("let n: ", stringify!($SelfT), " = 0b1011_1100;")]
+        ///
+        /// assert_eq!(n.gather_bits(0b0010_0100), 0b0000_0011);
+        /// assert_eq!(n.gather_bits(0xF0), 0b0000_1011);
+        /// ```
+        #[unstable(feature = "uint_gather_scatter_bits", issue = "149069")]
+        #[must_use = "this returns the result of the operation, \
+                      without modifying the original"]
+        #[inline]
+        pub const fn gather_bits(self, mut mask: Self) -> Self {
+            let mut bit_position = 1;
+            let mut result = 0;
+
+            // Iterate through the mask bits, unsetting the lowest bit after
+            // each iteration. We fill the bits in the result starting from the
+            // least significant bit.
+            while mask != 0 {
+                // Find the next lowest set bit in the mask
+                let next_mask_bit = mask.isolate_lowest_one();
+
+                // Retrieve the masked bit and if present, set it in the result
+                let src_bit = (self & next_mask_bit) != 0;
+                result |= if src_bit { bit_position } else { 0 };
+
+                // Unset lowest set bit in the mask, prepare next position to set
+                mask ^= next_mask_bit;
+                bit_position <<= 1;
+            }
+
+            result
+        }
+
+        /// Returns an integer with the least significant bits of `self`
+        /// distributed to the bit locations specified by `mask`.
+        /// ```
+        /// #![feature(uint_gather_scatter_bits)]
+        #[doc = concat!("let n: ", stringify!($SelfT), " = 0b1010_1101;")]
+        ///
+        /// assert_eq!(n.scatter_bits(0b0101_0101), 0b0101_0001);
+        /// assert_eq!(n.scatter_bits(0xF0), 0b1101_0000);
+        /// ```
+        #[unstable(feature = "uint_gather_scatter_bits", issue = "149069")]
+        #[must_use = "this returns the result of the operation, \
+                      without modifying the original"]
+        #[inline]
+        pub const fn scatter_bits(mut self, mut mask: Self) -> Self {
+            let mut result = 0;
+
+            // Iterate through the mask bits, unsetting the lowest bit after
+            // each iteration and right-shifting `self` by one to get the next
+            // bit into the least significant bit position.
+            while mask != 0 {
+                // Find the next bit position to potentially set
+                let next_mask_bit = mask.isolate_lowest_one();
+
+                // If bit is set, deposit it at the masked bit position
+                result |= if (self & 1) != 0 { next_mask_bit } else { 0 };
+
+                // Unset lowest set bit in the mask, shift in next `self` bit
+                mask ^= next_mask_bit;
+                self >>= 1;
+            }
+
+            result
+        }
+
         /// Reverses the order of bits in the integer. The least significant bit becomes the most significant bit,
         ///                 second least-significant bit becomes second most-significant bit, etc.
         ///

library/coretests/tests/lib.rs

@@ -114,6 +114,7 @@
 #![feature(try_find)]
 #![feature(try_trait_v2)]
 #![feature(uint_bit_width)]
+#![feature(uint_gather_scatter_bits)]
 #![feature(unsize)]
 #![feature(unwrap_infallible)]
 // tidy-alphabetical-end

library/coretests/tests/num/uint_macros.rs

@@ -127,6 +127,52 @@ macro_rules! uint_module {
                 assert_eq_const_safe!($T: _1.swap_bytes(), _1);
             }
 
+            fn test_gather_bits() {
+                assert_eq_const_safe!($T: $T::gather_bits(0b1010_0101, 0b0000_0011), 0b_0001);
+                assert_eq_const_safe!($T: $T::gather_bits(0b1010_0101, 0b0000_0110), 0b_0010);
+                assert_eq_const_safe!($T: $T::gather_bits(0b1010_0101, 0b0000_1100), 0b_0001);
+                assert_eq_const_safe!($T: $T::gather_bits(0b1010_0101, 0b0001_1000), 0b_0000);
+                assert_eq_const_safe!($T: $T::gather_bits(0b1010_0101, 0b0011_0000), 0b_0010);
+                assert_eq_const_safe!($T: $T::gather_bits(0b1010_0101, 0b0110_0000), 0b_0001);
+                assert_eq_const_safe!($T: $T::gather_bits(0b1010_0101, 0b1100_0000), 0b_0010);
+
+                assert_eq_const_safe!($T: A.gather_bits(_0), 0);
+                assert_eq_const_safe!($T: B.gather_bits(_0), 0);
+                assert_eq_const_safe!($T: C.gather_bits(_0), 0);
+                assert_eq_const_safe!($T: _0.gather_bits(A), 0);
+                assert_eq_const_safe!($T: _0.gather_bits(B), 0);
+                assert_eq_const_safe!($T: _0.gather_bits(C), 0);
+
+                assert_eq_const_safe!($T: A.gather_bits(_1), A);
+                assert_eq_const_safe!($T: B.gather_bits(_1), B);
+                assert_eq_const_safe!($T: C.gather_bits(_1), C);
+                assert_eq_const_safe!($T: _1.gather_bits(0b0010_0001), 0b0000_0011);
+                assert_eq_const_safe!($T: _1.gather_bits(0b0010_1100), 0b0000_0111);
+                assert_eq_const_safe!($T: _1.gather_bits(0b0111_1001), 0b0001_1111);
+            }
+
+            fn test_scatter_bits() {
+                assert_eq_const_safe!($T: $T::scatter_bits(0b1111, 0b1001_0110), 0b1001_0110);
+                assert_eq_const_safe!($T: $T::scatter_bits(0b0001, 0b1001_0110), 0b0000_0010);
+                assert_eq_const_safe!($T: $T::scatter_bits(0b0010, 0b1001_0110), 0b0000_0100);
+                assert_eq_const_safe!($T: $T::scatter_bits(0b0100, 0b1001_0110), 0b0001_0000);
+                assert_eq_const_safe!($T: $T::scatter_bits(0b1000, 0b1001_0110), 0b1000_0000);
+
+                assert_eq_const_safe!($T: A.scatter_bits(_0), 0);
+                assert_eq_const_safe!($T: B.scatter_bits(_0), 0);
+                assert_eq_const_safe!($T: C.scatter_bits(_0), 0);
+                assert_eq_const_safe!($T: _0.scatter_bits(A), 0);
+                assert_eq_const_safe!($T: _0.scatter_bits(B), 0);
+                assert_eq_const_safe!($T: _0.scatter_bits(C), 0);
+
+                assert_eq_const_safe!($T: A.scatter_bits(_1), A);
+                assert_eq_const_safe!($T: B.scatter_bits(_1), B);
+                assert_eq_const_safe!($T: C.scatter_bits(_1), C);
+                assert_eq_const_safe!($T: _1.scatter_bits(A), A);
+                assert_eq_const_safe!($T: _1.scatter_bits(B), B);
+                assert_eq_const_safe!($T: _1.scatter_bits(C), C);
+            }
+
             fn test_reverse_bits() {
                 assert_eq_const_safe!($T: A.reverse_bits().reverse_bits(), A);
                 assert_eq_const_safe!($T: B.reverse_bits().reverse_bits(), B);