Merge pull request #18 from yoshuawuyts/pch/attr_macro_main

Attribute macro for main

Author: pchickey

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "wstd"
-version = "0.4.0"
+version.workspace = true
 license.workspace = true
 repository = "https://github.com/yoshuawuyts/wstd"
 documentation = "https://docs.rs/wstd"
@@ -19,6 +19,7 @@ authors = [
 slab.workspace = true
 url.workspace = true
 wasi.workspace = true
+wstd-macro.workspace = true
 
 [dev-dependencies]
 anyhow.workspace = true
@@ -30,26 +31,31 @@ wasmtime-wasi-http.workspace = true
 
 [workspace]
 members = [
+    "macro",
     "test-programs",
     "test-programs/artifacts",
 ]
 resolver = "2"
 
 [workspace.package]
+version = "0.4.0"
 edition = "2021"
 license = "MIT OR Apache-2.0 OR Apache-2.0 WITH LLVM-exception"
 
 [workspace.dependencies]
 anyhow = "1"
 cargo_metadata = "0.18.1"
 heck = "0.5"
+quote = "1.0"
 serde_json = "1"
 slab = "0.4.9"
+syn = "2.0"
 test-programs = { path = "test-programs" }
 test-programs-artifacts = { path = "test-programs/artifacts" }
 url = "2.5.0"
 wasi = "0.13.1"
-wstd = { path = "." }
 wasmtime = "26"
 wasmtime-wasi = "26"
 wasmtime-wasi-http = "26"
+wstd = { path = "." }
+wstd-macro = { path = "macro" }

examples/http_get.rs

@@ -1,36 +1,34 @@
 use std::error::Error;
 use wstd::http::{Client, Method, Request, Url};
 use wstd::io::AsyncRead;
-use wstd::runtime::block_on;
 
-fn main() -> Result<(), Box<dyn Error>> {
-    block_on(async move {
-        let request = Request::new(Method::Get, Url::parse("https://postman-echo.com/get")?);
-        let mut response = Client::new().send(request).await?;
+#[wstd::main]
+async fn main() -> Result<(), Box<dyn Error>> {
+    let request = Request::new(Method::Get, Url::parse("https://postman-echo.com/get")?);
+    let mut response = Client::new().send(request).await?;
 
-        let content_type = response
-            .headers()
-            .get(&"content-type".into())
-            .ok_or_else(|| "response expected to have content-type header")?;
-        assert_eq!(content_type.len(), 1, "one header value for content-type");
-        assert_eq!(content_type[0], b"application/json; charset=utf-8");
+    let content_type = response
+        .headers()
+        .get(&"content-type".into())
+        .ok_or_else(|| "response expected to have content-type header")?;
+    assert_eq!(content_type.len(), 1, "one header value for content-type");
+    assert_eq!(content_type[0], b"application/json; charset=utf-8");
 
-        // Would much prefer read_to_end here:
-        let mut body_buf = vec![0; 4096];
-        let body_len = response.body().read(&mut body_buf).await?;
-        body_buf.truncate(body_len);
+    // Would much prefer read_to_end here:
+    let mut body_buf = vec![0; 4096];
+    let body_len = response.body().read(&mut body_buf).await?;
+    body_buf.truncate(body_len);
 
-        let val: serde_json::Value = serde_json::from_slice(&body_buf)?;
-        let body_url = val
-            .get("url")
-            .ok_or_else(|| "body json has url")?
-            .as_str()
-            .ok_or_else(|| "body json url is str")?;
-        assert!(
-            body_url.contains("postman-echo.com/get"),
-            "expected body url to contain the authority and path, got: {body_url}"
-        );
+    let val: serde_json::Value = serde_json::from_slice(&body_buf)?;
+    let body_url = val
+        .get("url")
+        .ok_or_else(|| "body json has url")?
+        .as_str()
+        .ok_or_else(|| "body json url is str")?;
+    assert!(
+        body_url.contains("postman-echo.com/get"),
+        "expected body url to contain the authority and path, got: {body_url}"
+    );
 
-        Ok(())
-    })
+    Ok(())
 }

examples/tcp_echo_server.rs

@@ -1,20 +1,18 @@
 use wstd::io;
 use wstd::iter::AsyncIterator;
 use wstd::net::TcpListener;
-use wstd::runtime::block_on;
 
-fn main() -> io::Result<()> {
-    block_on(async move {
-        let listener = TcpListener::bind("127.0.0.1:8080").await?;
-        println!("Listening on {}", listener.local_addr()?);
-        println!("type `nc localhost 8080` to create a TCP client");
+#[wstd::main]
+async fn main() -> io::Result<()> {
+    let listener = TcpListener::bind("127.0.0.1:8080").await?;
+    println!("Listening on {}", listener.local_addr()?);
+    println!("type `nc localhost 8080` to create a TCP client");
 
-        let mut incoming = listener.incoming();
-        while let Some(stream) = incoming.next().await {
-            let stream = stream?;
-            println!("Accepted from: {}", stream.peer_addr()?);
-            io::copy(&stream, &stream).await?;
-        }
-        Ok(())
-    })
+    let mut incoming = listener.incoming();
+    while let Some(stream) = incoming.next().await {
+        let stream = stream?;
+        println!("Accepted from: {}", stream.peer_addr()?);
+        io::copy(&stream, &stream).await?;
+    }
+    Ok(())
 }

macro/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "wstd-macro"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[lib]
+proc-macro = true
+
+[dependencies]
+syn = { workspace = true, features = ["full"] }
+quote.workspace = true

macro/src/lib.rs

@@ -0,0 +1,46 @@
+use proc_macro::TokenStream;
+use quote::{quote, quote_spanned};
+use syn::{parse_macro_input, spanned::Spanned, ItemFn};
+
+#[proc_macro_attribute]
+pub fn attr_macro_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(item as ItemFn);
+
+    if input.sig.asyncness.is_none() {
+        return quote_spanned! { input.sig.fn_token.span()=>
+            compile_error!("fn must be `async fn`");
+        }
+        .into();
+    }
+
+    if input.sig.ident != "main" {
+        return quote_spanned! { input.sig.ident.span()=>
+            compile_error!("only `async fn main` can be used for #[wstd::main]");
+        }
+        .into();
+    }
+
+    if !input.sig.inputs.is_empty() {
+        return quote_spanned! { input.sig.inputs.span()=>
+            compile_error!("arguments to main are not supported");
+        }
+        .into();
+    }
+    let attrs = input.attrs;
+    let output = input.sig.output;
+    let block = input.block;
+    quote! {
+        pub fn main() #output {
+
+            #(#attrs)*
+            async fn __run() #output {
+                #block
+            }
+
+            ::wstd::runtime::block_on(async {
+                __run().await
+            })
+        }
+    }
+    .into()
+}

src/lib.rs

@@ -51,3 +51,5 @@ pub mod net;
 pub mod rand;
 pub mod runtime;
 pub mod time;
+
+pub use wstd_macro::attr_macro_main as main;