some_executor_macros 0.7.2

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Attribute macros for [`some_executor`](https://docs.rs/some_executor).
//!
//! Nothing here is meant to be used directly; `some_executor` re-exports
//! [`main`], and `#[some_executor::main]` is the spelling to write. A
//! `proc-macro = true` crate can export nothing but macros, which is the only
//! reason this is a separate package.

use proc_macro::{Delimiter, TokenStream, TokenTree};

/// Turns an `async fn main` into a `fn main` that installs an executor and runs
/// it.
///
/// ```ignore
/// # // ignore because: this shows the macro's input, and naming a backend here
/// # // would make this proc-macro crate depend on an executor implementation.
/// #[some_executor::main(some_executor_tokio::TokioExecutor)]
/// async fn main() {
///     // ...
/// }
/// ```
///
/// expands to
///
/// ```ignore
/// # // ignore because: this is the macro's *output*, shown for reading. It names
/// # // types this crate cannot depend on, and is not meant to compile here.
/// fn main() {
///     <some_executor_tokio::TokioExecutor as ::some_executor::ExecutorMain>::main(
///         async move {
///             ::some_executor::entry_point::MainResult::report(async move { /* body */ }.await)
///         },
///     )
/// }
/// ```
///
/// The call goes through the [`ExecutorMain`] trait because this macro cannot
/// know an inherent method on a backend type it has never heard of.
///
/// # The backend argument is optional
///
/// With no argument the backend is
/// [`some_executor::entry_point::LastResort`], which is always available and
/// needs no dependency. It is meant for quick demos and examples: it warns on
/// first use, because it is not a production executor. Real programs name their
/// backend.
///
/// # A fallible `main`
///
/// `async fn main() -> Result<T, E>` works, as long as `E: Debug`. The error is
/// reported and the process is brought down, by
/// [`some_executor::entry_point::MainResult`] rather than by generated code —
/// so the behaviour is documented and tested in one place instead of being
/// whatever this macro happened to emit.
///
/// [`ExecutorMain`]: https://docs.rs/some_executor/latest/some_executor/entry_point/trait.ExecutorMain.html
/// [`some_executor::entry_point::LastResort`]: https://docs.rs/some_executor/latest/some_executor/entry_point/struct.LastResort.html
/// [`some_executor::entry_point::MainResult`]: https://docs.rs/some_executor/latest/some_executor/entry_point/trait.MainResult.html
#[proc_macro_attribute]
pub fn main(attribute: TokenStream, item: TokenStream) -> TokenStream {
    match expand(attribute, item) {
        Ok(stream) => stream,
        Err(message) => error(&message),
    }
}

/// A `compile_error!` the caller sees instead of a syntax error inside an
/// expansion they never wrote.
fn error(message: &str) -> TokenStream {
    // The trailing semicolon matters: this is spliced where an item was, and
    // without it the macro invocation runs into whatever follows.
    format!("compile_error!({:?});", message)
        .parse()
        .expect("a compile_error! invocation is always valid Rust")
}

fn expand(attribute: TokenStream, item: TokenStream) -> Result<TokenStream, String> {
    let backend = if attribute.is_empty() {
        "::some_executor::entry_point::LastResort".to_string()
    } else {
        attribute.to_string()
    };

    let function = parse_async_fn(item)?;

    // `async move` twice, not once: the outer block is what the backend is
    // handed, and the inner one is the user's body, so its output can be
    // awaited and passed to `report` without this macro having to look at the
    // declared return type at all. Which type `report` resolves to is the
    // compiler's problem, and a `main` returning something unsupported gets a
    // trait error naming `MainResult` rather than a mystery.
    let expanded = format!(
        "{attrs} {vis} fn {name}() {{
    <{backend} as ::some_executor::ExecutorMain>::main(async move {{
        let __some_executor_main_result = async move {body}.await;
        ::some_executor::entry_point::MainResult::report(__some_executor_main_result);
    }})
}}",
        attrs = function.attributes,
        vis = function.visibility,
        name = function.name,
        backend = backend,
        body = function.body,
    );

    expanded
        .parse()
        .map_err(|e| format!("could not build the expansion: {e}"))
}

