rustfmt: Format cfg_select!

src/tools/rustfmt/src/items.rs

@@ -1525,7 +1525,7 @@ fn get_bytepos_after_visibility(vis: &ast::Visibility, default_span: Span) -> By
 
 // Format tuple or struct without any fields. We need to make sure that the comments
 // inside the delimiters are preserved.
-fn format_empty_struct_or_tuple(
+pub(crate) fn format_empty_struct_or_tuple(
     context: &RewriteContext<'_>,
     span: Span,
     offset: Indent,

src/tools/rustfmt/src/macros.rs

@@ -9,14 +9,15 @@
 // List-like invocations with parentheses will be formatted as function calls,
 // and those with brackets will be formatted as array literals.
 
+use std::borrow::Cow;
 use std::collections::HashMap;
 use std::panic::{AssertUnwindSafe, catch_unwind};
 
 use rustc_ast::ast;
 use rustc_ast::token::{Delimiter, Token, TokenKind};
 use rustc_ast::tokenstream::{TokenStream, TokenStreamIter, TokenTree};
 use rustc_ast_pretty::pprust;
-use rustc_span::{BytePos, DUMMY_SP, Ident, Span, Symbol};
+use rustc_span::{BytePos, DUMMY_SP, Ident, Pos, Span, Symbol};
 use tracing::debug;
 
 use crate::comment::{
@@ -28,6 +29,7 @@ use crate::expr::{RhsAssignKind, rewrite_array, rewrite_assign_rhs};
 use crate::header::{HeaderPart, format_header};
 use crate::lists::{ListFormatting, itemize_list, write_list};
 use crate::overflow;
+use crate::parse::macros::cfg_select::{CfgSelectFormatPredicate, parse_cfg_select_arms};
 use crate::parse::macros::lazy_static::parse_lazy_static;
 use crate::parse::macros::{ParsedMacroArgs, parse_expr, parse_macro_args};
 use crate::rewrite::{
@@ -245,6 +247,26 @@ fn rewrite_macro_inner(
         }
     }
 
+    if macro_name.ends_with("cfg_select!") {
+        match format_cfg_select(context, shape, mac.span(), &macro_name, style, ts.clone()) {
+            Ok(rw) => return Ok(rw),
+            Err(err) => match err {
+                // We will move on to parsing macro args just like other macros
+                // if we could not parse cfg_select! with known syntax
+                RewriteError::MacroFailure { kind, span: _ }
+                    if kind == MacroErrorKind::ParseFailure => {}
+                // If formatting fails even though parsing succeeds, return the err early
+                other => return Err(other),
+            },
+        }
+    }
+
+    // If we're falling through to default macro handling check that the context is correct
+    debug_assert!(
+        context.inside_macro(),
+        "expect `context.inside_macro() == true`"
+    );
+
     let ParsedMacroArgs {
         args: arg_vec,
         vec_with_semi,
@@ -1530,3 +1552,117 @@ fn rewrite_macro_with_items(
     result.push_str(trailing_semicolon);
     Ok(result)
 }
+
+fn format_cfg_select(
+    context: &RewriteContext<'_>,
+    shape: Shape,
+    span: Span,
+    name: &str,
+    delim_token: Delimiter,
+    ts: TokenStream,
+) -> RewriteResult {
+    let mut rewrite = String::with_capacity((span.hi() - span.lo()).to_usize() * 2);
+    rewrite.push_str(name);
+
+    let (opening_delim, closing_delim) = match delim_token {
+        Delimiter::Brace => ("{", "}"),
+        Delimiter::Bracket => ("[", "]"),
+        Delimiter::Parenthesis => ("(", ")"),
+        Delimiter::Invisible(_) => {
+            unreachable!("cfg_select! macro will always have outer delimiters");
+        }
+    };
+
+    if matches!(delim_token, Delimiter::Brace) {
+        rewrite.push(' ');
+    };
+
+    let arms =
+        parse_cfg_select_arms(context.psess, ts).macro_error(MacroErrorKind::ParseFailure, span)?;
+
+    if arms.is_empty() {
+        let lo = context.snippet_provider.span_after(span, opening_delim);
+        let hi = context.snippet_provider.span_before(span, closing_delim);
+
+        // NOTE(ytmimi) reusing `format_empty_struct_or_tuple` since
+        // it handles proper indentation and recovering comments
+        crate::items::format_empty_struct_or_tuple(
+            context,
+            mk_sp(lo, hi),
+            shape.indent,
+            &mut rewrite,
+            opening_delim,
+            closing_delim,
+        );
+        return Ok(rewrite);
+    } else {
+        rewrite.push_str(opening_delim);
+    }
+
+    let nested_shape = shape.block_indent(context.config.tab_spaces());
+    rewrite.push_str(&nested_shape.indent.to_string_with_newline(context.config));
+
+    let last_arm = arms.last();
+
+    // We have to fib a little here and update the context to remove the `inside_macro` state.
+    // The code that flattens match arms will refuse to do so if it's inside a macro. Mostly
+    // this is done to prevent rustfmt from removing tokens in the context of a macro, but in
+    // this case it should be fine since we know that each `cfg_select!` arm must be a valid expr.
+    context.leave_macro();
+
+    let items = itemize_list(
+        context.snippet_provider,
+        arms.iter(),
+        closing_delim,
+        "}",
+        |arm| arm.span().lo(),
+        |arm| arm.span().hi(),
+        |arm| {
+            let predicate_str = match &arm.predicate {
+                CfgSelectFormatPredicate::Wildcard(_t) => Cow::Borrowed("_"),
+                CfgSelectFormatPredicate::Cfg(meta_item_inner) => {
+                    Cow::Owned(meta_item_inner.rewrite_result(context, nested_shape)?)
+                }
+            };
+
+            crate::matches::rewrite_match_body(
+                context,
+                &arm.expr,
+                &predicate_str,
+                nested_shape,
+                false,
+                arm.arrow.span,
+                last_arm.is_some_and(|la| la == arm),
+            )
+        },
+        // Start Span after the opening delimiter. For example,
+        // ```
+        // cfg_select! {
+        //              ^ start here
+        // }
+        // ```
+        context.snippet_provider.span_after(span, opening_delim),
+        // End on closing delimiter. For example,
+        // ```
+        // cfg_select! {
+        // }
+        // ^ end here
+        // ```
+        span.hi(),
+        false,
+    );
+    let arms_vec: Vec<_> = items.collect();
+
+    // We will add/remove commas inside `arm.rewrite()`, and hence no separator here.
+    let fmt = ListFormatting::new(nested_shape, context.config)
+        .separator("")
+        .align_comments(false)
+        .preserve_newline(true);
+
+    rewrite.push_str(&write_list(&arms_vec, &fmt)?);
+    rewrite.push('\n');
+    rewrite.push_str(&shape.indent.to_string(context.config));
+    rewrite.push_str(closing_delim);
+
+    Ok(rewrite)
+}

src/tools/rustfmt/src/matches.rs

@@ -394,7 +394,7 @@ fn flatten_arm_body<'a>(
     }
 }
 
-fn rewrite_match_body(
+pub(crate) fn rewrite_match_body(
     context: &RewriteContext<'_>,
     body: &Box<ast::Expr>,
     pats_str: &str,

src/tools/rustfmt/src/modules/visitor.rs

@@ -5,7 +5,7 @@ use tracing::debug;
 
 use crate::attr::MetaVisitor;
 use crate::parse::macros::cfg_if::parse_cfg_if;
-use crate::parse::macros::cfg_select::parse_cfg_select;
+use crate::parse::macros::cfg_select::parse_items_from_cfg_select;
 use crate::parse::session::ParseSess;
 
 pub(crate) struct ModItem {
@@ -123,7 +123,7 @@ impl<'a, 'ast: 'a> CfgSelectVisitor<'a> {
             }
         };
 
-        let items = parse_cfg_select(self.psess, mac)?;
+        let items = parse_items_from_cfg_select(self.psess, mac)?;
         self.mods
             .append(&mut items.into_iter().map(|item| ModItem { item }).collect());
 

src/tools/rustfmt/src/parse/macros/cfg_select.rs

@@ -1,25 +1,36 @@
+//! See [`cfg_select!` reference](
+//! https://doc.rust-lang.org/nightly/reference/conditional-compilation.html#the-cfg_select-macro
+//! ) for grammar.
+
 use std::panic::{AssertUnwindSafe, catch_unwind};
 
 use rustc_ast::ast;
-use rustc_ast::token::TokenKind;
+use rustc_ast::token;
+use rustc_ast::token::{Token, TokenKind};
+use rustc_ast::tokenstream::TokenStream;
 use rustc_parse::exp;
 use rustc_parse::parser::{AllowConstBlockItems, ForceCollect};
+use rustc_span::Span;
+use tracing::debug;
 
 use crate::parse::macros::build_stream_parser;
 use crate::parse::session::ParseSess;
+use crate::spanned::Spanned;
 
-pub(crate) fn parse_cfg_select<'a>(
+pub(crate) fn parse_items_from_cfg_select<'a>(
     psess: &'a ParseSess,
     mac: &'a ast::MacCall,
 ) -> Result<Vec<ast::Item>, &'static str> {
-    match catch_unwind(AssertUnwindSafe(|| parse_cfg_select_inner(psess, mac))) {
+    match catch_unwind(AssertUnwindSafe(|| {
+        parse_items_from_cfg_select_inner(psess, mac)
+    })) {
         Ok(Ok(items)) => Ok(items),
         Ok(err @ Err(_)) => err,
         Err(..) => Err("failed to parse cfg_select!"),
     }
 }
 
-fn parse_cfg_select_inner<'a>(
+fn parse_items_from_cfg_select_inner<'a>(
     psess: &'a ParseSess,
     mac: &'a ast::MacCall,
 ) -> Result<Vec<ast::Item>, &'static str> {
@@ -78,3 +89,113 @@ fn parse_cfg_select_inner<'a>(
 
     Ok(items)
 }
+
+/// LHS predicate of a `cfg_select!` arm.
+pub(crate) enum CfgSelectFormatPredicate {
+    /// Example: the `unix` in `unix => {}`. Notably, outer or inner attributes are not permitted.
+    Cfg(ast::MetaItemInner),
+    /// `_` in `_ => {}`.
+    Wildcard(Span),
+}
+
+impl Spanned for CfgSelectFormatPredicate {
+    fn span(&self) -> rustc_span::Span {
+        match self {
+            Self::Cfg(meta_item_inner) => meta_item_inner.span(),
+            Self::Wildcard(span) => *span,
+        }
+    }
+}
+
+/// Each `$predicate => $production` arm in `cfg_select!`.
+pub(crate) struct CfgSelectArm {
+    /// The `$predicate` part.
+    pub(crate) predicate: CfgSelectFormatPredicate,
+    /// Span of `=>`.
+    pub(crate) arrow: Token,
+    /// The RHS `$production` expression.
+    pub(crate) expr: Box<ast::Expr>,
+    /// `cfg_select!` arms `$production`s can be optionally `,` terminated, like `match` arms.
+    /// The `,` is not needed when `$production` is itself braced `{}`.
+    pub(crate) trailing_comma: Option<Span>,
+}
+
+impl PartialEq for &CfgSelectArm {
+    fn eq(&self, other: &Self) -> bool {
+        // consider the arms equal if they have the same span
+        self.span() == other.span()
+    }
+}
+
+impl Spanned for CfgSelectArm {
+    fn span(&self) -> Span {
+        self.predicate
+            .span()
+            .with_hi(if let Some(comma) = self.trailing_comma {
+                comma.hi()
+            } else {
+                self.expr.span.hi()
+            })
+    }
+}
+
+impl std::fmt::Debug for CfgSelectArm {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self.predicate {
+            CfgSelectFormatPredicate::Cfg(cfg_entry) => cfg_entry.fmt(f)?,
+            CfgSelectFormatPredicate::Wildcard(t) => t.fmt(f)?,
+        };
+        write!(f, "=> {:?}", self.expr)
+    }
+}
+
+// FIXME(ytmimi) would be nice if rustfmt didn't need to implement parsing logic on its own
+// and could instead just call rustc_attr_parsing::parse_cfg_select, but this is fine for now.
+pub(crate) fn parse_cfg_select_arms(
+    psess: &ParseSess,
+    ts: TokenStream,
+) -> Option<Vec<CfgSelectArm>> {
+    let mut cfg_select_predicates = vec![];
+    let mut parser = build_stream_parser(psess.inner(), ts);
+
+    while parser.token != token::Eof {
+        let predicate = if parser.eat_keyword(exp!(Underscore)) {
+            CfgSelectFormatPredicate::Wildcard(parser.prev_token.span)
+        } else {
+            let Ok(meta_item) = parser.parse_meta_item_inner().map_err(|e| e.cancel()) else {
+                debug!("Failed to parse cfg entry in cfg_select! predicate");
+                return None;
+            };
+            CfgSelectFormatPredicate::Cfg(meta_item)
+        };
+
+        if let Err(e) = parser.expect(exp!(FatArrow)) {
+            e.cancel();
+            debug!("Expected to find a `=>` after cfg_selec! predicate.");
+            return None;
+        };
+
+        let arrow = parser.prev_token;
+
+        let Ok(expr) = parser.parse_expr().map_err(|e| e.cancel()) else {
+            debug!("Couldn't parse cfg_select! arm body after `=>`.");
+            return None;
+        };
+
+        let trailing_comma = if parser.eat(exp!(Comma)) {
+            Some(parser.prev_token.span)
+        } else {
+            None
+        };
+
+        let arm = CfgSelectArm {
+            predicate,
+            arrow,
+            expr,
+            trailing_comma,
+        };
+
+        cfg_select_predicates.push(arm);
+    }
+    Some(cfg_select_predicates)
+}