verty 0.1.1

use proc_macro2::{Span, TokenStream};
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::visit_mut::VisitMut;
use syn::{Expr, ImplItem, Result, Stmt, Type};

use super::{Args, versioned_ident};
use crate::parse::{PunctTerminated, SpannedInterval, StagedParse};
use crate::util::error_sink::ErrorSink;

// --- preliminaries ---

trait PseudoMacro<T> {
	const IDENT: &'static str;

	fn expand(cx: ExpandPseudoMacros, span: Span, input: TokenStream) -> Result<T>;
}

trait PseudoMacroList<T> {
	/// `None` here *only* means the given macro wasn't one of the specified pseudo-macros.
	fn expand(cx: &mut ExpandPseudoMacros, m: &mut syn::Macro) -> Option<T>;
}
impl<T> PseudoMacroList<T> for () {
	#[inline]
	fn expand(_cx: &mut ExpandPseudoMacros, _m: &mut syn::Macro) -> Option<T> {
		None
	}
}
impl<T: Placeholder, P1: PseudoMacro<T>, P2: PseudoMacroList<T>> PseudoMacroList<T> for (P1, P2) {
	#[inline]
	fn expand(cx: &mut ExpandPseudoMacros, m: &mut syn::Macro) -> Option<T> {
		if m.path.is_ident(P1::IDENT) {
			let res = P1::expand(
				cx.reborrow(),
				syn::spanned::Spanned::span(m),
				std::mem::take(&mut m.tokens),
			);

			Some(cx.errs.eat_err(res).unwrap_or_else(T::placeholder))
		} else {
			P2::expand(cx, m)
		}
	}
}
macro_rules! list {
	() => { () };
	($T:ty $(, $($R:tt)*)?) => {
		($T, list!($($($R)*)?))
	}
}

trait Placeholder {
	fn placeholder() -> Self;
}

impl Placeholder for Expr {
	fn placeholder() -> Self {
		syn::parse_quote! { () }
	}
}

impl Placeholder for Stmt {
	fn placeholder() -> Self {
		syn::parse_quote! { ; }
	}
}

impl Placeholder for Type {
	fn placeholder() -> Self {
		syn::parse_quote! { () }
	}
}

impl Placeholder for ImplItem {
	fn placeholder() -> Self {
		syn::parse_quote! { const _: () = (); }
	}
}

// --- main impl ---

/// Match to Option
macro_rules! mopt {
	(let $p:pat = $x:expr => $r:expr) => {
		if let $p = $x { Some($r) } else { None }
	};
}

pub struct ExpandPseudoMacros<'a, 'e> {
	pub version: usize,
	pub validate_args: Option<&'a Args>,
	pub errs: &'e mut ErrorSink,
}

impl<'a, 'e> ExpandPseudoMacros<'a, 'e> {
	pub fn reborrow(&mut self) -> ExpandPseudoMacros<'_, '_> {
		ExpandPseudoMacros {
			version: self.version,
			validate_args: self.validate_args,
			errs: self.errs,
		}
	}
}

type ExprPseudoMacros = list![VerMatch];
type StmtPseudoMacros = list![VerMatch];
type TypePseudoMacros = list![VerMatch, VerType];
type ImplItemPseudoMacros = list![VerMatch];

impl<'a, 'e> VisitMut for ExpandPseudoMacros<'a, 'e> {
	fn visit_expr_mut(&mut self, e: &mut Expr) {
		let new_e = mopt!(let Expr::Macro(m) = e => &mut m.mac)
			.and_then(|m| ExprPseudoMacros::expand(self, m));

		if let Some(new_e) = new_e {
			*e = new_e;
		} else {
			syn::visit_mut::visit_expr_mut(self, e)
		}
	}

	// NOTE: a macro could be in visit_item_mut, but `versioned` is never invoked with anything that could contain that

	fn visit_stmt_mut(&mut self, s: &mut Stmt) {
		let new_s = mopt!(let Stmt::Macro(m) = s => &mut m.mac)
			.and_then(|m| StmtPseudoMacros::expand(self, m));

		if let Some(new_s) = new_s {
			*s = new_s;
		} else {
			syn::visit_mut::visit_stmt_mut(self, s)
		}
	}

	fn visit_type_mut(&mut self, t: &mut Type) {
		let new_t = mopt!(let Type::Macro(m) = t => &mut m.mac)
			.and_then(|m| TypePseudoMacros::expand(self, m));

		if let Some(new_t) = new_t {
			*t = new_t;
		} else {
			syn::visit_mut::visit_type_mut(self, t)
		}
	}

	fn visit_impl_item_mut(&mut self, i: &mut ImplItem) {
		let new_i = mopt!(let ImplItem::Macro(m) = i => &mut m.mac)
			.and_then(|m| ImplItemPseudoMacros::expand(self, m));

		if let Some(new_i) = new_i {
			*i = new_i;
		} else {
			syn::visit_mut::visit_impl_item_mut(self, i)
		}
	}

	// NOTE: a macro could be in visit_trait_mut, but `versioned` is never invoked with anything that could contain that

	fn visit_parenthesized_generic_arguments_mut(
		&mut self,
		args: &mut syn::ParenthesizedGenericArguments,
	) {
		ver_gen_visit_args_mut(self, &mut args.inputs, |cx, t| match t {
			Type::Macro(m) if m.mac.path.is_ident(VER_GEN_IDENT) => {
				Some(std::mem::take(&mut m.mac.tokens))
			}
			_ => {
				cx.visit_type_mut(t);
				None
			}
		});
		self.visit_return_type_mut(&mut args.output);
	}

