switchboard_starknet_macros/

lib.rs

extern crate proc_macro;

mod params;
mod utils;

use proc_macro::TokenStream;
use quote::quote;
use syn::{ FnArg, ItemFn, Result as SynResult, ReturnType, Type };

#[proc_macro_attribute]
pub fn switchboard_function(attr: TokenStream, item: TokenStream) -> TokenStream {
    // Parse the macro parameters to set a timeout
    let macro_params = match syn::parse::<params::SwitchboardStarknetFunctionArgs>(attr.clone()) {
        Ok(args) => args,
        Err(err) => {
            let e = syn::Error::new_spanned(
                err.to_compile_error(),
                format!("Failed to parse macro parameters: {:?}", err)
            );

            return e.to_compile_error().into();
        }
    };

    // Try to build the token stream, return errors if failed
    match build_token_stream(macro_params, item) {
        Ok(token_stream) => token_stream,
        Err(err) => err.to_compile_error().into(),
    }
}

/// Validates whether the first param is StarknetFunctionRunner
fn validate_function_runner_param(input: &ItemFn) -> SynResult<()> {
    // Extract the first parameter of the function.
    let first_param_type = input.sig.inputs
        .iter()
        .next()
        .ok_or_else(|| {
            syn::Error::new_spanned(
                &input.sig,
                "The switchboard_function must take at least one parameter"
            )
        })?;

    let typed_arg = match first_param_type {
        FnArg::Typed(typed) => { typed }
        _ => {
            return Err(syn::Error::new_spanned(first_param_type, "Expected a typed parameter"));
        }
    };

    let is_function_runner = if let Type::Path(type_path) = &*typed_arg.ty {
        &type_path.path.segments.last().unwrap().ident == "StarknetFunctionRunner"
    } else {
        false
    };

    if !is_function_runner {
        return Err(syn::Error::new_spanned(&typed_arg.ty, "Parameter must be StarknetFunctionRunner"));
    }

    Ok(())
}

/// Helper function to validate the return type is a Result with the correct arguments.
fn validate_function_return_type(input: &ItemFn) -> SynResult<()> {
    let ty = match &input.sig.output {
        ReturnType::Type(_, ty) => ty,
        ReturnType::Default => {
            return Err(
                syn::Error::new_spanned(&input.sig.output, "Function does not have a return type")
            );
        }
    };

    let (ok_type, err_type) = utils
        ::extract_result_args(ty)
        .ok_or_else(|| {
            syn::Error::new_spanned(&input.sig.output, "Return type must be a Result")
        })?;

    // Validate the inner Vec type
    let inner_vec_type = utils
        ::extract_inner_type_from_vec(ok_type)
        .ok_or_else(|| {
            syn::Error::new_spanned(
                &input.sig.output,
                "Ok variant of Result must be a Vec<Call>"
            )
        })?;

    if !matches!(inner_vec_type, Type::Path(t) if t.path.is_ident("Call")) {
        return Err(
            syn::Error::new_spanned(
                &input.sig.output,
                "Ok variant of Result must be a Vec<Call>"
            )
        );
    }

    // Validate the error type
    let error_type_path_segments = match err_type {
        Type::Path(type_path) => &type_path.path.segments,
        _ => {
            return Err(syn::Error::new_spanned(err_type, "Error type must be a path type"));
        }
    };

    // Check if the error type is SbFunctionError or switchboard_common::SbFunctionError
    let is_sb_function_error = match error_type_path_segments.last() {
        Some(last_segment) if last_segment.ident == "SbFunctionError" => true,
        Some(last_segment) if last_segment.ident == "Error" => {
            // If the last segment is "Error", check the preceding segment for "switchboard_common"
            error_type_path_segments.len() > 1 &&
                error_type_path_segments[error_type_path_segments.len() - 2].ident ==
                    "switchboard_common"
        }
        _ => false,
    };

    if !is_sb_function_error {
        return Err(
            syn::Error::new_spanned(
                &input.sig.output,
                "The error variant in the Result return type should be SbFunctionError"
            )
        );
    }

    Ok(())
}

