Fix Printer group modes (#6815)

crates/ruff_formatter/src/diagnostics.rs

@@ -1,4 +1,5 @@
 use crate::prelude::TagKind;
+use crate::GroupId;
 use ruff_text_size::TextRange;
 use std::error::Error;
 
@@ -86,7 +87,9 @@ pub enum InvalidDocumentError {
     /// Text
     /// EndGroup
     /// ```
-    StartTagMissing { kind: TagKind },
+    StartTagMissing {
+        kind: TagKind,
+    },
 
     /// Expected a specific start tag but instead is:
     /// - at the end of the document
@@ -96,6 +99,10 @@ pub enum InvalidDocumentError {
         expected_start: TagKind,
         actual: ActualStart,
     },
+
+    UnknownGroupId {
+        group_id: GroupId,
+    },
 }
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
@@ -148,6 +155,9 @@ impl std::fmt::Display for InvalidDocumentError {
                     }
                 }
             }
+            InvalidDocumentError::UnknownGroupId { group_id } => {
+                std::write!(f, "Encountered unknown group id {group_id:?}. Ensure that the group with the id {group_id:?} exists and that the group is a parent of or comes before the element referring to it.")
+            }
         }
     }
 }

crates/ruff_formatter/src/group_id.rs

@@ -1,9 +1,11 @@
 use std::num::NonZeroU32;
 use std::sync::atomic::{AtomicU32, Ordering};
 
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 #[derive(Clone, Copy, Eq, PartialEq, Hash)]
 pub struct DebugGroupId {
     value: NonZeroU32,
+    #[cfg_attr(feature = "serde", serde(skip))]
     name: &'static str,
 }
 
@@ -28,6 +30,7 @@ impl std::fmt::Debug for DebugGroupId {
 /// See [`crate::Formatter::group_id`] on how to get a unique id.
 #[repr(transparent)]
 #[derive(Clone, Copy, Eq, PartialEq, Hash)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub struct ReleaseGroupId {
     value: NonZeroU32,
 }

crates/ruff_formatter/src/printer/mod.rs

