From 0656fb8ea80811903c61946d267aa9c458c03efd Mon Sep 17 00:00:00 2001
From: Jim Blandy <jimb@red-bean.com>
Date: Sat, 22 Jun 2024 22:39:37 -0700
Subject: [PATCH] [naga] Use `HandleSet` in `naga::proc::index`.

Change `naga::proc::index::find_checked_indexes` to return a
`HandleSet<Expression>`, rather than an untyped `BitSet`.

Fix uses in the Metal Shading Language backend.
---
 naga/src/back/msl/writer.rs | 23 ++++++++++-------------
 naga/src/proc/index.rs      | 20 ++++++++++----------
 2 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/naga/src/back/msl/writer.rs b/naga/src/back/msl/writer.rs
index c2ad813921..8b86897007 100644
--- a/naga/src/back/msl/writer.rs
+++ b/naga/src/back/msl/writer.rs
@@ -1,12 +1,11 @@
 use super::{sampler as sm, Error, LocationMode, Options, PipelineOptions, TranslationInfo};
 use crate::{
-    arena::Handle,
+    arena::{Handle, HandleSet},
     back::{self, Baked},
     proc::index,
     proc::{self, NameKey, TypeResolution},
     valid, FastHashMap, FastHashSet,
 };
-use bit_set::BitSet;
 use std::{
     fmt::{Display, Error as FmtError, Formatter, Write},
     iter,
@@ -584,11 +583,10 @@ struct ExpressionContext<'a> {
     lang_version: (u8, u8),
     policies: index::BoundsCheckPolicies,
 
-    /// A bitset containing the `Expression` handle indexes of expressions used
-    /// as indices in `ReadZeroSkipWrite`-policy accesses. These may need to be
-    /// cached in temporary variables. See `index::find_checked_indexes` for
-    /// details.
-    guarded_indices: BitSet,
+    /// The set of expressions used as indices in `ReadZeroSkipWrite`-policy
+    /// accesses. These may need to be cached in temporary variables. See
+    /// `index::find_checked_indexes` for details.
+    guarded_indices: HandleSet<crate::Expression>,
 }
 
 impl<'a> ExpressionContext<'a> {
@@ -2873,12 +2871,11 @@ impl<W: Write> Writer<W> {
                             // If this expression is an index that we're going to first compare
                             // against a limit, and then actually use as an index, then we may
                             // want to cache it in a temporary, to avoid evaluating it twice.
-                            let bake =
-                                if context.expression.guarded_indices.contains(handle.index()) {
-                                    true
-                                } else {
-                                    self.need_bake_expressions.contains(&handle)
-                                };
+                            let bake = if context.expression.guarded_indices.contains(handle) {
+                                true
+                            } else {
+                                self.need_bake_expressions.contains(&handle)
+                            };
 
                             if bake {
                                 Some(Baked(handle).to_string())
diff --git a/naga/src/proc/index.rs b/naga/src/proc/index.rs
index e2c3de8eb0..48b987ce85 100644
--- a/naga/src/proc/index.rs
+++ b/naga/src/proc/index.rs
@@ -2,8 +2,8 @@
 Definitions for index bounds checking.
 */
 
-use crate::{valid, Handle, UniqueArena};
-use bit_set::BitSet;
+use crate::arena::{Handle, HandleSet, UniqueArena};
+use crate::valid;
 
 /// How should code generated by Naga do bounds checks?
 ///
@@ -196,7 +196,7 @@ pub enum GuardedIndex {
 /// Build a set of expressions used as indices, to cache in temporary variables when
 /// emitted.
 ///
-/// Given the bounds-check policies `policies`, construct a `BitSet` containing the handle
+/// Given the bounds-check policies `policies`, construct a `HandleSet` containing the handle
 /// indices of all the expressions in `function` that are ever used as guarded indices
 /// under the [`ReadZeroSkipWrite`] policy. The `module` argument must be the module to
 /// which `function` belongs, and `info` should be that function's analysis results.
@@ -241,10 +241,10 @@ pub fn find_checked_indexes(
     function: &crate::Function,
     info: &valid::FunctionInfo,
     policies: BoundsCheckPolicies,
-) -> BitSet {
+) -> HandleSet<crate::Expression> {
     use crate::Expression as Ex;
 
-    let mut guarded_indices = BitSet::new();
+    let mut guarded_indices = HandleSet::for_arena(&function.expressions);
 
     // Don't bother scanning if we never need `ReadZeroSkipWrite`.
     if policies.contains(BoundsCheckPolicy::ReadZeroSkipWrite) {
@@ -264,7 +264,7 @@ pub fn find_checked_indexes(
                         )
                         .is_some()
                     {
-                        guarded_indices.insert(index.index());
+                        guarded_indices.insert(index);
                     }
                 }
                 Ex::ImageLoad {
@@ -275,15 +275,15 @@ pub fn find_checked_indexes(
                     ..
                 } => {
                     if policies.image_load == BoundsCheckPolicy::ReadZeroSkipWrite {
-                        guarded_indices.insert(coordinate.index());
+                        guarded_indices.insert(coordinate);
                         if let Some(array_index) = array_index {
-                            guarded_indices.insert(array_index.index());
+                            guarded_indices.insert(array_index);
                         }
                         if let Some(sample) = sample {
-                            guarded_indices.insert(sample.index());
+                            guarded_indices.insert(sample);
                         }
                         if let Some(level) = level {
-                            guarded_indices.insert(level.index());
+                            guarded_indices.insert(level);
                         }
                     }
                 }