turso_macros 0.7.2

The Limbo database library
Documentation
use proc_macro::TokenStream;
use quote::quote;
use syn::{
    parse::{Parse, ParseStream},
    Expr, Ident, ImplItemFn, ItemFn, LitStr, Token,
};

enum TraceStackArgs {
    Inferred,
    InferredWithDetail(Expr),
    Label(LitStr),
    LabelWithDetail(LitStr, Expr),
}

impl Parse for TraceStackArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        if input.is_empty() {
            return Ok(Self::Inferred);
        }

        if input.peek(Ident) {
            let ident: Ident = input.parse()?;
            if ident != "detail" {
                return Err(syn::Error::new_spanned(ident, "expected `detail`"));
            }
            input.parse::<Token![=]>()?;
            let detail = input.parse()?;
            if !input.is_empty() {
                input.parse::<Token![,]>()?;
            }
            if !input.is_empty() {
                return Err(input.error("expected a single detail expression"));
            }
            return Ok(Self::InferredWithDetail(detail));
        }

        let label = input.parse()?;
        if input.is_empty() {
            return Ok(Self::Label(label));
        }

        input.parse::<Token![,]>()?;
        if !input.is_empty() {
            let detail = input.parse()?;
            if input.is_empty() {
                return Ok(Self::LabelWithDetail(label, detail));
            }
            input.parse::<Token![,]>()?;
        }
        if !input.is_empty() {
            return Err(
                input.error("expected a string literal label and optional detail expression")
            );
        }

        Ok(Self::Label(label))
    }
}

pub(crate) fn trace_stack_attribute(attr: TokenStream, input: TokenStream) -> TokenStream {
    let args = match syn::parse::<TraceStackArgs>(attr) {
        Ok(args) => args,
        Err(err) => return err.to_compile_error().into(),
    };
    let input: proc_macro2::TokenStream = input.into();

    if let Ok(mut function) = syn::parse2::<ItemFn>(input.clone()) {
        let guard = trace_guard(&args, &function.sig.ident);
        let body = &function.block;
        function.block = syn::parse_quote!({
            #guard
            #body
        });
        return quote!(#function).into();
    }

    if let Ok(mut function) = syn::parse2::<ImplItemFn>(input.clone()) {
        let guard = trace_guard(&args, &function.sig.ident);
        let body = &function.block;
        function.block = syn::parse_quote!({
            #guard
            #body
        });
        return quote!(#function).into();
    }

    syn::Error::new_spanned(input, "trace_stack can only be applied to functions")
        .to_compile_error()
        .into()
}

fn trace_guard(args: &TraceStackArgs, function_name: &syn::Ident) -> proc_macro2::TokenStream {
    let inferred_span_name = LitStr::new(&function_name.to_string(), function_name.span());

    let (span_name, guard) = match args {
        TraceStackArgs::Inferred => {
            let guard = quote! {
                crate::stack::trace_stack!(stringify!(#function_name));
            };
            (inferred_span_name, guard)
        }
        TraceStackArgs::InferredWithDetail(detail) => {
            let guard = quote! {
                crate::stack::trace_stack!(stringify!(#function_name), #detail);
            };
            (inferred_span_name, guard)
        }
        TraceStackArgs::Label(label) => {
            let guard = quote! {
                crate::stack::trace_stack!(#label);
            };
            (label.clone(), guard)
        }
        TraceStackArgs::LabelWithDetail(label, detail) => {
            let guard = quote! {
                crate::stack::trace_stack!(#label, #detail);
            };
            (label.clone(), guard)
        }
    };

    quote! {
        #[cfg(feature = "stacker")]
        let _stack_span = tracing::debug_span!(target: "turso_stack", #span_name, module_path = module_path!()).entered();
        #guard
    }
}