1 files changed, 193 insertions, 0 deletions
diff --git a/third_party/rust/naga/src/front/wgsl/index.rs b/third_party/rust/naga/src/front/wgsl/index.rs
new file mode 100644
index 0000000000..a5524fe8f1
--- /dev/null
+++ b/third_party/rust/naga/src/front/wgsl/index.rs
@@ -0,0 +1,193 @@
+use super::Error;
+use crate::front::wgsl::parse::ast;
+use crate::{FastHashMap, Handle, Span};
+
+/// A `GlobalDecl` list in which each definition occurs before all its uses.
+pub struct Index<'a> {
+    dependency_order: Vec<Handle<ast::GlobalDecl<'a>>>,
+}
+
+impl<'a> Index<'a> {
+    /// Generate an `Index` for the given translation unit.
+    ///
+    /// Perform a topological sort on `tu`'s global declarations, placing
+    /// referents before the definitions that refer to them.
+    ///
+    /// Return an error if the graph of references between declarations contains
+    /// any cycles.
+    pub fn generate(tu: &ast::TranslationUnit<'a>) -> Result<Self, Error<'a>> {
+        // Produce a map from global definitions' names to their `Handle<GlobalDecl>`s.
+        // While doing so, reject conflicting definitions.
+        let mut globals = FastHashMap::with_capacity_and_hasher(tu.decls.len(), Default::default());
+        for (handle, decl) in tu.decls.iter() {
+            let ident = decl_ident(decl);
+            let name = ident.name;
+            if let Some(old) = globals.insert(name, handle) {
+                return Err(Error::Redefinition {
+                    previous: decl_ident(&tu.decls[old]).span,
+                    current: ident.span,
+                });
+            }
+        }
+
+        let len = tu.decls.len();
+        let solver = DependencySolver {
+            globals: &globals,
+            module: tu,
+            visited: vec![false; len],
+            temp_visited: vec![false; len],
+            path: Vec::new(),
+            out: Vec::with_capacity(len),
+        };
+        let dependency_order = solver.solve()?;
+
+        Ok(Self { dependency_order })
+    }
+
+    /// Iterate over `GlobalDecl`s, visiting each definition before all its uses.
+    ///
+    /// Produce handles for all of the `GlobalDecl`s of the `TranslationUnit`
+    /// passed to `Index::generate`, ordered so that a given declaration is
+    /// produced before any other declaration that uses it.
+    pub fn visit_ordered(&self) -> impl Iterator<Item = Handle<ast::GlobalDecl<'a>>> + '_ {
+        self.dependency_order.iter().copied()
+    }
+}
+
+/// An edge from a reference to its referent in the current depth-first
+/// traversal.
+///
+/// This is like `ast::Dependency`, except that we've determined which
+/// `GlobalDecl` it refers to.
+struct ResolvedDependency<'a> {
+    /// The referent of some identifier used in the current declaration.
+    decl: Handle<ast::GlobalDecl<'a>>,
+
+    /// Where that use occurs within the current declaration.
+    usage: Span,
+}
+
+/// Local state for ordering a `TranslationUnit`'s module-scope declarations.
+///
+/// Values of this type are used temporarily by `Index::generate`
+/// to perform a depth-first sort on the declarations.
+/// Technically, what we want is a topological sort, but a depth-first sort
+/// has one key benefit - it's much more efficient in storing
+/// the path of each node for error generation.
+struct DependencySolver<'source, 'temp> {
+    /// A map from module-scope definitions' names to their handles.
+    globals: &'temp FastHashMap<&'source str, Handle<ast::GlobalDecl<'source>>>,
+
+    /// The translation unit whose declarations we're ordering.
+    module: &'temp ast::TranslationUnit<'source>,
+
+    /// For each handle, whether we have pushed it onto `out` yet.
+    visited: Vec<bool>,
+
+    /// For each handle, whether it is an predecessor in the current depth-first
+    /// traversal. This is used to detect cycles in the reference graph.
+    temp_visited: Vec<bool>,
+
+    /// The current path in our depth-first traversal. Used for generating
+    /// error messages for non-trivial reference cycles.
+    path: Vec<ResolvedDependency<'source>>,
+
+    /// The list of declaration handles, with declarations before uses.
+    out: Vec<Handle<ast::GlobalDecl<'source>>>,
+}
+
+impl<'a> DependencySolver<'a, '_> {
+    /// Produce the sorted list of declaration handles, and check for cycles.
+    fn solve(mut self) -> Result<Vec<Handle<ast::GlobalDecl<'a>>>, Error<'a>> {
+        for (id, _) in self.module.decls.iter() {
+            if self.visited[id.index()] {
+                continue;
+            }
+
+            self.dfs(id)?;
+        }
+
+        Ok(self.out)
+    }
+
+    /// Ensure that all declarations used by `id` have been added to the
+    /// ordering, and then append `id` itself.
+    fn dfs(&mut self, id: Handle<ast::GlobalDecl<'a>>) -> Result<(), Error<'a>> {
+        let decl = &self.module.decls[id];
+        let id_usize = id.index();
+
+        self.temp_visited[id_usize] = true;
+        for dep in decl.dependencies.iter() {
+            if let Some(&dep_id) = self.globals.get(dep.ident) {
+                self.path.push(ResolvedDependency {
+                    decl: dep_id,
+                    usage: dep.usage,
+                });
+                let dep_id_usize = dep_id.index();
+
+                if self.temp_visited[dep_id_usize] {
+                    // Found a cycle.
+                    return if dep_id == id {
+                        // A declaration refers to itself directly.
+                        Err(Error::RecursiveDeclaration {
+                            ident: decl_ident(decl).span,
+                            usage: dep.usage,
+                        })
+                    } else {
+                        // A declaration refers to itself indirectly, through
+                        // one or more other definitions. Report the entire path
+                        // of references.
+                        let start_at = self
+                            .path
+                            .iter()
+                            .rev()
+                            .enumerate()
+                            .find_map(|(i, dep)| (dep.decl == dep_id).then_some(i))
+                            .unwrap_or(0);
+
+                        Err(Error::CyclicDeclaration {
+                            ident: decl_ident(&self.module.decls[dep_id]).span,
+                            path: self.path[start_at..]
+                                .iter()
+                                .map(|curr_dep| {
+                                    let curr_id = curr_dep.decl;
+                                    let curr_decl = &self.module.decls[curr_id];
+
+                                    (decl_ident(curr_decl).span, curr_dep.usage)
+                                })
+                                .collect(),
+                        })
+                    };
+                } else if !self.visited[dep_id_usize] {
+                    self.dfs(dep_id)?;
+                }
+
+                // Remove this edge from the current path.
+                self.path.pop();
+            }
+
+            // Ignore unresolved identifiers; they may be predeclared objects.
+        }
+
+        // Remove this node from the current path.
+        self.temp_visited[id_usize] = false;
+
+        // Now everything this declaration uses has been visited, and is already
+        // present in `out`. That means we we can append this one to the
+        // ordering, and mark it as visited.
+        self.out.push(id);
+        self.visited[id_usize] = true;
+
+        Ok(())
+    }
+}
+
+const fn decl_ident<'a>(decl: &ast::GlobalDecl<'a>) -> ast::Ident<'a> {
+    match decl.kind {
+        ast::GlobalDeclKind::Fn(ref f) => f.name,
+        ast::GlobalDeclKind::Var(ref v) => v.name,
+        ast::GlobalDeclKind::Const(ref c) => c.name,
+        ast::GlobalDeclKind::Struct(ref s) => s.name,
+        ast::GlobalDeclKind::Type(ref t) => t.name,
+    }
+}