subako-function-sdk-macros 0.2.2

//! `#[function]` and `package!`, re-exported by `subako-function-sdk`.

#![cfg_attr(
    not(test),
    deny(
        clippy::unwrap_used,
        clippy::expect_used,
        clippy::panic,
        clippy::unreachable,
        clippy::todo,
        clippy::unimplemented,
        clippy::indexing_slicing,
    )
)]

use proc_macro2::{Span, TokenStream};
use quote::{format_ident, quote};
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::{Expr, Ident, ItemFn, LitBool, Path, Token, bracketed};

/// Makes a function a tool: its name is the tool's, its doc comment the
/// description the model reads, and its one argument and `Ok` type the input
/// and output schemas. The function itself is left as written.
///
/// ```ignore
/// /// Greet someone by name.
/// #[function]
/// fn greet(input: GreetInput) -> Result<Greeting, ToolError> { ... }
/// ```
#[proc_macro_attribute]
pub fn function(
    attr: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    expand_function(attr.into(), item.into())
        .unwrap_or_else(syn::Error::into_compile_error)
        .into()
}

/// Declares the package: `pub fn package()` building it and, unless
/// `main: false`, the program's `main` serving it.
///
/// ```ignore
/// package! {
///     name: "hello",
///     description: "Says hello.",
///     functions: [greet, tools::farewell],
/// }
/// ```
#[proc_macro]
pub fn package(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    expand_package(input.into())
        .unwrap_or_else(syn::Error::into_compile_error)
        .into()
}

/// The hidden function `#[function]` adds beside `name`, which `package!`
/// calls to get its [`Function`](subako_function_sdk::Function).
fn handle(name: &Ident) -> Ident {
    format_ident!("__subako_function_{}", name)
}

fn expand_function(attr: TokenStream, item: TokenStream) -> syn::Result<TokenStream> {
    if !attr.is_empty() {
        return Err(syn::Error::new_spanned(
            attr,
            "#[function] takes no arguments",
        ));
    }
    let function: ItemFn = syn::parse2(item)?;
    let name = &function.sig.ident;
    let description = description(&function)?;
    let tool = name.to_string();
    let vis = &function.vis;
    let handle = handle(name);
    Ok(quote! {
        #function

        #[doc(hidden)]
        #vis fn #handle() -> ::subako_function_sdk::Function {
            ::subako_function_sdk::Function::new(#tool, #description, #name)
        }
    })
}

/// The doc comment's lines, each without the one space `///` leaves, joined
/// and trimmed.
fn description(function: &ItemFn) -> syn::Result<String> {
    let lines: Vec<String> = function
        .attrs
        .iter()
        .filter(|a| a.path().is_ident("doc"))
        .filter_map(|a| match &a.meta {
            syn::Meta::NameValue(syn::MetaNameValue {
                value:
                    Expr::Lit(syn::ExprLit {
                        lit: syn::Lit::Str(text),
                        ..
                    }),
                ..
            }) => Some(text.value()),
            _ => None,
        })
        .map(|line| line.strip_prefix(' ').unwrap_or(&line).to_string())
        .collect();
    let description = lines.join("\n").trim().to_string();
    match description.is_empty() {
        false => Ok(description),
        true => Err(syn::Error::new_spanned(
            &function.sig.ident,
            "a function's doc comment is the description the model reads; add one",
        )),
    }
}

struct PackageInput {
    name: Expr,
    description: Expr,
    functions: Vec<Path>,
    main: bool,
}

impl Parse for PackageInput {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let (mut name, mut description, mut functions, mut main) = (None, None, None, None);
        while !input.is_empty() {
            let key: Ident = input.parse()?;
            input.parse::<Token![:]>()?;
            let duplicate = || syn::Error::new(key.span(), format!("duplicate key `{key}`"));
            match key.to_string().as_str() {
                "name" => name
                    .replace(input.parse()?)
                    .map_or(Ok(()), |_| Err(duplicate()))?,
                "description" => description
                    .replace(input.parse()?)
                    .map_or(Ok(()), |_| Err(duplicate()))?,
                "functions" => {
                    let list;
                    bracketed!(list in input);
                    let paths = Punctuated::<Path, Token![,]>::parse_terminated(&list)?;
                    functions
                        .replace(paths.into_iter().collect())
                        .map_or(Ok(()), |_| Err(duplicate()))?
                }
                "main" => main
                    .replace(input.parse::<LitBool>()?.value)
                    .map_or(Ok(()), |_| Err(duplicate()))?,
                _ => {
                    return Err(syn::Error::new(
                        key.span(),
                        format!(
                            "unknown key `{key}`; expected name, description, functions, or main"
                        ),
                    ));
                }
            }
            if !input.is_empty() {
                input.parse::<Token![,]>()?;
            }
        }
        let missing =
            |key: &str| syn::Error::new(Span::call_site(), format!("package! needs `{key}`"));
        Ok(PackageInput {
            name: name.ok_or_else(|| missing("name"))?,
            description: description.ok_or_else(|| missing("description"))?,
            functions: functions.ok_or_else(|| missing("functions"))?,
            main: main.unwrap_or(true),
        })
    }
}

