naga/proc/
namer.rs

1use alloc::{
2    borrow::Cow,
3    format,
4    string::{String, ToString},
5    vec::Vec,
6};
7
8use crate::{
9    arena::Handle,
10    proc::{keyword_set::CaseInsensitiveKeywordSet, KeywordSet},
11    FastHashMap,
12};
13
14pub type EntryPointIndex = u16;
15const SEPARATOR: char = '_';
16
17/// A component of a lowered external texture.
18///
19/// Whereas the WGSL backend implements [`ImageClass::External`]
20/// images directly, most other Naga backends lower them to a
21/// collection of ordinary textures that represent individual planes
22/// (as received from a video decoder, perhaps), together with a
23/// struct of parameters saying how they should be cropped, sampled,
24/// and color-converted.
25///
26/// This lowering means that individual globals and function
27/// parameters in Naga IR must be split out by the backends into
28/// collections of globals and parameters of simpler types.
29///
30/// A value of this enum serves as a name key for one specific
31/// component in the lowered representation of an external texture.
32/// That is, these keys are for variables/parameters that do not exist
33/// in the Naga IR, only in its lowered form.
34///
35/// [`ImageClass::External`]: crate::ir::ImageClass::External
36#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
37pub enum ExternalTextureNameKey {
38    Plane(usize),
39    Params,
40}
41
42impl ExternalTextureNameKey {
43    const ALL: &[(&str, ExternalTextureNameKey)] = &[
44        ("_plane0", ExternalTextureNameKey::Plane(0)),
45        ("_plane1", ExternalTextureNameKey::Plane(1)),
46        ("_plane2", ExternalTextureNameKey::Plane(2)),
47        ("_params", ExternalTextureNameKey::Params),
48    ];
49}
50
51#[derive(Debug, Eq, Hash, PartialEq)]
52pub enum NameKey {
53    Constant(Handle<crate::Constant>),
54    GlobalVariable(Handle<crate::GlobalVariable>),
55    Type(Handle<crate::Type>),
56    StructMember(Handle<crate::Type>, u32),
57    Function(Handle<crate::Function>),
58    FunctionArgument(Handle<crate::Function>, u32),
59    FunctionLocal(Handle<crate::Function>, Handle<crate::LocalVariable>),
60
61    /// A local variable used by ReadZeroSkipWrite bounds-check policy
62    /// when it needs to produce a pointer-typed result for an OOB access.
63    /// These are unique per accessed type, so the second element is a
64    /// type handle. See docs for [`crate::back::msl`].
65    FunctionOobLocal(Handle<crate::Function>, Handle<crate::Type>),
66
67    EntryPoint(EntryPointIndex),
68    EntryPointLocal(EntryPointIndex, Handle<crate::LocalVariable>),
69    EntryPointArgument(EntryPointIndex, u32),
70
71    /// Entry point version of `FunctionOobLocal`.
72    EntryPointOobLocal(EntryPointIndex, Handle<crate::Type>),
73
74    /// A global variable holding a component of a lowered external texture.
75    ///
76    /// See [`ExternalTextureNameKey`] for details.
77    ExternalTextureGlobalVariable(Handle<crate::GlobalVariable>, ExternalTextureNameKey),
78
79    /// A function argument holding a component of a lowered external
80    /// texture.
81    ///
82    /// See [`ExternalTextureNameKey`] for details.
83    ExternalTextureFunctionArgument(Handle<crate::Function>, u32, ExternalTextureNameKey),
84}
85
86/// This processor assigns names to all the things in a module
87/// that may need identifiers in a textual backend.
88#[derive(Default)]
89pub struct Namer {
90    /// The last numeric suffix used for each base name. Zero means "no suffix".
91    unique: FastHashMap<String, u32>,
92    keywords: &'static KeywordSet,
93    keywords_case_insensitive: &'static CaseInsensitiveKeywordSet,
94    reserved_prefixes: Vec<&'static str>,
95}
96
97impl Namer {
98    /// Return a form of `string` suitable for use as the base of an identifier.
99    ///
100    /// - Drop leading digits.
101    /// - Retain only alphanumeric and `_` characters.
102    /// - Avoid prefixes in [`Namer::reserved_prefixes`].
103    /// - Replace consecutive `_` characters with a single `_` character.
104    ///
105    /// The return value is a valid identifier prefix in all of Naga's output languages,
106    /// and it never ends with a `SEPARATOR` character.
107    /// It is used as a key into the unique table.
108    fn sanitize<'s>(&self, string: &'s str) -> Cow<'s, str> {
109        let string = string
110            .trim_start_matches(|c: char| c.is_numeric())
111            .trim_end_matches(SEPARATOR);
112
113        let base = if !string.is_empty()
114            && !string.contains("__")
115            && string
116                .chars()
117                .all(|c: char| c.is_ascii_alphanumeric() || c == '_')
118        {
119            Cow::Borrowed(string)
120        } else {
121            let mut filtered = string
122                .chars()
123                .filter(|&c| c.is_ascii_alphanumeric() || c == '_')
124                .fold(String::new(), |mut s, c| {
125                    if s.ends_with('_') && c == '_' {
126                        return s;
127                    }
128                    s.push(c);
129                    s
130                });
131            let stripped_len = filtered.trim_end_matches(SEPARATOR).len();
132            filtered.truncate(stripped_len);
133            if filtered.is_empty() {
134                filtered.push_str("unnamed");
135            }
136            Cow::Owned(filtered)
137        };
138
139        for prefix in &self.reserved_prefixes {
140            if base.starts_with(prefix) {
141                return format!("gen_{base}").into();
142            }
143        }
144
145        base
146    }
147
148    /// Return a new identifier based on `label_raw`.
149    ///
150    /// The result:
151    /// - is a valid identifier even if `label_raw` is not
152    /// - conflicts with no keywords listed in `Namer::keywords`, and
153    /// - is different from any identifier previously constructed by this
154    ///   `Namer`.
155    ///
156    /// Guarantee uniqueness by applying a numeric suffix when necessary. If `label_raw`
157    /// itself ends with digits, separate them from the suffix with an underscore.
158    pub fn call(&mut self, label_raw: &str) -> String {
159        use core::fmt::Write as _; // for write!-ing to Strings
160
161        let base = self.sanitize(label_raw);
162        debug_assert!(!base.is_empty() && !base.ends_with(SEPARATOR));
163
164        // This would seem to be a natural place to use `HashMap::entry`. However, `entry`
165        // requires an owned key, and we'd like to avoid heap-allocating strings we're
166        // just going to throw away. The approach below double-hashes only when we create
167        // a new entry, in which case the heap allocation of the owned key was more
168        // expensive anyway.
169        match self.unique.get_mut(base.as_ref()) {
170            Some(count) => {
171                *count += 1;
172                // Add the suffix. This may fit in base's existing allocation.
173                let mut suffixed = base.into_owned();
174                write!(suffixed, "{}{}", SEPARATOR, *count).unwrap();
175                suffixed
176            }
177            None => {
178                let mut suffixed = base.to_string();
179                if base.ends_with(char::is_numeric)
180                    || self.keywords.contains(base.as_ref())
181                    || self.keywords_case_insensitive.contains(base.as_ref())
182                {
183                    suffixed.push(SEPARATOR);
184                }
185                debug_assert!(!self.keywords.contains(&suffixed));
186                // `self.unique` wants to own its keys. This allocates only if we haven't
187                // already done so earlier.
188                self.unique.insert(base.into_owned(), 0);
189                suffixed
190            }
191        }
192    }
193
194    pub fn call_or(&mut self, label: &Option<String>, fallback: &str) -> String {
195        self.call(match *label {
196            Some(ref name) => name,
197            None => fallback,
198        })
199    }
200
201    /// Enter a local namespace for things like structs.
202    ///
203    /// Struct member names only need to be unique amongst themselves, not
204    /// globally. This function temporarily establishes a fresh, empty naming
205    /// context for the duration of the call to `body`.
206    fn namespace(&mut self, capacity: usize, body: impl FnOnce(&mut Self)) {
207        let fresh = FastHashMap::with_capacity_and_hasher(capacity, Default::default());
208        let outer = core::mem::replace(&mut self.unique, fresh);
209        body(self);
210        self.unique = outer;
211    }
212
213    pub fn reset(
214        &mut self,
215        module: &crate::Module,
216        reserved_keywords: &'static KeywordSet,
217        reserved_keywords_case_insensitive: &'static CaseInsensitiveKeywordSet,
218        reserved_prefixes: &[&'static str],
219        output: &mut FastHashMap<NameKey, String>,
220    ) {
221        self.reserved_prefixes.clear();
222        self.reserved_prefixes.extend(reserved_prefixes.iter());
223
224        self.unique.clear();
225        self.keywords = reserved_keywords;
226        self.keywords_case_insensitive = reserved_keywords_case_insensitive;
227
228        // Choose fallback names for anonymous entry point return types.
229        let mut entrypoint_type_fallbacks = FastHashMap::default();
230        for ep in &module.entry_points {
231            if let Some(ref result) = ep.function.result {
232                if let crate::Type {
233                    name: None,
234                    inner: crate::TypeInner::Struct { .. },
235                } = module.types[result.ty]
236                {
237                    let label = match ep.stage {
238                        crate::ShaderStage::Vertex => "VertexOutput",
239                        crate::ShaderStage::Fragment => "FragmentOutput",
240                        crate::ShaderStage::Compute => "ComputeOutput",
241                        crate::ShaderStage::Task | crate::ShaderStage::Mesh => unreachable!(),
242                    };
243                    entrypoint_type_fallbacks.insert(result.ty, label);
244                }
245            }
246        }
247
248        let mut temp = String::new();
249
250        for (ty_handle, ty) in module.types.iter() {
251            // If the type is anonymous, check `entrypoint_types` for
252            // something better than just `"type"`.
253            let raw_label = match ty.name {
254                Some(ref given_name) => given_name.as_str(),
255                None => entrypoint_type_fallbacks
256                    .get(&ty_handle)
257                    .cloned()
258                    .unwrap_or("type"),
259            };
260            let ty_name = self.call(raw_label);
261            output.insert(NameKey::Type(ty_handle), ty_name);
262
263            if let crate::TypeInner::Struct { ref members, .. } = ty.inner {
264                // struct members have their own namespace, because access is always prefixed
265                self.namespace(members.len(), |namer| {
266                    for (index, member) in members.iter().enumerate() {
267                        let name = namer.call_or(&member.name, "member");
268                        output.insert(NameKey::StructMember(ty_handle, index as u32), name);
269                    }
270                })
271            }
272        }
273
274        for (ep_index, ep) in module.entry_points.iter().enumerate() {
275            let ep_name = self.call(&ep.name);
276            output.insert(NameKey::EntryPoint(ep_index as _), ep_name);
277            for (index, arg) in ep.function.arguments.iter().enumerate() {
278                let name = self.call_or(&arg.name, "param");
279                output.insert(
280                    NameKey::EntryPointArgument(ep_index as _, index as u32),
281                    name,
282                );
283            }
284            for (handle, var) in ep.function.local_variables.iter() {
285                let name = self.call_or(&var.name, "local");
286                output.insert(NameKey::EntryPointLocal(ep_index as _, handle), name);
287            }
288        }
289
290        for (fun_handle, fun) in module.functions.iter() {
291            let fun_name = self.call_or(&fun.name, "function");
292            output.insert(NameKey::Function(fun_handle), fun_name);
293            for (index, arg) in fun.arguments.iter().enumerate() {
294                let name = self.call_or(&arg.name, "param");
295                output.insert(NameKey::FunctionArgument(fun_handle, index as u32), name);
296
297                if matches!(
298                    module.types[arg.ty].inner,
299                    crate::TypeInner::Image {
300                        class: crate::ImageClass::External,
301                        ..
302                    }
303                ) {
304                    let base = arg.name.as_deref().unwrap_or("param");
305                    for &(suffix, ext_key) in ExternalTextureNameKey::ALL {
306                        let name = self.call(&format!("{base}_{suffix}"));
307                        output.insert(
308                            NameKey::ExternalTextureFunctionArgument(
309                                fun_handle,
310                                index as u32,
311                                ext_key,
312                            ),
313                            name,
314                        );
315                    }
316                }
317            }
318            for (handle, var) in fun.local_variables.iter() {
319                let name = self.call_or(&var.name, "local");
320                output.insert(NameKey::FunctionLocal(fun_handle, handle), name);
321            }
322        }
323
324        for (handle, var) in module.global_variables.iter() {
325            let name = self.call_or(&var.name, "global");
326            output.insert(NameKey::GlobalVariable(handle), name);
327
328            if matches!(
329                module.types[var.ty].inner,
330                crate::TypeInner::Image {
331                    class: crate::ImageClass::External,
332                    ..
333                }
334            ) {
335                let base = var.name.as_deref().unwrap_or("global");
336                for &(suffix, ext_key) in ExternalTextureNameKey::ALL {
337                    let name = self.call(&format!("{base}_{suffix}"));
338                    output.insert(
339                        NameKey::ExternalTextureGlobalVariable(handle, ext_key),
340                        name,
341                    );
342                }
343            }
344        }
345
346        for (handle, constant) in module.constants.iter() {
347            let label = match constant.name {
348                Some(ref name) => name,
349                None => {
350                    use core::fmt::Write;
351                    // Try to be more descriptive about the constant values
352                    temp.clear();
353                    write!(temp, "const_{}", output[&NameKey::Type(constant.ty)]).unwrap();
354                    &temp
355                }
356            };
357            let name = self.call(label);
358            output.insert(NameKey::Constant(handle), name);
359        }
360    }
361}
362
363#[test]
364fn test() {
365    let mut namer = Namer::default();
366    assert_eq!(namer.call("x"), "x");
367    assert_eq!(namer.call("x"), "x_1");
368    assert_eq!(namer.call("x1"), "x1_");
369    assert_eq!(namer.call("__x"), "_x");
370    assert_eq!(namer.call("1___x"), "_x_1");
371}