Bit set iterator

Author: maneatingape

src/lib.rs

@@ -13,6 +13,7 @@
 /// # Utility modules to handle common recurring Advent of Code patterns.
 pub mod util {
     pub mod ansi;
+    pub mod bitset;
     pub mod grid;
     pub mod hash;
     pub mod heap;

src/util/bitset.rs

@@ -0,0 +1,40 @@
+//! Add `biterator` method that treats an integer as a set, iterating over each element where
+//! the respective bit is set. For example `1101` would return 0, 2 and 3.
+use crate::util::integer::*;
+
+pub trait BitOps<T> {
+    fn biterator(self) -> Bitset<T>;
+}
+
+impl<T> BitOps<T> for T
+where
+    T: Integer<T>,
+{
+    fn biterator(self) -> Bitset<T> {
+        Bitset { t: self }
+    }
+}
+
+pub struct Bitset<T> {
+    t: T,
+}
+
+impl<T> Iterator for Bitset<T>
+where
+    T: Integer<T>,
+    T: TryInto<usize>,
+{
+    type Item = usize;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.t == T::ZERO {
+            return None;
+        }
+
+        let tz = self.t.trailing_zeros();
+        self.t = self.t ^ (T::ONE << tz);
+
+        tz.try_into().ok()
+    }
+}

src/util/integer.rs

@@ -1,6 +1,6 @@
 //! Combines common [operators](https://doc.rust-lang.org/book/appendix-02-operators.html)
 //! and constants `0`, `1` and `10` to enable generic methods on integer types.
-use std::ops::{Add, BitAnd, BitOr, Div, Mul, Neg, Rem, Shl, Shr, Sub};
+use std::ops::*;
 
 pub trait Integer<T>:
     Copy
@@ -10,6 +10,7 @@ pub trait Integer<T>:
     + Add<Output = T>
     + BitAnd<Output = T>
     + BitOr<Output = T>
+    + BitXor<Output = T>
     + Div<Output = T>
     + Mul<Output = T>
     + Rem<Output = T>
@@ -22,6 +23,7 @@ pub trait Integer<T>:
     const TEN: T;
 
     fn ilog2(self) -> T;
+    fn trailing_zeros(self) -> T;
 }
 
 pub trait Unsigned<T>: Integer<T> {}
@@ -38,7 +40,13 @@ macro_rules! integer {
             #[inline]
             #[allow(trivial_numeric_casts)]
             fn ilog2(self) -> $t {
-                self.ilog2() as $t
+                <$t>::ilog2(self) as $t
+            }
+
+            #[inline]
+            #[allow(trivial_numeric_casts)]
+            fn trailing_zeros(self) -> $t {
+                <$t>::trailing_zeros(self) as $t
             }
         }
     )*)

src/year2017/day24.rs

@@ -13,6 +13,7 @@
 //!
 //! If we can place such a component then there's no need to consider further components which
 //! reduces the total number of combination to consider.
+use crate::util::bitset::*;
 use crate::util::iter::*;
 use crate::util::parse::*;
 
