feat(macros): builder params

Author: Berrysoft

compio-macros/README.md

@@ -39,3 +39,12 @@ async fn main() {
     println!("Hello from compio!");
 }
 ```
+
+You can customize the runtime through params:
+
+```rust
+#[compio::main(event_interval = 4, with_proactor(capacity = 16))]
+async fn main() {
+    println!("Hello from compio!");
+}
+```

compio-macros/src/item_fn.rs

@@ -1,25 +1,114 @@
 use proc_macro2::TokenStream;
-use quote::quote;
+use quote::{ToTokens, TokenStreamExt, quote};
 use syn::{
-    Attribute, Expr, Lit, Meta, Signature, Visibility, parse::Parse, punctuated::Punctuated,
+    Attribute, Expr, ExprLit, Ident, Lit, Meta, Path, Signature, Token, Visibility, parse::Parse,
+    punctuated::Punctuated,
 };
 
-type AttributeArgs = Punctuated<syn::Meta, syn::Token![,]>;
+use crate::{retrieve_driver_mod, retrieve_runtime_mod};
 
+struct MetaPunctuated(Punctuated<Meta, Token![,]>);
+
+impl Parse for MetaPunctuated {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        Ok(Self(Punctuated::parse_terminated(input)?))
+    }
+}
+
+pub(crate) struct BuilderMethod {
+    pub name: Ident,
+    pub value: Expr,
+}
+
+#[derive(Default)]
 pub(crate) struct RawAttr {
-    pub inner_attrs: AttributeArgs,
+    pub runtime_methods: Vec<BuilderMethod>,
+    pub proactor_methods: Vec<BuilderMethod>,
+    pub crate_name: Option<TokenStream>,
+    pub with_proactor_call: Option<Path>,
 }
 
 impl Parse for RawAttr {
     fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
-        let inner_attrs = AttributeArgs::parse_terminated(input)?;
-        Ok(Self { inner_attrs })
+        let items = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;
+
+        let mut runtime_methods = Vec::new();
+        let mut proactor_methods = Vec::new();
+        let mut crate_name = None;
+        let mut with_proactor_call = None;
+
+        for meta in items {
+            match meta {
+                Meta::List(list) => {
+                    if list.path.is_ident("with_proactor") {
+                        if !list.tokens.is_empty() {
+                            let inner_items = syn::parse2::<MetaPunctuated>(list.tokens)?.0;
+                            for inner_meta in inner_items {
+                                if let Meta::NameValue(nv) = inner_meta {
+                                    let name = nv.path.require_ident()?.clone();
+                                    proactor_methods.push(BuilderMethod {
+                                        name,
+                                        value: nv.value,
+                                    });
+                                } else {
+                                    return Err(syn::Error::new_spanned(
+                                        inner_meta,
+                                        "expected `name = value` inside `with_proactor`",
+                                    ));
+                                }
+                            }
+                            with_proactor_call = Some(list.path.clone());
+                        }
+                    } else {
+                        return Err(syn::Error::new_spanned(
+                            list.path,
+                            "unknown key; use `name = value` for parameters or \
+                             `with_proactor(...)` for proactor config",
+                        ));
+                    }
+                }
+                Meta::NameValue(nv) => {
+                    if nv.path.is_ident("crate") {
+                        if let Expr::Lit(ExprLit {
+                            lit: Lit::Str(s), ..
+                        }) = &nv.value
+                        {
+                            crate_name = Some(s.parse::<TokenStream>().unwrap());
+                        } else {
+                            return Err(syn::Error::new_spanned(
+                                &nv.value,
+                                "expected a string literal for `crate`",
+                            ));
+                        }
+                    } else {
+                        let name = nv.path.require_ident()?.clone();
+                        runtime_methods.push(BuilderMethod {
+                            name,
+                            value: nv.value,
+                        });
+                    }
+                }
+                Meta::Path(path) => {
+                    return Err(syn::Error::new_spanned(
+                        path,
+                        "expected `name = value` or `with_proactor(...)`",
+                    ));
+                }
+            }
+        }
+
+        Ok(Self {
+            runtime_methods,
+            proactor_methods,
+            crate_name,
+            with_proactor_call,
+        })
     }
 }
 
 pub(crate) struct RawBodyItemFn {
     pub attrs: Vec<Attribute>,
-    pub args: AttributeArgs,
+    pub args: RawAttr,
     pub vis: Visibility,
     pub sig: Signature,
     pub body: TokenStream,
@@ -29,37 +118,88 @@ impl RawBodyItemFn {
     pub fn new(attrs: Vec<Attribute>, vis: Visibility, sig: Signature, body: TokenStream) -> Self {
         Self {
             attrs,
-            args: AttributeArgs::new(),
+            args: RawAttr::default(),
             vis,
             sig,
             body,
         }
     }
 
-    pub fn set_args(&mut self, args: AttributeArgs) {
+    pub fn set_args(&mut self, args: RawAttr) {
         self.args = args;
     }
 
-    pub fn crate_name(&self) -> Option<TokenStream> {
-        for attr in &self.args {
-            if let Meta::NameValue(name) = &attr {
-                let ident = name
-                    .path
-                    .get_ident()
-                    .map(|ident| ident.to_string().to_lowercase())
-                    .unwrap_or_default();
-                if ident == "crate" {
-                    if let Expr::Lit(lit) = &name.value
-                        && let Lit::Str(s) = &lit.lit
-                    {
-                        let crate_name = s.parse::<TokenStream>().unwrap();
-                        return Some(quote!(#crate_name::runtime));
-                    }
-                } else {
-                    panic!("Unsupported property {ident}");
-                }
+    pub fn emit_fn_to_tokens(&self, tokens: &mut TokenStream) {
+        tokens.append_all(
+            self.attrs
+                .iter()
+                .filter(|a| matches!(a.style, syn::AttrStyle::Outer)),
+        );
+        self.vis.to_tokens(tokens);
+        self.sig.to_tokens(tokens);
+        tokens.append_all(self.gen_runtime_block());
+    }
+
+    fn gen_runtime_block(&self) -> TokenStream {
+        let runtime_mod = match &self.args.crate_name {
+            Some(c) => {
+                let c = c.clone();
+                quote!(#c::runtime)
+            }
+            None => retrieve_runtime_mod(),
+        };
+
+        let driver_mod = match &self.args.crate_name {
+            Some(c) => {
+                let c = c.clone();
+                quote!(#c::driver)
             }
+            None => retrieve_driver_mod(),
+        };
+
+        let block = &self.body;
+
+        let mut builder = quote! {
+            #runtime_mod::Runtime::builder()
+        };
+
+        for method in &self.args.runtime_methods {
+            let name = &method.name;
+            let value = &method.value;
+            builder = quote! {
+                #builder.#name(#value)
+            };
         }
-        None
+
+        if !self.args.proactor_methods.is_empty() {
+            let mut proactor_stmts: Vec<TokenStream> = Vec::new();
+            proactor_stmts.push(quote! {
+                let mut __compio_proactor_builder = #driver_mod::Proactor::builder();
+            });
+            for method in &self.args.proactor_methods {
+                let name = &method.name;
+                let value = &method.value;
+                proactor_stmts.push(quote! {
+                    __compio_proactor_builder.#name(#value);
+                });
+            }
+            // Preserve the original token for the `with_proactor` call to make the language
+            // server work better.
+            let with_proactor_call = if let Some(path) = &self.args.with_proactor_call {
+                quote!(#path)
+            } else {
+                quote!(with_proactor)
+            };
+            builder = quote! {
+                #builder.#with_proactor_call({
+                    #(#proactor_stmts)*
+                    __compio_proactor_builder
+                })
+            };
+        }
+
+        quote!({
+            #builder.build().expect("cannot create runtime").block_on(async move #block)
+        })
     }
 }

compio-macros/src/lib.rs

@@ -50,3 +50,24 @@ fn retrieve_runtime_mod() -> proc_macro2::TokenStream {
         },
     }
 }
+
+fn retrieve_driver_mod() -> proc_macro2::TokenStream {
+    match crate_name("compio-driver") {
+        Ok(FoundCrate::Itself) => quote!(crate),
+        Ok(FoundCrate::Name(name)) => {
+            let ident = Ident::new(&name, Span::call_site());
+            quote!(::#ident)
+        }
+        Err(_) => match crate_name("compio") {
+            Ok(FoundCrate::Itself) => quote!(crate::driver),
+            Ok(FoundCrate::Name(name)) => {
+                let ident = Ident::new(&name, Span::call_site());
+                quote!(::#ident::driver)
+            }
+            Err(_) => {
+                let ident = Ident::new("compio_driver", Span::call_site());
+                quote!(::#ident)
+            }
+        },
+    }
+}

compio-macros/src/main_fn.rs

@@ -1,17 +1,14 @@
 use proc_macro2::TokenStream;
-use quote::{ToTokens, TokenStreamExt, quote};
-use syn::{AttrStyle, Attribute, Signature, Visibility, parse::Parse};
+use quote::ToTokens;
+use syn::{Attribute, Signature, Visibility, parse::Parse};
 
-use crate::{
-    item_fn::{RawAttr, RawBodyItemFn},
-    retrieve_runtime_mod,
-};
+use crate::item_fn::{RawAttr, RawBodyItemFn};
 
 pub(crate) struct CompioMain(pub RawBodyItemFn);
 
 impl CompioMain {
     pub fn with_args(mut self, args: RawAttr) -> Self {
-        self.0.set_args(args.inner_attrs);
+        self.0.set_args(args);
         self
     }
 }
@@ -37,18 +34,6 @@ impl Parse for CompioMain {
 
 impl ToTokens for CompioMain {
     fn to_tokens(&self, tokens: &mut TokenStream) {
-        tokens.append_all(
-            self.0
-                .attrs
-                .iter()
-                .filter(|a| matches!(a.style, AttrStyle::Outer)),
-        );
-        self.0.vis.to_tokens(tokens);
-        self.0.sig.to_tokens(tokens);
-        let block = &self.0.body;
-        let runtime_mod = self.0.crate_name().unwrap_or_else(retrieve_runtime_mod);
-        tokens.append_all(quote!({
-            #runtime_mod::Runtime::new().expect("cannot create runtime").block_on(async move #block)
-        }));
+        self.0.emit_fn_to_tokens(tokens);
     }
 }

compio-macros/src/test_fn.rs

@@ -1,17 +1,14 @@
 use proc_macro2::TokenStream;
 use quote::{ToTokens, TokenStreamExt, quote};
-use syn::{AttrStyle, Attribute, Signature, Visibility, parse::Parse};
+use syn::{Attribute, Signature, Visibility, parse::Parse};
 
-use crate::{
-    item_fn::{RawAttr, RawBodyItemFn},
-    retrieve_runtime_mod,
-};
+use crate::item_fn::{RawAttr, RawBodyItemFn};
 
 pub(crate) struct CompioTest(pub RawBodyItemFn);
 
 impl CompioTest {
     pub fn with_args(mut self, args: RawAttr) -> Self {
-        self.0.set_args(args.inner_attrs);
+        self.0.set_args(args);
         self
     }
 }
@@ -37,18 +34,6 @@ impl Parse for CompioTest {
 impl ToTokens for CompioTest {
     fn to_tokens(&self, tokens: &mut TokenStream) {
         tokens.append_all(quote!(#[test]));
-        tokens.append_all(
-            self.0
-                .attrs
-                .iter()
-                .filter(|a| matches!(a.style, AttrStyle::Outer)),
-        );
-        self.0.vis.to_tokens(tokens);
-        self.0.sig.to_tokens(tokens);
-        let block = &self.0.body;
-        let runtime_mod = self.0.crate_name().unwrap_or_else(retrieve_runtime_mod);
-        tokens.append_all(quote!({
-            #runtime_mod::Runtime::new().expect("cannot create runtime").block_on(async move #block)
-        }));
+        self.0.emit_fn_to_tokens(tokens);
     }
 }

compio/tests/accept.rs

@@ -1,27 +1,18 @@
 use compio::{
-    driver::{DriverType, ProactorBuilder},
+    driver::DriverType,
     net::{TcpListener, TcpStream},
-    runtime::Runtime,
 };
 use compio_runtime::ResumeUnwind;
 
-#[test]
-fn accept() {
-    let mut proactor_builder = ProactorBuilder::new();
-    proactor_builder.driver_type(DriverType::Poll);
-    let runtime = Runtime::builder()
-        .with_proactor(proactor_builder)
-        .build()
-        .unwrap();
-    runtime.block_on(async {
-        let listener = TcpListener::bind("localhost:0").await.unwrap();
-        let addr = listener.local_addr().unwrap();
-        let task = compio_runtime::spawn(async move {
-            let (socket, _) = listener.accept().await.unwrap();
-            socket
-        });
-        let cli = TcpStream::connect(&addr).await.unwrap();
-        let srv = task.await.resume_unwind().expect("shouldn't be cancelled");
-        assert_eq!(cli.local_addr().unwrap(), srv.peer_addr().unwrap());
-    })
+#[compio_macros::test(with_proactor(driver_type = DriverType::Poll))]
+async fn accept() {
+    let listener = TcpListener::bind("localhost:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let task = compio_runtime::spawn(async move {
+        let (socket, _) = listener.accept().await.unwrap();
+        socket
+    });
+    let cli = TcpStream::connect(&addr).await.unwrap();
+    let srv = task.await.resume_unwind().expect("shouldn't be cancelled");
+    assert_eq!(cli.local_addr().unwrap(), srv.peer_addr().unwrap());
 }