wstd_macro/
lib.rs

1use proc_macro::TokenStream;
2use quote::{quote, quote_spanned};
3use syn::{ItemFn, parse_macro_input, spanned::Spanned};
4
5#[proc_macro_attribute]
6pub fn attr_macro_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(item as ItemFn);
8
9    if input.sig.asyncness.is_none() {
10        return quote_spanned! { input.sig.fn_token.span()=>
11            compile_error!("fn must be `async fn`");
12        }
13        .into();
14    }
15
16    if input.sig.ident != "main" {
17        return quote_spanned! { input.sig.ident.span()=>
18            compile_error!("only `async fn main` can be used for #[wstd::main]");
19        }
20        .into();
21    }
22
23    if !input.sig.inputs.is_empty() {
24        return quote_spanned! { input.sig.inputs.span()=>
25            compile_error!("arguments to main are not supported");
26        }
27        .into();
28    }
29    let attrs = input.attrs;
30    let output = input.sig.output;
31    let block = input.block;
32    quote! {
33        pub fn main() #output {
34
35            #(#attrs)*
36            async fn __run() #output {
37                #block
38            }
39
40            ::wstd::runtime::block_on(async {
41                __run().await
42            })
43        }
44    }
45    .into()
46}
47
48#[proc_macro_attribute]
49pub fn attr_macro_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
50    let input = parse_macro_input!(item as ItemFn);
51
52    if input.sig.asyncness.is_none() {
53        return quote_spanned! { input.sig.fn_token.span()=>
54            compile_error!("fn must be `async fn`");
55        }
56        .into();
57    }
58
59    let name = input.sig.ident;
60
61    if !input.sig.inputs.is_empty() {
62        return quote_spanned! { input.sig.inputs.span()=>
63            compile_error!("arguments to main are not supported");
64        }
65        .into();
66    }
67    let attrs = input.attrs;
68    let output = input.sig.output;
69    let block = input.block;
70    quote! {
71        #[::core::prelude::v1::test]
72        pub fn #name() #output {
73
74            #(#attrs)*
75            async fn __run() #output {
76                #block
77            }
78
79            ::wstd::runtime::block_on(async {
80                __run().await
81            })
82        }
83    }
84    .into()
85}
86
87/// Enables a HTTP server main function, for creating [HTTP servers].
88///
89/// [HTTP servers]: https://docs.rs/wstd/latest/wstd/http/server/index.html
90///
91/// # Examples
92///
93/// ```ignore
94/// #[wstd::http_server]
95/// async fn main(request: Request<Body>) -> Result<Response<Body>> {
96///     Ok(Response::new("Hello!\n".into()))
97/// }
98/// ```
99#[proc_macro_attribute]
100pub fn attr_macro_http_server(_attr: TokenStream, item: TokenStream) -> TokenStream {
101    let input = parse_macro_input!(item as ItemFn);
102
103    let (run_async, run_await) = if input.sig.asyncness.is_some() {
104        (quote!(async), quote!(.await))
105    } else {
106        (quote!(), quote!())
107    };
108
109    let output = &input.sig.output;
110    let inputs = &input.sig.inputs;
111    let name = &input.sig.ident;
112    let body = &input.block;
113    let attrs = &input.attrs;
114    let vis = &input.vis;
115
116    if name != "main" {
117        return quote_spanned! { input.sig.ident.span()=>
118            compile_error!("only `async fn main` can be used for #[wstd::http_server]");
119        }
120        .into();
121    }
122
123    quote! {
124        struct TheServer;
125
126        impl ::wstd::__internal::wasip2::exports::http::incoming_handler::Guest for TheServer {
127            fn handle(
128                request: ::wstd::__internal::wasip2::http::types::IncomingRequest,
129                response_out: ::wstd::__internal::wasip2::http::types::ResponseOutparam
130            ) {
131                #(#attrs)*
132                #vis #run_async fn __run(#inputs) #output {
133                    #body
134                }
135
136                let responder = ::wstd::http::server::Responder::new(response_out);
137                ::wstd::runtime::block_on(async move {
138                    match ::wstd::http::request::try_from_incoming(request) {
139                        Ok(request) => match __run(request) #run_await {
140                            Ok(response) => { responder.respond(response).await.unwrap() },
141                            Err(err) => responder.fail(err).unwrap(),
142                        }
143                        Err(err) => responder.fail(err).unwrap(),
144                    }
145                })
146            }
147        }
148
149        ::wstd::__internal::wasip2::http::proxy::export!(TheServer with_types_in ::wstd::__internal::wasip2);
150
151        // Provide an actual function named `main`.
152        //
153        // WASI HTTP server components don't use a traditional `main` function.
154        // They export a function named `handle` which takes a `Request`
155        // argument, and which may be called multiple times on the same
156        // instance. To let users write a familiar `fn main` in a file
157        // named src/main.rs, we provide this `wstd::http_server` macro, which
158        // transforms the user's `fn main` into the appropriate `handle`
159        // function.
160        //
161        // However, when the top-level file is named src/main.rs, rustc
162        // requires there to be a function named `main` somewhere in it. This
163        // requirement can be disabled using `#![no_main]`, however we can't
164        // use that automatically because macros can't contain inner
165        // attributes, and we don't want to require users to add `#![no_main]`
166        // in their own code.
167        //
168        // So, we include a definition of a function named `main` here, which
169        // isn't intended to ever be called, and exists just to satify the
170        // requirement for a `main` function.
171        //
172        // Users could use `#![no_main]` if they want to. Or, they could name
173        // their top-level file src/lib.rs and add
174        // ```toml
175        // [lib]
176        // crate-type = ["cdylib"]
177        // ```
178        // to their Cargo.toml. With either of these, this "main" function will
179        // be ignored as dead code.
180        fn main() {
181            unreachable!("HTTP server components should be run with `handle` rather than `run`")
182        }
183    }
184    .into()
185}