perf: Reduce best fitting allocations (#7411)

crates/ruff_formatter/src/builders.rs

@@ -9,9 +9,7 @@ use Tag::*;
 use crate::format_element::tag::{Condition, Tag};
 use crate::prelude::tag::{DedentMode, GroupMode, LabelId};
 use crate::prelude::*;
-use crate::{
-    format_element, write, Argument, Arguments, FormatContext, FormatOptions, GroupId, TextSize,
-};
+use crate::{write, Argument, Arguments, FormatContext, FormatOptions, GroupId, TextSize};
 use crate::{Buffer, VecBuffer};
 
 /// A line break that only gets printed if the enclosing `Group` doesn't fit on a single line.
@@ -2543,25 +2541,20 @@ impl<Context> Format<Context> for BestFitting<'_, Context> {
     fn fmt(&self, f: &mut Formatter<Context>) -> FormatResult<()> {
         let variants = self.variants.items();
 
-        let mut formatted_variants = Vec::with_capacity(variants.len());
+        let mut buffer = VecBuffer::with_capacity(variants.len() * 8, f.state_mut());
 
         for variant in variants {
-            let mut buffer = VecBuffer::with_capacity(8, f.state_mut());
-            buffer.write_element(FormatElement::Tag(StartEntry));
+            buffer.write_element(FormatElement::Tag(StartBestFittingEntry));
             buffer.write_fmt(Arguments::from(variant))?;
-            buffer.write_element(FormatElement::Tag(EndEntry));
-
-            formatted_variants.push(buffer.into_vec().into_boxed_slice());
+            buffer.write_element(FormatElement::Tag(EndBestFittingEntry));
         }
 
         // SAFETY: The constructor guarantees that there are always at least two variants. It's, therefore,
         // safe to call into the unsafe `from_vec_unchecked` function
         #[allow(unsafe_code)]
         let element = unsafe {
             FormatElement::BestFitting {
-                variants: format_element::BestFittingVariants::from_vec_unchecked(
-                    formatted_variants,
-                ),
+                variants: BestFittingVariants::from_vec_unchecked(buffer.into_vec()),
                 mode: self.mode,
             }
         };

crates/ruff_formatter/src/format_element.rs

@@ -3,6 +3,7 @@ pub mod tag;
 
 use std::borrow::Cow;
 use std::hash::{Hash, Hasher};
+use std::iter::FusedIterator;
 use std::num::NonZeroU32;
 use std::ops::Deref;
 use std::rc::Rc;
@@ -67,6 +68,16 @@ pub enum FormatElement {
     Tag(Tag),
 }
 
+impl FormatElement {
+    pub fn tag_kind(&self) -> Option<TagKind> {
+        if let FormatElement::Tag(tag) = self {
+            Some(tag.kind())
+        } else {
+            None
+        }
+    }
+}
+
 impl std::fmt::Debug for FormatElement {
     fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
@@ -318,7 +329,7 @@ pub enum BestFittingMode {
 /// The first element is the one that takes up the most space horizontally (the most flat),
 /// The last element takes up the least space horizontally (but most horizontal space).
 #[derive(Clone, Eq, PartialEq, Debug)]
-pub struct BestFittingVariants(Box<[Box<[FormatElement]>]>);
+pub struct BestFittingVariants(Box<[FormatElement]>);
 
 impl BestFittingVariants {
     /// Creates a new best fitting IR with the given variants. The method itself isn't unsafe
@@ -331,9 +342,13 @@ impl BestFittingVariants {
     /// The slice must contain at least two variants.
     #[doc(hidden)]
     #[allow(unsafe_code)]
-    pub unsafe fn from_vec_unchecked(variants: Vec<Box<[FormatElement]>>) -> Self {
+    pub unsafe fn from_vec_unchecked(variants: Vec<FormatElement>) -> Self {
         debug_assert!(
-            variants.len() >= 2,
+            variants
+                .iter()
+                .filter(|element| matches!(element, FormatElement::Tag(Tag::StartBestFittingEntry)))
+                .count()
+                >= 2,
             "Requires at least the least expanded and most expanded variants"
         );
 
@@ -342,40 +357,85 @@ impl BestFittingVariants {
 
     /// Returns the most expanded variant
     pub fn most_expanded(&self) -> &[FormatElement] {
-        self.0.last().expect(
-            "Most contain at least two elements, as guaranteed by the best fitting builder.",
-        )
+        self.into_iter().last().unwrap()
     }
 
-    pub fn as_slice(&self) -> &[Box<[FormatElement]>] {
+    pub fn as_slice(&self) -> &[FormatElement] {
         &self.0
     }
 
     /// Returns the least expanded variant
     pub fn most_flat(&self) -> &[FormatElement] {
-        self.0.first().expect(
-            "Most contain at least two elements, as guaranteed by the best fitting builder.",
-        )
+        self.into_iter().next().unwrap()
     }
 }
 
 impl Deref for BestFittingVariants {
-    type Target = [Box<[FormatElement]>];
+    type Target = [FormatElement];
 
     fn deref(&self) -> &Self::Target {
         self.as_slice()
     }
 }
 
+pub struct BestFittingVariantsIter<'a> {
+    elements: &'a [FormatElement],
+}
+
 impl<'a> IntoIterator for &'a BestFittingVariants {
-    type Item = &'a Box<[FormatElement]>;
-    type IntoIter = std::slice::Iter<'a, Box<[FormatElement]>>;
+    type Item = &'a [FormatElement];
+    type IntoIter = BestFittingVariantsIter<'a>;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.as_slice().iter()
+        BestFittingVariantsIter { elements: &self.0 }
+    }
+}
+
+impl<'a> Iterator for BestFittingVariantsIter<'a> {
+    type Item = &'a [FormatElement];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.elements.first()? {
+            FormatElement::Tag(Tag::StartBestFittingEntry) => {
+                let end = self
+                    .elements
+                    .iter()
+                    .position(|element| {
+                        matches!(element, FormatElement::Tag(Tag::EndBestFittingEntry))
+                    })
+                    .map_or(self.elements.len(), |position| position + 1);
+
+                let (variant, rest) = self.elements.split_at(end);
+                self.elements = rest;
+
+                Some(variant)
+            }
+            _ => None,
+        }
+    }
+
+    fn last(mut self) -> Option<Self::Item>
+    where
+        Self: Sized,
+    {
+        self.next_back()
     }
 }
 
+impl<'a> DoubleEndedIterator for BestFittingVariantsIter<'a> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        let start_position = self.elements.iter().rposition(|element| {
+            matches!(element, FormatElement::Tag(Tag::StartBestFittingEntry))
+        })?;
+
+        let (rest, variant) = self.elements.split_at(start_position);
+        self.elements = rest;
+        Some(variant)
+    }
+}
+
+impl FusedIterator for BestFittingVariantsIter<'_> {}
+
 pub trait FormatElements {
     /// Returns true if this [`FormatElement`] is guaranteed to break across multiple lines by the printer.
     /// This is the case if this format element recursively contains a:

crates/ruff_formatter/src/format_element/document.rs

@@ -35,15 +35,18 @@ impl Document {
         enum Enclosing<'a> {
             Group(&'a tag::Group),
             ConditionalGroup(&'a tag::ConditionalGroup),
-            FitsExpanded(&'a tag::FitsExpanded),
+            FitsExpanded {
+                tag: &'a tag::FitsExpanded,
+                expands_before: bool,
+            },
             BestFitting,
         }
 
         fn expand_parent(enclosing: &[Enclosing]) {
             match enclosing.last() {
                 Some(Enclosing::Group(group)) => group.propagate_expand(),
                 Some(Enclosing::ConditionalGroup(group)) => group.propagate_expand(),
-                Some(Enclosing::FitsExpanded(fits_expanded)) => fits_expanded.propagate_expand(),
+                Some(Enclosing::FitsExpanded { tag, .. }) => tag.propagate_expand(),
                 _ => {}
             }
         }
@@ -85,23 +88,24 @@ impl Document {
                     FormatElement::BestFitting { variants, mode: _ } => {
                         enclosing.push(Enclosing::BestFitting);
 
-                        for variant in variants {
-                            propagate_expands(variant, enclosing, checked_interned);
-                        }
-
-                        // Best fitting acts as a boundary
-                        expands = false;
+                        propagate_expands(variants, enclosing, checked_interned);
                         enclosing.pop();
                         continue;
                     }
                     FormatElement::Tag(Tag::StartFitsExpanded(fits_expanded)) => {
-                        enclosing.push(Enclosing::FitsExpanded(fits_expanded));
+                        enclosing.push(Enclosing::FitsExpanded {
+                            tag: fits_expanded,
+                            expands_before: expands,
+                        });
                         false
                     }
                     FormatElement::Tag(Tag::EndFitsExpanded) => {
-                        enclosing.pop();
-                        // Fits expanded acts as a boundary
-                        expands = false;
+                        if let Some(Enclosing::FitsExpanded { expands_before, .. }) =
+                            enclosing.pop()
+                        {
+                            expands = expands_before;
+                        }
+
                         continue;
                     }
                     FormatElement::Text {
@@ -338,28 +342,28 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                 }
 
                 FormatElement::BestFitting { variants, mode } => {
-                    write!(f, [token("best_fitting([")])?;
+                    write!(f, [token("best_fitting(")])?;
+
+                    if *mode != BestFittingMode::FirstLine {
+                        write!(f, [text(&std::format!("mode: {mode:?}, "), None)])?;
+                    }
+
+                    write!(f, [token("[")])?;
                     f.write_elements([
                         FormatElement::Tag(StartIndent),
                         FormatElement::Line(LineMode::Hard),
                     ]);
 
                     for variant in variants {
-                        write!(f, [&**variant, hard_line_break()])?;
+                        write!(f, [variant, hard_line_break()])?;
                     }
 
                     f.write_elements([
                         FormatElement::Tag(EndIndent),
                         FormatElement::Line(LineMode::Hard),
                     ]);
 
-                    write!(f, [token("]")])?;
-
-                    if *mode != BestFittingMode::FirstLine {
-                        write!(f, [text(&std::format!(", mode: {mode:?}"), None),])?;
-                    }
-
-                    write!(f, [token(")")])?;
+                    write!(f, [token("])")])?;
                 }
 
                 FormatElement::Interned(interned) => {
@@ -594,10 +598,10 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                             }
                         }
 
-                        StartEntry => {
+                        StartEntry | StartBestFittingEntry { .. } => {
                             // handled after the match for all start tags
                         }
-                        EndEntry => write!(f, [ContentArrayEnd])?,
+                        EndEntry | EndBestFittingEntry => write!(f, [ContentArrayEnd])?,
 
                         EndFill
                         | EndLabelled

crates/ruff_formatter/src/format_element/tag.rs

@@ -83,6 +83,9 @@ pub enum Tag {
 
     StartFitsExpanded(FitsExpanded),
     EndFitsExpanded,
+
+    StartBestFittingEntry,
+    EndBestFittingEntry,
 }
 
 impl Tag {
@@ -103,6 +106,7 @@ impl Tag {
                 | Tag::StartVerbatim(_)
                 | Tag::StartLabelled(_)
                 | Tag::StartFitsExpanded(_)
+                | Tag::StartBestFittingEntry,
         )
     }
 
@@ -129,6 +133,7 @@ impl Tag {
             StartVerbatim(_) | EndVerbatim => TagKind::Verbatim,
             StartLabelled(_) | EndLabelled => TagKind::Labelled,
             StartFitsExpanded { .. } | EndFitsExpanded => TagKind::FitsExpanded,
+            StartBestFittingEntry { .. } | EndBestFittingEntry => TagKind::BestFittingEntry,
         }
     }
 }
@@ -152,6 +157,7 @@ pub enum TagKind {
     Verbatim,
     Labelled,
     FitsExpanded,
+    BestFittingEntry,
 }
 
 #[derive(Debug, Copy, Default, Clone, Eq, PartialEq)]