wesl/
mangler.rs

1//! [`Mangler`] trait and implementations.
2
3use std::cell::RefCell;
4use std::collections::HashMap;
5use std::fmt;
6use std::hash::DefaultHasher;
7use std::hash::Hash;
8use std::hash::Hasher;
9
10use itertools::Itertools;
11use wgsl_parse::syntax::Expression;
12use wgsl_parse::syntax::PathOrigin;
13use wgsl_parse::syntax::TypeExpression;
14
15use super::ModulePath;
16
17/// A name mangler is responsible for renaming import-qualified identifiers into valid
18/// WGSL identifiers.
19///
20/// `Mangler` implementations must respect the following constraints:
21/// * A pair {path, item} (aka. fully-qualified name) must be associated with a unique mangled name.
22/// * A mangled name must be associated with a unique pair {path, item} (or at least, the risk of
23///   a collision must be negligible).
24/// * The mangled name must be a valid WGSL identifier.
25///
26/// Calls to `Mangler::mangle` must respect these preconditions:
27/// * the item must be a valid WGSL identifier.
28///
29/// # WESL Reference
30/// spec: [NameMangling.md](https://github.com/webgpu-tools/wesl-spec/blob/main/NameMangling.md)
31pub trait Mangler {
32    /// Turn an import path and item name into a mangled WGSL identifier.
33    fn mangle(&self, path: &ModulePath, item: &str) -> String;
34    /// Reverse the [`Mangler::mangle`] operation. Implementing this is optional.
35    fn unmangle(&self, _mangled: &str) -> Option<(ModulePath, String)> {
36        None
37    }
38    /// Used for generics. Is experimental. Implementing is optional.
39    fn mangle_types(&self, item: &str, variant: u32, _types: &[TypeExpression]) -> String {
40        format!("{item}_{variant}")
41    }
42}
43
44impl<T: Mangler + ?Sized> Mangler for Box<T> {
45    fn mangle(&self, path: &ModulePath, item: &str) -> String {
46        (**self).mangle(path, item)
47    }
48    fn unmangle(&self, mangled: &str) -> Option<(ModulePath, String)> {
49        (**self).unmangle(mangled)
50    }
51    fn mangle_types(&self, item: &str, variant: u32, types: &[TypeExpression]) -> String {
52        (**self).mangle_types(item, variant, types)
53    }
54}
55
56impl<T: Mangler + ?Sized> Mangler for &T {
57    fn mangle(&self, path: &ModulePath, item: &str) -> String {
58        (**self).mangle(path, item)
59    }
60    fn unmangle(&self, mangled: &str) -> Option<(ModulePath, String)> {
61        (**self).unmangle(mangled)
62    }
63    fn mangle_types(&self, item: &str, variant: u32, types: &[TypeExpression]) -> String {
64        (**self).mangle_types(item, variant, types)
65    }
66}
67
68/// A mangler that hashes the module path.
69/// e.g. `foo::bar::baz item => item_32938483840293402930392`
70#[derive(Default, Clone, Debug)]
71pub struct HashMangler;
72
73impl Mangler for HashMangler {
74    fn mangle(&self, path: &ModulePath, item: &str) -> String {
75        let mut hasher = DefaultHasher::new();
76        path.hash(&mut hasher);
77        item.hash(&mut hasher);
78        let hash = hasher.finish();
79        format!("{item}_{hash}")
80    }
81}
82
83/// A mangler that replaces `::` with `_` and prefixes components with the number of `_`
84/// they contain.
85/// e.g. `foo::bar_baz item => foo__1bar_baz_item`
86///
87/// This is WESL's default mangler.
88///
89/// If a component contains `_`, it is prefixed by `_`, followed by the number of `_`.
90/// This scheme guarantees that paths can be unambiguously mangled and unmangled.
91#[derive(Default, Clone, Debug)]
92pub struct EscapeMangler;
93
94impl EscapeMangler {
95    fn escape_component(comp: &str) -> String {
96        if comp.contains('/') {
97            // This can exist only for dependencies of dependencies. see ModulePath::join_path.
98            let count = comp.chars().filter(|c| *c == '/').count() + 1;
99            format!(
100                "import_{count}_{}",
101                comp.split('/').map(Self::escape_component).format("_")
102            )
103        } else {
104            let underscores = comp.chars().filter(|c| *c == '_').count();
105            if underscores > 0 {
106                format!("_{underscores}{comp}")
107            } else {
108                comp.to_string()
109            }
110        }
111    }
112
113    fn unescape_component<'a>(
114        part: &str,
115        parts: &mut impl Iterator<Item = &'a str>,
116    ) -> Option<String> {
117        fn extract_count(part: &str) -> Option<(usize, &str)> {
118            let digits = part.chars().take_while(|c| c.is_ascii_digit()).count();
119
120            if digits == 0 {
121                return None;
122            }
123
124            let count = part[..digits].parse().ok()?;
125            Some((count, &part[digits..]))
126        }
127
128        let component = if part.is_empty() {
129            let first = parts.next()?;
130            let (n, rem) = extract_count(first)?;
131
132            std::iter::once(Some(rem))
133                .chain((0..n).map(|_| parts.next()))
134                .collect::<Option<Vec<_>>>()?
135                .join("_")
136        } else {
137            part.to_string()
138        };
139
140        Some(component)
141    }
142}
143
144impl Mangler for EscapeMangler {
145    fn mangle(&self, path: &ModulePath, item: &str) -> String {
146        let origin = match &path.origin {
147            PathOrigin::Absolute => "package".to_string(),
148            PathOrigin::Relative(0) => "self".to_string(),
149            PathOrigin::Relative(n) => format!("{}", (0..*n).map(|_| "super").format("_")),
150            PathOrigin::Package(name) => Self::escape_component(name).to_string(),
151        };
152
153        let item = Self::escape_component(item);
154
155        if path.components.is_empty() {
156            format!("{origin}_{item}",)
157        } else {
158            let path = path
159                .components
160                .iter()
161                .map(|comp| Self::escape_component(comp))
162                .format("_");
163
164            format!("{origin}_{path}_{item}",)
165        }
166    }
167
168    /// This function does not check that the module names are valid. Unmangling invalid
169    /// names (or names not mangled with the EscapeMangler) may produce erroneous results.
170    fn unmangle(&self, mangled: &str) -> Option<(ModulePath, String)> {
171        let mut components = Vec::new();
172        let mut parts = mangled.split('_').peekable();
173
174        let origin = match parts.next() {
175            Some("package") => PathOrigin::Absolute,
176            Some("self") => PathOrigin::Relative(0),
177            Some("super") => {
178                let mut n = 1;
179                while let Some(&"super") = parts.peek() {
180                    n += 1;
181                    parts.next().unwrap();
182                }
183                PathOrigin::Relative(n)
184            }
185            Some("import") => {
186                let count: u32 = parts.next()?.parse().ok()?;
187                let pkg_name = (0..count)
188                    .map(|_| Self::unescape_component(parts.next()?, &mut parts))
189                    .collect::<Option<Vec<_>>>()?
190                    .join("/");
191                PathOrigin::Package(pkg_name)
192            }
193            Some(name) => PathOrigin::Package(Self::unescape_component(name, &mut parts)?),
194            None => return None,
195        };
196
197        while let Some(part) = parts.next() {
198            components.push(Self::unescape_component(part, &mut parts)?);
199        }
200
201        let item = components.pop()?;
202
203        let path = ModulePath::new(origin, components);
204        Some((path, item))
205    }
206}
207
208#[test]
209fn test_escape_mangler() {
210    let tests = [
211        (
212            "package::bevy_pbr::lighting",
213            "package__1bevy_pbr_lighting_item",
214        ),
215        ("a_b::c_d::e", "_1a_b__1c_d_e_item"),
216        ("self::a", "self_a_item"),
217        ("super::super", "super_super_item"),
218        ("package::_a_b__c___", "package__7_a_b__c____item"),
219        // test package dependencies with slashes
220        ("bevy/pbr", "import_2_bevy_pbr_item"),
221        ("bevy/pbr::lighting", "import_2_bevy_pbr_lighting_item"),
222        ("bevy_pbr::lighting", "_1bevy_pbr_lighting_item"),
223    ];
224
225    for (p, m) in tests {
226        let p = p.parse().expect("failed to parse module path");
227        println!("testing {p}::item -> {m}");
228        assert_eq!(EscapeMangler.mangle(&p, "item"), m);
229        assert_eq!(EscapeMangler.unmangle(m), Some((p, "item".to_string())));
230    }
231}
232
233/// A mangler that just returns the identifier as-is (no mangling).
234/// e.g. `foo::bar::baz item => item`
235///
236/// Warning: this mangler is not spec-compliant. It can cause name collisions.
237#[derive(Default, Clone, Debug)]
238pub struct NoMangler;
239
240impl Mangler for NoMangler {
241    fn mangle(&self, _path: &ModulePath, item: &str) -> String {
242        item.to_string()
243    }
244}
245
246/// A mangler that remembers and can unmangle.
247pub struct CacheMangler<'a, T: Mangler> {
248    cache: RefCell<HashMap<String, (ModulePath, String)>>,
249    mangler: &'a T,
250}
251
252impl<'a, T: Mangler> CacheMangler<'a, T> {
253    pub fn new(mangler: &'a T) -> Self {
254        Self {
255            cache: Default::default(),
256            mangler,
257        }
258    }
259}
260
261impl<T: Mangler> Mangler for CacheMangler<'_, T> {
262    fn mangle(&self, path: &ModulePath, item: &str) -> String {
263        let res = self.mangler.mangle(path, item);
264        let mut cache = self.cache.borrow_mut();
265        cache.insert(res.clone(), (path.clone(), item.to_string()));
266        res
267    }
268    fn unmangle(&self, mangled: &str) -> Option<(ModulePath, String)> {
269        {
270            let cache = self.cache.borrow();
271            if let Some(res) = cache.get(mangled).cloned() {
272                return Some(res);
273            }
274        }
275
276        self.mangler.unmangle(mangled)
277    }
278}
279
280/// A mangler that uses cryptic unicode symbols that look like :, < and >
281/// e.g. `foo::bar::baz array<f32,2> => foo::bar::baz::arrayᐸf32ˏ2ᐳ`
282///
283/// Very unlikely to collide unless using U+02D0 characters.
284///
285/// # Panics
286/// if the TypeExpression is not normalized (i.e. contains only identifiers and literals)
287#[derive(Default, Clone, Debug)]
288pub struct UnicodeMangler;
289
290impl UnicodeMangler {
291    const LT: char = 'ᐸ'; // U+1438
292    const GT: char = 'ᐳ'; // U+02CF
293    const SEP: &'static str = "::"; // <-- these are NOT colons, they are U+02D0
294    const TY_SEP: &'static str = "ˏ"; // U+02CF
295
296    fn display_ty(ty: &TypeExpression) -> impl fmt::Display + '_ {
297        format_args!(
298            "{}{}",
299            ty.ident,
300            ty.template_args
301                .iter()
302                .format_with(Self::TY_SEP, |tplt, f| {
303                    f(&format_args!(
304                        "{}{}{}",
305                        Self::LT,
306                        tplt.iter().format_with(Self::TY_SEP, |tplt, f| {
307                            match tplt.expression.node() {
308                                Expression::Literal(lit) => f(lit),
309                                Expression::TypeOrIdentifier(ty) => f(&Self::display_ty(ty)),
310                                _ => panic!("only type names can be mangled"),
311                            }
312                        }),
313                        Self::GT
314                    ))
315                })
316        )
317        .to_string()
318    }
319}
320
321impl Mangler for UnicodeMangler {
322    fn mangle(&self, path: &ModulePath, item: &str) -> String {
323        let sep = Self::SEP;
324        let origin = match &path.origin {
325            PathOrigin::Absolute => format!("package{sep}"),
326            PathOrigin::Relative(0) => format!("self{sep}"),
327            PathOrigin::Relative(n) => format!("{}{sep}", (1..*n).map(|_| "super").format(sep)),
328            PathOrigin::Package(name) => format!("{name}{sep}"),
329        };
330        let path = path.components.iter().format(sep);
331        format!("{origin}{path}{sep}{item}")
332    }
333
334    fn unmangle(&self, mangled: &str) -> Option<(ModulePath, String)> {
335        let mut components = mangled.split(Self::SEP).map(str::to_string).collect_vec();
336        let mut origin = PathOrigin::Absolute;
337        while let Some(comp) = components.first() {
338            match comp.as_str() {
339                "package" => {
340                    components.remove(0);
341                    origin = PathOrigin::Absolute;
342                    break;
343                }
344                "self" => {
345                    components.remove(0);
346                    origin = PathOrigin::Relative(0);
347                    break;
348                }
349                "super" => {
350                    components.remove(0);
351                    origin = match origin {
352                        PathOrigin::Relative(n) => PathOrigin::Relative(n + 1),
353                        _ => PathOrigin::Relative(1),
354                    };
355                }
356                _ => break,
357            }
358        }
359        let name = components.pop()?;
360        Some((ModulePath::new(origin, components), name))
361    }
362    fn mangle_types(&self, item: &str, _variant: u32, types: &[TypeExpression]) -> String {
363        // these are NOT chevrons and comma!
364        format!(
365            "{item}{}{}{}",
366            Self::LT,
367            types
368                .iter()
369                .format_with(Self::TY_SEP, |ty, f| f(&Self::display_ty(ty))),
370            Self::GT
371        )
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    use super::*;
378
379    #[test]
380    fn mangle_textures3d() {
381        let mangler = EscapeMangler;
382        let module_path = ModulePath::new_root();
383        let mangled = mangler.mangle(&module_path, "textures_3d");
384        assert_eq!("package__1textures_3d", mangled);
385    }
386
387    #[test]
388    fn unmangle_textures3d() {
389        let mangler = EscapeMangler;
390        let unmangled = mangler.unmangle("package__1textures_3d").map(|x| x.1);
391        assert_eq!(Some("textures_3d"), unmangled.as_deref());
392    }
393}
wesl/mangler.rs

wesl/
mangler.rs
