rustfmt: Format cfg_select!

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::{
@@ -27,6 +28,7 @@ use crate::config::lists::*;
 use crate::expr::{RhsAssignKind, rewrite_array, rewrite_assign_rhs};
 use crate::lists::{ListFormatting, itemize_list, write_list};
 use crate::overflow;
+use crate::parse::macros::cfg_select::{CfgSelectFormatPredicate, parse_cfg_select};
 use crate::parse::macros::lazy_static::parse_lazy_static;
 use crate::parse::macros::{ParsedMacroArgs, parse_expr, parse_macro_args};
 use crate::rewrite::{
@@ -244,6 +246,20 @@ fn rewrite_macro_inner(
         }
     }
 
+    if macro_name.ends_with("cfg_select!") {
+        match format_cfg_select(&macro_name, style, context, shape, ts.clone(), mac.span()) {
+            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),
+            },
+        }
+    }
+
     let ParsedMacroArgs {
         args: arg_vec,
         vec_with_semi,
@@ -397,16 +413,22 @@ fn rewrite_empty_macro_def_body(
     context: &RewriteContext<'_>,
     span: Span,
     shape: Shape,
+    delim_token: Delimiter,
 ) -> RewriteResult {
-    // Create an empty, dummy `ast::Block` representing an empty macro body
-    let block = ast::Block {
-        stmts: vec![].into(),
-        id: rustc_ast::node_id::DUMMY_NODE_ID,
-        rules: ast::BlockCheckMode::Default,
-        span,
-        tokens: None,
-    };
-    block.rewrite_result(context, shape)
+    match delim_token {
+        Delimiter::Brace => {
+            // Create an empty, dummy `ast::Block` representing an empty macro body
+            let block = ast::Block {
+                stmts: vec![].into(),
+                id: rustc_ast::node_id::DUMMY_NODE_ID,
+                rules: ast::BlockCheckMode::Default,
+                span,
+                tokens: None,
+            };
+            block.rewrite_result(context, shape)
+        }
+        _ => Ok(context.snippet(span).to_owned()),
+    }
 }
 
 pub(crate) fn rewrite_macro_def(
@@ -452,7 +474,8 @@ pub(crate) fn rewrite_macro_def(
     if parsed_def.branches.len() == 0 {
         let lo = context.snippet_provider.span_before(span, "{");
         result += " ";
-        result += &rewrite_empty_macro_def_body(context, span.with_lo(lo), shape)?;
+        result +=
+            &rewrite_empty_macro_def_body(context, span.with_lo(lo), shape, Delimiter::Brace)?;
         return Ok(result);
     }
 
@@ -1523,3 +1546,114 @@ fn rewrite_macro_with_items(
     result.push_str(trailing_semicolon);
     Ok(result)
 }
+
+fn format_cfg_select(
+    name: &str,
+    delim_token: Delimiter,
+    context: &RewriteContext<'_>,
+    shape: Shape,
+    ts: TokenStream,
+    span: Span,
+) -> RewriteResult {
+    let mut rewrite = String::with_capacity((span.hi() - span.lo()).to_usize() * 2);
+    rewrite.push_str(name);
+
+    let opening_delim = match delim_token {
+        Delimiter::Brace => "{",
+        Delimiter::Bracket => "[",
+        Delimiter::Parenthesis => "(",
+        Delimiter::Invisible(_) => {
+            unreachable!("cfg_selec! macro will always have outer delimiters");
+        }
+    };
+
+    if matches!(delim_token, Delimiter::Brace) {
+        rewrite.push(' ');
+    };
+
+    let cfg_select_body_start = context.snippet_provider.span_before(span, opening_delim);
+    let arms =
+        parse_cfg_select(context.psess, ts).macro_error(MacroErrorKind::ParseFailure, span)?;
+
+    if arms.is_empty() {
+        rewrite.push_str(&rewrite_empty_macro_def_body(
+            context,
+            span.with_lo(cfg_select_body_start),
+            shape,
+            delim_token,
+        )?);
+        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();
+
+    let closing_delim = match delim_token {
+        Delimiter::Brace => "}",
+        Delimiter::Bracket => "]",
+        Delimiter::Parenthesis => ")",
+        Delimiter::Invisible(_) => {
+            unreachable!("cfg_selec! macro will always have outer delimiters");
+        }
+    };
+
+    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/parse/macros/cfg_select.rs

@@ -0,0 +1,108 @@
+use rustc_ast::tokenstream::TokenStream;
+use rustc_ast::{
+    ast,
+    token::{self, Token},
+};
+use rustc_parse::exp;
+use rustc_span::Span;
+use tracing::debug;
+
+use crate::parse::session::ParseSess;
+use crate::spanned::Spanned;
+
+pub(crate) enum CfgSelectFormatPredicate {
+    Cfg(ast::MetaItemInner),
+    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,
+        }
+    }
+}
+
+pub(crate) struct CfgSelectArm {
+    pub(crate) predicate: CfgSelectFormatPredicate,
+    pub(crate) arrow: Token,
+    pub(crate) expr: Box<ast::Expr>,
+    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(psess: &ParseSess, ts: TokenStream) -> Option<Vec<CfgSelectArm>> {
+    let mut cfg_select_predicates = vec![];
+    let mut parser = super::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(_) = parser.expect(exp!(FatArrow)) {
+            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)
+}

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

@@ -11,6 +11,7 @@ use crate::rewrite::RewriteContext;
 
 pub(crate) mod cfg_if;
 pub(crate) mod cfg_match;
+pub(crate) mod cfg_select;
 pub(crate) mod lazy_static;
 
 fn build_stream_parser<'a>(psess: &'a ParseSess, tokens: TokenStream) -> Parser<'a> {