fn validate_second_parameter(input: &ItemFn) -> SynResult<()> {
    let second_param = input.sig.inputs
        .iter()
        .nth(1)
        .ok_or_else(|| {
            syn::Error::new_spanned(&input.sig, "The switchboard_function must take two parameters")
        })?;

    let typed_arg = match second_param {
        FnArg::Typed(typed) => typed,
        _ => {
            return Err(syn::Error::new_spanned(second_param, "Expected a typed second parameter"));
        }
    };

    // Use the utility function to extract the inner type from a Vec
    let inner_type = utils
        ::extract_inner_type_from_vec(&typed_arg.ty)
        .ok_or_else(||
            syn::Error::new_spanned(&typed_arg.ty, "The second parameter must be of type Vec<FieldElement>")
        )?;

    // Ensure the inner type of the Vec is u8
    if let Type::Path(type_path) = inner_type {
        if !type_path.path.is_ident("FieldElement") {
            return Err(
                syn::Error::new_spanned(
                    &typed_arg.ty,
                    "The second parameter must be of type Vec<FieldElement>"
                )
            );
        }
    } else {
        return Err(
            syn::Error::new_spanned(&typed_arg.ty, "The second parameter must be of type Vec<FieldElement>")
        );
    }

    Ok(())
}

fn build_token_stream(
    _params: params::SwitchboardStarknetFunctionArgs,
    item: TokenStream
) -> SynResult<TokenStream> {
    let input: ItemFn = syn::parse(item.clone())?;
    let function_name = &input.sig.ident;

    // Validate that there's exactly one input of the correct type
    if input.sig.inputs.len() != 2 {
        return Err(
            syn::Error::new_spanned(
                &input.sig,
                "The switchboard_function must take exactly one parameter of type 'Arc<StarknetFunctionRunner>' and 'Vec<FieldElement>'"
            )
        );
    }

    validate_function_return_type(&input)?;

    // Validate input parameters
    validate_function_runner_param(&input)?;
    validate_second_parameter(&input)?;

    let expanded =
        quote! {

            // Include the original function definition
            #input

            pub type SwitchboardFunctionResult<T> = std::result::Result<T, SbFunctionError>;

            /// Run an async function and catch any panics
            pub async fn run_switchboard_function<F, T>(
                logic: F,
            ) -> SwitchboardFunctionResult<()>
            where
                F: Fn(StarknetFunctionRunner, Vec<FieldElement>) -> T + Send + 'static,
                T: futures::Future<Output = SwitchboardFunctionResult<Vec<Call>>>
                    + Send,
            {
                // Initialize the runner
                let mut runner = StarknetFunctionRunner::new().map_err(|_e| SbFunctionError::FunctionResultEmitError)?;
                match logic(runner.clone(), runner.params.clone()).await {
                    Ok(calls) => {
                        runner
                            .emit(calls)
                            .map_err(|_e| SbFunctionError::FunctionResultEmitError)?;
                        Ok(())
                    }
                    Err(e) => {
                        println!("Error: Switchboard function failed with error code: {:?}", e);
                        let mut err_code = 199;
                        if let SbFunctionError::FunctionError(code) = e {
                            err_code = code;
                        }
                        runner
                            .emit_error(err_code)
                            .map_err(|_e| SbFunctionError::FunctionResultEmitError)?;
                        Ok(())
                    }
                }
            }

            #[tokio::main(worker_threads = 12)]
            async fn main() -> SwitchboardFunctionResult<()> {
                run_switchboard_function(#function_name).await?;
                Ok(())
            }
    };

    Ok(TokenStream::from(expanded))
}

#[proc_macro_attribute]
pub fn sb_error(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(item as syn::DeriveInput);

    let name = &input.ident;
    let expanded = quote! {
        #[derive(Clone, Copy, Debug, PartialEq)]
        #[repr(u8)]
        #input

        impl From<#name> for SbFunctionError {
            fn from(item: #name) -> Self {
                SbFunctionError::FunctionError(item as u8 + 1)
            }
        }

        impl From<#name> for u8 {
            fn from(item: #name) -> Self {
                item as u8 + 1
            }
        }

        impl std::fmt::Display for #name {
            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                write!(f, "{:?}", self)
            }
        }

        impl std::error::Error for #name {}
    };

    TokenStream::from(expanded)
}