@@ -144,23 +144,40 @@ impl<'a> Printer<'a> {
             }
 
             FormatElement::Tag(StartGroup(group)) => {
-                let print_mode =
-                    self.print_group(TagKind::Group, group.mode(), args, queue, stack)?;
+                let print_mode = match group.mode() {
+                    GroupMode::Expand | GroupMode::Propagated => PrintMode::Expanded,
+                    GroupMode::Flat => {
+                        self.flat_group_print_mode(TagKind::Group, group.id(), args, queue, stack)?
+                    }
+                };
 
                 if let Some(id) = group.id() {
                     self.state.group_modes.insert_print_mode(id, print_mode);
                 }
+
+                stack.push(TagKind::Group, args.with_print_mode(print_mode));
             }
 
             FormatElement::Tag(StartConditionalGroup(group)) => {
                 let condition = group.condition();
                 let expected_mode = match condition.group_id {
                     None => args.mode(),
-                    Some(id) => self.state.group_modes.unwrap_print_mode(id, element),
+                    Some(id) => self.state.group_modes.get_print_mode(id)?,
                 };
 
                 if expected_mode == condition.mode {
-                    self.print_group(TagKind::ConditionalGroup, group.mode(), args, queue, stack)?;
+                    let print_mode = match group.mode() {
+                        GroupMode::Expand | GroupMode::Propagated => PrintMode::Expanded,
+                        GroupMode::Flat => self.flat_group_print_mode(
+                            TagKind::ConditionalGroup,
+                            None,
+                            args,
+                            queue,
+                            stack,
+                        )?,
+                    };
+
+                    stack.push(TagKind::ConditionalGroup, args.with_print_mode(print_mode));
                 } else {
                     // Condition isn't met, render as normal content
                     stack.push(TagKind::ConditionalGroup, args);
@@ -193,7 +210,7 @@ impl<'a> Printer<'a> {
             FormatElement::Tag(StartConditionalContent(Condition { mode, group_id })) => {
                 let group_mode = match group_id {
                     None => args.mode(),
-                    Some(id) => self.state.group_modes.unwrap_print_mode(*id, element),
+                    Some(id) => self.state.group_modes.get_print_mode(*id)?,
                 };
 
                 if *mode == group_mode {
@@ -204,7 +221,7 @@ impl<'a> Printer<'a> {
             }
 
             FormatElement::Tag(StartIndentIfGroupBreaks(group_id)) => {
-                let group_mode = self.state.group_modes.unwrap_print_mode(*group_id, element);
+                let group_mode = self.state.group_modes.get_print_mode(*group_id)?;
 
                 let args = match group_mode {
                     PrintMode::Flat => args,
@@ -237,9 +254,7 @@ impl<'a> Printer<'a> {
                 let condition_met = match condition {
                     Some(condition) => {
                         let group_mode = match condition.group_id {
-                            Some(group_id) => {
-                                self.state.group_modes.unwrap_print_mode(group_id, element)
-                            }
+                            Some(group_id) => self.state.group_modes.get_print_mode(group_id)?,
                             None => args.mode(),
                         };
 
@@ -289,47 +304,46 @@ impl<'a> Printer<'a> {
         result
     }
 
-    fn print_group(
+    fn flat_group_print_mode(
         &mut self,
         kind: TagKind,
-        mode: GroupMode,
+        id: Option<GroupId>,
         args: PrintElementArgs,
         queue: &PrintQueue<'a>,
         stack: &mut PrintCallStack,
     ) -> PrintResult<PrintMode> {
-        let group_mode = match mode {
-            GroupMode::Expand | GroupMode::Propagated => PrintMode::Expanded,
-            GroupMode::Flat => {
-                match args.mode() {
-                    PrintMode::Flat if self.state.measured_group_fits => {
-                        // A parent group has already verified that this group fits on a single line
-                        // Thus, just continue in flat mode
-                        PrintMode::Flat
-                    }
-                    // The printer is either in expanded mode or it's necessary to re-measure if the group fits
-                    // because the printer printed a line break
-                    _ => {
-                        self.state.measured_group_fits = true;
+        let print_mode = match args.mode() {
+            PrintMode::Flat if self.state.measured_group_fits => {
+                // A parent group has already verified that this group fits on a single line
+                // Thus, just continue in flat mode
+                PrintMode::Flat
+            }
+            // The printer is either in expanded mode or it's necessary to re-measure if the group fits
+            // because the printer printed a line break
+            _ => {
+                self.state.measured_group_fits = true;
 
-                        // Measure to see if the group fits up on a single line. If that's the case,
-                        // print the group in "flat" mode, otherwise continue in expanded mode
-                        stack.push(kind, args.with_print_mode(PrintMode::Flat));
-                        let fits = self.fits(queue, stack)?;
-                        stack.pop(kind)?;
+                if let Some(id) = id {
+                    self.state
+                        .group_modes
+                        .insert_print_mode(id, PrintMode::Flat);
+                }
 
-                        if fits {
-                            PrintMode::Flat
-                        } else {
-                            PrintMode::Expanded
-                        }
-                    }
+                // Measure to see if the group fits up on a single line. If that's the case,
+                // print the group in "flat" mode, otherwise continue in expanded mode
+                stack.push(kind, args.with_print_mode(PrintMode::Flat));
+                let fits = self.fits(queue, stack)?;
+                stack.pop(kind)?;
+
+                if fits {
+                    PrintMode::Flat
+                } else {
+                    PrintMode::Expanded
                 }
             }
         };
 
-        stack.push(kind, args.with_print_mode(group_mode));
-
-        Ok(group_mode)
+        Ok(print_mode)
     }
 
     fn print_text(&mut self, text: &str, source_range: Option<TextRange>) {
@@ -785,17 +799,15 @@ impl GroupModes {
         self.0[index] = Some(mode);
     }
 
-    fn get_print_mode(&self, group_id: GroupId) -> Option<PrintMode> {
+    fn get_print_mode(&self, group_id: GroupId) -> PrintResult<PrintMode> {
         let index = u32::from(group_id) as usize;
-        self.0
-            .get(index)
-            .and_then(|option| option.as_ref().copied())
-    }
 
-    fn unwrap_print_mode(&self, group_id: GroupId, next_element: &FormatElement) -> PrintMode {
-        self.get_print_mode(group_id).unwrap_or_else(|| {
-            panic!("Expected group with id {group_id:?} to exist but it wasn't present in the document. Ensure that a group with such a document appears in the document before the element {next_element:?}.")
-        })
+        match self.0.get(index) {
+            Some(Some(print_mode)) => Ok(*print_mode),
+            None | Some(None) => Err(PrintError::InvalidDocument(
+                InvalidDocumentError::UnknownGroupId { group_id },
+            )),
+        }
     }
 }
 
@@ -1123,10 +1135,7 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
 
                 let print_mode = match condition.group_id {
                     None => args.mode(),
-                    Some(group_id) => self
-                        .group_modes()
-                        .get_print_mode(group_id)
-                        .unwrap_or_else(|| args.mode()),
+                    Some(group_id) => self.group_modes().get_print_mode(group_id)?,
                 };
 
                 if condition.mode == print_mode {
@@ -1143,10 +1152,7 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
             FormatElement::Tag(StartConditionalContent(condition)) => {
                 let print_mode = match condition.group_id {
                     None => args.mode(),
-                    Some(group_id) => self
-                        .group_modes()
-                        .get_print_mode(group_id)
-                        .unwrap_or_else(|| args.mode()),
+                    Some(group_id) => self.group_modes().get_print_mode(group_id)?,
                 };
 
                 if condition.mode == print_mode {
@@ -1157,10 +1163,7 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
             }
 
             FormatElement::Tag(StartIndentIfGroupBreaks(id)) => {
-                let print_mode = self
-                    .group_modes()
-                    .get_print_mode(*id)
-                    .unwrap_or_else(|| args.mode());
+                let print_mode = self.group_modes().get_print_mode(*id)?;
 
                 match print_mode {
                     PrintMode::Flat => {
@@ -1191,10 +1194,7 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
                 let condition_met = match condition {
                     Some(condition) => {
                         let group_mode = match condition.group_id {
-                            Some(group_id) => self
-                                .group_modes()
-                                .get_print_mode(group_id)
-                                .unwrap_or_else(|| args.mode()),
+                            Some(group_id) => self.group_modes().get_print_mode(group_id)?,
                             None => args.mode(),
                         };
 
@@ -1250,17 +1250,17 @@ impl<'a, 'print> FitsMeasurer<'a, 'print> {
     fn fits_group(
         &mut self,
         kind: TagKind,
-        mode: GroupMode,
+        group_mode: GroupMode,
         id: Option<GroupId>,
         args: PrintElementArgs,
     ) -> Fits {
-        if self.must_be_flat && !mode.is_flat() {
+        if self.must_be_flat && !group_mode.is_flat() {
             return Fits::No;
         }
 
         // Continue printing groups in expanded mode if measuring a `best_fitting` element where
         // a group expands.
-        let print_mode = if mode.is_flat() {
+        let print_mode = if group_mode.is_flat() {
             args.mode()
         } else {
             PrintMode::Expanded