//! Provides a `#[safety]` attribute for generating a corresponding section in the documentation
//! and, if provided, checks for a constraint in debug builds.

extern crate proc_macro;

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

/// Adds a `# Safety` section to the documentation of the function and tests the given constraints
/// in debug builds.
///
/// The attribute has four different forms:
/// - `#[safety("Description")]`: Only the description is added to the documentation
/// - `#[safety(assert(constraint), "Description")]`: `constraint` must evaluate to `true`.
/// - `#[safety(eq(lhs, rhs), "Description")]`: `lhs` and `rhs` needs to be equal
/// - `#[safety(ne(lhs, rhs), "Description")]`: `lhs` and `rhs` must not be equal
///
/// A function with a `#[safety]` attribute must be marked as `unsafe`. Otherwise a compile error
/// is generated.
///
/// If `# Safety` already exists in the documentation, the heading is not added.
///
/// # Examples
///
/// ```
/// use safety_guard::safety;
///
/// #[safety(assert(lhs.checked_add(rhs).is_some()), "`lhs` + `rhs` must not overflow")]
/// unsafe fn add_unchecked(lhs: usize, rhs: usize) -> usize {
/// 	lhs + rhs
/// }
/// ```
///
/// generates
///
/// ```
/// /// # Safety
/// /// - `lhs` + `rhs` must not overflow
/// unsafe fn add_unchecked(lhs: usize, rhs: usize) -> usize {
/// 	debug_assert!(lhs.checked_add(rhs).is_some(), "`lhs` + `rhs` must not overflow");
/// 	lhs + rhs
/// }
/// ```
///
/// Without a constraint, only the documentation is added:
///
/// ```
/// use safety_guard::safety;
///
/// #[safety("`hash` must correspond to the `string`s hash value")]
/// unsafe fn add_string_with_hash(string: &str, hash: u64) -> u64 {
/// 	# unimplemented!()
/// 	// ...
/// }
/// ```
///
/// generates
///
/// ```
/// /// # Safety
/// /// - `hash` must correspond to the `string`s hash value
/// unsafe fn add_string_with_hash(string: &str, hash: u64) -> u64 {
/// 	# unimplemented!()
/// 	// ...
/// }
/// ```
///
/// It is also possible to use multiple `#[safety]` attributes:
///
/// ```
/// # use core::alloc::Layout;
/// use safety_guard::safety;
///
///	#[safety(eq(ptr as usize % layout.align(), 0), "`layout` must *fit* the `ptr`")]
///	#[safety(assert(new_size > 0), "`new_size` must be greater than zero")]
///	#[safety(
///		"`new_size`, when rounded up to the nearest multiple of `layout.align()`, must not \
///		 overflow (i.e., the rounded value must be less than `usize::MAX`)."
///	)]
///	unsafe fn realloc(
///		ptr: *mut u8,
///		layout: Layout,
///		new_size: usize,
///	) -> *mut u8 {
/// 	# unimplemented!()
/// 	// ...
/// }
/// ```
///
/// However, the documentation is generated in reversed order:
///
///
/// ```
/// # use core::alloc::Layout;
/// /// # Safety
/// /// - `new_size`, when rounded up to the nearest multiple of `layout.align()`, must not
/// ///   overflow (i.e., the rounded value must be less than `usize::MAX`).
/// /// - `new_size` must be greater than zero
/// /// - `layout` must *fit* the `ptr`
///	unsafe fn realloc(
///		ptr: *mut u8,
///		layout: Layout,
///		new_size: usize,
///	) -> *mut u8 {
/// 	debug_assert!(new_size > 0, "`new_size` must be greater than zero");
/// 	debug_assert_eq!(ptr as usize % layout.align(), 0, "`layout` must *fit* the `ptr`");
/// 	# unimplemented!()
/// 	// ...
/// }
/// ```
#[proc_macro_attribute]
pub fn safety(args: TokenStream, input: TokenStream) -> TokenStream {
	let input = parse_macro_input!(input);
	let mut args = parse_macro_input!(args as Args);
	let output = args.fold_item_fn(input);
	TokenStream::from(quote!(#output))
}

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))
	}
}

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

#[allow(clippy::large_enum_variant)]
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
	}
}