Tell OutputSink what encoding is used

Author: kornelski

src/rewriter/mod.rs

@@ -343,7 +343,7 @@ mod tests {
     use crate::html::TextType;
     use crate::html_content::ContentType;
     use crate::test_utils::{ASCII_COMPATIBLE_ENCODINGS, NON_ASCII_COMPATIBLE_ENCODINGS, Output};
-    use encoding_rs::Encoding;
+    use encoding_rs::{Encoding, WINDOWS_1252};
     use itertools::Itertools;
     use static_assertions::assert_impl_all;
     use std::convert::TryInto;
@@ -835,8 +835,28 @@ mod tests {
     #[test]
     fn test_charset_switch_latency() {
         let html = b"<title>\xC3\xB0</title>\xC3\xB0<meta charset=latin1 attr='\xC3\xB0'>\xF0<meta attr='\xF0'>\xF0";
-        let rewritten = rewrite_html_bytes(
-            html,
+
+        struct Sink {
+            out: Vec<u8>,
+            charsets: Vec<(usize, AsciiCompatibleEncoding)>,
+        }
+
+        impl OutputSink for &mut Sink {
+            fn handle_chunk(&mut self, chunk: &[u8]) {
+                self.out.extend_from_slice(chunk);
+            }
+
+            fn set_encoding(&mut self, enc: AsciiCompatibleEncoding) {
+                self.charsets.push((self.out.len(), enc));
+            }
+        }
+
+        let mut sink = Sink {
+            out: Vec::with_capacity(html.len()),
+            charsets: vec![],
+        };
+
+        let mut rewriter = HtmlRewriter::new(
             Settings {
                 element_content_handlers: vec![element!("[attr]", |el| {
                     assert_eq!(el.get_attribute("attr").unwrap(), "ð");
@@ -850,8 +870,20 @@ mod tests {
                 adjust_charset_on_meta_tag: true,
                 ..Default::default()
             },
+            &mut sink,
+        );
+
+        rewriter.write(html).unwrap();
+        rewriter.end().unwrap();
+
+        assert_eq!(html, sink.out.as_slice());
+        assert_eq!(
+            &[
+                (0, AsciiCompatibleEncoding::utf_8()),
+                (50, WINDOWS_1252.try_into().unwrap())
+            ],
+            sink.charsets.as_slice()
         );
-        assert_eq!(html, rewritten.as_slice());
     }
 
     #[test]

src/transform_stream/dispatcher.rs

@@ -59,6 +59,11 @@ pub trait OutputSink {
     /// # Note
     /// The last chunk of the output has zero length.
     fn handle_chunk(&mut self, chunk: &[u8]);
+
+    /// Called before the first `handle_chunk` and once when `<meta charset>` is applied.
+    ///
+    /// Implementations for closures won't be able to access this method
+    fn set_encoding(&mut self, _new_encoding: AsciiCompatibleEncoding) {}
 }
 
 impl<F: FnMut(&[u8])> OutputSink for F {
@@ -188,11 +193,12 @@ where
     #[inline]
     pub fn new(
         transform_controller: C,
-        output_sink: O,
+        mut output_sink: O,
         encoding: AsciiCompatibleEncoding,
         next_encoding: SharedEncoding,
     ) -> Self {
         let capture_flags = transform_controller.initial_capture_flags();
+        output_sink.set_encoding(encoding);
 
         Self {
             delegate: DispatcherDelegate {
@@ -217,6 +223,7 @@ where
         {
             self.encoding = next_encoding;
             self.text_decoder.set_encoding(next_encoding);
+            self.delegate.output_sink.set_encoding(next_encoding);
         }
     }