Add optimized best_fit_parenthesize IR

crates/ruff_formatter/src/builders.rs

@@ -1451,6 +1451,150 @@ impl<Context> std::fmt::Debug for Group<'_, Context> {
     }
 }
 
+/// Content that may get parenthesized if it exceeds the configured line width but only if the parenthesized
+/// layout doesn't exceed the line width too, in which case it falls back to the flat layout.
+///
+/// This IR is identical to the following [`best_fitting`] layout but is implemented as custom IR for
+/// best performance.
+///
+/// ```rust
+/// # use ruff_formatter::prelude::*;
+/// # use ruff_formatter::format_args;
+///
+/// let format_expression = format_with(|f: &mut Formatter<SimpleFormatContext>| token("A long string").fmt(f));
+/// let _ = best_fitting![
+///     // ---------------------------------------------------------------------
+///     // Variant 1:
+///     // Try to fit the expression without any parentheses
+///     group(&format_expression),
+///     // ---------------------------------------------------------------------
+///     // Variant 2:
+///     // Try to fit the expression by adding parentheses and indenting the expression.
+///     group(&format_args![
+///         token("("),
+///         soft_block_indent(&format_expression),
+///         token(")")
+///     ])
+///     .should_expand(true),
+///     // ---------------------------------------------------------------------
+///     // Variant 3: Fallback, no parentheses
+///     // Expression doesn't fit regardless of adding the parentheses. Remove the parentheses again.
+///     group(&format_expression).should_expand(true)
+/// ]
+/// // Measure all lines, to avoid that the printer decides that this fits right after hitting
+/// // the `(`.
+/// .with_mode(BestFittingMode::AllLines)        ;
+/// ```
+///
+/// The element breaks from left-to-right because it uses the unintended version as *expanded* layout, the same as the above showed best fitting example.
+///
+/// ## Examples
+///
+/// ### Content that fits into the configured line width.
+///
+/// ```rust
+/// # use ruff_formatter::prelude::*;
+/// # use ruff_formatter::{format, PrintResult, write};
+///
+/// # fn main() -> FormatResult<()> {
+///     let formatted = format!(SimpleFormatContext::default(), [format_with(|f| {
+///         write!(f, [
+///             token("aLongerVariableName = "),
+///             best_fit_parenthesize(&token("'a string that fits into the configured line width'"))
+///         ])
+///     })])?;
+///
+///     assert_eq!(formatted.print()?.as_code(), "aLongerVariableName = 'a string that fits into the configured line width'");
+///     # Ok(())
+/// # }
+/// ```
+///
+/// ### Content that fits parenthesized
+///
+/// ```rust
+/// # use ruff_formatter::prelude::*;
+/// # use ruff_formatter::{format, PrintResult, write};
+///
+/// # fn main() -> FormatResult<()> {
+///     let formatted = format!(SimpleFormatContext::default(), [format_with(|f| {
+///         write!(f, [
+///             token("aLongerVariableName = "),
+///             best_fit_parenthesize(&token("'a string that exceeds configured line width but fits parenthesized'"))
+///         ])
+///     })])?;
+///
+///     assert_eq!(formatted.print()?.as_code(), "aLongerVariableName = (\n\t'a string that exceeds configured line width but fits parenthesized'\n)");
+///     # Ok(())
+/// # }
+/// ```
+///
+/// ### Content that exceeds the line width, parenthesized or not
+///
+/// ```rust
+/// # use ruff_formatter::prelude::*;
+/// # use ruff_formatter::{format, PrintResult, write};
+///
+/// # fn main() -> FormatResult<()> {
+///     let formatted = format!(SimpleFormatContext::default(), [format_with(|f| {
+///         write!(f, [
+///             token("aLongerVariableName = "),
+///             best_fit_parenthesize(&token("'a string that exceeds the configured line width and even parenthesizing doesn't make it fit'"))
+///         ])
+///     })])?;
+///
+///     assert_eq!(formatted.print()?.as_code(), "aLongerVariableName = 'a string that exceeds the configured line width and even parenthesizing doesn't make it fit'");
+///     # Ok(())
+/// # }
+/// ```
+#[inline]
+pub fn best_fit_parenthesize<Context>(
+    content: &impl Format<Context>,
+) -> BestFitParenthesize<Context> {
+    BestFitParenthesize {
+        content: Argument::new(content),
+        group_id: None,
+    }
+}
+
+#[derive(Copy, Clone)]
+pub struct BestFitParenthesize<'a, Context> {
+    content: Argument<'a, Context>,
+    group_id: Option<GroupId>,
+}
+
+impl<Context> BestFitParenthesize<'_, Context> {
+    /// Optional ID that can be used in conditional content that supports [`Condition`] to gate content
+    /// depending on whether the parentheses are rendered (flat: no parentheses, expanded: parentheses).
+    #[must_use]
+    pub fn with_group_id(mut self, group_id: Option<GroupId>) -> Self {
+        self.group_id = group_id;
+        self
+    }
+}
+
+impl<Context> Format<Context> for BestFitParenthesize<'_, Context> {
+    fn fmt(&self, f: &mut Formatter<Context>) -> FormatResult<()> {
+        f.write_element(FormatElement::Tag(StartBestFitParenthesize {
+            id: self.group_id,
+        }));
+
+        Arguments::from(&self.content).fmt(f)?;
+
+        f.write_element(FormatElement::Tag(EndBestFitParenthesize));
+
+        Ok(())
+    }
+}
+
+impl<Context> std::fmt::Debug for BestFitParenthesize<'_, Context> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BestFitParenthesize")
+            .field("group_id", &self.group_id)
+            .field("content", &"{{content}}")
+            .finish()
+    }
+}
+
 /// Sets the `condition` for the group. The element will behave as a regular group if `condition` is met,
 /// and as *ungrouped* content if the condition is not met.
 ///

crates/ruff_formatter/src/format_element/document.rs

@@ -40,6 +40,9 @@ impl Document {
                 expands_before: bool,
             },
             BestFitting,
+            BestFitParenthesize {
+                expanded: bool,
+            },
         }
 
         fn expand_parent(enclosing: &[Enclosing]) {
@@ -67,6 +70,18 @@ impl Document {
                         Some(Enclosing::Group(group)) => !group.mode().is_flat(),
                         _ => false,
                     },
+                    FormatElement::Tag(Tag::StartBestFitParenthesize { .. }) => {
+                        enclosing.push(Enclosing::BestFitParenthesize { expanded: expands });
+                        expands = false;
+                        continue;
+                    }
+
+                    FormatElement::Tag(Tag::EndBestFitParenthesize) => {
+                        if let Some(Enclosing::BestFitParenthesize { expanded }) = enclosing.pop() {
+                            expands = expanded;
+                        }
+                        continue;
+                    }
                     FormatElement::Tag(Tag::StartConditionalGroup(group)) => {
                         enclosing.push(Enclosing::ConditionalGroup(group));
                         false
@@ -503,6 +518,21 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                             }
                         }
 
+                        StartBestFitParenthesize { id } => {
+                            write!(f, [token("best_fit_parenthesize(")])?;
+
+                            if let Some(group_id) = id {
+                                write!(
+                                    f,
+                                    [
+                                        text(&std::format!("\"{group_id:?}\""), None),
+                                        token(","),
+                                        space(),
+                                    ]
+                                )?;
+                            }
+                        }
+
                         StartConditionalGroup(group) => {
                             write!(
                                 f,
@@ -611,6 +641,7 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                         | EndIndent
                         | EndGroup
                         | EndConditionalGroup
+                        | EndBestFitParenthesize
                         | EndLineSuffix
                         | EndDedent
                         | EndFitsExpanded

crates/ruff_formatter/src/format_element/tag.rs

@@ -86,6 +86,13 @@ pub enum Tag {
 
     StartBestFittingEntry,
     EndBestFittingEntry,
+
+    /// Parenthesizes the content but only if adding the parentheses and indenting the content
+    /// makes the content fit in the configured line width.
+    StartBestFitParenthesize {
+        id: Option<GroupId>,
+    },
+    EndBestFitParenthesize,
 }
 
 impl Tag {
@@ -102,11 +109,12 @@ impl Tag {
                 | Tag::StartIndentIfGroupBreaks(_)
                 | Tag::StartFill
                 | Tag::StartEntry
-                | Tag::StartLineSuffix { reserved_width: _ }
+                | Tag::StartLineSuffix { .. }
                 | Tag::StartVerbatim(_)
                 | Tag::StartLabelled(_)
                 | Tag::StartFitsExpanded(_)
-                | Tag::StartBestFittingEntry,
+                | Tag::StartBestFittingEntry
+                | Tag::StartBestFitParenthesize { .. }
         )
     }
 
@@ -134,6 +142,9 @@ impl Tag {
             StartLabelled(_) | EndLabelled => TagKind::Labelled,
             StartFitsExpanded { .. } | EndFitsExpanded => TagKind::FitsExpanded,
             StartBestFittingEntry { .. } | EndBestFittingEntry => TagKind::BestFittingEntry,
+            StartBestFitParenthesize { .. } | EndBestFitParenthesize => {
+                TagKind::BestFitParenthesize
+            }
         }
     }
 }
@@ -158,6 +169,7 @@ pub enum TagKind {
     Labelled,
     FitsExpanded,
     BestFittingEntry,
+    BestFitParenthesize,
 }
 
 #[derive(Debug, Copy, Default, Clone, Eq, PartialEq)]

