perf: Reduce best fitting allocations (#7411)

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: