naga/front/
atomic_upgrade.rs1use alloc::{format, sync::Arc};
35use core::sync::atomic::AtomicUsize;
36
37use crate::{GlobalVariable, Handle, Module, Type, TypeInner};
38
39#[derive(Clone, Debug, thiserror::Error)]
40pub enum Error {
41 #[error("encountered an unsupported expression")]
42 Unsupported,
43 #[error("unexpected end of struct field access indices")]
44 UnexpectedEndOfIndices,
45 #[error("encountered unsupported global initializer in an atomic variable")]
46 GlobalInitUnsupported,
47 #[error("expected to find a global variable")]
48 GlobalVariableMissing,
49 #[error("atomic compare exchange requires a scalar base type")]
50 CompareExchangeNonScalarBaseType,
51}
52
53#[derive(Clone, Default)]
54struct Padding(Arc<AtomicUsize>);
55
56impl core::fmt::Display for Padding {
57 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
58 for _ in 0..self.0.load(core::sync::atomic::Ordering::Relaxed) {
59 f.write_str(" ")?;
60 }
61 Ok(())
62 }
63}
64
65impl Drop for Padding {
66 fn drop(&mut self) {
67 let _ = self.0.fetch_sub(1, core::sync::atomic::Ordering::Relaxed);
68 }
69}
70
71impl Padding {
72 fn trace(&self, msg: impl core::fmt::Display, t: impl core::fmt::Debug) {
73 format!("{msg} {t:#?}")
74 .split('\n')
75 .for_each(|ln| log::trace!("{self}{ln}"));
76 }
77
78 fn debug(&self, msg: impl core::fmt::Display, t: impl core::fmt::Debug) {
79 format!("{msg} {t:#?}")
80 .split('\n')
81 .for_each(|ln| log::debug!("{self}{ln}"));
82 }
83
84 fn inc_padding(&self) -> Padding {
85 let _ = self.0.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
86 self.clone()
87 }
88}
89
90#[derive(Debug, Default)]
91pub struct Upgrades {
92 globals: crate::arena::HandleSet<GlobalVariable>,
97
98 fields: crate::FastHashMap<Handle<Type>, bit_set::BitSet>,
110}
111
112impl Upgrades {
113 pub fn insert_global(&mut self, global: Handle<GlobalVariable>) {
114 self.globals.insert(global);
115 }
116
117 pub fn insert_field(&mut self, struct_type: Handle<Type>, field: usize) {
118 self.fields.entry(struct_type).or_default().insert(field);
119 }
120
121 pub fn is_empty(&self) -> bool {
122 self.globals.is_empty()
123 }
124}
125
126struct UpgradeState<'a> {
127 padding: Padding,
128 module: &'a mut Module,
129
130 upgraded_types: crate::FastHashMap<Handle<Type>, Handle<Type>>,
134}
135
136impl UpgradeState<'_> {
137 fn inc_padding(&self) -> Padding {
138 self.padding.inc_padding()
139 }
140
141 fn upgrade_type(
161 &mut self,
162 ty: Handle<Type>,
163 upgrades: &Upgrades,
164 ) -> Result<Handle<Type>, Error> {
165 let padding = self.inc_padding();
166 padding.trace("visiting type: ", ty);
167
168 if let Some(&new) = self.upgraded_types.get(&ty) {
171 return Ok(new);
172 }
173
174 let inner = match self.module.types[ty].inner {
175 TypeInner::Scalar(scalar) => {
176 log::trace!("{padding}hit the scalar leaf, replacing with an atomic");
177 TypeInner::Atomic(scalar)
178 }
179 TypeInner::Pointer { base, space } => TypeInner::Pointer {
180 base: self.upgrade_type(base, upgrades)?,
181 space,
182 },
183 TypeInner::Array { base, size, stride } => TypeInner::Array {
184 base: self.upgrade_type(base, upgrades)?,
185 size,
186 stride,
187 },
188 TypeInner::Struct { ref members, span } => {
189 let Some(fields) = upgrades.fields.get(&ty) else {
192 unreachable!("global or field incorrectly flagged as atomically accessed");
193 };
194
195 let mut new_members = members.clone();
196 for field in fields {
197 new_members[field].ty = self.upgrade_type(new_members[field].ty, upgrades)?;
198 }
199
200 TypeInner::Struct {
201 members: new_members,
202 span,
203 }
204 }
205 TypeInner::BindingArray { base, size } => TypeInner::BindingArray {
206 base: self.upgrade_type(base, upgrades)?,
207 size,
208 },
209 _ => return Ok(ty),
210 };
211
212 let r#type = &self.module.types[ty];
216 let span = self.module.types.get_span(ty);
217 let new_type = Type {
218 name: r#type.name.clone(),
219 inner,
220 };
221 padding.debug("ty: ", ty);
222 padding.debug("from: ", r#type);
223 padding.debug("to: ", &new_type);
224 let new_handle = self.module.types.insert(new_type, span);
225 self.upgraded_types.insert(ty, new_handle);
226 Ok(new_handle)
227 }
228
229 fn upgrade_all(&mut self, upgrades: &Upgrades) -> Result<(), Error> {
230 for handle in upgrades.globals.iter() {
231 let padding = self.inc_padding();
232
233 let global = &self.module.global_variables[handle];
234 padding.trace("visiting global variable: ", handle);
235 padding.trace("var: ", global);
236
237 if global.init.is_some() {
238 return Err(Error::GlobalInitUnsupported);
239 }
240
241 let var_ty = global.ty;
242 let new_ty = self.upgrade_type(var_ty, upgrades)?;
243 if new_ty != var_ty {
244 padding.debug("upgrading global variable: ", handle);
245 padding.debug("from ty: ", var_ty);
246 padding.debug("to ty: ", new_ty);
247 self.module.global_variables[handle].ty = new_ty;
248 }
249 }
250
251 Ok(())
252 }
253}
254
255impl Module {
256 pub(crate) fn upgrade_atomics(&mut self, upgrades: &Upgrades) -> Result<(), Error> {
260 let mut state = UpgradeState {
261 padding: Default::default(),
262 module: self,
263 upgraded_types: crate::FastHashMap::with_capacity_and_hasher(
264 upgrades.fields.len(),
265 Default::default(),
266 ),
267 };
268
269 state.upgrade_all(upgrades)?;
270
271 Ok(())
272 }
273}