	fn visit_angle_bracketed_generic_arguments_mut(
		&mut self,
		args: &mut syn::AngleBracketedGenericArguments,
	) {
		ver_gen_visit_args_mut(self, &mut args.args, |cx, t| match t {
			// NOTE: exprs can be macros too, but those will take the form `{ name!(...) }`, which I'll explicitly disallow,
			// since `ver_param` is explicitly NOT self-contained and may expand to nothing, making the arg list shorter.
			syn::GenericArgument::Type(Type::Macro(m)) if m.mac.path.is_ident(VER_GEN_IDENT) => {
				Some(std::mem::take(&mut m.mac.tokens))
			}
			_ => {
				cx.visit_generic_argument_mut(t);
				None
			}
		});
	}
}

// --- pseudo-macros ---

struct VerMatch;

impl<T: Parse> PseudoMacro<T> for VerMatch {
	const IDENT: &'static str = "ver_match";

	fn expand(cx: ExpandPseudoMacros, span: Span, input: TokenStream) -> Result<T> {
		let out_tokens = cx
			.errs
			.eat_err(crate::parse::staged_parse2::<
				PunctTerminated<VerMatchArm, syn::Token![;]>,
			>(input, cx.validate_args.map(|a| (a.start, a.end))))
			.and_then(|arms| {
				let mut res = None;

				for VerMatchArm(SpannedInterval(span, iv), out) in arms.0 {
					if iv.contains(cx.version) {
						match res {
							None => res = Some(out),
							Some(_) => cx.errs.push(syn::Error::new(
								span,
								format!("more than one option for version {}", cx.version),
							)),
						}
					}
				}

				if res.is_none() {
					cx.errs.push(syn::Error::new(
						span,
						format!("version {} not covered", cx.version),
					));
				}

				res
			})
			.unwrap_or_default();

		syn::parse2(out_tokens)
	}
}

struct VerType;

impl PseudoMacro<Type> for VerType {
	const IDENT: &'static str = "VerType";

	#[inline]
	fn expand(mut cx: ExpandPseudoMacros, _span: Span, input: TokenStream) -> Result<Type> {
		use syn::spanned::Spanned;

		let mut path = syn::parse2::<syn::Path>(input)?;

		cx.visit_path_mut(&mut path);

		let name = &mut path.segments.last_mut().unwrap().ident;
		*name = versioned_ident(name, cx.version);

		Ok(syn::parse_quote_spanned! { path.span() => #path })
	}
}

const VER_GEN_IDENT: &str = "ver_gen";
// This one is special since it has to remove itself from its parent if necessary,
// so it can't just operate on its input itself; It needs to be handled externally.
fn ver_gen_visit_args_mut<T: Parse>(
	cx: &mut ExpandPseudoMacros,
	args: &mut Punctuated<T, syn::Token![,]>,
	// Descend into `T` and return `None`, or,
	// if the `T` is an instance of the `ver_param!` pseudo-macro,
	// return its input inside `Some`.
	mut subvisit_extract_ver_param: impl FnMut(&mut ExpandPseudoMacros, &mut T) -> Option<TokenStream>,
) {
	// since `Punctuated` doesn't have neat in-place modification,
	// this just builds a new one.
	let mut new_args = Punctuated::new();
	for mut t in std::mem::take(args).into_iter() {
		if let Some(input) = subvisit_extract_ver_param(cx, &mut t) {
			if let Some(VerParamInput(SpannedInterval(_span, iv), new_t)) =
				cx.errs
					.eat_err(crate::parse::staged_parse2::<VerParamInput<T>>(
						input,
						cx.validate_args.map(|a| (a.start, a.end)),
					)) {
				if iv.contains(cx.version) {
					new_args.push(new_t);
				}
			}
			// all other cases mean the arg is removed
		} else {
			// it was subvisited; nothing else needs to be done, so just append it
			new_args.push(t);
		}
	}
	*args = new_args;
}

// --- parsing types ---

#[cfg_attr(test, derive(Debug))]
struct VerMatchArm<V = SpannedInterval>(V, TokenStream);

impl Parse for VerMatchArm<Expr> {
	fn parse(input: ParseStream) -> Result<Self> {
		use quote::TokenStreamExt;
		// since `x.. => ...` would get parsed as `x ..= >`, producing an error,
		// we get the arrow FIRST with a bit more effort before parsing what comes before it.
		let mut before = TokenStream::new();
		while !input.is_empty() {
			if input.parse::<syn::Token![=>]>().is_ok() {
				break;
			} else {
				before.append(input.parse::<proc_macro2::TokenTree>().unwrap());
			}
		}
		let inner;
		Ok(Self(syn::parse2(before)?, {
			syn::braced!(inner in input);
			inner
				.parse::<TokenStream>()
				.expect("parse TokenStream failed?!")
		}))
	}
}
impl StagedParse for VerMatchArm {
	type Stage1 = VerMatchArm<Expr>;
	type Cx = <SpannedInterval as StagedParse>::Cx;

	fn stage2(from: Self::Stage1, cx: Self::Cx) -> Result<Self> {
		Ok(Self(SpannedInterval::stage2(from.0, cx)?, from.1))
	}
}

#[cfg_attr(test, derive(Debug))]
struct VerParamInput<T, V = SpannedInterval>(V, T);

impl<T: Parse> Parse for VerParamInput<T, Expr> {
	fn parse(input: ParseStream) -> Result<Self> {
		Ok(Self(input.parse()?, {
			let _: syn::Token![,] = input.parse()?;
			input.parse()?
		}))
	}
}
impl<T: Parse> StagedParse for VerParamInput<T> {
	type Stage1 = VerParamInput<T, Expr>;
	type Cx = <SpannedInterval as StagedParse>::Cx;

	fn stage2(from: Self::Stage1, cx: Self::Cx) -> Result<Self> {
		Ok(Self(SpannedInterval::stage2(from.0, cx)?, from.1))
	}
}