distributed_slice_elements (multiple) plus tests

Author: jdonszelmann

compiler/rustc_ast/src/ast.rs

@@ -3624,6 +3624,10 @@ pub enum DistributedSlice {
     /// This const (we never do this to statics) represents an addition to a global registry
     /// declared somewhere else.
     Addition { declaration: Path, id: NodeId },
+    /// This const (we never do this to statics) represents an addition of an array
+    /// to a global registry declared somewhere else. All elements are added, though not necessarily
+    /// in the same order as in the original slice.
+    AdditionMany { declaration: Path, id: NodeId },
     /// Applied to an invalid item, error guaranteed to have be emitted
     Err(ErrorGuaranteed),
 }

compiler/rustc_ast/src/visit.rs

@@ -467,6 +467,10 @@ macro_rules! common_visitor_and_walkers {
                                 try_visit!(vis.visit_path(declaration$(${ignore($lt)}, *id)?));
                                 try_visit!(visit_id(vis, id));
                             }
+                            DistributedSlice::AdditionMany { declaration, id } => {
+                                try_visit!(vis.visit_path(declaration$(${ignore($lt)}, *id)?));
+                                try_visit!(visit_id(vis, id));
+                            }
                         }
                         walk_define_opaques(vis, define_opaque)
                     }
@@ -645,6 +649,10 @@ macro_rules! common_visitor_and_walkers {
                     try_visit!(vis.visit_path(declaration$(${ignore($lt)}, *id)?));
                     try_visit!(visit_id(vis, id));
                 }
+                DistributedSlice::AdditionMany { declaration, id } => {
+                    try_visit!(vis.visit_path(declaration$(${ignore($lt)}, *id)?));
+                    try_visit!(visit_id(vis, id));
+                }
             }
 
 

compiler/rustc_ast_lowering/messages.ftl

@@ -56,6 +56,9 @@ ast_lowering_coroutine_too_many_parameters =
 ast_lowering_default_field_in_tuple = default fields are not supported in tuple structs
     .label = default fields are only supported on structs
 
+ast_lowering_distributed_slice_elements_wrong_expr =
+    `distributed_slice_elements!()` only accepts a path or array literal
+    .note = arbitrary expressions are not supported
 ast_lowering_distributed_slice_with_initializer =
     distributed slice elements are added with `distributed_slice_element!(...)`
 ast_lowering_does_not_support_modifiers =

compiler/rustc_ast_lowering/src/errors.rs

@@ -482,3 +482,11 @@ pub(crate) struct DistributedSliceWithInitializer {
     #[primary_span]
     pub span: Span,
 }
+
+#[derive(Diagnostic)]
+#[diag(ast_lowering_distributed_slice_elements_wrong_expr)]
+#[note]
+pub(crate) struct DistributedSliceElementsWrongExpr {
+    #[primary_span]
+    pub span: Span,
+}

compiler/rustc_ast_lowering/src/item.rs

@@ -6,8 +6,8 @@ use rustc_errors::ErrorGuaranteed;
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::{CRATE_DEF_ID, LocalDefId};
 use rustc_hir::{
-    self as hir, DistributedSlice, HirId, InvalidDistributedSliceDeclaration, LifetimeSource,
-    PredicateOrigin,
+    self as hir, DistributedSlice, DistributedSliceAdditionManyKind, HirId,
+    InvalidDistributedSliceDeclaration, LifetimeSource, PredicateOrigin,
 };
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::ty::{ResolverAstLowering, TyCtxt};
@@ -25,7 +25,7 @@ use super::{
     AstOwner, FnDeclKind, ImplTraitContext, ImplTraitPosition, LoweringContext, ParamMode,
     ResolverAstLoweringExt,
 };
