wgsl-parser 0.5.0

use std::sync::Arc;

use gramatika::{Spanned, Substr, Token as _};

use crate::{
	decl::{
		Decl, ImportDecl, ImportPath, ImportPathBlock, ImportPathDecl, ImportPathLeaf,
		NamespacedImportPath,
	},
	token::Token,
	traversal::{FlowControl, Visitor, Walk},
	utils, SyntaxTree,
};

#[derive(Clone, DebugLisp, PartialEq, Eq, Hash)]
pub struct ImportedSymbol {
	pub qualified_path: Arc<[Token]>,
	pub local_binding: Token,
}

impl SyntaxTree {
	pub fn find_exported_symbol<'a>(&'a self, needle: &Token) -> Option<&'a Decl> {
		self.inner.iter().find(|decl| {
			decl.name().lexeme() == needle.lexeme()
				&& (needle.kind() == decl.name().kind() || !matches!(decl, Decl::ImportPath(_)))
		})
	}

	pub fn import_path_decl(&self) -> Option<&Decl> {
		self.inner
			.iter()
			.find(|decl| matches!(decl, Decl::ImportPath(_)))
	}
}

impl ImportPathDecl {
	/// Returns the fully-qualified path defined by this declaration as an array
	/// of tokens.
	///
	/// I.e., for the following import path declaration:
	/// ```wgsl
	/// #define_import_path foo::bar::baz
	/// ```
	/// This function will return the tokens `[foo, bar, baz]`.
	pub fn qualified_path(&self) -> Arc<[Token]> {
		let mut result = vec![];
		let mut import_path = &self.path;
		loop {
			match import_path {
				ImportPath::Namespaced(NamespacedImportPath { namespace, path }) => {
					result.push(namespace.clone());
					import_path = path.as_ref();
				}
				ImportPath::Leaf(ImportPathLeaf { name, .. }) => {
					result.push(name.clone());
					break result.into();
				}
				ImportPath::Block(_) => unreachable!(),
			}
		}
	}

	/// Returns an ad-hoc `Token::Path` that represents a combination of the
	/// tokens that make up the path defined by this declaration.
	///
	/// I.e., for the following import path declaration:
	/// ```wgsl
	/// #define_import_path foo::bar::baz
	/// ```
	/// This function will return a `Token::Path("foo::bar::baz", (1:21..1:34))`.
	///
	/// The token returned by this function is _created by this function_ by
	/// "joining" the tokens actually contained in the [`ImportPath`] into a
	/// single token. This is in contrast to [`qualified_path`](Self::qualified_path),
	/// which returns the actual sequence of tokens contained in the [`ImportPath`].
	pub fn qualified_name(&self) -> Token {
		let mut result = Substr::new();
		let mut import_path = &self.path;
		loop {
			match import_path {
				ImportPath::Namespaced(NamespacedImportPath { namespace, path }) => {
					if result.is_empty() {
						(result, _) = namespace.as_inner();
					} else {
						let (next, _) = namespace.as_inner();
						result = utils::join_substrs(&result, &next);
					}

					import_path = path.as_ref();
				}
				ImportPath::Leaf(ImportPathLeaf { name, .. }) => {
					if result.is_empty() {
						(result, _) = name.as_inner();
					} else {
						let (next, _) = name.as_inner();
						result = utils::join_substrs(&result, &next);
					}

					break;
				}
				ImportPath::Block(_) => unreachable!(),
			}
		}

		Token::Path(result, self.path.span())
	}
}

impl ImportDecl {
	pub fn build_imported_symbols(&self) -> Vec<ImportedSymbol> {
		let mut builder = ImportedSymbolsBuilder::default();
		self.walk(&mut builder);

		let mut result = builder.build();
		result.sort_unstable_by_key(|sym| sym.local_binding.span());

		result
	}
}

#[derive(Default)]
struct ImportedSymbolsBuilder {
	path: Vec<Token>,
	binding: Option<Token>,
	forks: Vec<ImportedSymbolsBuilder>,
}

impl ImportedSymbolsBuilder {
	fn build(self) -> Vec<ImportedSymbol> {
		std::iter::once(ImportedSymbol {
			qualified_path: self.path.into(),
			local_binding: self.binding.unwrap(),
		})
		.chain(self.forks.into_iter().flat_map(|builder| builder.build()))
		.collect()
	}

	fn fork(&self) -> Self {
		Self {
			path: self.path.clone(),
			binding: None,
			forks: vec![],
		}
	}
}

impl Visitor for ImportedSymbolsBuilder {
	fn visit_import_path(&mut self, path: &ImportPath) -> FlowControl {
		match path {
			ImportPath::Leaf(ImportPathLeaf { name, as_binding }) => {
				self.path.push(name.clone());
				self.binding = Some(as_binding.as_ref().cloned().unwrap_or(name.clone()));

				FlowControl::Break
			}
			ImportPath::Namespaced(NamespacedImportPath { namespace, .. }) => {
				self.path.push(namespace.clone());

				FlowControl::Continue
			}
			ImportPath::Block(ImportPathBlock { paths, .. }) => {
				for (idx, path) in paths.iter().enumerate().rev() {
					if idx == 0 {
						path.walk(self);
					} else {
						let mut fork = self.fork();
						path.walk(&mut fork);
						self.forks.push(fork);
					}
				}

				FlowControl::Break
			}
		}
	}
}