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#[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 FunctionOobLocal(Handle<crate::Function>, Handle<crate::Type>),
66
67 EntryPoint(EntryPointIndex),
68 EntryPointLocal(EntryPointIndex, Handle<crate::LocalVariable>),
69 EntryPointArgument(EntryPointIndex, u32),
70
71 EntryPointOobLocal(EntryPointIndex, Handle<crate::Type>),
73
74 ExternalTextureGlobalVariable(Handle<crate::GlobalVariable>, ExternalTextureNameKey),
78
79 ExternalTextureFunctionArgument(Handle<crate::Function>, u32, ExternalTextureNameKey),
84}
85
86#[derive(Default)]
89pub struct Namer {
90 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 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 pub fn call(&mut self, label_raw: &str) -> String {
159 use core::fmt::Write as _; let base = self.sanitize(label_raw);
162 debug_assert!(!base.is_empty() && !base.ends_with(SEPARATOR));
163
164 match self.unique.get_mut(base.as_ref()) {
170 Some(count) => {
171 *count += 1;
172 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.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 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 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 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 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 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}