Make mutations struct smaller

cloudflare · Nov 6, 2024 · 8a7d8b7 · 8a7d8b7
1 parent bfd7b80
commit 8a7d8b7
Show file tree

Hide file tree

Showing 8 changed files with 118 additions and 60 deletions.
diff --git a/src/rewritable_units/element.rs b/src/rewritable_units/element.rs
@@ -1,3 +1,4 @@
+use super::mutations::MutationsInner;
 use super::{Attribute, AttributeNameError, ContentType, EndTag, Mutations, StartTag, StringChunk};
 use crate::base::Bytes;
 use crate::rewriter::{HandlerTypes, LocalHandlerTypes};
@@ -88,19 +89,18 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
 
     #[inline]
     fn remove_content(&mut self) {
-        self.start_tag.mutations.content_after.clear();
-        if let Some(end) = &mut self.end_tag_mutations {
+        self.start_tag.mutations.mutate().content_after.clear();
+        if let Some(end) = self.end_tag_mutations.as_mut().and_then(|m| m.if_mutated()) {
             end.content_before.clear();
         }
         self.should_remove_content = true;
     }
 
     #[inline]
-    fn end_tag_mutations_mut(&mut self) -> &mut Mutations {
-        let encoding = self.encoding;
-
+    fn end_tag_mutations_mut(&mut self) -> &mut MutationsInner {
         self.end_tag_mutations
-            .get_or_insert_with(|| Mutations::new(encoding))
+            .get_or_insert_with(Mutations::new)
+            .mutate()
     }
 
     /// Returns the tag name of the element.
@@ -236,6 +236,7 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
     pub fn before(&mut self, content: &str, content_type: ContentType) {
         self.start_tag
             .mutations
+            .mutate()
             .content_before
             .push_back((content, content_type).into());
     }
@@ -277,7 +278,7 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
         if self.can_have_content {
             &mut self.end_tag_mutations_mut().content_after
         } else {
-            &mut self.start_tag.mutations.content_after
+            &mut self.start_tag.mutations.mutate().content_after
         }
         .push_front(chunk);
     }
@@ -324,7 +325,11 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
 
     fn prepend_chunk(&mut self, chunk: StringChunk) {
         if self.can_have_content {
-            self.start_tag.mutations.content_after.push_front(chunk);
+            self.start_tag
+                .mutations
+                .mutate()
+                .content_after
+                .push_front(chunk);
         }
     }
 
@@ -416,7 +421,11 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
     fn set_inner_content_chunk(&mut self, chunk: StringChunk) {
         if self.can_have_content {
             self.remove_content();
-            self.start_tag.mutations.content_after.push_front(chunk);
+            self.start_tag
+                .mutations
+                .mutate()
+                .content_after
+                .push_front(chunk);
         }
     }
 
@@ -453,7 +462,7 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
     }
 
     fn replace_chunk(&mut self, chunk: StringChunk) {
-        self.start_tag.mutations.replace(chunk);
+        self.start_tag.mutations.mutate().replace(chunk);
 
         if self.can_have_content {
             self.remove_content();
@@ -464,7 +473,7 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
     /// Removes the element and its inner content.
     #[inline]
     pub fn remove(&mut self) {
-        self.start_tag.mutations.remove();
+        self.start_tag.mutations.mutate().remove();
 
         if self.can_have_content {
             self.remove_content();
@@ -497,7 +506,7 @@ impl<'r, 't, H: HandlerTypes> Element<'r, 't, H> {
     /// ```
     #[inline]
     pub fn remove_and_keep_content(&mut self) {
-        self.start_tag.mutations.remove();
+        self.start_tag.remove();
 
         if self.can_have_content {
             self.end_tag_mutations_mut().remove();

diff --git a/src/rewritable_units/mutations.rs b/src/rewritable_units/mutations.rs
@@ -1,5 +1,4 @@
 use super::text_encoder::StreamingHandlerSink;
-use encoding_rs::Encoding;
 use std::error::Error as StdError;
 
 type BoxResult = Result<(), Box<dyn StdError + Send + Sync>>;
@@ -16,27 +15,14 @@ pub enum ContentType {
     Text,
 }
 
-pub(crate) struct Mutations {
+pub(crate) struct MutationsInner {
     pub content_before: DynamicString,
     pub replacement: DynamicString,
     pub content_after: DynamicString,
     pub removed: bool,
-    pub encoding: &'static Encoding,
 }
 
-impl Mutations {
-    #[inline]
-    #[must_use]
-    pub const fn new(encoding: &'static Encoding) -> Self {
-        Self {
-            content_before: DynamicString::new(),
-            replacement: DynamicString::new(),
-            content_after: DynamicString::new(),
-            removed: false,
-            encoding,
-        }
-    }
-
+impl MutationsInner {
     #[inline]
     pub fn replace(&mut self, chunk: StringChunk) {
         self.remove();
@@ -48,10 +34,52 @@ impl Mutations {
     pub fn remove(&mut self) {
         self.removed = true;
     }
+}
+
+pub(crate) struct Mutations {
+    inner: Option<Box<MutationsInner>>,
+}
+
+impl Mutations {
+    #[inline]
+    #[must_use]
+    pub const fn new() -> Self {
+        Self { inner: None }
+    }
+
+    #[inline]
+    pub fn take(&mut self) -> Option<Box<MutationsInner>> {
+        self.inner.take()
+    }
+
+    #[inline]
+    pub fn if_mutated(&mut self) -> Option<&mut MutationsInner> {
+        self.inner.as_deref_mut()
+    }
+
+    #[inline]
+    pub fn mutate(&mut self) -> &mut MutationsInner {
+        #[inline(never)]
+        fn alloc_content(inner: &mut Option<Box<MutationsInner>>) -> &mut MutationsInner {
+            inner.get_or_insert_with(move || {
+                Box::new(MutationsInner {
+                    content_before: DynamicString::new(),
+                    replacement: DynamicString::new(),
+                    content_after: DynamicString::new(),
+                    removed: false,
+                })
+            })
+        }
+
+        match &mut self.inner {
+            Some(inner) => inner,
+            uninit => alloc_content(uninit),
+        }
+    }
 
     #[inline]
-    pub const fn removed(&self) -> bool {
-        self.removed
+    pub fn removed(&self) -> bool {
+        self.inner.as_ref().is_some_and(|inner| inner.removed)
     }
 }
 

diff --git a/src/rewritable_units/tokens/comment.rs b/src/rewritable_units/tokens/comment.rs
@@ -43,7 +43,7 @@ impl<'i> Comment<'i> {
             text,
             raw: Some(raw),
             encoding,
-            mutations: Mutations::new(encoding),
+            mutations: Mutations::new(),
             user_data: Box::new(()),
         })
     }
@@ -107,6 +107,7 @@ impl<'i> Comment<'i> {
     #[inline]
     pub fn before(&mut self, content: &str, content_type: crate::rewritable_units::ContentType) {
         self.mutations
+            .mutate()
             .content_before
             .push_back((content, content_type).into());
     }
@@ -141,6 +142,7 @@ impl<'i> Comment<'i> {
     #[inline]
     pub fn after(&mut self, content: &str, content_type: crate::rewritable_units::ContentType) {
         self.mutations
+            .mutate()
             .content_after
             .push_front((content, content_type).into());
     }
@@ -174,13 +176,15 @@ impl<'i> Comment<'i> {
     /// ```
     #[inline]
     pub fn replace(&mut self, content: &str, content_type: crate::rewritable_units::ContentType) {
-        self.mutations.replace((content, content_type).into());
+        self.mutations
+            .mutate()
+            .replace((content, content_type).into());
     }
 
     /// Removes the comment.
     #[inline]
     pub fn remove(&mut self) {
-        self.mutations.remove();
+        self.mutations.mutate().remove();
     }
 
     /// Returns `true` if the comment has been replaced or removed.

diff --git a/src/rewritable_units/tokens/end_tag.rs b/src/rewritable_units/tokens/end_tag.rs
@@ -27,7 +27,7 @@ impl<'i> EndTag<'i> {
             name,
             raw: Some(raw),
             encoding,
-            mutations: Mutations::new(encoding),
+            mutations: Mutations::new(),
         })
     }
 
@@ -70,6 +70,7 @@ impl<'i> EndTag<'i> {
     #[inline]
     pub fn before(&mut self, content: &str, content_type: ContentType) {
         self.mutations
+            .mutate()
             .content_before
             .push_back((content, content_type).into());
     }
@@ -80,6 +81,7 @@ impl<'i> EndTag<'i> {
     #[inline]
     pub fn after(&mut self, content: &str, content_type: ContentType) {
         self.mutations
+            .mutate()
             .content_after
             .push_front((content, content_type).into());
     }
@@ -89,13 +91,15 @@ impl<'i> EndTag<'i> {
     /// Consequent calls to the method overwrite previous replacement content.
     #[inline]
     pub fn replace(&mut self, content: &str, content_type: ContentType) {
-        self.mutations.replace((content, content_type).into());
+        self.mutations
+            .mutate()
+            .replace((content, content_type).into());
     }
 
     /// Removes the end tag.
     #[inline]
     pub fn remove(&mut self) {
-        self.mutations.remove();
+        self.mutations.mutate().remove();
     }
 
     #[inline]

diff --git a/src/rewritable_units/tokens/mod.rs b/src/rewritable_units/tokens/mod.rs
@@ -22,27 +22,32 @@ macro_rules! impl_serialize {
                 output_handler: &mut dyn FnMut(&[u8]),
             ) -> Result<(), crate::errors::RewritingError> {
                 let mut encoder = crate::rewritable_units::text_encoder::StreamingHandlerSink::new(
-                    self.mutations.encoding,
+                    self.encoding,
                     output_handler,
                 );
-                let content_before = ::std::mem::take(&mut self.mutations.content_before);
-                content_before
-                    .encode(&mut encoder)
-                    .map_err(crate::errors::RewritingError::ContentHandlerError)?;
+                match self.mutations.take() {
+                    None => self.serialize_self(encoder.output_handler()),
+                    Some(mutations) => {
+                        mutations
+                            .content_before
+                            .encode(&mut encoder)
+                            .map_err(crate::errors::RewritingError::ContentHandlerError)?;
 
-                if !self.mutations.removed {
-                    self.serialize_self(encoder.output_handler())?;
-                } else {
-                    self.mutations
-                        .replacement
-                        .encode(&mut encoder)
-                        .map_err(crate::errors::RewritingError::ContentHandlerError)?;
-                }
+                        if !mutations.removed {
+                            self.serialize_self(encoder.output_handler())?;
+                        } else {
+                            mutations
+                                .replacement
+                                .encode(&mut encoder)
+                                .map_err(crate::errors::RewritingError::ContentHandlerError)?;
+                        }
 
-                self.mutations
-                    .content_after
-                    .encode(&mut encoder)
-                    .map_err(crate::errors::RewritingError::ContentHandlerError)
+                        mutations
+                            .content_after
+                            .encode(&mut encoder)
+                            .map_err(crate::errors::RewritingError::ContentHandlerError)
+                    }
+                }
             }
         }
     };

diff --git a/src/rewritable_units/tokens/start_tag.rs b/src/rewritable_units/tokens/start_tag.rs
@@ -38,7 +38,7 @@ impl<'i> StartTag<'i> {
             self_closing,
             raw: Some(raw),
             encoding,
-            mutations: Mutations::new(encoding),
+            mutations: Mutations::new(),
         })
     }
 
@@ -113,6 +113,7 @@ impl<'i> StartTag<'i> {
     #[inline]
     pub fn before(&mut self, content: &str, content_type: ContentType) {
         self.mutations
+            .mutate()
             .content_before
             .push_back((content, content_type).into());
     }
@@ -123,6 +124,7 @@ impl<'i> StartTag<'i> {
     #[inline]
     pub fn after(&mut self, content: &str, content_type: ContentType) {
         self.mutations
+            .mutate()
             .content_after
             .push_front((content, content_type).into());
     }
@@ -132,13 +134,15 @@ impl<'i> StartTag<'i> {
     /// Consequent calls to the method overwrite previous replacement content.
     #[inline]
     pub fn replace(&mut self, content: &str, content_type: ContentType) {
-        self.mutations.replace((content, content_type).into());
+        self.mutations
+            .mutate()
+            .replace((content, content_type).into());
     }
 
     /// Removes the start tag.
     #[inline]
     pub fn remove(&mut self) {
-        self.mutations.remove();
+        self.mutations.mutate().remove();
     }
 
     fn serialize_self(&self, output_handler: &mut dyn FnMut(&[u8])) -> Result<(), RewritingError> {