Add runtime check to avoid overwrite arg easily in diag and store and restore snapshot when set subdiag arg

Signed-off-by: xizheyin <[email protected]>

Author: xizheyin

compiler/rustc_borrowck/src/diagnostics/mod.rs

@@ -1254,8 +1254,14 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                         && !spans.is_empty()
                     {
                         let mut span: MultiSpan = spans.clone().into();
+                        err.arg("ty", param_ty.to_string());
+                        let msg = err.dcx.eagerly_translate_to_string(
+                            fluent::borrowck_moved_a_fn_once_in_call_def,
+                            err.args.iter(),
+                        );
+                        err.remove_arg("ty");
                         for sp in spans {
-                            span.push_span_label(sp, fluent::borrowck_moved_a_fn_once_in_call_def);
+                            span.push_span_label(sp, msg.clone());
                         }
                         span.push_span_label(
                             fn_call_span,

compiler/rustc_builtin_macros/src/errors.rs

@@ -672,6 +672,7 @@ impl Subdiagnostic for FormatUnusedArg {
     fn add_to_diag<G: EmissionGuarantee>(self, diag: &mut Diag<'_, G>) {
         diag.arg("named", self.named);
         let msg = diag.eagerly_translate(crate::fluent_generated::builtin_macros_format_unused_arg);
+        diag.remove_arg("named");
         diag.span_label(self.span, msg);
     }
 }

compiler/rustc_const_eval/src/errors.rs

@@ -293,6 +293,9 @@ impl Subdiagnostic for FrameNote {
             span.push_span_label(self.span, fluent::const_eval_frame_note_last);
         }
         let msg = diag.eagerly_translate(fluent::const_eval_frame_note);
+        diag.remove_arg("times");
+        diag.remove_arg("where_");
+        diag.remove_arg("instance");
         diag.span_note(span, msg);
     }
 }

compiler/rustc_errors/src/diagnostic.rs

