Store a reference to the OutputSink in DocumentEnd.

src/lib.rs

@@ -51,7 +51,7 @@ pub mod errors {
 /// HTML content descriptors that can be produced and modified by a rewriter.
 pub mod html_content {
     pub use super::rewritable_units::{
-        Attribute, Comment, ContentType, Doctype, Element, TextChunk, UserData,
+        Attribute, Comment, ContentType, Doctype, DocumentEnd, Element, TextChunk, UserData,
     };
 
     pub use super::html::TextType;

src/rewritable_units/document_end.rs

@@ -1,25 +1,56 @@
 use super::mutations::content_to_bytes;
 use super::ContentType;
+
 use encoding_rs::Encoding;
 
+use crate::transform_stream::OutputSink;
+
+/// A rewritable unit that represents the end of the document.
+///
+/// This exposes the [append](#method.append) function that can be used to append content at the
+/// end of the document. The content will only be appended after the rewriter has finished processing
+/// the final chunk.
 pub struct DocumentEnd<'a> {
-    output_handler: &'a mut dyn FnMut(&[u8]),
+    output_sink: &'a mut dyn OutputSink,
     encoding: &'static Encoding,
 }
 
 impl<'a> DocumentEnd<'a> {
-    pub(crate) fn new(
-        output_handler: &'a mut dyn FnMut(&[u8]),
-        encoding: &'static Encoding,
-    ) -> Self {
+    pub(crate) fn new(output_sink: &'a mut dyn OutputSink, encoding: &'static Encoding) -> Self {
         DocumentEnd {
-            output_handler,
+            output_sink,
             encoding,
         }
     }
 
+    /// Appends `content` at the end of the document.
+    ///
+    /// Subsequent calls to this method append `content` to the previously inserted content.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use cool_thing::{end, rewrite_str, RewriteStrSettings};
+    /// use cool_thing::html_content::{ContentType, DocumentEnd};
+    ///
+    /// let html = rewrite_str(
+    ///     r#"<div id="foo"><!-- content --></div><img>"#,
+    ///     RewriteStrSettings {
+    ///         document_content_handlers: vec![end!(|end| {
+    ///             end.append("<bar>", ContentType::Html);
+    ///             end.append("<baz>", ContentType::Text);
+    ///             Ok(())
+    ///         })],
+    ///         ..RewriteStrSettings::default()
+    ///     }
+    /// ).unwrap();
+    ///
+    /// assert_eq!(html, r#"<div id="foo"><!-- content --></div><img><bar>&lt;baz&gt;"#);
+    /// ```
     #[inline]
     pub fn append(&mut self, content: &str, content_type: ContentType) {
-        content_to_bytes(content, content_type, self.encoding, self.output_handler);
+        content_to_bytes(content, content_type, self.encoding, &mut |c: &[u8]| {
+            self.output_sink.handle_chunk(c)
+        });
     }
 }

src/rewriter/settings.rs

@@ -94,6 +94,7 @@ impl<'h> DocumentContentHandlers<'h> {
         self
     }
 
+    /// Sets a handler for the document end, which is called after the last chunk is processed.
     #[inline]
     pub fn end(mut self, handler: impl FnMut(&mut DocumentEnd) -> HandlerResult + 'h) -> Self {
         self.end = Some(Box::new(handler));
@@ -308,6 +309,38 @@ macro_rules! doc_comments {
     };
 }
 
+/// A convenience macro to construct a rewriting handler for the end of the document.
+///
+/// This handler will only be called after the rewriter has finished processing the final chunk.
+///
+/// # Example
+/// ```
+/// use cool_thing::{rewrite_str, element, end, RewriteStrSettings};
+/// use cool_thing::html_content::ContentType;
+///
+/// let html = rewrite_str(
+///     r#"<span>foo</span>"#,
+///     RewriteStrSettings {
+///         element_content_handlers: vec![
+///             element!("span", |el| {
+///                 el.append("bar", ContentType::Text);
+///
+///                 Ok(())
+///             })
+///         ],
+///         document_content_handlers: vec![
+///             end!(|end| {
+///                 end.append("<div>baz</div>", ContentType::Html);
+///
+///                 Ok(())
+///             })
+///         ],
+///         ..RewriteStrSettings::default()
+///     }
+/// ).unwrap();
+///
+/// assert_eq!(html, r#"<span>foobar</span><div>baz</div>"#);
+/// ```
 #[macro_export(local_inner_macros)]
 macro_rules! end {
     ($handler:expr) => {

src/selectors_vm/mod.rs

@@ -530,7 +530,7 @@ mod tests {
     use crate::errors::RewritingError;
     use crate::html::Namespace;
     use crate::memory::MemoryLimiter;
-    use crate::rewritable_units::{Token, TokenCaptureFlags};
+    use crate::rewritable_units::{DocumentEnd, Token, TokenCaptureFlags};
     use crate::transform_stream::{
         StartTagHandlingResult, TransformController, TransformStream, TransformStreamSettings,
     };
@@ -591,6 +591,10 @@ mod tests {
                 TokenCaptureFlags::all()
             }
 
+            fn handle_end(&mut self, _document_end: &mut DocumentEnd) -> Result<(), RewritingError> {
+                Ok(())
+            }
+
             fn handle_token(&mut self, token: &mut Token) -> Result<(), RewritingError> {
                 (self.0)(token);
                 Ok(())

src/transform_stream/dispatcher.rs

@@ -108,9 +108,7 @@ where
     pub fn finish(&mut self, input: &[u8]) -> Result<(), RewritingError> {
         self.flush_remaining_input(input, input.len());
 
-        let output_sink = &mut self.output_sink;
-        let mut output_sink_handle_chunk = |c: &[u8]| output_sink.handle_chunk(c);
-        let mut document_end = DocumentEnd::new(&mut output_sink_handle_chunk, self.encoding);
+        let mut document_end = DocumentEnd::new(&mut self.output_sink, self.encoding);
 
         self.transform_controller.handle_end(&mut document_end)?;