crates/ruff_formatter/src/printer/mod.rs

@@ -174,6 +174,75 @@ impl<'a> Printer<'a> {
                 stack.push(TagKind::Group, args.with_print_mode(print_mode));
             }
 
+            FormatElement::Tag(StartBestFitParenthesize { id }) => {
+                const OPEN_PAREN: FormatElement = FormatElement::Token { text: "(" };
+                const INDENT: FormatElement = FormatElement::Tag(Tag::StartIndent);
+                const HARD_LINE_BREAK: FormatElement = FormatElement::Line(LineMode::Hard);
+
+                let fits_flat = self.flat_group_print_mode(
+                    TagKind::BestFitParenthesize,
+                    *id,
+                    args,
+                    queue,
+                    stack,
+                )? == PrintMode::Flat;
+
+                let print_mode = if fits_flat {
+                    PrintMode::Flat
+                } else {
+                    // Test if the content fits in expanded mode. If not, prefer avoiding the parentheses
+                    // over parenthesizing the expression.
+                    if let Some(id) = id {
+                        self.state
+                            .group_modes
+                            .insert_print_mode(*id, PrintMode::Expanded);
+                    }
+
+                    stack.push(
+                        TagKind::BestFitParenthesize,
+                        args.with_measure_mode(MeasureMode::AllLines),
+                    );
+
+                    queue.extend_back(&[OPEN_PAREN, INDENT, HARD_LINE_BREAK]);
+                    let fits_expanded = self.fits(queue, stack)?;
+                    queue.pop_slice();
+                    stack.pop(TagKind::BestFitParenthesize)?;
+
+                    if fits_expanded {
+                        PrintMode::Expanded
+                    } else {
+                        PrintMode::Flat
+                    }
+                };
+
+                if let Some(id) = id {
+                    self.state.group_modes.insert_print_mode(*id, print_mode);
+                }
+
+                if print_mode.is_expanded() {
+                    // Parenthesize the content. The `EndIndent` is handled inside of the `EndBestFitParenthesize`
+                    queue.extend_back(&[OPEN_PAREN, INDENT, HARD_LINE_BREAK]);
+                }
+
+                stack.push(
+                    TagKind::BestFitParenthesize,
+                    args.with_print_mode(print_mode),
+                );
+            }
+
+            FormatElement::Tag(EndBestFitParenthesize) => {
+                if args.mode().is_expanded() {
+                    const HARD_LINE_BREAK: FormatElement = FormatElement::Line(LineMode::Hard);
+                    const CLOSE_PAREN: FormatElement = FormatElement::Token { text: ")" };
+
+                    // Finish the indent and print the hardline break and closing parentheses.
+                    stack.pop(TagKind::Indent)?;
+                    queue.extend_back(&[HARD_LINE_BREAK, CLOSE_PAREN]);
+                }
+
+                stack.pop(TagKind::BestFitParenthesize)?;
+            }
+
             FormatElement::Tag(StartConditionalGroup(group)) => {
                 let condition = group.condition();
                 let expected_mode = match condition.group_id {
@@ -1204,6 +1273,38 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
                 return Ok(self.fits_group(TagKind::Group, group.mode(), group.id(), args));
             }
 
+            FormatElement::Tag(StartBestFitParenthesize { id }) => {
+                if let Some(id) = id {
+                    self.printer
+                        .state
+                        .group_modes
+                        .insert_print_mode(*id, args.mode());
+                }
+
+                // Don't use the parenthesized with indent layout even when measuring expanded mode similar to `BestFitting`.
+                // This is to expand the left and not right after the `(` parentheses (it is okay to expand after the content that it wraps).
+                self.stack.push(TagKind::BestFitParenthesize, args);
+            }
+
+            FormatElement::Tag(EndBestFitParenthesize) => {
+                // If this is the end tag of the outer most parentheses for which we measure if it fits,
+                // pop the indent.
+                if args.mode().is_expanded() && self.stack.top_kind() == Some(TagKind::Indent) {
+                    self.stack.pop(TagKind::Indent).unwrap();
+                    let unindented = self.stack.pop(TagKind::BestFitParenthesize)?;
+
+                    // There's a hard line break after the indent but don't return `Fits::Yes` here
+                    // to ensure any trailing comments (that, unfortunately, are attached to the statement and not the expression)
+                    // fit too.
+                    self.state.line_width = 0;
+                    self.state.pending_indent = unindented.indention();
+
+                    return Ok(self.fits_text(Text::Token(")"), unindented));
+                }
+
+                self.stack.pop(TagKind::BestFitParenthesize)?;
+            }
+
             FormatElement::Tag(StartConditionalGroup(group)) => {
                 let condition = group.condition();