@@ -289,6 +289,9 @@ pub struct DiagInner {
     pub suggestions: Suggestions,
     pub args: DiagArgMap,
 
+    // This is used to store args and restore them after a subdiagnostic is rendered.
+    pub reserved_args: DiagArgMap,
+
     /// This is not used for highlighting or rendering any error message. Rather, it can be used
     /// as a sort key to sort a buffer of diagnostics. By default, it is the primary span of
     /// `span` if there is one. Otherwise, it is `DUMMY_SP`.
@@ -319,6 +322,7 @@ impl DiagInner {
             children: vec![],
             suggestions: Suggestions::Enabled(vec![]),
             args: Default::default(),
+            reserved_args: Default::default(),
             sort_span: DUMMY_SP,
             is_lint: None,
             long_ty_path: None,
@@ -390,7 +394,27 @@ impl DiagInner {
     }
 
     pub(crate) fn arg(&mut self, name: impl Into<DiagArgName>, arg: impl IntoDiagArg) {
-        self.args.insert(name.into(), arg.into_diag_arg(&mut self.long_ty_path));
+        let name = name.into();
+        let value = arg.into_diag_arg(&mut self.long_ty_path);
+        // This assertion is to avoid subdiagnostics overwriting an existing diagnostic arg.
+        debug_assert!(
+            !self.args.contains_key(&name) || self.args.get(&name) == Some(&value),
+            "arg {} already exists",
+            name
+        );
+        self.args.insert(name, value);
+    }
+
+    pub fn remove_arg(&mut self, name: &str) {
+        self.args.swap_remove(name);
+    }
+
+    pub fn store_args(&mut self) {
+        self.reserved_args = self.args.clone();
+    }
+
+    pub fn restore_args(&mut self) {
+        self.args = std::mem::take(&mut self.reserved_args);
     }
 
     /// Fields used for Hash, and PartialEq trait.
@@ -1423,6 +1447,12 @@ impl<'a, G: EmissionGuarantee> Diag<'a, G> {
         self.downgrade_to_delayed_bug();
         self.emit()
     }
+
+    pub fn remove_arg(&mut self, name: &str) {
+        if let Some(diag) = self.diag.as_mut() {
+            diag.remove_arg(name);
+        }
+    }
 }
 
 /// Destructor bomb: every `Diag` must be consumed (emitted, cancelled, etc.)

compiler/rustc_hir_analysis/src/errors.rs

@@ -127,6 +127,7 @@ pub(crate) enum AssocItemNotFoundSugg<'a> {
     SimilarInOtherTrait {
         #[primary_span]
         span: Span,
+        trait_name: &'a str,
         assoc_kind: &'static str,
         suggested_name: Symbol,
     },

compiler/rustc_hir_analysis/src/hir_ty_lowering/errors.rs

@@ -309,6 +309,7 @@ impl<'tcx> dyn HirTyLowerer<'tcx> + '_ {
                         // change the associated item.
                         err.sugg = Some(errors::AssocItemNotFoundSugg::SimilarInOtherTrait {
                             span: assoc_ident.span,
+                            trait_name: &trait_name,
                             assoc_kind: assoc_kind_str,
                             suggested_name,
                         });

compiler/rustc_lint/src/levels.rs

@@ -482,7 +482,8 @@ impl<'s, P: LintLevelsProvider> LintLevelsBuilder<'s, P> {
                     let name = lint_name.as_str();
                     let suggestion = RenamedLintSuggestion::WithoutSpan { replace };
                     let requested_level = RequestedLevel { level, lint_name };
-                    let lint = RenamedLintFromCommandLine { name, suggestion, requested_level };
+                    let lint =
+                        RenamedLintFromCommandLine { name, replace, suggestion, requested_level };
                     self.emit_lint(RENAMED_AND_REMOVED_LINTS, lint);
                 }
                 CheckLintNameResult::Removed(ref reason) => {
@@ -824,7 +825,7 @@ impl<'s, P: LintLevelsProvider> LintLevelsBuilder<'s, P> {
                                 RenamedLintSuggestion::WithSpan { suggestion: sp, replace };
                             let name =
                                 tool_ident.map(|tool| format!("{tool}::{name}")).unwrap_or(name);
-                            let lint = RenamedLint { name: name.as_str(), suggestion };
+                            let lint = RenamedLint { name: name.as_str(), replace, suggestion };
                             self.emit_span_lint(RENAMED_AND_REMOVED_LINTS, sp.into(), lint);
                         }
 

compiler/rustc_lint/src/lifetime_syntax.rs

@@ -422,12 +422,12 @@ fn build_mismatch_suggestion(
     lifetime_name: &str,
     infos: &[&Info<'_>],
 ) -> lints::MismatchedLifetimeSyntaxesSuggestion {
-    let lifetime_name = lifetime_name.to_owned();
+    let lifetime_name_sugg = lifetime_name.to_owned();
 
     let suggestions = infos.iter().map(|info| info.suggestion(&lifetime_name)).collect();
 
     lints::MismatchedLifetimeSyntaxesSuggestion::Explicit {
-        lifetime_name,
+        lifetime_name_sugg,
         suggestions,
         tool_only: false,
     }

compiler/rustc_lint/src/lints.rs

@@ -1089,6 +1089,7 @@ pub(crate) struct DeprecatedLintNameFromCommandLine<'a> {
 #[diag(lint_renamed_lint)]
 pub(crate) struct RenamedLint<'a> {
     pub name: &'a str,
+    pub replace: &'a str,
     #[subdiagnostic]
     pub suggestion: RenamedLintSuggestion<'a>,
 }
@@ -1109,6 +1110,7 @@ pub(crate) enum RenamedLintSuggestion<'a> {
 #[diag(lint_renamed_lint)]
 pub(crate) struct RenamedLintFromCommandLine<'a> {
     pub name: &'a str,
+    pub replace: &'a str,
     #[subdiagnostic]
     pub suggestion: RenamedLintSuggestion<'a>,
     #[subdiagnostic]
@@ -3227,7 +3229,7 @@ pub(crate) enum MismatchedLifetimeSyntaxesSuggestion {
     },
 
     Explicit {
-        lifetime_name: String,
+        lifetime_name_sugg: String,
         suggestions: Vec<(Span, String)>,
         tool_only: bool,
     },
@@ -3281,13 +3283,12 @@ impl Subdiagnostic for MismatchedLifetimeSyntaxesSuggestion {
                 );
             }
 
-            Explicit { lifetime_name, suggestions, tool_only } => {
-                diag.arg("lifetime_name", lifetime_name);
-
+            Explicit { lifetime_name_sugg, suggestions, tool_only } => {
+                diag.arg("lifetime_name_sugg", lifetime_name_sugg);
                 let msg = diag.eagerly_translate(
                     fluent::lint_mismatched_lifetime_syntaxes_suggestion_explicit,
                 );
-
+                diag.remove_arg("lifetime_name_sugg");
                 diag.multipart_suggestion_with_style(
                     msg,
                     suggestions,

compiler/rustc_macros/src/diagnostics/subdiagnostic.rs

@@ -220,7 +220,7 @@ impl<'parent, 'a> SubdiagnosticDeriveVariantBuilder<'parent, 'a> {
     }
 
     /// Generates the code for a field with no attributes.
-    fn generate_field_arg(&mut self, binding_info: &BindingInfo<'_>) -> TokenStream {
+    fn generate_field_arg(&mut self, binding_info: &BindingInfo<'_>) -> (TokenStream, TokenStream) {
         let diag = &self.parent.diag;
 
         let field = binding_info.ast();
@@ -230,12 +230,16 @@ impl<'parent, 'a> SubdiagnosticDeriveVariantBuilder<'parent, 'a> {
         let ident = field.ident.as_ref().unwrap();
         let ident = format_ident!("{}", ident); // strip `r#` prefix, if present
 
-        quote! {
+        let args = quote! {
             #diag.arg(
                 stringify!(#ident),
                 #field_binding
             );
-        }
+        };
+        let remove_args = quote! {
+            #diag.remove_arg(stringify!(#ident));
+        };
+        (args, remove_args)
     }
 
     /// Generates the necessary code for all attributes on a field.
@@ -600,8 +604,13 @@ impl<'parent, 'a> SubdiagnosticDeriveVariantBuilder<'parent, 'a> {
 
             calls.extend(call);
         }
-
-        let plain_args: TokenStream = self
+        let store_args = quote! {
+            #diag.store_args();
+        };
+        let restore_args = quote! {
+            #diag.restore_args();
+        };
+        let (plain_args, remove_args): (TokenStream, TokenStream) = self
             .variant
             .bindings()
             .iter()
@@ -610,12 +619,23 @@ impl<'parent, 'a> SubdiagnosticDeriveVariantBuilder<'parent, 'a> {
             .collect();
 
         let formatting_init = &self.formatting_init;
+
+        // For #[derive(Subdiagnostic)]
+        //
+        // - Store args of the main diagnostic for later restore.
+        // - add args of subdiagnostic.
+        // - Generate the calls, such as note, label, etc.
+        // - Remove the arguments for allowing Vec<Subdiagnostic> to be used.
+        // - Restore the arguments for allowing main and subdiagnostic share the same fields.
         Ok(quote! {
             #init
             #formatting_init
             #attr_args
+            #store_args
             #plain_args
             #calls
+            #remove_args
+            #restore_args
         })
     }
 }