Streaming content mutations

fuzz/test_case/src/lib.rs

@@ -11,7 +11,8 @@ use std::ffi::{CStr, CString};
 
 use encoding_rs::*;
 use lol_html::html_content::ContentType;
-use lol_html::{comments, doc_comments, doc_text, element, text, HtmlRewriter, MemorySettings, Settings};
+use lol_html::{comments, doc_comments, doc_text, element, streaming, text};
+use lol_html::{HtmlRewriter, MemorySettings, Settings};
 
 include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
 
@@ -111,10 +112,12 @@ fn run_rewriter_iter(data: &[u8], selector: &str, encoding: &'static Encoding) {
                         &format!("<!--[/ELEMENT('{selector}')]-->"),
                         ContentType::Html,
                     );
-                    el.set_inner_content(
-                        &format!("<!--Replaced ({selector}) -->"),
-                        ContentType::Html,
-                    );
+
+                    let replaced = format!("<!--Replaced ({selector}) -->");
+                    el.streaming_set_inner_content(streaming!(move |sink| {
+                        sink.write_str(&replaced, ContentType::Html);
+                        Ok(())
+                    }));
 
                     Ok(())
                 }),

src/lib.rs

@@ -86,7 +86,9 @@ pub mod send {
 pub mod errors {
     pub use super::memory::MemoryLimitExceededError;
     pub use super::parser::ParsingAmbiguityError;
-    pub use super::rewritable_units::{AttributeNameError, CommentTextError, TagNameError};
+    pub use super::rewritable_units::{
+        AttributeNameError, CommentTextError, TagNameError, Utf8Error,
+    };
     pub use super::rewriter::RewritingError;
     pub use super::selectors_vm::SelectorError;
 }
@@ -95,7 +97,7 @@ pub mod errors {
 pub mod html_content {
     pub use super::rewritable_units::{
         Attribute, Comment, ContentType, Doctype, DocumentEnd, Element, EndTag, StartTag,
-        TextChunk, UserData,
+        StreamingHandler, StreamingHandlerSink, TextChunk, UserData,
     };
 
     pub use super::html::TextType;

src/rewritable_units/document_end.rs

@@ -1,5 +1,4 @@
-use super::text_encoder::StreamingHandlerSink;
-use super::ContentType;
+use super::{ContentType, StreamingHandlerSink};
 use encoding_rs::Encoding;
 
 use crate::transform_stream::OutputSink;
@@ -53,7 +52,7 @@ impl<'a> DocumentEnd<'a> {
         StreamingHandlerSink::new(self.encoding, &mut |c| {
             self.output_sink.handle_chunk(c);
         })
-        .write_str_chunk(content, content_type);
+        .write_str(content, content_type);
     }
 }
 

src/rewritable_units/element.rs

@@ -1,4 +1,7 @@
-use super::{Attribute, AttributeNameError, ContentType, EndTag, Mutations, StartTag, StringChunk};
+use super::{
+    Attribute, AttributeNameError, ContentType, EndTag, Mutations, StartTag, StreamingHandler,
+    StringChunk,
+};
 use crate::base::Bytes;
 use crate::rewriter::{HandlerTypes, LocalHandlerTypes};
 use encoding_rs::Encoding;
@@ -240,6 +243,18 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
             .push_back((content, content_type).into());
     }
 
+    /// Inserts  content from a [`StreamingHandler`] before the element.
+    ///
+    /// Consequent calls to the method append to the previously inserted content.
+    ///
+    /// Use the [`streaming!`] macro to make a `StreamingHandler` from a closure.
+    pub fn streaming_before(&mut self, string_writer: Box<dyn StreamingHandler>) {
+        self.start_tag
+            .mutations
+            .content_before
+            .push_back(string_writer.into());
+    }
+
     /// Inserts `content` after the element.
     ///
     /// Consequent calls to the method prepend `content` to the previously inserted content.
@@ -282,6 +297,16 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
         .push_front(chunk);
     }
 
+    /// Inserts content from a [`StreamingHandler`] after the element.
+    ///
+    /// Consequent calls to the method prepend to the previously inserted content.
+    ///
+    ///
+    /// Use the [`streaming!`] macro to make a `StreamingHandler` from a closure.
+    pub fn streaming_after(&mut self, string_writer: Box<dyn StreamingHandler>) {
+        self.after_chunk(string_writer.into());
+    }
+
     /// Prepends `content` to the element's inner content, i.e. inserts content right after
     /// the element's start tag.
     ///
