distributed_slice_elements (multiple) plus tests

Author: jdonszelmann

compiler/rustc_ast/src/ast.rs

@@ -3611,6 +3611,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

@@ -193,3 +193,7 @@ ast_lowering_yield_in_closure =
 
 ast_lowering_distributed_slice_with_initializer =
     distributed slice elements are added with `distributed_slice_element!(...)`
+
+ast_lowering_distributed_slice_elements_wrong_expr =
+    `distributed_slice_elements!()` only accepts a path or array literal
+    .note = arbitrary expressions are not supported

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

@@ -5,7 +5,7 @@ use rustc_ast::*;
 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};
+use rustc_hir::{self as hir, DistributedSlice, DistributedSliceAdditionManyKind, HirId, InvalidDistributedSliceDeclaration, LifetimeSource, PredicateOrigin};
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::ty::{ResolverAstLowering, TyCtxt};
 use rustc_span::edit_distance::find_best_match_for_name;
@@ -14,7 +14,7 @@ use smallvec::{SmallVec, smallvec};
 use thin_vec::ThinVec;
 use tracing::instrument;
 
-use crate::errors::DistributedSliceWithInitializer;
+use crate::errors::{DistributedSliceElementsWrongExpr, DistributedSliceWithInitializer};
 
 use super::errors::{
     InvalidAbi, InvalidAbiSuggestion, MisplacedRelaxTraitBound, TupleStructWithDefault,
@@ -153,6 +153,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,
@@ -177,6 +178,49 @@ 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)
+                        }
+                    }
+                )
+            }
         }
     }
 
@@ -223,7 +267,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     ty,
                     *m,
                     body_id,
-                    self.lower_distributed_slice(distributed_slice),
+                    self.lower_distributed_slice(distributed_slice, e.as_deref()),
                 )
             }
             ItemKind::Const(box ast::ConstItem {
@@ -256,7 +300,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     ty,
                     generics,
                     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

@@ -126,3 +126,42 @@ 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

@@ -91,6 +91,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,24 @@ 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/messages.ftl

@@ -631,3 +631,4 @@ hir_analysis_distributed_slice_non_array =
     expected this type to be an array
     .note = a distributed slice must have an array type with an inferred length
     .suggestion = change the type to an array
+

compiler/rustc_hir_analysis/src/collect/distributed_slice.rs

@@ -1,25 +1,32 @@
 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))