328 lines · 12.7 KB · Rust
RawDownload
use proc_macro::TokenStream;use proc_macro2::Ident;use quote::{format_ident, quote};use std::mem;use syn::{    self, Expr, ExprLit, FnArg, ItemFn, Lit, Meta, MetaList, PathSegment, Token, Type,    parse::{Parse, ParseStream},    parse_quote,    punctuated::Punctuated,    spanned::Spanned,}; struct Args {    seeds: Vec<u64>,    max_retries: usize,    max_iterations: usize,    on_failure_fn_name: proc_macro2::TokenStream,} impl Parse for Args {    fn parse(input: ParseStream) -> Result<Self, syn::Error> {        let mut seeds = Vec::<u64>::new();        let mut max_retries = 0;        let mut max_iterations = 1;        let mut on_failure_fn_name = quote!(None);         let metas = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;         for meta in metas {            let ident = {                let meta_path = match &meta {                    Meta::NameValue(meta) => &meta.path,                    Meta::List(list) => &list.path,                    Meta::Path(path) => {                        return Err(syn::Error::new(path.span(), "invalid path argument"));                    }                };                let Some(ident) = meta_path.get_ident() else {                    return Err(syn::Error::new(meta_path.span(), "unexpected path"));                };                ident.to_string()            };             match (&meta, ident.as_str()) {                (Meta::NameValue(meta), "retries") => {                    max_retries = parse_usize_from_expr(&meta.value)?                }                (Meta::NameValue(meta), "iterations") => {                    max_iterations = parse_usize_from_expr(&meta.value)?                }                (Meta::NameValue(meta), "on_failure") => {                    let Expr::Lit(ExprLit {                        lit: Lit::Str(name),                        ..                    }) = &meta.value                    else {                        return Err(syn::Error::new(                            meta.value.span(),                            "on_failure argument must be a string",                        ));                    };                    let segments = name                        .value()                        .split("::")                        .map(|part| PathSegment::from(Ident::new(part, name.span())))                        .collect();                    let path = syn::Path {                        leading_colon: None,                        segments,                    };                    on_failure_fn_name = quote!(Some(#path));                }                (Meta::NameValue(meta), "seed") => {                    seeds = vec![parse_usize_from_expr(&meta.value)? as u64]                }                (Meta::List(list), "seeds") => seeds = parse_u64_array(&list)?,                (Meta::Path(_), _) => {                    return Err(syn::Error::new(meta.span(), "invalid path argument"));                }                (_, _) => {                    return Err(syn::Error::new(meta.span(), "invalid argument name"));                }            }        }         Ok(Args {            seeds,            max_retries,            max_iterations: max_iterations,            on_failure_fn_name,        })    }} pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {    let args = syn::parse_macro_input!(args as Args);    let mut inner_fn = match syn::parse::<ItemFn>(function) {        Ok(f) => f,        Err(err) => return error_to_stream(err),    };     let inner_fn_attributes = mem::take(&mut inner_fn.attrs);    let inner_fn_name = format_ident!("_{}", inner_fn.sig.ident);    let outer_fn_name = mem::replace(&mut inner_fn.sig.ident, inner_fn_name.clone());     let result = generate_test_function(        args,        inner_fn,        inner_fn_attributes,        inner_fn_name,        outer_fn_name,    );    match result {        Ok(tokens) => tokens,        Err(tokens) => tokens,    }} fn generate_test_function(    args: Args,    inner_fn: ItemFn,    inner_fn_attributes: Vec<syn::Attribute>,    inner_fn_name: Ident,    outer_fn_name: Ident,) -> Result<TokenStream, TokenStream> {    let seeds = &args.seeds;    let max_retries = args.max_retries;    let num_iterations = args.max_iterations;    let on_failure_fn_name = &args.on_failure_fn_name;    let seeds = quote!( #(#seeds),* );     let mut outer_fn: ItemFn = if inner_fn.sig.asyncness.is_some() {        // Pass to the test function the number of app contexts that it needs,        // based on its parameter list.        let mut cx_vars = proc_macro2::TokenStream::new();        let mut cx_teardowns = proc_macro2::TokenStream::new();        let mut inner_fn_args = proc_macro2::TokenStream::new();        for (ix, arg) in inner_fn.sig.inputs.iter().enumerate() {            if let FnArg::Typed(arg) = arg {                if let Type::Path(ty) = &*arg.ty {                    let last_segment = ty.path.segments.last();                    match last_segment.map(|s| s.ident.to_string()).as_deref() {                        Some("StdRng") => {                            inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(_seed),));                            continue;                        }                        Some("BackgroundExecutor") => {                            inner_fn_args.extend(quote!(gpui::BackgroundExecutor::new(                                std::sync::Arc::new(dispatcher.clone()),                            ),));                            continue;                        }                        _ => {}                    }                } else if let Type::Reference(ty) = &*arg.ty {                    if let Type::Path(ty) = &*ty.elem {                        let last_segment = ty.path.segments.last();                        if let Some("TestAppContext") =                            last_segment.map(|s| s.ident.to_string()).as_deref()                        {                            let cx_varname = format_ident!("cx_{}", ix);                            cx_vars.extend(quote!(                                let mut #cx_varname = gpui::TestAppContext::build(                                    dispatcher.clone(),                                    Some(stringify!(#outer_fn_name)),                                );                            ));                            cx_teardowns.extend(quote!(                                dispatcher.run_until_parked();                                #cx_varname.quit();                                dispatcher.run_until_parked();                            ));                            inner_fn_args.extend(quote!(&mut #cx_varname,));                            continue;                        }                    }                }            }             return Err(error_with_message("invalid function signature", arg));        }         parse_quote! {            #[test]            fn #outer_fn_name() {                #inner_fn                 gpui::run_test(                    #num_iterations,                    &[#seeds],                    #max_retries,                    &mut |dispatcher, _seed| {                        let executor = gpui::BackgroundExecutor::new(std::sync::Arc::new(dispatcher.clone()));                        #cx_vars                        executor.block_test(#inner_fn_name(#inner_fn_args));                        #cx_teardowns                    },                    #on_failure_fn_name                );            }        }    } else {        // Pass to the test function the number of app contexts that it needs,        // based on its parameter list.        let mut cx_vars = proc_macro2::TokenStream::new();        let mut cx_teardowns = proc_macro2::TokenStream::new();        let mut inner_fn_args = proc_macro2::TokenStream::new();        for (ix, arg) in inner_fn.sig.inputs.iter().enumerate() {            if let FnArg::Typed(arg) = arg {                if let Type::Path(ty) = &*arg.ty {                    let last_segment = ty.path.segments.last();                     if let Some("StdRng") = last_segment.map(|s| s.ident.to_string()).as_deref() {                        inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(_seed),));                        continue;                    }                } else if let Type::Reference(ty) = &*arg.ty {                    if let Type::Path(ty) = &*ty.elem {                        let last_segment = ty.path.segments.last();                        match last_segment.map(|s| s.ident.to_string()).as_deref() {                            Some("App") => {                                let cx_varname = format_ident!("cx_{}", ix);                                let cx_varname_lock = format_ident!("cx_{}_lock", ix);                                cx_vars.extend(quote!(                                    let mut #cx_varname = gpui::TestAppContext::build(                                       dispatcher.clone(),                                       Some(stringify!(#outer_fn_name))                                    );                                    let mut #cx_varname_lock = #cx_varname.app.borrow_mut();                                ));                                inner_fn_args.extend(quote!(&mut #cx_varname_lock,));                                cx_teardowns.extend(quote!(                                    drop(#cx_varname_lock);                                    dispatcher.run_until_parked();                                    #cx_varname.update(|cx| { cx.quit() });                                    dispatcher.run_until_parked();                                ));                                continue;                            }                            Some("TestAppContext") => {                                let cx_varname = format_ident!("cx_{}", ix);                                cx_vars.extend(quote!(                                    let mut #cx_varname = gpui::TestAppContext::build(                                        dispatcher.clone(),                                        Some(stringify!(#outer_fn_name))                                    );                                ));                                cx_teardowns.extend(quote!(                                    dispatcher.run_until_parked();                                    #cx_varname.quit();                                    dispatcher.run_until_parked();                                ));                                inner_fn_args.extend(quote!(&mut #cx_varname,));                                continue;                            }                            _ => {}                        }                    }                }            }             return Err(error_with_message("invalid function signature", arg));        }         parse_quote! {            #[test]            fn #outer_fn_name() {                #inner_fn                 gpui::run_test(                    #num_iterations,                    &[#seeds],                    #max_retries,                    &mut |dispatcher, _seed| {                        #cx_vars                        #inner_fn_name(#inner_fn_args);                        #cx_teardowns                    },                    #on_failure_fn_name,                );            }        }    };    outer_fn.attrs.extend(inner_fn_attributes);     Ok(TokenStream::from(quote!(#outer_fn)))} fn parse_usize_from_expr(expr: &Expr) -> Result<usize, syn::Error> {    let Expr::Lit(ExprLit {        lit: Lit::Int(int), ..    }) = expr    else {        return Err(syn::Error::new(expr.span(), "expected an integer"));    };    int.base10_parse()        .map_err(|_| syn::Error::new(int.span(), "failed to parse integer"))} fn parse_u64_array(meta_list: &MetaList) -> Result<Vec<u64>, syn::Error> {    let mut result = Vec::new();    let tokens = &meta_list.tokens;    let parser = |input: ParseStream| {        let exprs = Punctuated::<Expr, Token![,]>::parse_terminated(input)?;        for expr in exprs {            if let Expr::Lit(ExprLit {                lit: Lit::Int(int), ..            }) = expr            {                let value: usize = int.base10_parse()?;                result.push(value as u64);            } else {                return Err(syn::Error::new(expr.span(), "expected an integer"));            }        }        Ok(())    };    syn::parse::Parser::parse2(parser, tokens.clone())?;    Ok(result)} fn error_with_message(message: &str, spanned: impl Spanned) -> TokenStream {    error_to_stream(syn::Error::new(spanned.span(), message))} fn error_to_stream(err: syn::Error) -> TokenStream {    TokenStream::from(err.into_compile_error())}