anodized_core/instrument/traits/

mod.rs

#[cfg(test)]
mod tests;

use quote::quote;
use syn::{
    Attribute, Block, FnArg, ImplItem, ImplItemFn, Pat, TraitItem, TraitItemFn, Visibility,
    parse_quote,
};

use crate::{
    DataSpec, Spec,
    instrument::{Config, find_spec_attr, make_item_error},
};

impl Config {
    /// Expand trait items by mangling each method and adding a wrapper default impl.
    ///
    /// Mangling a function involves the following:
    /// 1. Rename the function following the pattern: `fn add` -> `fn __anodized_add`.
    /// 2. Make a new function with the original name that has a default impl; the
    ///    default impl performs runtime validation and calls the mangled function.
    pub fn instrument_trait(
        &self,
        spec: DataSpec,
        mut the_trait: syn::ItemTrait,
    ) -> syn::Result<syn::ItemTrait> {
        // Currently we don't support any spec arguments for traits themselves.
        if !spec.is_empty() {
            return Err(spec.spec_err(
                "Unsupported spec element on trait. Try placing it on an item inside the trait",
            ));
        }
        let _ = move || spec;

        let mut new_trait_items = Vec::with_capacity(the_trait.items.len() * 5);

        for item in the_trait.items.into_iter() {
            match item {
                TraitItem::Fn(mut func) => {
                    let (spec_attr, other_attrs) = find_spec_attr(func.attrs)?;
                    // NOTE: We have no way of knowing which attributes are
                    //   "external" - meant for the interface and belong on the wrapper,
                    //   "internal" - meant for the mangled implementation.
                    //   Right now we put all attribs on both functions, but that's certainly
                    //   not going to work in every situation.
                    func.attrs = other_attrs.clone();

                    let fn_spec: Spec = match spec_attr {
                        Some(spec_attr) => spec_attr.parse_args()?,
                        None => Spec::empty(),
                    };

                    if self.embed_spec {
                        let attrs: [Attribute; 2] = [
                            parse_quote!(#[doc(hidden)]),
                            parse_quote!(#[allow(warnings)]),
                        ];

                        // Embed `spec` elements as `__anodized_fn_*` items.
                        let spec_trait_qualifiers_const = Self::build_qualifier_const_item(
                            &attrs,
                            "__anodized_fn_qualifiers_trait",
                            fn_spec.qualifiers,
                            &func.sig.ident,
                        );
                        let spec_qualifiers_const = Self::build_qualifier_const_item(
                            &attrs,
                            "__anodized_fn_qualifiers",
                            fn_spec.qualifiers,
                            &func.sig.ident,
                        );
                        let spec_requires_fn = TraitItemFn {
                            attrs: attrs.to_vec(),
                            sig: Self::build_spec_fn_sig("__anodized_fn_requires", &func.sig),
                            default: Some(Self::build_precondition_fn_body(&fn_spec.requires)),
                            semi_token: None,
                        };
                        let spec_maintains_fn = TraitItemFn {
                            attrs: attrs.to_vec(),
                            sig: Self::build_spec_fn_sig("__anodized_fn_maintains", &func.sig),
                            default: Some(Self::build_precondition_fn_body(&fn_spec.maintains)),
                            semi_token: None,
                        };
                        let spec_ensures_fn = TraitItemFn {
                            attrs: attrs.to_vec(),
                            sig: Self::build_spec_fn_sig("__anodized_fn_ensures", &func.sig),
                            default: Some(Self::build_poscondition_fn_body(
                                &fn_spec.captures,
                                &fn_spec.ensures,
                                &func.sig.output,
                            )?),
                            semi_token: None,
                        };

                        new_trait_items.push(TraitItem::Const(spec_trait_qualifiers_const));
                        new_trait_items.push(TraitItem::Const(spec_qualifiers_const));
                        new_trait_items.push(TraitItem::Fn(spec_requires_fn));
                        new_trait_items.push(TraitItem::Fn(spec_maintains_fn));
                        new_trait_items.push(TraitItem::Fn(spec_ensures_fn));
                    }

                    if self.emit_anything() {
                        let mangled_ident = mangle_ident(&func.sig.ident);

                        let mut mangled_fn = func.clone();
                        mangled_fn.sig.ident = mangled_ident.clone();
                        mangled_fn.attrs.retain(|attr| !attr.path().is_ident("doc"));
                        mangled_fn.attrs.push(parse_quote!(#[doc(hidden)]));
                        new_trait_items.push(TraitItem::Fn(mangled_fn));

                        let call_args = build_call_args(&func.sig.inputs)?;
                        let forwarding_body: Block = parse_quote!({
                            Self::#mangled_ident(#(#call_args),*)
                        });
                        func.default = Some(forwarding_body);
                        func.semi_token = None;
                    }

                    func.attrs = other_attrs;

                    if let Some(default_body) = &mut func.default {
                        // NOTE: Needed to handle loop specs in the body of the default impl.
                        self.instrument_fn(&fn_spec, &func.sig, default_body)?;
                    }
                    new_trait_items.push(TraitItem::Fn(func));
                }
                TraitItem::Const(mut const_item) => {
                    let (spec, attrs) = find_spec_attr(const_item.attrs)?;
                    if let Some(ref spec_attr) = spec {
                        return Err(make_item_error(&spec_attr, "trait const"));
                    }
                    const_item.attrs = attrs;
                    new_trait_items.push(TraitItem::Const(const_item));
                }
                TraitItem::Type(mut type_item) => {
                    let (spec, attrs) = find_spec_attr(type_item.attrs)?;
                    if let Some(ref spec_attr) = spec {
                        return Err(make_item_error(&spec_attr, "trait type"));
                    }
                    type_item.attrs = attrs;
                    new_trait_items.push(TraitItem::Type(type_item));
                }
                TraitItem::Macro(mut macro_item) => {
                    let (spec, attrs) = find_spec_attr(macro_item.attrs)?;
                    if let Some(ref spec_attr) = spec {
                        return Err(make_item_error(&spec_attr, "trait macro"));
                    }
                    macro_item.attrs = attrs;
                    new_trait_items.push(TraitItem::Macro(macro_item));
                }
                TraitItem::Verbatim(token_stream) => {
                    new_trait_items.push(TraitItem::Verbatim(token_stream));
                }
                _ => unimplemented!(),
            }
        }
        the_trait.items = new_trait_items;
        Ok(the_trait)
    }

    /// Expand impl items by mangling methods for trait impls.
    ///
    /// The `#[spec]` attribute on an impl `fn` must narrow the `#[spec]` of the trait `fn`:
    /// - The impl's preconditions must follow from the trait's preconditions.
    /// - The impl's postconditions must entail the trait's postconditions.
    pub fn instrument_trait_impl(
        &self,
        spec: DataSpec,
        mut the_impl: syn::ItemImpl,
    ) -> syn::Result<syn::ItemImpl> {
        let Some((trait_bang, ref trait_path, _trait_for)) = the_impl.trait_ else {
            return Err(make_item_error(&the_impl, "inherent impl"));
        };

        if trait_bang.is_some() {
            return Err(make_item_error(&the_impl, "negative trait impl"));
        }

        if !spec.is_empty() {
            return Err(spec.spec_err("Unsupported spec element on trait impl."));
        }

        let mut new_items = Vec::with_capacity(the_impl.items.len() * 4);

        for item in the_impl.items.into_iter() {
            match item {
                ImplItem::Fn(mut func) => {
                    let (spec_attr, func_attrs) = find_spec_attr(func.attrs)?;
                    func.attrs = func_attrs;

                    if func.sig.ident.to_string().starts_with("__anodized_") {
                        return Err(syn::Error::new_spanned(
                            func.sig.ident,
                            r#"An item with the `__anodized_` prefix is internal. Do not implement it directly.
Instead, ensure that both the trait and the impl fn have a `#[spec]` annotation."#,
                        ));
                    }

                    let fn_spec: Spec = match spec_attr {
                        Some(spec_attr) => spec_attr.parse_args()?,
                        None => Spec::empty(),
                    };

                    if self.embed_spec {
                        let attrs: [Attribute; 2] = [
                            parse_quote!(#[doc(hidden)]),
                            parse_quote!(#[allow(warnings)]),
                        ];

                        // Embed `spec` elements as `__anodized_fn_*` items.
                        let spec_qualifiers_const = Self::build_qualifier_const_item(
                            &attrs,
                            "__anodized_fn_qualifiers",
                            fn_spec.qualifiers,
                            &func.sig.ident,
                        );
                        let spec_requires_fn = ImplItemFn {
                            attrs: attrs.to_vec(),
                            sig: Self::build_spec_fn_sig("__anodized_fn_requires", &func.sig),
                            block: Self::build_precondition_fn_body(&fn_spec.requires),
                            vis: Visibility::Inherited,
                            defaultness: None,
                        };
                        let spec_maintains_fn = ImplItemFn {
                            attrs: attrs.to_vec(),
                            sig: Self::build_spec_fn_sig("__anodized_fn_maintains", &func.sig),
                            block: Self::build_precondition_fn_body(&fn_spec.maintains),
                            vis: Visibility::Inherited,
                            defaultness: None,
                        };
                        let spec_ensures_fn = ImplItemFn {
                            attrs: attrs.to_vec(),
                            sig: Self::build_spec_fn_sig("__anodized_fn_ensures", &func.sig),
                            block: Self::build_poscondition_fn_body(
                                &fn_spec.captures,
                                &fn_spec.ensures,
                                &func.sig.output,
                            )?,
                            vis: Visibility::Inherited,
                            defaultness: None,
                        };

                        new_items.push(ImplItem::Const(spec_qualifiers_const));
                        new_items.push(ImplItem::Fn(spec_requires_fn));
                        new_items.push(ImplItem::Fn(spec_maintains_fn));
                        new_items.push(ImplItem::Fn(spec_ensures_fn));
                    }

                    self.instrument_fn(&fn_spec, &func.sig, &mut func.block)?;

                    if self.embed_spec {
                        // Add a compile-time check to the body.
                        func.block.stmts.insert(
                            0,
                            Self::build_qualifier_check_stmt(
                                &func.sig.ident,
                                &the_impl.self_ty,
                                trait_path,
                            ),
                        );
                    }

                    if self.emit_anything() {
                        func.sig.ident = mangle_ident(&func.sig.ident);

                        // Add a default `#[inline]` attribute unless one is already there.
                        // The caller can supress this with `#[inline(never)]`
                        if !has_inline_attr(&func.attrs) {
                            func.attrs.push(parse_quote!(#[inline]));
                        }
                    }
                    new_items.push(ImplItem::Fn(func));
                }
                ImplItem::Const(mut const_item) => {
                    let (spec, attrs) = find_spec_attr(const_item.attrs)?;
                    if let Some(ref spec_attr) = spec {
                        return Err(make_item_error(&spec_attr, "trait impl const"));
                    }
                    const_item.attrs = attrs;
                    new_items.push(ImplItem::Const(const_item));
                }
                ImplItem::Type(mut type_item) => {
                    let (spec, attrs) = find_spec_attr(type_item.attrs)?;
                    if let Some(ref spec_attr) = spec {
                        return Err(make_item_error(&spec_attr, "trait impl type"));
                    }
                    type_item.attrs = attrs;
                    new_items.push(ImplItem::Type(type_item));
                }
                ImplItem::Macro(mut macro_item) => {
                    let (spec, attrs) = find_spec_attr(macro_item.attrs)?;
                    if let Some(ref spec_attr) = spec {
                        return Err(make_item_error(&spec_attr, "trait impl macro"));
                    }
                    macro_item.attrs = attrs;
                    new_items.push(ImplItem::Macro(macro_item));
                }
                ImplItem::Verbatim(token_stream) => {
                    new_items.push(ImplItem::Verbatim(token_stream))
                }
                _ => unimplemented!(),
            };
        }

        the_impl.items = new_items;
        Ok(the_impl)
    }
}

/// Build argument tokens for calling the mangled trait method from the wrapper.
///
/// Purpose: the wrapper method needs to forward its arguments to the mangled
/// implementation, so this extracts a usable token for each input.
///
/// Examples (inputs -> output tokens):
/// - `fn f(&self, x: i32)` -> `self, x`
/// - `fn f(self, a: u8, b: u8)` -> `self, a, b`
///
/// The caller is responsible for ensuring these tokens are used in a call
/// expression like `Self::__anodized_f(#(#args),*)`.
///
/// Callers: only `instrument_trait` in this module should use this; it is not
/// part of the public API.
fn build_call_args(
    inputs: &syn::punctuated::Punctuated<FnArg, syn::Token![,]>,
) -> syn::Result<Vec<proc_macro2::TokenStream>> {
    let mut args = Vec::new();
    for input in inputs.iter() {
        match input {
            FnArg::Receiver(_) => {
                args.push(quote! { self });
            }
            FnArg::Typed(pat) => match pat.pat.as_ref() {
                Pat::Ident(pat_ident) => {
                    let ident = &pat_ident.ident;
                    args.push(quote! { #ident });
                }
                _ => {
                    return Err(syn::Error::new_spanned(
                        &pat.pat,
                        "unsupported pattern in trait method arguments",
                    ));
                }
            },
        }
    }
    Ok(args)
}

/// Prefix an identifier with `__anodized_`, preserving the original span.
/// Used when generating mangled method names in trait and impl expansion.
fn mangle_ident(original_ident: &syn::Ident) -> syn::Ident {
    syn::Ident::new(
        &format!("__anodized_{original_ident}"),
        original_ident.span(),
    )
}

/// Checks to see if any `#[inline]` (with or without arg) exists in the function's attribs.
fn has_inline_attr(attrs: &[syn::Attribute]) -> bool {
    attrs.iter().any(|attr| attr.path().is_ident("inline"))
}