[naga] Use HandleSet in Validator::needs_visit.

Change the type of `naga::valid::Validator::needs_visit` from `BitSet` to `HandleSet`. Adjust uses. Add `HandleSet` method `iter`.
gfx-rs · Jun 25, 2024 · 29e3b98 · 29e3b98
1 parent 71a52bd
commit 29e3b98
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 9 deletions.
diff --git a/naga/src/arena/handle_set.rs b/naga/src/arena/handle_set.rs
@@ -76,6 +76,11 @@ impl<T> HandleSet<T> {
     pub fn contains(&self, handle: Handle<T>) -> bool {
         self.members.contains(handle.index())
     }
+
+    /// Return an iterator over all handles in `self`.
+    pub fn iter(&self) -> impl '_ + Iterator<Item = Handle<T>> {
+        self.members.iter().map(Handle::from_usize)
+    }
 }
 
 pub trait ArenaType<T> {

diff --git a/naga/src/valid/function.rs b/naga/src/valid/function.rs
@@ -323,7 +323,7 @@ impl super::Validator {
                 crate::Expression::CallResult(callee)
                     if fun.result.is_some() && callee == function =>
                 {
-                    if !self.needs_visit.remove(expr.index()) {
+                    if !self.needs_visit.remove(expr) {
                         return Err(CallError::ResultAlreadyPopulated(expr)
                             .with_span_handle(expr, context.expressions));
                     }
@@ -462,7 +462,7 @@ impl super::Validator {
 
                 // Note that this expression has been visited by the proper kind
                 // of statement.
-                if !self.needs_visit.remove(result.index()) {
+                if !self.needs_visit.remove(result) {
                     return Err(AtomicError::ResultAlreadyPopulated(result)
                         .with_span_handle(result, context.expressions)
                         .into_other());
@@ -1429,7 +1429,7 @@ impl super::Validator {
 
         self.valid_expression_set.clear_for_arena(&fun.expressions);
         self.valid_expression_list.clear();
-        self.needs_visit.clear();
+        self.needs_visit.clear_for_arena(&fun.expressions);
         for (handle, expr) in fun.expressions.iter() {
             if expr.needs_pre_emit() {
                 self.valid_expression_set.insert(handle);
@@ -1440,7 +1440,7 @@ impl super::Validator {
                 if let crate::Expression::CallResult(_) | crate::Expression::AtomicResult { .. } =
                     *expr
                 {
-                    self.needs_visit.insert(handle.index());
+                    self.needs_visit.insert(handle);
                 }
 
                 match self.validate_expression(
@@ -1471,9 +1471,7 @@ impl super::Validator {
             info.available_stages &= stages;
 
             if self.flags.contains(super::ValidationFlags::EXPRESSIONS) {
-                if let Some(unvisited) = self.needs_visit.iter().next() {
-                    let index = crate::non_max_u32::NonMaxU32::new(unvisited as u32).unwrap();
-                    let handle = Handle::new(index);
+                if let Some(handle) = self.needs_visit.iter().next() {
                     return Err(FunctionError::UnvisitedExpression(handle)
                         .with_span_handle(handle, &fun.expressions));
                 }

diff --git a/naga/src/valid/mod.rs b/naga/src/valid/mod.rs
@@ -281,7 +281,7 @@ pub struct Validator {
     /// [`Atomic`]: crate::Statement::Atomic
     /// [`Expression`]: crate::Expression
     /// [`Statement`]: crate::Statement
-    needs_visit: BitSet,
+    needs_visit: HandleSet<crate::Expression>,
 }
 
 #[derive(Clone, Debug, thiserror::Error)]
@@ -451,7 +451,7 @@ impl Validator {
             valid_expression_set: HandleSet::new(),
             override_ids: FastHashSet::default(),
             allow_overrides: true,
-            needs_visit: BitSet::new(),
+            needs_visit: HandleSet::new(),
         }
     }