fn expand_package(input: TokenStream) -> syn::Result<TokenStream> {
    let PackageInput {
        name,
        description,
        functions,
        main,
    } = syn::parse2(input)?;
    let handles = functions
        .into_iter()
        .map(|mut path| {
            // COVERAGE(never): a parsed `Path` always has at least one segment.
            let last = path
                .segments
                .last_mut()
                .ok_or_else(|| syn::Error::new(Span::call_site(), "an empty function path"))?;
            last.ident = handle(&last.ident);
            Ok(path)
        })
        // COVERAGE(never): the map above never fails, as its marker says.
        .collect::<syn::Result<Vec<_>>>()?;
    let main = main.then(|| {
        quote! {
            fn main() -> ::std::process::ExitCode {
                package().run()
            }
        }
    });
    Ok(quote! {
        /// This package: its functions and the manifest they describe.
        pub fn package() -> ::subako_function_sdk::Package {
            ::subako_function_sdk::Package::new(#name, #description, ::std::vec![#(#handles()),*])
                .with_version(::core::env!("CARGO_PKG_VERSION"))
        }

        #main
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    fn error(result: syn::Result<TokenStream>) -> String {
        result.unwrap_err().to_string()
    }

    #[test]
    fn a_function_keeps_itself_and_gains_its_handle() {
        let out = expand_function(
            quote!(),
            quote! {
                /// Greet someone.
                ///
                /// By name.
                pub fn greet(input: In) -> Result<Out, ToolError> { todo!() }
            },
        )
        .unwrap()
        .to_string();
        assert!(out.contains("pub fn greet"));
        assert!(out.contains("pub fn __subako_function_greet"));
        assert!(out.contains(r#""Greet someone.\n\nBy name.""#));
    }

    #[test]
    fn a_function_needs_a_doc_comment_and_no_arguments() {
        let item = quote!(
            fn greet(input: In) -> Result<Out, ToolError> {
                todo!()
            }
        );
        assert!(error(expand_function(quote!(), item.clone())).contains("doc comment"));
        assert!(error(expand_function(quote!(name = "x"), item)).contains("no arguments"));
    }

    #[test]
    fn a_package_names_its_handles_and_main() {
        let out = expand_package(quote! {
            name: "hello",
            description: "Says hello.",
            functions: [greet, tools::farewell],
        })
        .unwrap()
        .to_string();
        assert!(out.contains("__subako_function_greet ()"));
        assert!(out.contains("tools :: __subako_function_farewell ()"));
        assert!(out.contains("fn main"));
    }

    #[test]
    fn main_false_leaves_main_out() {
        let out = expand_package(quote! {
            name: "hello", description: "d", functions: [], main: false
        })
        .unwrap()
        .to_string();
        assert!(out.contains("fn package"));
        assert!(!out.contains("fn main"));
    }

    #[test]
    fn a_package_refuses_missing_duplicate_and_unknown_keys() {
        for (input, message) in [
            (quote!(name: "a", functions: []), "needs `description`"),
            (
                quote!(name: "a", name: "b", description: "d", functions: []),
                "duplicate key `name`",
            ),
            (
                quote!(name: "a", description: "d", functions: [], extra: 1),
                "unknown key `extra`",
            ),
            (quote!(description: "d", functions: []), "needs `name`"),
            (quote!(name: "a", description: "d"), "needs `functions`"),
            (
                quote!(name: "a", description: "d", description: "e", functions: []),
                "duplicate key `description`",
            ),
            (
                quote!(name: "a", description: "d", functions: [], functions: []),
                "duplicate key `functions`",
            ),
            (
                quote!(name: "a", description: "d", functions: [], main: true, main: false),
                "duplicate key `main`",
            ),
        ] {
            assert!(error(expand_package(input)).contains(message));
        }
    }

    #[test]
    fn a_package_refuses_malformed_syntax() {
        for input in [
            quote!(1: "a"),
            quote!(name "a"),
            quote!(name: , description: "d", functions: []),
            quote!(name: "a", description: , functions: []),
            quote!(name: "a", description: "d", functions: [1]),
            quote!(name: "a", description: "d", functions: [], main: 1),
            quote!(name: "a" description: "d"),
        ] {
            assert!(expand_package(input).is_err());
        }
    }

    #[test]
    fn a_function_refuses_a_non_function_and_skips_non_literal_docs() {
        assert!(
            expand_function(
                quote!(),
                quote!(
                    struct Greet;
                )
            )
            .is_err()
        );
        let out = expand_function(
            quote!(),
            quote! {
                /// Greet.
                #[doc = concat!("ignored")]
                fn greet(input: In) -> Result<Out, ToolError> { todo!() }
            },
        )
        .unwrap()
        .to_string();
        assert!(out.contains(r#""Greet.""#));
    }
}