@@ -328,6 +353,20 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
         }
     }
 
+    /// Prepends content from a [`StreamingHandler`] to the element's inner content,
+    /// i.e. inserts content right after the element's start tag.
+    ///
+    /// Consequent calls to the method prepend to the previously inserted content.
+    /// A call to the method doesn't make any effect if the element is an [empty element].
+    ///
+    /// [empty element]: https://developer.mozilla.org/en-US/docs/Glossary/Empty_element
+    ///
+    ///
+    /// Use the [`streaming!`] macro to make a `StreamingHandler` from a closure.
+    pub fn streaming_prepend(&mut self, string_writer: Box<dyn StreamingHandler>) {
+        self.prepend_chunk(string_writer.into());
+    }
+
     /// Appends `content` to the element's inner content, i.e. inserts content right before
     /// the element's end tag.
     ///
@@ -374,6 +413,19 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
         }
     }
 
+    /// Appends content from a [`StreamingHandler`] to the element's inner content,
+    /// i.e. inserts content right before the element's end tag.
+    ///
+    /// Consequent calls to the method append to the previously inserted content.
+    /// A call to the method doesn't make any effect if the element is an [empty element].
+    ///
+    /// [empty element]: https://developer.mozilla.org/en-US/docs/Glossary/Empty_element
+    ///
+    /// Use the [`streaming!`] macro to make a `StreamingHandler` from a closure.
+    pub fn streaming_append(&mut self, string_writer: Box<dyn StreamingHandler>) {
+        self.append_chunk(string_writer.into());
+    }
+
     /// Replaces inner content of the element with `content`.
     ///
     /// Consequent calls to the method overwrite previously inserted content.
@@ -420,6 +472,19 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
         }
     }
 
+    /// Replaces inner content of the element with content from a [`StreamingHandler`].
+    ///
+    /// Consequent calls to the method overwrite previously inserted content.
+    /// A call to the method doesn't make any effect if the element is an [empty element].
+    ///
+    /// [empty element]: https://developer.mozilla.org/en-US/docs/Glossary/Empty_element
+    ///
+    ///
+    /// Use the [`streaming!`] macro to make a `StreamingHandler` from a closure.
+    pub fn streaming_set_inner_content(&mut self, string_writer: Box<dyn StreamingHandler>) {
+        self.set_inner_content_chunk(string_writer.into());
+    }
+
     /// Replaces the element and its inner content with `content`.
     ///
     /// Consequent calls to the method overwrite previously inserted content.
@@ -461,6 +526,16 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
         }
     }
 
+    /// Replaces the element and its inner content with content from a [`StreamingHandler`].
+    ///
+    /// Consequent calls to the method overwrite previously inserted content.
+    ///
+    ///
+    /// Use the [`streaming!`] macro to make a `StreamingHandler` from a closure.
+    pub fn streaming_replace(&mut self, string_writer: Box<dyn StreamingHandler>) {
+        self.replace_chunk(string_writer.into());
+    }
+
     /// Removes the element and its inner content.
     #[inline]
     pub fn remove(&mut self) {
@@ -629,6 +704,7 @@ mod tests {
     use crate::rewritable_units::test_utils::*;
     use crate::*;
     use encoding_rs::{Encoding, EUC_JP, UTF_8};
+    use rewritable_units::StreamingHandlerSink;
 
     fn rewrite_element(
         html: &[u8],
@@ -651,7 +727,11 @@ mod tests {
                     el.before("[before: should be removed]", ContentType::Text);
                     el.after("[after: should be removed]", ContentType::Text);
                     el.append("[append: should be removed]", ContentType::Text);
-                    el.before("[before: should be removed]", ContentType::Text);
+                    el.streaming_before(Box::new(|sink: &mut StreamingHandlerSink<'_>| {
+                        sink.write_str("[before:", ContentType::Text);
+                        sink.write_str(" should be removed]", ContentType::Text);
+                        Ok(())
+                    }));
                     Ok(())
                 }),
             ],
