perf: optimise embedded formatter (#7567)

Author: arendjr

crates/biome_formatter/src/format_element.rs

@@ -317,6 +317,10 @@ impl BestFittingElement {
         &self.variants
     }
 
+    pub(crate) fn variants_mut(&mut self) -> &mut [Box<[FormatElement]>] {
+        &mut self.variants
+    }
+
     /// Returns the least expanded variant
     pub fn most_flat(&self) -> &[FormatElement] {
         self.variants.first().expect(

crates/biome_formatter/src/format_element/document.rs

@@ -145,66 +145,52 @@ impl Document {
     pub fn as_elements(&self) -> &[FormatElement] {
         &self.elements
     }
-}
-
-pub trait DocumentVisitor {
-    /// Visit an element and optionally return a replacement
-    fn visit_element(&mut self, element: &FormatElement) -> Option<FormatElement>;
-}
 
-/// Applies a visitor to transform elements in the document
-pub struct ElementTransformer;
-
-impl ElementTransformer {
-    /// Visits a mutable [Document] and replaces its internal elements.
-    pub fn transform_document<V: DocumentVisitor>(document: &mut Document, visitor: &mut V) {
-        let elements = std::mem::take(&mut document.elements);
-        document.elements = Self::transform_elements(elements, visitor);
+    /// Transforms the document by visiting every element, optionally replacing
+    /// them.
+    ///
+    /// Accepts a `visitor` that will be called to visit each element, and which
+    /// may optionally return a replacement.
+    ///
+    /// Elements that contain nested elements, such as [FormatElement::Interned]
+    /// and [FormatElement::BestFitting], have the visitor called on their
+    /// nested elements, but not on the elements themselves.
+    pub(crate) fn transform(
+        &mut self,
+        mut visitor: impl FnMut(&FormatElement) -> Option<FormatElement>,
+    ) {
+        transform_elements(&mut self.elements, &mut visitor);
     }
+}
 
-    /// Iterates over each element of the document and map each element to a new element. The new element is
-    /// optionally crated using [DocumentVisitor::visit_element]. If no element is returned, the original element is kept.
-    ///
-    /// Nested data structures such as [FormatElement::Interned] and [FormatElement::BestFitting] use recursion and call
-    /// [Self::transform_elements] again.
-    fn transform_elements<V: DocumentVisitor>(
-        elements: Vec<FormatElement>,
-        visitor: &mut V,
-    ) -> Vec<FormatElement> {
-        elements
-            .into_iter()
-            .map(|element| {
-                // Transform nested elements first
-                let transformed_element = match element {
-                    FormatElement::Interned(interned) => {
-                        let nested_elements = interned.deref().to_vec();
-                        let transformed_nested = Self::transform_elements(nested_elements, visitor);
-                        FormatElement::Interned(Interned::new(transformed_nested))
-                    }
-                    FormatElement::BestFitting(best_fitting) => {
-                        let variants: Vec<Box<[FormatElement]>> = best_fitting
-                            .variants()
-                            .iter()
-                            .map(|variant| {
-                                Self::transform_elements(variant.to_vec(), visitor)
-                                    .into_boxed_slice()
-                            })
-                            .collect();
-                        // SAFETY: Safe because the number of variants is the same after the transformation
-                        unsafe {
-                            FormatElement::BestFitting(BestFittingElement::from_vec_unchecked(
-                                variants,
-                            ))
-                        }
-                    }
-                    other => other,
-                };
-                // Then apply a visitor to the element itself
-                visitor
-                    .visit_element(&transformed_element)
-                    .unwrap_or(transformed_element)
-            })
-            .collect()
+/// Iterates over each of the given `elements` and optionally replaces each
+/// element with a new one.
+///
+/// Nested data structures such as [FormatElement::Interned] and
+/// [FormatElement::BestFitting] use recursion and call [transform_elements()]
+/// again. The visitor is *not* invoked on these elements.
+fn transform_elements(
+    elements: &mut [FormatElement],
+    visitor: &mut impl FnMut(&FormatElement) -> Option<FormatElement>,
+) {
+    for element in elements {
+        match element {
+            FormatElement::Interned(interned) => {
+                let mut nested_elements = interned.deref().to_vec();
+                transform_elements(&mut nested_elements, visitor);
+                *element = FormatElement::Interned(Interned::new(nested_elements));
+            }
+            FormatElement::BestFitting(best_fitting) => {
+                for variant in best_fitting.variants_mut() {
+                    transform_elements(variant, visitor);
+                }
+            }
+            _ => {
+                if let Some(replacement) = visitor(element) {
+                    *element = replacement;
+                }
+            }
+        }
     }
 }
 

crates/biome_formatter/src/lib.rs

@@ -50,9 +50,8 @@ use std::fmt::{Debug, Display};
 use crate::builders::syntax_token_cow_slice;
 use crate::comments::{CommentStyle, Comments, SourceComment};
 pub use crate::diagnostics::{ActualStart, FormatError, InvalidDocumentError, PrintError};
-use crate::format_element::document::{Document, ElementTransformer};
+use crate::format_element::document::Document;
 use crate::format_element::{Interned, LineMode};
-use crate::prelude::document::DocumentVisitor;
 #[cfg(debug_assertions)]
 use crate::printed_tokens::PrintedTokens;
 use crate::printer::{Printer, PrinterOptions};
@@ -892,41 +891,20 @@ impl<Context> Formatted<Context> {
 
     /// Visits each embedded element and replaces it with elements contained inside the [Document]
     /// emitted by `fn_format_embedded`
-    pub fn format_embedded<F>(&mut self, fn_format_embedded: F)
+    pub fn format_embedded<F>(&mut self, mut fn_format_embedded: F)
     where
         F: FnMut(TextRange) -> Option<Document>,
     {
-        let document = &mut self.document;
-
-        struct EmbeddedVisitor<F> {
-            fn_format_embedded: F,
-        }
-
-        impl<F> DocumentVisitor for EmbeddedVisitor<F>
-        where
-            F: FnMut(TextRange) -> Option<Document>,
-        {
-            fn visit_element(&mut self, element: &FormatElement) -> Option<FormatElement> {
-                match element {
-                    FormatElement::Tag(Tag::StartEmbedded(range)) => {
-                        (self.fn_format_embedded)(*range).map(|document| {
-                            FormatElement::Interned(Interned::new(document.into_elements()))
-                        })
-                    }
-                    FormatElement::Tag(Tag::EndEmbedded) => {
-                        // FIXME: this might not play well for all cases, so we need to figure out
-                        // a nicer way to replace the tag
-                        Some(FormatElement::Line(LineMode::Hard))
-                    }
-                    _ => None,
-                }
+        self.document.transform(move |element| match element {
+            FormatElement::Tag(Tag::StartEmbedded(range)) => fn_format_embedded(*range)
+                .map(|document| FormatElement::Interned(Interned::new(document.into_elements()))),
+            FormatElement::Tag(Tag::EndEmbedded) => {
+                // FIXME: this might not play well for all cases, so we need to figure out
+                // a nicer way to replace the tag
+                Some(FormatElement::Line(LineMode::Hard))
             }
-        }
-
-        ElementTransformer::transform_document(
-            document,
-            &mut EmbeddedVisitor { fn_format_embedded },
-        );
+            _ => None,
+        });
     }
 
     /// Returns the formatted document.
@@ -938,10 +916,6 @@ impl<Context> Formatted<Context> {
     pub fn into_document(self) -> Document {
         self.document
     }
-
-    pub fn swap_document(&mut self, document: Document) {
-        self.document = document;
-    }
 }
 
 impl<Context> Formatted<Context>