struct AsyncFn {
    /// Attributes that were below this one, re-emitted verbatim.
    attributes: String,
    visibility: String,
    name: String,
    /// The braced body, braces included, re-emitted verbatim.
    body: String,
}

/// Recognises `#[..]* $vis async fn NAME() [-> T] { .. }`.
///
/// Deliberately shallow. The body and any attributes are opaque token trees
/// that are passed straight through, and the return type is discarded because
/// the expansion never names it. What has to be checked is what produces a bad
/// error message if it is wrong: that the function is `async`, that it takes no
/// arguments, and that it is not generic.
fn parse_async_fn(item: TokenStream) -> Result<AsyncFn, String> {
    let tokens: Vec<TokenTree> = item.into_iter().collect();
    let mut index = 0;

    let mut attributes = String::new();
    while index < tokens.len() {
        match (&tokens[index], tokens.get(index + 1)) {
            (TokenTree::Punct(punct), Some(TokenTree::Group(group)))
                if punct.as_char() == '#' && group.delimiter() == Delimiter::Bracket =>
            {
                attributes.push_str(&tokens[index].to_string());
                attributes.push_str(&group.to_string());
                attributes.push(' ');
                index += 2;
            }
            _ => break,
        }
    }

    let mut visibility = String::new();
    if let Some(TokenTree::Ident(ident)) = tokens.get(index)
        && ident.to_string() == "pub"
    {
        visibility.push_str("pub");
        index += 1;
        if let Some(TokenTree::Group(group)) = tokens.get(index)
            && group.delimiter() == Delimiter::Parenthesis
        {
            visibility.push_str(&group.to_string());
            index += 1;
        }
        visibility.push(' ');
    }

    let keyword = |index: usize| match tokens.get(index) {
        Some(TokenTree::Ident(ident)) => ident.to_string(),
        _ => String::new(),
    };

    if keyword(index) != "async" {
        return Err(
            "#[some_executor::main] expects an `async fn`; a synchronous one has nothing \
             to run on an executor"
                .to_string(),
        );
    }
    index += 1;

    if keyword(index) != "fn" {
        return Err("#[some_executor::main] expects an `async fn`".to_string());
    }
    index += 1;

    let name = match tokens.get(index) {
        Some(TokenTree::Ident(ident)) => ident.to_string(),
        _ => return Err("#[some_executor::main] expects a named function".to_string()),
    };
    index += 1;

    if let Some(TokenTree::Punct(punct)) = tokens.get(index)
        && punct.as_char() == '<'
    {
        return Err(
            "#[some_executor::main] cannot be applied to a generic function: an entry point \
             has nothing to infer its parameters from"
                .to_string(),
        );
    }

    match tokens.get(index) {
        Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Parenthesis => {
            if !group.stream().is_empty() {
                return Err(
                    "#[some_executor::main] cannot be applied to a function with arguments: \
                     nothing would supply them"
                        .to_string(),
                );
            }
        }
        _ => return Err("#[some_executor::main] expects a function signature".to_string()),
    }
    index += 1;

    // Everything between the arguments and the body is the return type, which
    // the expansion never names -- `MainResult` is selected by the body's own
    // type. Skipped rather than parsed.
    let body = tokens
        .iter()
        .skip(index)
        .find_map(|token| match token {
            TokenTree::Group(group) if group.delimiter() == Delimiter::Brace => {
                Some(group.to_string())
            }
            _ => None,
        })
        .ok_or_else(|| "#[some_executor::main] expects a function body".to_string())?;

    Ok(AsyncFn {
        attributes,
        visibility,
        name,
        body,
    })
}