-use crate::errors::DistributedSliceWithInitializer;
+use crate::errors::{DistributedSliceElementsWrongExpr, DistributedSliceWithInitializer};
 
 pub(super) struct ItemLowerer<'a, 'hir> {
     pub(super) tcx: TyCtxt<'hir>,
@@ -155,6 +155,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     fn lower_distributed_slice(
         &mut self,
         distributed_slice: &ast::DistributedSlice,
+        expr: Option<&Expr>,
     ) -> DistributedSlice {
         match distributed_slice {
             ast::DistributedSlice::None => DistributedSlice::None,
@@ -179,6 +180,58 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
                 DistributedSlice::Addition(local)
             }
+            ast::DistributedSlice::AdditionMany { declaration, id } => {
+                let Some(res) = self.resolver.get_partial_res(*id) else {
+                    self.dcx().span_delayed_bug(declaration.span, "should have errored in resolve");
+                    return DistributedSlice::None;
+                };
+
+                let Some(did) = res.expect_full_res().opt_def_id() else {
+                    self.dcx().span_delayed_bug(declaration.span, "should have errored in resolve");
+                    return DistributedSlice::None;
+                };
+
+                let Some(local) = did.as_local() else {
+                    panic!("adding to slice outside local crate");
+                };
+
+                let initializer = expr
+                    .expect("generated by `distributed_slice_elements!` and always has an expr");
+
+                DistributedSlice::AdditionMany(
+                    local,
+                    match &initializer.kind {
+                        ExprKind::Array(elems) => {
+                            DistributedSliceAdditionManyKind::ArrayLit { length: elems.len() }
+                        }
+                        ExprKind::Path(_, _) => {
+                            let Some(res) = self.resolver.get_partial_res(initializer.id) else {
+                                self.dcx().span_delayed_bug(
+                                    declaration.span,
+                                    "should have errored in resolve",
+                                );
+                                return DistributedSlice::None;
+                            };
+
+                            let Some(did) = res.expect_full_res().opt_def_id() else {
+                                self.dcx().span_delayed_bug(
+                                    declaration.span,
+                                    "should have errored in resolve",
+                                );
+                                return DistributedSlice::None;
+                            };
+
+                            DistributedSliceAdditionManyKind::Path { res: did }
+                        }
+                        _ => {
+                            let eg = self.dcx().emit_err(DistributedSliceElementsWrongExpr {
+                                span: initializer.span,
+                            });
+                            DistributedSliceAdditionManyKind::Err(eg)
+                        }
+                    },
+                )
+            }
         }
     }
 
@@ -225,7 +278,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     ident,
                     ty,
                     body_id,
-                    self.lower_distributed_slice(distributed_slice),
+                    self.lower_distributed_slice(distributed_slice, e.as_deref()),
                 )
             }
             ItemKind::Const(box ast::ConstItem {
@@ -258,7 +311,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     generics,
                     ty,
                     body_id,
-                    self.lower_distributed_slice(distributed_slice),
+                    self.lower_distributed_slice(distributed_slice, expr.as_deref()),
                 )
             }
             ItemKind::Fn(box Fn {

compiler/rustc_builtin_macros/src/distributed_slice.rs

@@ -143,3 +143,45 @@ pub(crate) fn distributed_slice_element(
         tokens: None
     })]))
 }
+
+/// ```rust
+/// distributed_slice_elements!(MEOWS, ["mrow"]);
+/// ```
+pub(crate) fn distributed_slice_elements(
+    cx: &mut ExtCtxt<'_>,
+    span: Span,
+    tts: TokenStream,
+) -> MacroExpanderResult<'static> {
+    let (path, expr) = match parse_element(cx.new_parser_from_tts(tts)) {
+        Ok((ident, expr)) => (ident, expr),
+        Err(mut err) => {
+            if err.span.is_dummy() {
+                err.span(span);
+            }
+            let guar = err.emit();
+            return ExpandResult::Ready(DummyResult::any(span, guar));
+        }
+    };
+
+    ExpandResult::Ready(MacEager::items(smallvec![P(Item {
+        attrs: ThinVec::new(),
+        id: DUMMY_NODE_ID,
+        span,
+        vis: ast::Visibility { kind: ast::VisibilityKind::Inherited, span, tokens: None },
+        kind: ItemKind::Const(Box::new(ConstItem {
+            defaultness: Defaultness::Final,
+            ident: Ident { name: kw::Underscore, span },
+            generics: Generics::default(),
+            // leave out the ty, we discover it when
+            // when name-resolving to the registry definition
+            ty: P(Ty { id: DUMMY_NODE_ID, kind: TyKind::Infer, span, tokens: None }),
+            expr: Some(expr),
+            define_opaque: None,
+            distributed_slice: DistributedSlice::AdditionMany {
+                declaration: path,
+                id: DUMMY_NODE_ID
+            }
+        })),
+        tokens: None
+    })]))
+}

