Feat: added streaming feature

Added a streaming version of `Markup`, with support for Axum

Author: giuseppe-zanichelli

maud/Cargo.toml

@@ -17,6 +17,9 @@ include.workspace = true
 [features]
 default = []
 
+# Streaming documents
+streaming = ["dep:futures", "dep:async-stream", "maud_macros/streaming"]
+
 # Web framework integrations
 actix-web = ["actix-web-dep", "futures-util"]
 axum = ["axum-core", "http"]
@@ -35,6 +38,8 @@ http = { version = "1", optional = true }
 warp = { version = "0.4", optional = true }
 poem = { version = "3", optional = true }
 salvo_core = { version = "0.78.0", optional = true }
+futures = { version = "0.3.31", default-features = false, optional = true }
+async-stream = { version = "0.3.6", default-features = false, optional = true }
 
 [dev-dependencies]
 trybuild = { version = "1.0.33", features = ["diff"] }

maud/src/lib.rs

@@ -18,6 +18,9 @@ pub use maud_macros::html;
 
 mod escape;
 
+#[cfg(feature = "streaming")]
+pub mod streaming;
+
 /// An adapter that escapes HTML special characters.
 ///
 /// The following characters are escaped:

maud/src/streaming.rs

@@ -0,0 +1,42 @@
+//! # Streaming generation
+//!
+//! Generation of streaming responses
+
+pub use async_stream;
+
+/// A streaming document.
+///
+/// This is what the [`streaming_html`] macro extends to, with `S` implementing
+/// [`futures::Stream<Item=Markup>`]
+#[derive(Debug, Clone, Copy)]
+pub struct StreamingMarkup<S>(pub S);
+
+#[cfg(feature = "axum")]
+mod axum_support {
+    use core::convert::Infallible;
+
+    use crate::{PreEscaped, streaming::StreamingMarkup};
+    use alloc::string::String;
+    use axum_core::{
+        body::Body,
+        response::{IntoResponse, Response},
+    };
+    use futures::{Stream, StreamExt};
+    use http::{HeaderValue, header};
+
+    impl<S> IntoResponse for StreamingMarkup<S>
+    where
+        S: Stream<Item = PreEscaped<String>> + Send + 'static,
+    {
+        fn into_response(self) -> Response {
+            let headers = [(
+                header::CONTENT_TYPE,
+                HeaderValue::from_static("text/html; charset=utf-8"),
+            )];
+
+            let body = Body::from_stream(self.0.map(|PreEscaped(s)| Ok::<_, Infallible>(s)));
+
+            (headers, body).into_response()
+        }
+    }
+}

maud_macros/Cargo.toml

@@ -12,6 +12,11 @@ repository.workspace = true
 edition.workspace = true
 include.workspace = true
 
+[features]
+
+# Streaming documents
+streaming = []
+
 [dependencies]
 syn = { version = "2", features = ["extra-traits", "full"] }
 quote = "1.0.7"

maud_macros/src/ast.rs

@@ -12,7 +12,7 @@ use syn::{
     spanned::Spanned,
     token::{
         At, Brace, Bracket, Colon, Comma, Dot, Else, Eq, FatArrow, For, If, In, Let, Match, Minus,
-        Paren, Pound, Question, Semi, Slash, While,
+        Paren, Pound, Question, Semi, Slash, While, Yield,
     },
 };
 
@@ -810,6 +810,15 @@ impl<E: MaybeElement> DiagnosticParse for ControlFlow<E> {
                     };
 
                     ControlFlowKind::Let(local)
+                } else if lookahead.peek(Yield) {
+                    #[cfg(feature = "streaming")]
+                    {
+                        ControlFlowKind::Yield(input.diagnostic_parse(diagnostics)?)
+                    }
+                    #[cfg(not(feature = "streaming"))]
+                    {
+                        return Err(lookahead.error());
+                    }
                 } else {
                     return Err(lookahead.error());
                 }
@@ -827,6 +836,8 @@ impl<E: ToTokens> ToTokens for ControlFlow<E> {
             ControlFlowKind::For(for_) => for_.to_tokens(tokens),
             ControlFlowKind::While(while_) => while_.to_tokens(tokens),
             ControlFlowKind::Match(match_) => match_.to_tokens(tokens),
+            #[cfg(feature = "streaming")]
+            ControlFlowKind::Yield(yield_) => yield_.to_tokens(tokens),
         }
     }
 }
@@ -838,6 +849,8 @@ pub enum ControlFlowKind<E> {
     For(ForExpr<E>),
     While(WhileExpr<E>),
     Match(MatchExpr<E>),
+    #[cfg(feature = "streaming")]
+    Yield(YieldExpr),
 }
 
 #[derive(Debug, Clone)]
@@ -1073,6 +1086,31 @@ impl<E: ToTokens> ToTokens for MatchArm<E> {
     }
 }
 