@@ -953,7 +1033,10 @@ mod tests {
             encoded("<div><span>Hi<inner-remove-me>RemoveŴ</inner-remove-me></span></div>")
         {
             let output = rewrite_element(&html, enc, "span", |el| {
-                el.prepend("<prepended>", ContentType::Html);
+                el.streaming_prepend(streaming!(|s| {
+                    s.write_utf8_chunk(b"<prepended>", ContentType::Html)?;
+                    Ok(())
+                }));
                 el.append("<appended>", ContentType::Html);
                 el.set_inner_content("<imgŵ>", ContentType::Html);
                 el.set_inner_content("<imgŵ>", ContentType::Text);
@@ -1087,7 +1170,17 @@ mod tests {
     #[test]
     fn self_closing_element() {
         let output = rewrite_element(b"<svg><foo/>Hi</foo></svg>", UTF_8, "foo", |el| {
-            el.after("<!--after-->", ContentType::Html);
+            el.after("->", ContentType::Html);
+            el.streaming_after(streaming!(|sink| {
+                sink.write_str("er-", ContentType::Html);
+                Ok(())
+            }));
+            el.after("t", ContentType::Html);
+            el.streaming_after(streaming!(|sink| {
+                sink.write_str("af", ContentType::Html);
+                Ok(())
+            }));
+            el.after("<!--", ContentType::Html);
             el.set_tag_name("bar").unwrap();
         });
 

src/rewritable_units/mod.rs

@@ -2,9 +2,9 @@ use std::any::Any;
 
 pub use self::document_end::*;
 pub use self::element::*;
-pub use self::mutations::ContentType;
+pub use self::mutations::{ContentType, StreamingHandler};
 pub(crate) use self::mutations::{Mutations, StringChunk};
-pub use self::text_encoder::StreamingHandlerSink;
+pub use self::text_encoder::{StreamingHandlerSink, Utf8Error};
 pub use self::tokens::*;
 
 /// Data that can be attached to a rewritable unit by a user and shared between content handler

src/rewritable_units/mutations.rs

@@ -1,6 +1,7 @@
-use super::text_encoder::StreamingHandlerSink;
+use super::StreamingHandlerSink;
 use encoding_rs::Encoding;
 use std::error::Error as StdError;
+use std::panic::{RefUnwindSafe, UnwindSafe};
 
 type BoxResult = Result<(), Box<dyn StdError + Send + Sync>>;
 
@@ -64,6 +65,7 @@ impl From<(&str, ContentType)> for StringChunk {
 
 pub(crate) enum StringChunk {
     Buffer(Box<str>, ContentType),
+    Stream(Box<dyn StreamingHandler>),
 }
 
 #[derive(Default)]
@@ -96,10 +98,59 @@ impl DynamicString {
         for chunk in self.chunks {
             match chunk {
                 StringChunk::Buffer(content, content_type) => {
-                    sink.write_str_chunk(&content, content_type);
+                    sink.write_str(&content, content_type);
+                }
+                StringChunk::Stream(handler) => {
+                    handler.write_all(sink)?;
                 }
             };
         }
         Ok(())
     }
 }
+
+/// A callback used to write content asynchronously.
+pub trait StreamingHandler: Send {
+    /// This method is called only once, and is expected to write content
+    /// by calling the [`sink.write_str()`](StreamingHandlerSink::write_str) one or more times.
+    ///
+    /// Multiple calls to `sink.write_str()` append more content to the output.
+    ///
+    /// See [`StreamingHandlerSink`].
+    fn write_all(self: Box<Self>, sink: &mut StreamingHandlerSink<'_>) -> BoxResult;
+
+    // Safety: due to lack of Sync, this trait must not have `&self` methods
+}
+
+/// Avoid requring `StreamingHandler` to be `Sync`.
+/// It only has a method taking exclusive ownership, so there's no sharing possible.
+unsafe impl Sync for StringChunk {}
+impl RefUnwindSafe for StringChunk {}
+impl UnwindSafe for StringChunk {}
+
+impl<F> From<F> for Box<dyn StreamingHandler>
+where
+    F: FnOnce(&mut StreamingHandlerSink<'_>) -> BoxResult + Send + 'static,
+{
+    #[inline]
+    fn from(f: F) -> Self {
+        Box::new(f)
+    }
+}
+
+impl<F> StreamingHandler for F
+where
+    F: FnOnce(&mut StreamingHandlerSink<'_>) -> BoxResult + Send + 'static,
+{
+    #[inline]
+    fn write_all(self: Box<F>, sink: &mut StreamingHandlerSink<'_>) -> BoxResult {
+        (self)(sink)
+    }
+}
+
+impl From<Box<dyn StreamingHandler>> for StringChunk {
+    #[inline]
+    fn from(writer: Box<dyn StreamingHandler>) -> Self {
+        Self::Stream(writer)
+    }
+}