Extend BestFitting with mode (#6814)

crates/ruff_benchmark/benches/formatter.rs

@@ -1,8 +1,9 @@
 use std::path::Path;
+use std::time::Duration;
 
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
 
-use ruff_benchmark::{TestCase, TestFile, TestFileDownloadError};
+use ruff_benchmark::{TestCase, TestCaseSpeed, TestFile, TestFileDownloadError};
 use ruff_python_formatter::{format_node, PyFormatOptions};
 use ruff_python_index::CommentRangesBuilder;
 use ruff_python_parser::lexer::lex;
@@ -46,6 +47,12 @@ fn benchmark_formatter(criterion: &mut Criterion) {
     for case in test_cases {
         group.throughput(Throughput::Bytes(case.code().len() as u64));
 
+        group.measurement_time(match case.speed() {
+            TestCaseSpeed::Fast => Duration::from_secs(5),
+            TestCaseSpeed::Normal => Duration::from_secs(10),
+            TestCaseSpeed::Slow => Duration::from_secs(20),
+        });
+
         group.bench_with_input(
             BenchmarkId::from_parameter(case.name()),
             &case,

crates/ruff_formatter/src/builders.rs

@@ -2399,6 +2399,7 @@ impl<'a, 'buf, Context> FillBuilder<'a, 'buf, Context> {
 #[derive(Copy, Clone)]
 pub struct BestFitting<'a, Context> {
     variants: Arguments<'a, Context>,
+    mode: BestFittingMode,
 }
 
 impl<'a, Context> BestFitting<'a, Context> {
@@ -2418,7 +2419,116 @@ impl<'a, Context> BestFitting<'a, Context> {
             "Requires at least the least expanded and most expanded variants"
         );
 
-        Self { variants }
+        Self {
+            variants,
+            mode: BestFittingMode::default(),
+        }
+    }
+
+    /// Changes the mode used by this best fitting element to determine whether a variant fits.
+    ///
+    /// ## Examples
+    ///
+    /// ### All Lines
+    ///
+    /// ```
+    /// use ruff_formatter::{Formatted, LineWidth, format, format_args, SimpleFormatOptions};
+    /// use ruff_formatter::prelude::*;
+    ///
+    /// # fn main() -> FormatResult<()> {
+    /// let formatted = format!(
+    ///     SimpleFormatContext::default(),
+    ///     [
+    ///         best_fitting!(
+    ///             // Everything fits on a single line
+    ///             format_args!(
+    ///                 group(&format_args![
+    ///                     text("["),
+    ///                         soft_block_indent(&format_args![
+    ///                         text("1,"),
+    ///                         soft_line_break_or_space(),
+    ///                         text("2,"),
+    ///                         soft_line_break_or_space(),
+    ///                         text("3"),
+    ///                     ]),
+    ///                     text("]")
+    ///                 ]),
+    ///                 space(),
+    ///                 text("+"),
+    ///                 space(),
+    ///                 text("aVeryLongIdentifier")
+    ///             ),
+    ///
+    ///             // Breaks after `[` and prints each elements on a single line
+    ///             // The group is necessary because the variant, by default is printed in flat mode and a
+    ///             // hard line break indicates that the content doesn't fit.
+    ///             format_args!(
+    ///                 text("["),
+    ///                 group(&block_indent(&format_args![text("1,"), hard_line_break(), text("2,"), hard_line_break(), text("3")])).should_expand(true),
+    ///                 text("]"),
+    ///                 space(),
+    ///                 text("+"),
+    ///                 space(),
+    ///                 text("aVeryLongIdentifier")
+    ///             ),
+    ///
+    ///             // Adds parentheses and indents the body, breaks after the operator
+    ///             format_args!(
+    ///                 text("("),
+    ///                 block_indent(&format_args![
+    ///                     text("["),
+    ///                     block_indent(&format_args![
+    ///                         text("1,"),
+    ///                         hard_line_break(),
+    ///                         text("2,"),
+    ///                         hard_line_break(),
+    ///                         text("3"),
+    ///                     ]),
+    ///                     text("]"),
+    ///                     hard_line_break(),
+    ///                     text("+"),
+    ///                     space(),
+    ///                     text("aVeryLongIdentifier")
+    ///                 ]),
+    ///                 text(")")
+    ///             )
+    ///         ).with_mode(BestFittingMode::AllLines)
+    ///     ]
+    /// )?;
+    ///
+    /// let document = formatted.into_document();
+    ///
+    /// // Takes the first variant if everything fits on a single line
+    /// assert_eq!(
+    ///     "[1, 2, 3] + aVeryLongIdentifier",
+    ///     Formatted::new(document.clone(), SimpleFormatContext::default())
+    ///         .print()?
+    ///         .as_code()
+    /// );
+    ///
+    /// // It takes the second if the first variant doesn't fit on a single line. The second variant
+    /// // has some additional line breaks to make sure inner groups don't break
+    /// assert_eq!(
+    ///     "[\n\t1,\n\t2,\n\t3\n] + aVeryLongIdentifier",
+    ///     Formatted::new(document.clone(), SimpleFormatContext::new(SimpleFormatOptions { line_width: 23.try_into().unwrap(), ..SimpleFormatOptions::default() }))
+    ///         .print()?
+    ///         .as_code()
+    /// );
+    ///
+    /// // Prints the last option as last resort
+    /// assert_eq!(
+    ///     "(\n\t[\n\t\t1,\n\t\t2,\n\t\t3\n\t]\n\t+ aVeryLongIdentifier\n)",
+    ///     Formatted::new(document.clone(), SimpleFormatContext::new(SimpleFormatOptions { line_width: 22.try_into().unwrap(), ..SimpleFormatOptions::default() }))
+    ///         .print()?
+    ///         .as_code()
+    /// );
+    /// # Ok(())
+    /// # }
+    /// ```
+    #[must_use]
+    pub fn with_mode(mut self, mode: BestFittingMode) -> Self {
+        self.mode = mode;
+        self
     }
 }
 
@@ -2445,6 +2555,7 @@ impl<Context> Format<Context> for BestFitting<'_, Context> {
                 variants: format_element::BestFittingVariants::from_vec_unchecked(
                     formatted_variants,
                 ),
+                mode: self.mode,
             }
         };
 

crates/ruff_formatter/src/format_element.rs

@@ -57,7 +57,10 @@ pub enum FormatElement {
 
     /// A list of different variants representing the same content. The printer picks the best fitting content.
     /// Line breaks inside of a best fitting don't propagate to parent groups.
-    BestFitting { variants: BestFittingVariants },
+    BestFitting {
+        variants: BestFittingVariants,
+        mode: BestFittingMode,
+    },
 
     /// A [Tag] that marks the start/end of some content to which some special formatting is applied.
     Tag(Tag),
@@ -84,9 +87,10 @@ impl std::fmt::Debug for FormatElement {
                 .field(contains_newlines)
                 .finish(),
             FormatElement::LineSuffixBoundary => write!(fmt, "LineSuffixBoundary"),
-            FormatElement::BestFitting { variants } => fmt
+            FormatElement::BestFitting { variants, mode } => fmt
                 .debug_struct("BestFitting")
                 .field("variants", variants)
+                .field("mode", &mode)
                 .finish(),
             FormatElement::Interned(interned) => fmt.debug_list().entries(&**interned).finish(),
             FormatElement::Tag(tag) => fmt.debug_tuple("Tag").field(tag).finish(),
@@ -295,6 +299,29 @@ impl FormatElements for FormatElement {
     }
 }
 
+/// Mode used to determine if any variant (except the most expanded) fits for [`BestFittingVariants`].
+#[repr(u8)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
+pub enum BestFittingMode {
+    /// The variant fits if the content up to the first hard or a soft line break inside a [`Group`] with
+    /// [`PrintMode::Expanded`] fits on the line. The default mode.
+    ///
+    /// [`Group`]: tag::Group
+    #[default]
+    FirstLine,
+
+    /// A variant fits if all lines fit into the configured print width. A line ends if by any
+    /// hard or a soft line break inside a [`Group`] with [`PrintMode::Expanded`].
+    /// The content doesn't fit if there's any hard line break  outside a [`Group`] with [`PrintMode::Expanded`]
+    /// (a hard line break in content that should be considered in [`PrintMode::Flat`].
+    ///
+    /// Use this mode with caution as it requires measuring all content of the variant which is more
+    /// expensive than using [`BestFittingMode::FirstLine`].
+    ///
+    /// [`Group`]: tag::Group
+    AllLines,
+}
+
 /// The different variants for this element.
 /// 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).

crates/ruff_formatter/src/format_element/document.rs

@@ -81,7 +81,7 @@ impl Document {
                             interned_expands
                         }
                     }
-                    FormatElement::BestFitting { variants } => {
+                    FormatElement::BestFitting { variants, mode: _ } => {
                         enclosing.push(Enclosing::BestFitting);
 
                         for variant in variants {
@@ -326,7 +326,7 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                     write!(f, [text("line_suffix_boundary")])?;
                 }
 
-                FormatElement::BestFitting { variants } => {
+                FormatElement::BestFitting { variants, mode } => {
                     write!(f, [text("best_fitting([")])?;
                     f.write_elements([
                         FormatElement::Tag(StartIndent),
@@ -342,7 +342,13 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                         FormatElement::Line(LineMode::Hard),
                     ]);
 
-                    write!(f, [text("])")])?;
+                    write!(f, [text("]")])?;
+
+                    if *mode != BestFittingMode::FirstLine {
+                        write!(f, [dynamic_text(&std::format!(", mode: {mode:?}"), None),])?;
+                    }
+
+                    write!(f, [text(")")])?;
                 }
 
                 FormatElement::Interned(interned) => {

crates/ruff_formatter/src/printer/mod.rs

@@ -6,7 +6,7 @@ mod stack;
 
 use crate::format_element::document::Document;
 use crate::format_element::tag::{Condition, GroupMode};
-use crate::format_element::{BestFittingVariants, LineMode, PrintMode};
+use crate::format_element::{BestFittingMode, BestFittingVariants, LineMode, PrintMode};
 use crate::prelude::tag;
 use crate::prelude::tag::{DedentMode, Tag, TagKind, VerbatimKind};
 use crate::printer::call_stack::{
@@ -135,8 +135,8 @@ impl<'a> Printer<'a> {
                 self.flush_line_suffixes(queue, stack, Some(HARD_BREAK));
             }
 
-            FormatElement::BestFitting { variants } => {
-                self.print_best_fitting(variants, queue, stack)?;
+            FormatElement::BestFitting { variants, mode } => {
+                self.print_best_fitting(variants, *mode, queue, stack)?;
             }
 
             FormatElement::Interned(content) => {
@@ -429,6 +429,7 @@ impl<'a> Printer<'a> {
     fn print_best_fitting(
         &mut self,
         variants: &'a BestFittingVariants,
+        mode: BestFittingMode,
         queue: &mut PrintQueue<'a>,
         stack: &mut PrintCallStack,
     ) -> PrintResult<()> {
@@ -454,7 +455,9 @@ impl<'a> Printer<'a> {
                 // args must be popped from the stack as soon as it sees the matching end entry.
                 let content = &variant[1..];
 
-                let entry_args = args.with_print_mode(PrintMode::Flat);
+                let entry_args = args
+                    .with_print_mode(PrintMode::Flat)
+                    .with_measure_mode(MeasureMode::from(mode));
 
                 queue.extend_back(content);
                 stack.push(TagKind::Entry, entry_args);
@@ -1072,10 +1075,13 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
 
             FormatElement::SourcePosition(_) => {}
 
-            FormatElement::BestFitting { variants } => {
-                let slice = match args.mode() {
-                    PrintMode::Flat => variants.most_flat(),
-                    PrintMode::Expanded => variants.most_expanded(),
+            FormatElement::BestFitting { variants, mode } => {
+                let (slice, args) = match args.mode() {
+                    PrintMode::Flat => (
+                        variants.most_flat(),
+                        args.with_measure_mode(MeasureMode::from(*mode)),
+                    ),
+                    PrintMode::Expanded => (variants.most_expanded(), args),
                 };
 
                 if !matches!(slice.first(), Some(FormatElement::Tag(Tag::StartEntry))) {
@@ -1402,6 +1408,15 @@ enum MeasureMode {
     AllLines,
 }
 
+impl From<BestFittingMode> for MeasureMode {
+    fn from(value: BestFittingMode) -> Self {
+        match value {
+            BestFittingMode::FirstLine => Self::FirstLine,
+            BestFittingMode::AllLines => Self::AllLines,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::prelude::*;