@@ -107,16 +108,12 @@ pub fn part2(input: &[usize]) -> usize {
 
 fn build(state: &mut State, current: usize, used: usize, strength: usize, length: usize) {
     // Bitset of all unused components that have a matching port.
-    let mut remaining = state.possible[current] & !used;
-
-    while remaining > 0 {
-        // Extract the index of each component from the bitset.
-        let index = remaining.trailing_zeros() as usize;
-        let mask = 1 << index;
-        remaining ^= mask;
+    let remaining = state.possible[current] & !used;
 
+    // Extract the index of each component from the bitset.
+    for index in remaining.biterator() {
         let next = current ^ state.both[index];
-        let used = used | mask;
+        let used = used | (1 << index);
         let strength = strength + state.weight[index];
         let length = length + state.length[index];
 

src/year2019/day18.rs

@@ -50,6 +50,7 @@
 #![allow(clippy::needless_range_loop)]
 #![allow(clippy::unnecessary_lazy_evaluations)]
 
+use crate::util::bitset::*;
 use crate::util::grid::*;
 use crate::util::hash::*;
 use crate::util::heap::*;
@@ -196,27 +197,17 @@ fn explore(width: i32, bytes: &[u8]) -> u32 {
             return total;
         }
 
-        let mut robots = position;
-
-        while robots != 0 {
-            // The set of robots is stored as bits in a `u32` shifted by the index of the location.
-            let from = robots.trailing_zeros() as usize;
-            let from_mask = 1 << from;
-            robots ^= from_mask;
-
-            let mut keys = remaining;
-
-            while keys != 0 {
-                // The set of keys still needed is also stored as bits in a `u32` similar as robots.
-                let to = keys.trailing_zeros() as usize;
-                let to_mask = 1 << to;
-                keys ^= to_mask;
-
+        // The set of robots is stored as bits in a `u32` shifted by the index of the location.
+        for from in position.biterator() {
+            // The set of keys still needed is also stored as bits in a `u32` similar as robots.
+            for to in remaining.biterator() {
                 let Door { distance, needed } = maze[from][to];
 
                 // u32::MAX indicates that two nodes are not connected. Only possible in part two.
                 if distance != u32::MAX && remaining & needed == 0 {
                     let next_total = total + distance;
+                    let from_mask = 1 << from;
+                    let to_mask = 1 << to;
                     let next_state = State {
                         position: position ^ from_mask ^ to_mask,
                         remaining: remaining ^ to_mask,

src/year2021/day12.rs

@@ -12,6 +12,7 @@
 //!   cave connections as an [adjacency matrix](https://en.wikipedia.org/wiki/Adjacency_matrix)
 //!   and the list of visited caves compressed into a single `u32` is the low level strategy to
 //!   quickly and efficiently store the small cardinality set of caves.
+use crate::util::bitset::*;
 use crate::util::hash::*;
 use crate::util::iter::*;
 
@@ -139,19 +140,17 @@ fn paths(input: &Input, state: &State, cache: &mut [u32]) -> u32 {
 
     let mut caves = input.edges[from];
     let mut total = 0;
-    let mut mask = 1 << END;
+    let end = 1 << END;
 
-    if caves & mask != 0 {
-        caves ^= mask;
+    if caves & end != 0 {
+        caves ^= end;
         total += 1;
     }
 
-    while caves != 0 {
-        let to = caves.trailing_zeros() as usize;
-        mask = 1 << to;
-        caves ^= mask;
-
+    for to in caves.biterator() {
+        let mask = 1 << to;
         let once = input.small & mask == 0 || visited & mask == 0;
+
         if once || twice {
             let next = State { from: to, visited: visited | mask, twice: once && twice };
             total += paths(input, &next, cache);

src/year2022/day16.rs

@@ -38,6 +38,7 @@
 //! Then we check every possible pair formed by those values, considering only the pairs
 //! where the sets of valves are [disjoint](https://en.wikipedia.org/wiki/Disjoint_sets),
 //! which is when you and the elephant have visited different sets of valves.
+use crate::util::bitset::*;
 use crate::util::hash::*;
 use crate::util::parse::*;
 use std::cmp::Ordering;
@@ -264,24 +265,19 @@ pub fn part2(input: &Input) -> u32 {
 fn explore(input: &Input, state: &State, high_score: &mut impl FnMut(usize, u32) -> u32) {
     let State { todo, from, time, pressure } = *state;
     let score = high_score(todo, pressure);
-    let mut valves = todo;
-
-    while valves > 0 {
-        // Stores the set of unopened valves in a single integer as a bit mask with a 1
-        // for each unopened valve. This code iterates over each valve by finding the lowest
-        // 1 bit then removing it from the set.
-        let to = valves.trailing_zeros() as usize;
-        let mask = 1 << to;
-        valves ^= mask;
 
+    // Stores the set of unopened valves in a single integer as a bit mask with a 1
+    // for each unopened valve. This code iterates over each valve by finding the lowest
+    // 1 bit then removing it from the set.
+    for to in todo.biterator() {
         // Check if there's enough time to reach the valve.
         let needed = input.distance[from * input.size + to];
         if needed >= time {
             continue;
         }
 
         // Calculate the total pressure released by a valve up front.
-        let todo = todo ^ mask;
+        let todo = todo ^ (1 << to);
         let time = time - needed;
         let pressure = pressure + time * input.flow[to];