Support cfg_select!

Author: ChayimFriedman2

crates/cfg/src/cfg_expr.rs

@@ -68,6 +68,11 @@ impl CfgExpr {
         next_cfg_expr(&mut tt.iter()).unwrap_or(CfgExpr::Invalid)
     }
 
+    #[cfg(feature = "tt")]
+    pub fn parse_from_iter<S: Copy>(tt: &mut tt::iter::TtIter<'_, S>) -> CfgExpr {
+        next_cfg_expr(tt).unwrap_or(CfgExpr::Invalid)
+    }
+
     /// Fold the cfg by querying all basic `Atom` and `KeyValue` predicates.
     pub fn fold(&self, query: &dyn Fn(&CfgAtom) -> bool) -> Option<bool> {
         match self {
@@ -96,7 +101,14 @@ fn next_cfg_expr<S: Copy>(it: &mut tt::iter::TtIter<'_, S>) -> Option<CfgExpr> {
     };
 
     let ret = match it.peek() {
-        Some(TtElement::Leaf(tt::Leaf::Punct(punct))) if punct.char == '=' => {
+        Some(TtElement::Leaf(tt::Leaf::Punct(punct)))
+            // Don't consume on e.g. `=>`.
+            if punct.char == '='
+                && (punct.spacing == tt::Spacing::Alone
+                    || it.remaining().flat_tokens().get(1).is_none_or(|peek2| {
+                        !matches!(peek2, tt::TokenTree::Leaf(tt::Leaf::Punct(_)))
+                    })) =>
+        {
             match it.remaining().flat_tokens().get(1) {
                 Some(tt::TokenTree::Leaf(tt::Leaf::Literal(literal))) => {
                     it.next();

crates/hir-def/src/macro_expansion_tests/builtin_fn_macro.rs

@@ -550,3 +550,51 @@ fn main() { "\"hello\""; }
 "##]],
     );
 }
+
+#[test]
+fn cfg_select() {
+    check(
+        r#"
+#[rustc_builtin_macro]
+pub macro cfg_select($($tt:tt)*) {}
+
+cfg_select! {
+    false => { fn false_1() {} }
+    any(false, true) => { fn true_1() {} }
+}
+
+cfg_select! {
+    false => { fn false_2() {} }
+    _ => { fn true_2() {} }
+}
+
+cfg_select! {
+    false => { fn false_3() {} }
+}
+
+cfg_select! {
+    false
+}
+
+cfg_select! {
+    false =>
+}
+
+    "#,
+        expect![[r#"
+#[rustc_builtin_macro]
+pub macro cfg_select($($tt:tt)*) {}
+
+fn true_1() {}
+
+fn true_2() {}
+
+/* error: none of the predicates in this `cfg_select` evaluated to true */
+
+/* error: expected `=>` after cfg expression */
+
+/* error: expected a token tree after `=>` */
+
+    "#]],
+    );
+}

crates/hir-expand/src/builtin/fn_macro.rs

@@ -127,6 +127,7 @@ register_builtin! {
     (asm, Asm) => asm_expand,
     (global_asm, GlobalAsm) => global_asm_expand,
     (naked_asm, NakedAsm) => naked_asm_expand,
+    (cfg_select, CfgSelect) => cfg_select_expand,
     (cfg, Cfg) => cfg_expand,
     (core_panic, CorePanic) => panic_expand,
     (std_panic, StdPanic) => panic_expand,
@@ -355,6 +356,71 @@ fn naked_asm_expand(
     ExpandResult::ok(expanded)
 }
 
+fn cfg_select_expand(
+    db: &dyn ExpandDatabase,
+    id: MacroCallId,
+    tt: &tt::TopSubtree,
+    span: Span,
+) -> ExpandResult<tt::TopSubtree> {
+    let loc = db.lookup_intern_macro_call(id);
+    let cfg_options = loc.krate.cfg_options(db);
+
+    let mut iter = tt.iter();
+    let mut expand_to = None;
+    while let Some(next) = iter.peek() {
+        let active = if let tt::TtElement::Leaf(tt::Leaf::Ident(ident)) = next
+            && ident.sym == sym::underscore
+        {
+            iter.next();
+            true
+        } else {
+            cfg_options.check(&CfgExpr::parse_from_iter(&mut iter)) != Some(false)
+        };
+        match iter.expect_glued_punct() {
+            Ok(it) if it.len() == 2 && it[0].char == '=' && it[1].char == '>' => {}
+            _ => {
+                let err_span = iter.peek().map(|it| it.first_span()).unwrap_or(span);
+                return ExpandResult::new(
+                    tt::TopSubtree::empty(tt::DelimSpan::from_single(span)),
+                    ExpandError::other(err_span, "expected `=>` after cfg expression"),
+                );
+            }
+        }
+        let expand_to_if_active = match iter.next() {
+            Some(tt::TtElement::Subtree(_, tt)) => tt.remaining(),
+            _ => {
+                let err_span = iter.peek().map(|it| it.first_span()).unwrap_or(span);
+                return ExpandResult::new(
+                    tt::TopSubtree::empty(tt::DelimSpan::from_single(span)),
+                    ExpandError::other(err_span, "expected a token tree after `=>`"),
+                );
+            }
+        };
+
+        if expand_to.is_none() && active {
+            expand_to = Some(expand_to_if_active);
+        }
+    }
+    match expand_to {
+        Some(expand_to) => {
+            let mut builder = tt::TopSubtreeBuilder::new(tt::Delimiter {
+                kind: tt::DelimiterKind::Invisible,
+                open: span,
+                close: span,
+            });
+            builder.extend_with_tt(expand_to);
+            ExpandResult::ok(builder.build())
+        }
+        None => ExpandResult::new(
+            tt::TopSubtree::empty(tt::DelimSpan::from_single(span)),
+            ExpandError::other(
+                span,
+                "none of the predicates in this `cfg_select` evaluated to true",
+            ),
+        ),
+    }
+}
+
 fn cfg_expand(
     db: &dyn ExpandDatabase,
     id: MacroCallId,

crates/intern/src/symbol/symbols.rs

@@ -156,6 +156,7 @@ define_symbols! {
     cfg_attr,
     cfg_eval,
     cfg,
+    cfg_select,
     char,
     clone,
     Clone,

crates/tt/src/iter.rs

@@ -217,6 +217,17 @@ pub enum TtElement<'a, S> {
     Subtree(&'a Subtree<S>, TtIter<'a, S>),
 }
 
+impl<S: Copy + fmt::Debug> fmt::Debug for TtElement<'_, S> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Leaf(leaf) => f.debug_tuple("Leaf").field(leaf).finish(),
+            Self::Subtree(subtree, inner) => {
+                f.debug_tuple("Subtree").field(subtree).field(inner).finish()
+            }
+        }
+    }
+}
+
 impl<S: Copy> TtElement<'_, S> {
     #[inline]
     pub fn first_span(&self) -> S {