safety-guard 0.1.8

use syn::{
	fold::Fold,
	parenthesized,
	parse::{Parse, ParseBuffer},
	parse_quote,
	Attribute,
	Error,
	Expr,
	Ident,
	ItemFn,
	LitStr,
	Result,
	Stmt,
	Token,
};

pub(in crate) struct Binary(Expr, Expr);

impl Parse for Binary {
	fn parse(input: &ParseBuffer<'_>) -> Result<Self> {
		let lhs = input.parse()?;
		input.parse::<Token![,]>()?;
		let rhs = input.parse()?;
		Ok(Binary(lhs, rhs))
	}
}

pub(in crate) enum Condition {
	Eq(Binary),
	Ne(Binary),
	Assert(Expr),
}

#[allow(clippy::large_enum_variant)]
pub(in crate) enum Args {
	Literal(LitStr),
	Condition(Condition, LitStr),
}

impl Condition {
	fn stmt(&self, text: &LitStr) -> Stmt {
		match self {
			Condition::Eq(Binary(lhs, rhs)) => parse_quote!(debug_assert_eq!(#lhs, #rhs, #text);),
			Condition::Ne(Binary(lhs, rhs)) => parse_quote!(debug_assert_ne!(#lhs, #rhs, #text);),
			Condition::Assert(expr) => parse_quote!(debug_assert!(#expr, #text);),
		}
	}
}

impl Parse for Condition {
	fn parse(input: &ParseBuffer<'_>) -> Result<Self> {
		let ident: Ident = input.parse()?;
		let parens;
		parenthesized!(parens in input);
		match ident {
			ref i if i == "eq" => Ok(Condition::Eq(parens.parse()?)),
			ref i if i == "ne" => Ok(Condition::Ne(parens.parse()?)),
			ref i if i == "assert" => Ok(Condition::Assert(parens.parse()?)),
			_ => {
				let message = format!(
					"expected string literal, `assert()`, `eq()`, or `ne()`. Found `{}`",
					ident
				);
				Err(Error::new_spanned(ident, message))
			}
		}
	}
}

impl Args {
	fn attribute(&self) -> Attribute {
		let text = match self {
			Args::Literal(text) | Args::Condition(_, text) => text,
		};
		let text = format!("- {}", text.value());
		parse_quote!(#[doc = #text])
	}

	fn stmt(&self) -> Option<Stmt> {
		if let Args::Condition(condition, text) = self {
			Some(condition.stmt(text))
		} else {
			None
		}
	}
}

impl Parse for Args {
	fn parse(input: &ParseBuffer<'_>) -> Result<Self> {
		if input.peek(LitStr) {
			Ok(Args::Literal(input.parse()?))
		} else {
			let condition = input.parse()?;
			input.parse::<Token![,]>()?;
			let literal = input.parse()?;
			Ok(Args::Condition(condition, literal))
		}
	}
}

impl Fold for Args {
	fn fold_item_fn(&mut self, mut item: ItemFn) -> ItemFn {
		let safety_attr = parse_quote!(#[doc = " # Safety"]);
		let mut heading_inserted = None;
		for (i, attr) in item.attrs.iter().enumerate() {
			if *attr == safety_attr {
				heading_inserted = Some(i + 1);
				break;
			}
		}

		if item.unsafety.is_none() {
			let error = Error::new_spanned(
				&item.ident,
				"A function with safety attributes has to be marked as `unsafe`",
			)
			.to_compile_error();
			item.block.stmts.insert(0, parse_quote!(#error;));
		};

		if let Some(i) = heading_inserted {
			item.attrs.insert(i, self.attribute());
		} else {
			item.attrs.push(safety_attr);
			item.attrs.push(self.attribute());
		}

		if let Some(stmt) = self.stmt() {
			item.block.stmts.insert(0, parse_quote!(#stmt));
		};

		item
	}
}