Skip to main content

wesl/pass/
import.rs

1use itertools::Itertools;
2use std::collections::HashMap;
3use wgsl_parse::syntax::*;
4
5#[derive(Clone, Debug)]
6pub struct ImportedItem {
7    pub path: ModulePath,
8    pub ident: Ident, // this is the ident's original name before `as` renaming.
9    pub visibility: Visibility,
10}
11
12pub type Imports = HashMap<Ident, ImportedItem>;
13
14/// Flatten imports to a list.
15pub fn flatten_imports(imports: &[ImportStatement], path: &ModulePath) -> Imports {
16    fn rec(content: &ImportContent, path: ModulePath, visibility: Visibility, res: &mut Imports) {
17        match content {
18            ImportContent::Item(item) => {
19                let ident = item.rename.as_ref().unwrap_or(&item.ident).clone();
20                res.insert(
21                    ident,
22                    ImportedItem {
23                        path,
24                        ident: item.ident.clone(),
25                        visibility,
26                    },
27                );
28            }
29            ImportContent::Collection(coll) => {
30                for import in coll {
31                    let path = path.clone().join(import.path.iter().cloned());
32                    rec(&import.content, path, visibility, res);
33                }
34            }
35        }
36    }
37
38    let mut res = Imports::default();
39
40    for import in imports {
41        let visibility = import.visibility;
42        match &import.path {
43            Some(import_path) => {
44                let path = path.join_path(import_path);
45                rec(&import.content, path, visibility, &mut res);
46            }
47            None => {
48                // this covers two cases: `import foo;` and `import {foo, ..};`.
49                // COMBAK: these edge-cases smell
50                match &import.content {
51                    ImportContent::Item(_) => {
52                        // `import foo`, this import statement does nothing currently.
53                        // In the future, it may become a visibility/re-export mechanism.
54                    }
55                    ImportContent::Collection(coll) => {
56                        for import in coll {
57                            let mut components = import.path.iter().cloned();
58                            if let Some(pkg_name) = components.next() {
59                                // `import {foo::bar}`, foo becomes the package name.
60                                let path = ModulePath::new(
61                                    PathOrigin::Package(pkg_name),
62                                    components.collect_vec(),
63                                );
64                                rec(&import.content, path, visibility, &mut res);
65                            }
66                        }
67                    }
68                }
69            }
70        }
71    }
72
73    res
74}
75
76/// Find the normalized module path for an identifier in source, if it refers to an external declaration.
77///
78/// Inline imports differ from import statements only in case of package imports:
79/// the package component may refer to a local import shadowing the package name.
80pub fn imported_item_path(
81    ty_expr: &TypeExpression,
82    parent_path: &ModulePath,
83    imports: &Imports,
84) -> Option<(ModulePath, Ident)> {
85    if let Some(path) = &ty_expr.path {
86        match &path.origin {
87            PathOrigin::Package(pkg_name) => {
88                // the path could be either a package, or referencing an imported module alias.
89                let imported_item = imports.iter().find(|(ident, _)| *ident.name() == *pkg_name);
90
91                if let Some((_, ext_item)) = imported_item {
92                    // this inline path references an imported item. Example:
93                    // import a::b::c as foo; foo::bar::baz() => a::b::c::bar::baz()
94                    let mut res = ext_item.path.clone(); // a::b
95                    res.push(&ext_item.ident.name()); // c (not foo)
96                    let import_path = res.join(path.components.iter().cloned()); // a::b::c::bar
97                    Some((import_path, ty_expr.ident.clone()))
98                } else {
99                    let import_path = parent_path.join_path(path);
100                    Some((import_path, ty_expr.ident.clone()))
101                }
102            }
103            _ => {
104                let import_path = parent_path.join_path(path);
105                Some((import_path, ty_expr.ident.clone()))
106            }
107        }
108    } else {
109        imports
110            .get(&ty_expr.ident)
111            .map(|item| (item.path.clone(), item.ident.clone()))
112    }
113}