sausage/sdk/macros/src/lib.rs

//! Attribute macros for documented schemas and plugin dispatch.
use proc_macro::TokenStream;
use quote::quote;
use syn::{
    Attribute, Fields, FnArg, ImplItem, Item, ItemImpl, ReturnType, Type, ext::IdentExt,
    parse_macro_input, spanned::Spanned,
};
fn docs(attrs: &[Attribute]) -> String {
    attrs
        .iter()
        .filter_map(|a| {
            if a.path().is_ident("doc")
                && let syn::Meta::NameValue(n) = &a.meta
                && let syn::Expr::Lit(e) = &n.value
                && let syn::Lit::Str(s) = &e.lit
            {
                return Some(s.value().trim().to_string());
            }
            None
        })
        .collect::<Vec<_>>()
        .join("\n")
}
fn documented(attrs: &[Attribute], span: proc_macro2::Span) -> syn::Result<()> {
    if docs(attrs).is_empty() {
        Err(syn::Error::new(span, "a doc comment is required"))
    } else {
        Ok(())
    }
}
fn serde_subset(attrs: &[Attribute]) -> syn::Result<()> {
    for a in attrs {
        if a.path().is_ident("schemars") {
            return Err(syn::Error::new(
                a.span(),
                "schemars attributes are outside the schema subset",
            ));
        }
        if a.path().is_ident("serde") {
            a.parse_nested_meta(|m| {
                if !["rename", "rename_all", "default", "skip_serializing_if"]
                    .iter()
                    .any(|s| m.path.is_ident(s))
                {
                    return Err(m.error("serde attribute is outside the schema subset"));
                }
                if m.input.peek(syn::Token![=]) {
                    let _: syn::Expr = m.value()?.parse()?;
                }
                Ok(())
            })?;
        }
    }
    Ok(())
}
/// Derive serde and Draft 7 schemas for a documented type in the supported subset.
#[proc_macro_attribute]
pub fn schema(_: TokenStream, input: TokenStream) -> TokenStream {
    let item = parse_macro_input!(input as Item);
    let result = (|| -> syn::Result<proc_macro2::TokenStream> {
        let (name, attrs, fields, is_params) = match &item {
            Item::Struct(s) => {
                if !matches!(s.fields, Fields::Named(_)) {
                    return Err(syn::Error::new(
                        s.span(),
                        "schema structs require named fields",
                    ));
                }
                (
                    &s.ident,
                    &s.attrs,
                    s.fields.iter().collect::<Vec<_>>(),
                    true,
                )
            }
            Item::Enum(e) => {
                for v in &e.variants {
                    if !matches!(v.fields, Fields::Unit) {
                        return Err(syn::Error::new(
                            v.span(),
                            "schema enums require unit variants",
                        ));
                    }
                    serde_subset(&v.attrs)?;
                }
                (&e.ident, &e.attrs, vec![], false)
            }
            _ => {
                return Err(syn::Error::new(
                    item.span(),
                    "schema requires a struct or unit enum",
                ));
            }
        };
        documented(attrs, name.span())?;
        serde_subset(attrs)?;
        let mut types = vec![];
        for f in fields {
            documented(&f.attrs, f.span())?;
            serde_subset(&f.attrs)?;
            types.push(&f.ty);
        }
        let params = if is_params {
            quote!(impl ::sausage_plugin::Params for #name {})
        } else {
            quote!()
        };
        Ok(quote! {
            #[derive(::sausage_plugin::serde::Serialize,::sausage_plugin::serde::Deserialize,::sausage_plugin::schemars::JsonSchema)]
            #[serde(crate="::sausage_plugin::serde")]
            #[schemars(crate="::sausage_plugin::schemars")]
            #item
            impl ::sausage_plugin::SchemaType for #name {}
            #params
            const _: () = { fn check() { fn supported<T: ::sausage_plugin::SchemaType>() {} #(supported::<#types>();)* } };
        })
    })();
    result.unwrap_or_else(syn::Error::into_compile_error).into()
}
/// Generate interface dispatch and discovery from documented methods.
#[proc_macro_attribute]
pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream {
    let mut item = parse_macro_input!(input as ItemImpl);
    let args = proc_macro2::TokenStream::from(args);
    let result = (|| -> syn::Result<proc_macro2::TokenStream> {
        let mut grace = None;
        let mut endpoints: Vec<(syn::Ident, syn::Type)> = vec![];
        let parser = syn::meta::parser(|m| {
            if m.path.is_ident("grace_seconds") {
                grace = Some(m.value()?.parse::<syn::LitInt>()?);
                Ok(())
            } else if m.path.is_ident("endpoints") {
                let content;
                syn::parenthesized!(content in m.input);
                while !content.is_empty() {
                    let name = content.parse()?;
                    let _: syn::Token![=] = content.parse()?;
                    let ty = content.parse()?;
                    endpoints.push((name, ty));
                    if !content.is_empty() {
                        let _: syn::Token![,] = content.parse()?;
                    }
                }
                Ok(())
            } else {
                Err(m.error("unknown RPC option"))
            }
        });
        syn::parse::Parser::parse2(parser, args)?;
        let ty = &item.self_ty;
        let mut arms = vec![];
        let mut entries = vec![];
        let mut stop = quote!();
        let mut has_stop = false;
        let mut has_retired = false;
        for member in &mut item.items {
            let ImplItem::Fn(f) = member else { continue };
            let retired = f.attrs.iter().any(|a| {
                a.path()
                    .segments
                    .last()
                    .is_some_and(|s| s.ident == "on_retired")
            });
            let stopping = f.attrs.iter().any(|a| {
                a.path()
                    .segments
                    .last()
                    .is_some_and(|s| s.ident == "on_stop")
            });
            if stopping && std::mem::replace(&mut has_stop, true) {
                return Err(syn::Error::new(
                    f.span(),
                    "only one on_stop method is allowed",
                ));
            }
            if retired && std::mem::replace(&mut has_retired, true) {
                return Err(syn::Error::new(
                    f.span(),
                    "only one on_retired method is allowed",
                ));
            }
            f.attrs.retain(|a| {
                !a.path()
                    .segments
                    .last()
                    .is_some_and(|s| s.ident == "on_retired" || s.ident == "on_stop")
            });
            documented(&f.attrs, f.span())?;
            if f.sig.inputs.is_empty()
                || !matches!(f.sig.inputs.first(),Some(FnArg::Receiver(r)) if r.reference.is_some() && r.mutability.is_none())
            {
                return Err(syn::Error::new(f.sig.span(), "RPC methods require &self"));
            }
            let name = &f.sig.ident;
            let mut args = vec![];
            let mut param = None;
            for arg in f.sig.inputs.iter().skip(1) {
                let FnArg::Typed(p) = arg else { unreachable!() };
                let is_ctx = matches!(p.ty.as_ref(),Type::Reference(r) if matches!(r.elem.as_ref(),Type::Path(t) if t.path.segments.last().is_some_and(|s|s.ident=="Ctx")));
                if is_ctx {
                    if !args.is_empty() {
                        return Err(syn::Error::new(p.span(), "ctx must follow &self"));
                    }
                    args.push(quote!(ctx));
                } else {
                    if param.is_some() {
                        return Err(syn::Error::new(
                            p.span(),
                            "RPC methods take at most one params struct",
                        ));
                    }
                    param = Some(p.ty.as_ref());
                    args.push(quote!(
                        ::sausage_plugin::serde_json::from_value(params)
                            .map_err(|e| ::sausage_plugin::Error::new(-32602, e.to_string()))?
                    ));
                }
            }
            let await_ = f.sig.asyncness.map(|_| quote!(.await));
            if stopping {
                if !args.is_empty() {
                    return Err(syn::Error::new(f.span(), "stop hook takes only &self"));
                }
                stop = quote!(fn stop(&self)->::std::pin::Pin<Box<dyn ::std::future::Future<Output=()> + '_>> {Box::pin(async move {let _=self.#name() #await_;})});
                continue;
            }
            let wire = if retired {
                "rpc.retired".to_string()
            } else {
                name.unraw().to_string()
            };
            let p = param
                .map(|t| quote!(#t))
                .unwrap_or(quote!(::sausage_plugin::Empty));
            let output = match &f.sig.output {
                ReturnType::Type(_, t) => match t.as_ref() {
                    Type::Path(t)
                        if t.path.segments.last().is_some_and(|s| s.ident == "Result") =>
                    {
                        let syn::PathArguments::AngleBracketed(a) =
                            &t.path.segments.last().unwrap().arguments
                        else {
                            return Err(syn::Error::new(t.span(), "expected Result<T>"));
                        };
                        let Some(syn::GenericArgument::Type(t)) = a.args.first() else {
                            return Err(syn::Error::new(a.span(), "expected Result<T>"));
                        };
                        quote!(#t)
                    }
                    _ => return Err(syn::Error::new(t.span(), "RPC methods return Result<T>")),
                },
                _ => return Err(syn::Error::new(f.span(), "RPC methods return Result<T>")),
            };
            let doc = docs(&f.attrs);
            arms.push(quote!(#wire => { fn params_type<T: ::sausage_plugin::Params>() {} params_type::<#p>(); let result=self.#name(#(#args),*) #await_ ?; Ok(::sausage_plugin::serde_json::to_value(result)?) }));
            if !retired {
                entries.push(quote!(::sausage_plugin::method::<#p,#output>(&mut doc,#wire,#doc);));
            }
        }
        let grace = grace
            .map(|n| quote!(doc.value["x-sausage-grace"]=::sausage_plugin::serde_json::json!(#n);));
        let names: Vec<_> = endpoints
            .iter()
            .map(|(n, _)| n.unraw().to_string())
            .collect();
        let types: Vec<_> = endpoints.iter().map(|(_, t)| t).collect();
        Ok(quote! {#item impl ::sausage_plugin::Interface for #ty {
            fn endpoint_types(&self)->Vec<(::std::any::TypeId,String)>{vec![#((::std::any::TypeId::of::<#types>(),#names.into())),*]}
            fn add_methods(mut doc:&mut ::sausage_plugin::Discovery) {#(#entries)*}
            fn methods_document()->::sausage_plugin::serde_json::Value {let mut doc=::sausage_plugin::document(env!("CARGO_PKG_NAME"),env!("CARGO_PKG_VERSION")); Self::add_methods(&mut doc); #grace ::sausage_plugin::finish_document(doc)}
            fn describe()->::sausage_plugin::serde_json::Value {let mut doc=::sausage_plugin::document(env!("CARGO_PKG_NAME"),env!("CARGO_PKG_VERSION")); Self::add_methods(&mut doc); #grace #(doc.endpoint::<#types>(#names);)* ::sausage_plugin::finish_document(doc)}
            fn discover(&self)->::sausage_plugin::serde_json::Value {<Self as ::sausage_plugin::Interface>::describe()}
            fn dispatch<'a>(&'a self, method:&'a str, params: ::sausage_plugin::serde_json::Value,ctx:&'a ::sausage_plugin::Ctx)->::std::pin::Pin<Box<dyn ::std::future::Future<Output=::sausage_plugin::Result<::sausage_plugin::serde_json::Value>>+'a>> {Box::pin(async move {if !params.is_object() {return Err(::sausage_plugin::Error::new(-32602,"params must be an object"));} match method {#(#arms,)* _=>Err(::sausage_plugin::Error::new(-32601,"method not found"))}})}
            #stop
        }})
    })();
    result.unwrap_or_else(syn::Error::into_compile_error).into()
}
/// Generate a private client from `imports/<provider>.json`.
#[proc_macro]
pub fn import(input: TokenStream) -> TokenStream {
    let name = parse_macro_input!(input as syn::LitStr);
    let result = (|| -> Result<proc_macro2::TokenStream, String> {
        let path = std::path::PathBuf::from(
            std::env::var("CARGO_MANIFEST_DIR").map_err(|e| e.to_string())?,
        )
        .join("imports")
        .join(format!("{}.json", name.value()));
        let contents =
            std::fs::read_to_string(&path).map_err(|e| format!("{}: {e}", path.display()))?;
        let doc = serde_json::from_str(&contents).map_err(|e| e.to_string())?;
        let code = sausage_interface::client_rs(&name.value(), &doc)?;
        let generated: proc_macro2::TokenStream = code
            .parse()
            .map_err(|e: proc_macro2::LexError| e.to_string())?;
        fn source_spans(
            tokens: proc_macro2::TokenStream,
            span: proc_macro2::Span,
        ) -> proc_macro2::TokenStream {
            tokens
                .into_iter()
                .map(|mut token| {
                    if let proc_macro2::TokenTree::Group(group) = &token {
                        let mut replacement = proc_macro2::Group::new(
                            group.delimiter(),
                            source_spans(group.stream(), span),
                        );
                        replacement.set_span(span);
                        token = proc_macro2::TokenTree::Group(replacement);
                    }
                    token.set_span(span);
                    token
                })
                .collect()
        }
        let generated = source_spans(generated, name.span());
        let path = path.to_str().ok_or("non-UTF8 import path")?;
        Ok(
            quote! {#generated ::sausage_plugin::inventory::submit! {::sausage_plugin::Import {name:#name,document:include_str!(#path)}}},
        )
    })();
    result
        .unwrap_or_else(|e| syn::Error::new(name.span(), e).into_compile_error())
        .into()
}
/// Mark the method that handles `rpc.retired` after an upgrade.
#[proc_macro_attribute]
pub fn on_retired(_: TokenStream, input: TokenStream) -> TokenStream {
    input
}
/// Mark the hook that runs at stdin EOF without a live call transport.
#[proc_macro_attribute]
pub fn on_stop(_: TokenStream, input: TokenStream) -> TokenStream {
    input
}
/// Run an async test with the SDK runtime and the available core transport.
#[proc_macro_attribute]
pub fn test(_: TokenStream, input: TokenStream) -> TokenStream {
    let mut f = parse_macro_input!(input as syn::ItemFn);
    f.sig.asyncness = None;
    let body = f.block;
    f.block = Box::new(syn::parse_quote!({::sausage_plugin::test_run(async #body)}));
    quote!(#[::core::prelude::v1::test] #f).into()
}