compiler/rustc_builtin_macros/src/lib.rs

@@ -89,6 +89,7 @@ pub fn register_builtin_macros(resolver: &mut dyn ResolverExpand) {
         const_format_args: format::expand_format_args,
         core_panic: edition_panic::expand_panic,
         distributed_slice_element: distributed_slice::distributed_slice_element,
+        distributed_slice_elements: distributed_slice::distributed_slice_elements,
         env: env::expand_env,
         file: source_util::expand_file,
         format_args: format::expand_format_args,

compiler/rustc_hir/src/hir.rs

@@ -4271,12 +4271,20 @@ impl FnHeader {
     }
 }
 
+#[derive(Debug, Clone, Copy, HashStable_Generic)]
+pub enum DistributedSliceAdditionManyKind {
+    ArrayLit { length: usize },
+    Path { res: DefId },
+    Err(ErrorGuaranteed),
+}
+
 #[derive(Debug, Clone, Copy, HashStable_Generic, Default)]
 pub enum DistributedSlice {
     #[default]
     None,
     Declaration(Span),
     Addition(LocalDefId),
+    AdditionMany(LocalDefId, DistributedSliceAdditionManyKind),
 }
 
 #[derive(Debug, Clone, Copy, HashStable_Generic)]

compiler/rustc_hir_analysis/src/collect/distributed_slice.rs

@@ -1,25 +1,38 @@
 use rand::SeedableRng;
 use rand::seq::SliceRandom;
-use rustc_hir::def_id::{DefIdMap, LocalDefId};
+use rustc_hir::def_id::DefIdMap;
 use rustc_hir::{DistributedSlice, ItemKind};
+use rustc_middle::middle::distributed_slice::DistributedSliceAddition;
 use rustc_middle::ty::TyCtxt;
 
 pub(super) fn distributed_slice_elements<'tcx>(
     tcx: TyCtxt<'tcx>,
     _: (),
-) -> DefIdMap<Vec<LocalDefId>> {
-    let mut slice_elements = DefIdMap::<Vec<LocalDefId>>::default();
+) -> DefIdMap<Vec<DistributedSliceAddition>> {
+    let mut slice_elements = DefIdMap::<Vec<DistributedSliceAddition>>::default();
 
     for i in tcx.hir_free_items() {
         let addition_def_id = i.owner_id.def_id;
         if let ItemKind::Const(.., DistributedSlice::Addition(declaration_def_id)) =
             tcx.hir_expect_item(addition_def_id).kind
         {
-            slice_elements.entry(declaration_def_id.to_def_id()).or_default().push(addition_def_id);
+            slice_elements
+                .entry(declaration_def_id.to_def_id())
+                .or_default()
+                .push(DistributedSliceAddition::Single(addition_def_id));
+        }
+
+        if let ItemKind::Const(.., DistributedSlice::AdditionMany(declaration_def_id, _)) =
+            tcx.hir_expect_item(addition_def_id).kind
+        {
+            slice_elements
+                .entry(declaration_def_id.to_def_id())
+                .or_default()
+                .push(DistributedSliceAddition::Many(addition_def_id));
         }
     }
 
-    let mut res = DefIdMap::<Vec<LocalDefId>>::default();
+    let mut res = DefIdMap::<Vec<DistributedSliceAddition>>::default();
 
     for (key, mut registered_values) in
         tcx.with_stable_hashing_context(|hcx| slice_elements.into_sorted(&hcx, true))