Upgrade structs only once.

In `atomic_upgrade`, update each struct type only once, regardless of
how many of its fields/subfields are accessed atomically.

naga/src/arena/handle_set.rs

@@ -25,6 +25,10 @@ impl<T> HandleSet<T> {
         }
     }
 
+    pub fn is_empty(&self) -> bool {
+        self.members.is_empty()
+    }
+
     /// Return a new, empty `HandleSet`, sized to hold handles from `arena`.
     pub fn for_arena(arena: &impl ArenaType<T>) -> Self {
         let len = arena.len();

naga/src/front/atomic_upgrade.rs

@@ -85,15 +85,50 @@ impl Padding {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Hash)]
-pub struct ContainedGlobalVariable {
-    pub field_indices: Vec<u32>,
-    pub handle: Handle<GlobalVariable>,
+#[derive(Debug, Default)]
+pub struct Upgrades {
+    /// Global variables that we've accessed using atomic operations.
+    ///
+    /// This includes globals with composite types (arrays, structs) where we've
+    /// only accessed some components (elements, fields) atomically.
+    globals: crate::arena::HandleSet<GlobalVariable>,
+
+    /// Struct fields that we've accessed using atomic operations.
+    ///
+    /// Each key refers to some [`Struct`] type, and each value is a set of
+    /// the indices of the fields in that struct that have been accessed
+    /// atomically.
+    ///
+    /// This includes fields with composite types (arrays, structs)
+    /// of which we've only accessed some components (elements, fields)
+    /// atomically.
+    ///
+    /// [`Struct`]: crate::TypeInner::Struct
+    fields: crate::FastHashMap<Handle<Type>, bit_set::BitSet>,
+}
+
+impl Upgrades {
+    pub fn insert_global(&mut self, global: Handle<GlobalVariable>) {
+        self.globals.insert(global);
+    }
+
+    pub fn insert_field(&mut self, struct_type: Handle<Type>, field: usize) {
+        self.fields.entry(struct_type).or_default().insert(field);
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.globals.is_empty()
+    }
 }
 
 struct UpgradeState<'a> {
     padding: Padding,
     module: &'a mut Module,
+
+    /// A map from old types to their upgraded versions.
+    ///
+    /// This ensures we never try to rebuild a type more than once.
+    upgraded_types: crate::FastHashMap<Handle<Type>, Handle<Type>>,
 }
 
 impl UpgradeState<'_> {
@@ -123,37 +158,50 @@ impl UpgradeState<'_> {
     fn upgrade_type(
         &mut self,
         ty: Handle<Type>,
-        field_indices: &mut Vec<u32>,
+        upgrades: &Upgrades,
     ) -> Result<Handle<Type>, Error> {
         let padding = self.inc_padding();
         padding.trace("visiting type: ", ty);
 
+        // If we've already upgraded this type, return the handle we produced at
+        // the time.
+        if let Some(&new) = self.upgraded_types.get(&ty) {
+            return Ok(new);
+        }
+
         let inner = match self.module.types[ty].inner {
             TypeInner::Scalar(scalar) => {
                 log::trace!("{padding}hit the scalar leaf, replacing with an atomic");
                 TypeInner::Atomic(scalar)
             }
             TypeInner::Pointer { base, space } => TypeInner::Pointer {
-                base: self.upgrade_type(base, field_indices)?,
+                base: self.upgrade_type(base, upgrades)?,
                 space,
             },
             TypeInner::Array { base, size, stride } => TypeInner::Array {
-                base: self.upgrade_type(base, field_indices)?,
+                base: self.upgrade_type(base, upgrades)?,
                 size,
                 stride,
             },
             TypeInner::Struct { ref members, span } => {
-                let index = field_indices.pop().ok_or(Error::UnexpectedEndOfIndices)? as usize;
+                // If no field or subfield of this struct was ever accessed
+                // atomically, no change is needed. We should never have arrived here.
+                let Some(fields) = upgrades.fields.get(&ty) else {
+                    unreachable!("global or field incorrectly flagged as atomically accessed");
+                };
+
                 let mut new_members = members.clone();
-                new_members[index].ty = self.upgrade_type(new_members[index].ty, field_indices)?;
+                for field in fields {
+                    new_members[field].ty = self.upgrade_type(new_members[field].ty, upgrades)?;
+                }
 
                 TypeInner::Struct {
                     members: new_members,
                     span,
                 }
             }
             TypeInner::BindingArray { base, size } => TypeInner::BindingArray {
-                base: self.upgrade_type(base, field_indices)?,
+                base: self.upgrade_type(base, upgrades)?,
                 size,
             },
             _ => return Ok(ty),
@@ -172,31 +220,32 @@ impl UpgradeState<'_> {
         padding.debug("from: ", r#type);
         padding.debug("to:   ", &new_type);
         let new_handle = self.module.types.insert(new_type, span);
+        self.upgraded_types.insert(ty, new_handle);
         Ok(new_handle)
     }
 
-    fn upgrade_global_variable(
-        &mut self,
-        mut global: ContainedGlobalVariable,
-    ) -> Result<(), Error> {
-        let padding = self.inc_padding();
-        padding.trace("visiting global variable: ", &global);
+    fn upgrade_all(&mut self, upgrades: &Upgrades) -> Result<(), Error> {
+        for handle in upgrades.globals.iter() {
+            let padding = self.inc_padding();
 
-        let var = &self.module.global_variables[global.handle];
-        padding.trace("var: ", var);
+            let global = &self.module.global_variables[handle];
+            padding.trace("visiting global variable: ", handle);
+            padding.trace("var: ", global);
 
-        if var.init.is_some() {
-            return Err(Error::GlobalInitUnsupported);
-        }
+            if global.init.is_some() {
+                return Err(Error::GlobalInitUnsupported);
+            }
 
-        let var_ty = var.ty;
-        let new_ty = self.upgrade_type(var.ty, &mut global.field_indices)?;
-        if new_ty != var_ty {
-            padding.debug("upgrading global variable: ", global.handle);
-            padding.debug("from ty: ", var_ty);
-            padding.debug("to ty:   ", new_ty);
-            self.module.global_variables[global.handle].ty = new_ty;
+            let var_ty = global.ty;
+            let new_ty = self.upgrade_type(var_ty, upgrades)?;
+            if new_ty != var_ty {
+                padding.debug("upgrading global variable: ", handle);
+                padding.debug("from ty: ", var_ty);
+                padding.debug("to ty:   ", new_ty);
+                self.module.global_variables[handle].ty = new_ty;
+            }
         }
+
         Ok(())
     }
 }
@@ -205,18 +254,17 @@ impl Module {
     /// Upgrade `global_var_handles` to have [`Atomic`] leaf types.
     ///
     /// [`Atomic`]: TypeInner::Atomic
-    pub(crate) fn upgrade_atomics(
-        &mut self,
-        globals: impl IntoIterator<Item = ContainedGlobalVariable>,
-    ) -> Result<(), Error> {
+    pub(crate) fn upgrade_atomics(&mut self, upgrades: &Upgrades) -> Result<(), Error> {
         let mut state = UpgradeState {
             padding: Default::default(),
             module: self,
+            upgraded_types: crate::FastHashMap::with_capacity_and_hasher(
+                upgrades.fields.len(),
+                Default::default(),
+            ),
         };
 
-        for global in globals {
-            state.upgrade_global_variable(global)?;
-        }
+        state.upgrade_all(upgrades)?;
 
         Ok(())
     }