wgsl-parser 0.5.0

use std::marker::PhantomData;

use gramatika::{Span, Spanned, Token};

use crate::{
	decl::Decl,
	expr::Expr,
	stmt::Stmt,
	traversal::{FlowControl, Visitor, Walk},
	SyntaxTree,
};

use super::{GetSyntaxKind, SyntaxNode};

pub fn find_parent<T>(tree: &SyntaxTree, target: Span) -> Option<T>
where T: GetSyntaxKind + Clone + From<SyntaxNode> {
	let mut finder = ParentFinder::<T>::new(target);
	tree.walk(&mut finder);

	finder.parents.pop().map(T::from)
}

pub fn find_parents<T>(tree: &SyntaxTree, target: Span) -> Vec<T>
where T: GetSyntaxKind + Clone + From<SyntaxNode> {
	let mut finder = ParentFinder::<T>::new(target);
	tree.walk(&mut finder);

	finder.parents.reverse();
	finder.parents.into_iter().map(T::from).collect()
}

struct ParentFinder<T>
where T: GetSyntaxKind + Clone
{
	target: Span,
	parents: Vec<SyntaxNode>,
	_marker: PhantomData<T>,
}

impl<T> ParentFinder<T>
where T: GetSyntaxKind + Clone
{
	fn new(target: Span) -> Self {
		Self {
			target,
			parents: vec![],
			_marker: Default::default(),
		}
	}
}

impl<T> Visitor for ParentFinder<T>
where T: GetSyntaxKind + Clone
{
	fn visit_decl(&mut self, decl: &Decl) -> FlowControl {
		let span = decl.span();

		if span != self.target && span.contains(self.target) {
			if T::KIND.contains(Decl::KIND) {
				self.parents.push(SyntaxNode::Decl(decl.clone()));
			}

			FlowControl::Continue
		} else {
			FlowControl::Break
		}
	}

	fn visit_stmt(&mut self, stmt: &Stmt) -> FlowControl {
		let span = stmt.span();

		if span != self.target && span.contains(self.target) {
			// FIXME: This is pretty awkward. Consider updating `Stmt::Var(...)`
			//        so that it wraps a `Decl` instead of a `VarDecl`.
			if T::KIND == Decl::KIND && matches!(stmt, Stmt::Var(_)) {
				let Stmt::Var(var_decl) = stmt else {
					unreachable!();
				};
				match var_decl.storage.lexeme().as_str() {
					"let" | "const" => {
						self.parents
							.push(SyntaxNode::Decl(Decl::Const(var_decl.clone())));
					}
					"var" => {
						self.parents
							.push(SyntaxNode::Decl(Decl::Var(var_decl.clone())));
					}
					_ => unreachable!(),
				}
			} else if T::KIND.contains(Stmt::KIND) {
				self.parents.push(SyntaxNode::Stmt(stmt.clone()));
			}

			FlowControl::Continue
		} else {
			FlowControl::Break
		}
	}

	fn visit_expr(&mut self, expr: &Expr) -> FlowControl {
		if !T::KIND.contains(Expr::KIND) {
			return FlowControl::Break;
		}

		let span = expr.span();

		if span != self.target && span.contains(self.target) {
			self.parents.push(SyntaxNode::Expr(expr.clone()));

			FlowControl::Continue
		} else {
			FlowControl::Break
		}
	}
}