naga/proc/
namer.rs

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

naga/proc/
namer.rs
