websockets and upgrade mechanism

core/codegen/src/bang/mod.rs

@@ -4,6 +4,7 @@ mod test_guide;
 mod export;
 
 pub mod typed_stream;
+pub mod websocket;
 
 use devise::Result;
 use syn::{Path, punctuated::Punctuated, parse::Parser, Token};
@@ -77,3 +78,8 @@ pub fn typed_stream(input: proc_macro::TokenStream) -> TokenStream {
     typed_stream::_macro(input)
         .unwrap_or_else(|diag| diag.emit_as_item_tokens())
 }
+
+pub fn websocket(input: proc_macro::TokenStream) -> TokenStream {
+    websocket::_macro(input)
+        .unwrap_or_else(|diag| diag.emit_as_item_tokens())
+}

core/codegen/src/bang/websocket.rs

@@ -0,0 +1,56 @@
+use devise::{Level, Diagnostic, Spanned};
+use proc_macro2::TokenStream;
+use syn::{parse::{Parse, ParseStream, discouraged::Speculative}, Pat};
+
+pub fn _macro(input: proc_macro::TokenStream) -> devise::Result<TokenStream> {
+    let closure: syn::ExprClosure = syn::parse(input)?;
+
+    if closure.inputs.len() != 1 {
+        return Err(Diagnostic::spanned(
+            closure.inputs.span(),
+            Level::Error,
+            "rocket::response::websocket::CreateWebsocket! needs exactly one closure input"
+        ));
+    }
+
+    if closure.capture.is_none() {
+        return Err(
+            Diagnostic::spanned(
+                closure.span(),
+                Level::Error,
+                "rocket::response::websocket::CreateWebsocket! needs an closure that captures it's inputs"
+            )
+                .span_help(closure.or1_token.span(), "add the `move` keyword to the closure")
+        );
+    }
+
+    let inp = closure.inputs.first().unwrap();
+    match inp {
+        Pat::Ident(_) => {}
+        Pat::Type(_) => {}
+        _ => {
+            return Err(
+                Diagnostic::spanned(
+                    inp.span(),
+                    Level::Error,
+                    "rocket::response::websocket::CreateWebsocket! can only accept an identifier or a type ascription for closure input"
+                )
+            )
+        }
+    }
+
+    let body = closure.body;
+    let capture = closure.capture;
+    let tokens = quote!(
+        Websocket::create(|#inp| {
+            ::std::boxed::Box::new(
+                ::std::boxed::Box::pin(
+                    async #capture {
+                        #body
+                    }
+                )
+            )
+        })
+    );
+    Ok(tokens)
+}

core/codegen/src/lib.rs

@@ -1490,18 +1490,7 @@ pub fn __typed_stream(input: TokenStream) -> TokenStream {
 #[proc_macro]
 #[doc(hidden)]
 pub fn __websocket(input: TokenStream) -> TokenStream {
-    let stmts = syn::Block::parse_within.parse(input).expect("Input to __websocket! should be statements");
-    quote!(
-        Websocket::create(|mut ch: WebsocketChannel| {
-            ::std::boxed::Box::new(
-                ::std::boxed::Box::pin(
-                    async move {
-                        #(#stmts)*
-                    }
-                )
-            )
-        })
-    ).into()
+    emit!(bang::websocket(input))
 }
 
 /// Private Rocket internal macro: `internal_guide_tests!`.