Extend formatter IR to support Black's expression formatting (#5596)

crates/ruff_formatter/src/format_element/document.rs

@@ -1,5 +1,5 @@
 use super::tag::Tag;
-use crate::format_element::tag::DedentMode;
+use crate::format_element::tag::{Condition, DedentMode};
 use crate::prelude::tag::GroupMode;
 use crate::prelude::*;
 use crate::printer::LineEnding;
@@ -33,12 +33,17 @@ impl Document {
         #[derive(Debug)]
         enum Enclosing<'a> {
             Group(&'a tag::Group),
+            ConditionalGroup(&'a tag::ConditionalGroup),
+            FitsExpanded(&'a tag::FitsExpanded),
             BestFitting,
         }
 
         fn expand_parent(enclosing: &[Enclosing]) {
-            if let Some(Enclosing::Group(group)) = enclosing.last() {
-                group.propagate_expand();
+            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(),
+                _ => {}
             }
         }
 
@@ -58,6 +63,14 @@ impl Document {
                         Some(Enclosing::Group(group)) => !group.mode().is_flat(),
                         _ => false,
                     },
+                    FormatElement::Tag(Tag::StartConditionalGroup(group)) => {
+                        enclosing.push(Enclosing::ConditionalGroup(group));
+                        false
+                    }
+                    FormatElement::Tag(Tag::EndConditionalGroup) => match enclosing.pop() {
+                        Some(Enclosing::ConditionalGroup(group)) => !group.mode().is_flat(),
+                        _ => false,
+                    },
                     FormatElement::Interned(interned) => match checked_interned.get(interned) {
                         Some(interned_expands) => *interned_expands,
                         None => {
@@ -79,6 +92,16 @@ impl Document {
                         enclosing.pop();
                         continue;
                     }
+                    FormatElement::Tag(Tag::StartFitsExpanded(fits_expanded)) => {
+                        enclosing.push(Enclosing::FitsExpanded(fits_expanded));
+                        false
+                    }
+                    FormatElement::Tag(Tag::EndFitsExpanded) => {
+                        enclosing.pop();
+                        // Fits expanded acts as a boundary
+                        expands = false;
+                        continue;
+                    }
                     FormatElement::StaticText { text } => text.contains('\n'),
                     FormatElement::DynamicText { text, .. } => text.contains('\n'),
                     FormatElement::SourceCodeSlice {
@@ -441,6 +464,29 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                             }
                         }
 
+                        StartConditionalGroup(group) => {
+                            write!(
+                                f,
+                                [
+                                    text("conditional_group(condition:"),
+                                    space(),
+                                    group.condition(),
+                                    text(","),
+                                    space()
+                                ]
+                            )?;
+
+                            match group.mode() {
+                                GroupMode::Flat => {}
+                                GroupMode::Expand => {
+                                    write!(f, [text("expand: true,"), space()])?;
+                                }
+                                GroupMode::Propagated => {
+                                    write!(f, [text("expand: propagated,"), space()])?;
+                                }
+                            }
+                        }
+
                         StartIndentIfGroupBreaks(id) => {
                             write!(
                                 f,
@@ -491,6 +537,28 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                             write!(f, [text("fill(")])?;
                         }
 
+                        StartFitsExpanded(tag::FitsExpanded {
+                            condition,
+                            propagate_expand,
+                        }) => {
+                            write!(f, [text("fits_expanded(propagate_expand:"), space()])?;
+
+                            if propagate_expand.get() {
+                                write!(f, [text("true")])?;
+                            } else {
+                                write!(f, [text("false")])?;
+                            }
+
+                            write!(f, [text(","), space()])?;
+
+                            if let Some(condition) = condition {
+                                write!(
+                                    f,
+                                    [text("condition:"), space(), condition, text(","), space()]
+                                )?;
+                            }
+                        }
+
                         StartEntry => {
                             // handled after the match for all start tags
                         }
@@ -503,8 +571,10 @@ impl Format<IrFormatContext<'_>> for &[FormatElement] {
                         | EndAlign
                         | EndIndent
                         | EndGroup
+                        | EndConditionalGroup
                         | EndLineSuffix
                         | EndDedent
+                        | EndFitsExpanded
                         | EndVerbatim => {
                             write!(f, [ContentArrayEnd, text(")")])?;
                         }
@@ -658,6 +728,31 @@ impl FormatElements for [FormatElement] {
     }
 }
 
+impl Format<IrFormatContext<'_>> for Condition {
+    fn fmt(&self, f: &mut Formatter<IrFormatContext>) -> FormatResult<()> {
+        match (self.mode, self.group_id) {
+            (PrintMode::Flat, None) => write!(f, [text("if_fits_on_line")]),
+            (PrintMode::Flat, Some(id)) => write!(
+                f,
+                [
+                    text("if_group_fits_on_line("),
+                    dynamic_text(&std::format!("\"{id:?}\""), None),
+                    text(")")
+                ]
+            ),
+            (PrintMode::Expanded, None) => write!(f, [text("if_breaks")]),
+            (PrintMode::Expanded, Some(id)) => write!(
+                f,
+                [
+                    text("if_group_breaks("),
+                    dynamic_text(&std::format!("\"{id:?}\""), None),
+                    text(")")
+                ]
+            ),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::prelude::*;

crates/ruff_formatter/src/format_element/tag.rs

@@ -33,6 +33,16 @@ pub enum Tag {
     StartGroup(Group),
     EndGroup,
 
+    /// Creates a logical group similar to [`Tag::StartGroup`] but only if the condition is met.
+    /// This is an optimized representation for (assuming the content should only be grouped if another group fits):
+    ///
+    /// ```text
+    /// if_group_breaks(content, other_group_id),
+    /// if_group_fits_on_line(group(&content), other_group_id)
+    /// ```
+    StartConditionalGroup(ConditionalGroup),
+    EndConditionalGroup,
+
     /// Allows to specify content that gets printed depending on whatever the enclosing group
     /// is printed on a single line or multiple lines. See [crate::builders::if_group_breaks] for examples.
     StartConditionalContent(Condition),
@@ -67,6 +77,9 @@ pub enum Tag {
     /// See [crate::builders::labelled] for documentation.
     StartLabelled(LabelId),
     EndLabelled,
+
+    StartFitsExpanded(FitsExpanded),
+    EndFitsExpanded,
 }
 
 impl Tag {
@@ -77,7 +90,8 @@ impl Tag {
             Tag::StartIndent
                 | Tag::StartAlign(_)
                 | Tag::StartDedent(_)
-                | Tag::StartGroup { .. }
+                | Tag::StartGroup(_)
+                | Tag::StartConditionalGroup(_)
                 | Tag::StartConditionalContent(_)
                 | Tag::StartIndentIfGroupBreaks(_)
                 | Tag::StartFill
@@ -85,6 +99,7 @@ impl Tag {
                 | Tag::StartLineSuffix
                 | Tag::StartVerbatim(_)
                 | Tag::StartLabelled(_)
+                | Tag::StartFitsExpanded(_)
         )
     }
 
@@ -101,6 +116,7 @@ impl Tag {
             StartAlign(_) | EndAlign => TagKind::Align,
             StartDedent(_) | EndDedent => TagKind::Dedent,
             StartGroup(_) | EndGroup => TagKind::Group,
+            StartConditionalGroup(_) | EndConditionalGroup => TagKind::ConditionalGroup,
             StartConditionalContent(_) | EndConditionalContent => TagKind::ConditionalContent,
             StartIndentIfGroupBreaks(_) | EndIndentIfGroupBreaks => TagKind::IndentIfGroupBreaks,
             StartFill | EndFill => TagKind::Fill,
@@ -108,6 +124,7 @@ impl Tag {
             StartLineSuffix | EndLineSuffix => TagKind::LineSuffix,
             StartVerbatim(_) | EndVerbatim => TagKind::Verbatim,
             StartLabelled(_) | EndLabelled => TagKind::Labelled,
+            StartFitsExpanded { .. } | EndFitsExpanded => TagKind::FitsExpanded,
         }
     }
 }
@@ -122,6 +139,7 @@ pub enum TagKind {
     Align,
     Dedent,
     Group,
+    ConditionalGroup,
     ConditionalContent,
     IndentIfGroupBreaks,
     Fill,
@@ -129,6 +147,7 @@ pub enum TagKind {
     LineSuffix,
     Verbatim,
     Labelled,
+    FitsExpanded,
 }
 
 #[derive(Debug, Copy, Default, Clone, Eq, PartialEq)]
@@ -150,6 +169,27 @@ impl GroupMode {
     }
 }
 
+#[derive(Debug, Clone, Eq, PartialEq, Default)]
+pub struct FitsExpanded {
+    pub(crate) condition: Option<Condition>,
+    pub(crate) propagate_expand: Cell<bool>,
+}
+
+impl FitsExpanded {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn with_condition(mut self, condition: Option<Condition>) -> Self {
+        self.condition = condition;
+        self
+    }
+
+    pub fn propagate_expand(&self) {
+        self.propagate_expand.set(true)
+    }
+}
+
 #[derive(Debug, Clone, Eq, PartialEq, Default)]
 pub struct Group {
     id: Option<GroupId>,
@@ -189,6 +229,33 @@ impl Group {
     }
 }
 
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct ConditionalGroup {
+    mode: Cell<GroupMode>,
+    condition: Condition,
+}
+
+impl ConditionalGroup {
+    pub fn new(condition: Condition) -> Self {
+        Self {
+            mode: Cell::new(GroupMode::Flat),
+            condition,
+        }
+    }
+
+    pub fn condition(&self) -> Condition {
+        self.condition
+    }
+
+    pub fn propagate_expand(&self) {
+        self.mode.set(GroupMode::Propagated)
+    }
+
+    pub fn mode(&self) -> GroupMode {
+        self.mode.get()
+    }
+}
+
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub enum DedentMode {
     /// Reduces the indent by a level (if the current indent is > 0)
@@ -198,7 +265,7 @@ pub enum DedentMode {
     Root,
 }
 
-#[derive(Debug, Clone, Eq, PartialEq)]
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub struct Condition {
     /// - `Flat` -> Omitted if the enclosing group is a multiline group, printed for groups fitting on a single line
     /// - `Expanded` -> Omitted if the enclosing group fits on a single line, printed if the group breaks over multiple lines.
@@ -217,6 +284,34 @@ impl Condition {
         }
     }
 
+    pub fn if_fits_on_line() -> Self {
+        Self {
+            mode: PrintMode::Flat,
+            group_id: None,
+        }
+    }
+
+    pub fn if_group_fits_on_line(group_id: GroupId) -> Self {
+        Self {
+            mode: PrintMode::Flat,
+            group_id: Some(group_id),
+        }
+    }
+
+    pub fn if_breaks() -> Self {
+        Self {
+            mode: PrintMode::Expanded,
+            group_id: None,
+        }
+    }
+
+    pub fn if_group_breaks(group_id: GroupId) -> Self {
+        Self {
+            mode: PrintMode::Expanded,
+            group_id: Some(group_id),
+        }
+    }
+
     pub fn with_group_id(mut self, id: Option<GroupId>) -> Self {
         self.group_id = id;
         self