From 04658a2a11ffa8d554935aca635a8b959b58f374 Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Wed, 27 Sep 2023 10:40:10 -0500
Subject: [PATCH] nak/bitset: Add a helper for modifying in-place

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24998>
---
 src/nouveau/compiler/bitset.rs | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/nouveau/compiler/bitset.rs b/src/nouveau/compiler/bitset.rs
index a6034cb49f1..7a21995f0c7 100644
--- a/src/nouveau/compiler/bitset.rs
+++ b/src/nouveau/compiler/bitset.rs
@@ -7,7 +7,7 @@
 use crate::util::DivCeil;
 
 use std::ops::{
-    BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not,
+    BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not, Range,
 };
 
 #[derive(Clone)]
@@ -47,6 +47,10 @@ impl BitSet {
         }
     }
 
+    pub fn get_word(&self, word: usize) -> u32 {
+        self.words.get(word).cloned().unwrap_or(0)
+    }
+
     #[allow(dead_code)]
     pub fn next_set(&self, start: usize) -> Option<usize> {
         if start >= self.words.len() * 32 {
@@ -103,6 +107,33 @@ impl BitSet {
         exists
     }
 
+    pub fn set_words(
+        &mut self,
+        bits: Range<usize>,
+        mut f: impl FnMut(usize) -> u32,
+    ) {
+        if bits.is_empty() {
+            return;
+        }
+
+        let first_word = bits.start / 32;
+        let last_word = (bits.end - 1) / 32;
+        let start_mask = !0_u32 << (bits.start % 32);
+        let end_mask = !0_u32 >> (31 - ((bits.end - 1) % 32));
+
+        self.reserve(last_word + 1);
+
+        if first_word == last_word {
+            self.words[first_word] = f(first_word) & start_mask & end_mask;
+        } else {
+            self.words[first_word] = f(first_word) & start_mask;
+            for w in (first_word + 1)..last_word {
+                self.words[w] = f(w);
+            }
+            self.words[last_word] = f(last_word) & end_mask;
+        }
+    }
+
     pub fn union_with(&mut self, other: &BitSet) -> bool {
         let mut added_bits = false;
         self.reserve_words(other.words.len());