+#[cfg(feature = "streaming")]
+#[derive(Debug, Clone)]
+pub struct YieldExpr {
+    pub yield_token: Yield,
+}
+
+#[cfg(feature = "streaming")]
+impl DiagnosticParse for YieldExpr {
+    fn diagnostic_parse(
+        input: ParseStream,
+        _diagnostics: &mut Vec<Diagnostic>,
+    ) -> syn::Result<Self> {
+        Ok(Self {
+            yield_token: input.parse()?,
+        })
+    }
+}
+
+#[cfg(feature = "streaming")]
+impl ToTokens for YieldExpr {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
+        self.yield_token.to_tokens(tokens);
+    }
+}
+
 pub trait DiagnosticParse: Sized {
     fn diagnostic_parse(input: ParseStream, diagnostics: &mut Vec<Diagnostic>)
     -> syn::Result<Self>;

maud_macros/src/generate.rs

@@ -10,13 +10,34 @@ pub fn generate(markups: Markups<Element>, output_ident: Ident) -> TokenStream {
     build.finish()
 }
 
+#[cfg(feature = "streaming")]
+pub fn generate_streaming(markups: Markups<Element>, output_ident: Ident) -> TokenStream {
+    let mut build = Builder::new(output_ident.clone());
+    Generator::new_streaming(output_ident).markups(markups, &mut build);
+    build.finish()
+}
+
 struct Generator {
     output_ident: Ident,
+    #[cfg(feature = "streaming")]
+    streaming: bool,
 }
 
 impl Generator {
     fn new(output_ident: Ident) -> Generator {
-        Generator { output_ident }
+        Generator {
+            output_ident,
+            #[cfg(feature = "streaming")]
+            streaming: false,
+        }
+    }
+
+    #[cfg(feature = "streaming")]
+    fn new_streaming(output_ident: Ident) -> Generator {
+        Generator {
+            output_ident,
+            streaming: true,
+        }
     }
 
     fn builder(&self) -> Builder {
@@ -191,6 +212,8 @@ impl Generator {
             ControlFlowKind::For(for_) => self.control_flow_for(for_, build),
             ControlFlowKind::While(while_) => self.control_flow_while(while_, build),
             ControlFlowKind::Match(match_) => self.control_flow_match(match_, build),
+            #[cfg(feature = "streaming")]
+            ControlFlowKind::Yield(yield_) => self.control_flow_yield(yield_, build),
         }
     }
 
@@ -303,6 +326,23 @@ impl Generator {
 
         build.push_tokens(quote!(#match_token #expr #arm_block));
     }
+
+    #[cfg(feature = "streaming")]
+    fn control_flow_yield(&self, YieldExpr { yield_token }: YieldExpr, build: &mut Builder) {
+        if !self.streaming {
+            build.push_tokens(
+                syn::Error::new_spanned(
+                    yield_token,
+                    "`yield` instructions are permitted only in `streaming_html`",
+                )
+                .into_compile_error(),
+            );
+        }
+        let output_ident = &self.output_ident;
+        build.push_tokens(
+            quote! { #yield_token maud::PreEscaped( ::core::mem::take( #output_ident ) ) ; },
+        );
+    }
 }
 
 ////////////////////////////////////////////////////////

maud_macros/src/lib.rs

@@ -14,6 +14,9 @@ use proc_macro2_diagnostics::Diagnostic;
 use quote::quote;
 use syn::parse::{ParseStream, Parser};
 
+#[cfg(feature = "streaming")]
+mod streaming;
+
 #[proc_macro]
 pub fn html(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     expand(input.into()).into()
@@ -54,3 +57,9 @@ fn expand(input: TokenStream) -> TokenStream {
         maud::PreEscaped(#output_ident)
     }}
 }
+
+#[cfg(feature = "streaming")]
+#[proc_macro]
+pub fn streaming_html(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    streaming::expand(input.into()).into()
+}

maud_macros/src/streaming.rs

@@ -0,0 +1,46 @@
+use crate::ast::DiagnosticParse;
+use proc_macro2::{Ident, Span, TokenStream};
+use proc_macro2_diagnostics::Diagnostic;
+use quote::quote;
+use syn::parse::{ParseStream, Parser};
+
+pub fn expand(input: TokenStream) -> TokenStream {
+    // TODO: How to replace the size hint? Maybe measure the distance between
+    // two yields?
+
+    let mut diagnostics = Vec::new();
+    let markups = match Parser::parse2(
+        |input: ParseStream| crate::ast::Markups::diagnostic_parse(input, &mut diagnostics),
+        input,
+    ) {
+        Ok(data) => data,
+        Err(err) => {
+            let err = err.to_compile_error();
+            let diag_tokens = diagnostics.into_iter().map(Diagnostic::emit_as_expr_tokens);
+
+            return quote! {{
+                #err
+                #(#diag_tokens)*
+            }};
+        }
+    };
+
+    let diag_tokens = diagnostics.into_iter().map(Diagnostic::emit_as_expr_tokens);
+
+    let output_ident = Ident::new("__maud_output", Span::mixed_site());
+    let stmts = crate::generate::generate_streaming(markups, output_ident.clone());
+    quote! {{
+        extern crate alloc;
+        extern crate maud;
+
+        use maud::streaming::async_stream::stream;
+
+        maud::streaming::StreamingMarkup(stream!{
+            let mut #output_ident = alloc::string::String::new();
+
+            #stmts
+            #(#diag_tokens)*
+            yield maud::PreEscaped(#output_ident);
+        })
+    }}
+}