feat: add expansion_grow_limit attr as another expansion limit

Signed-off-by: Vincenzo Palazzo <[email protected]>

Author: vincenzopalazzo

compiler/rustc_expand/src/base.rs

@@ -1008,6 +1008,7 @@ pub struct ExtCtxt<'a> {
     pub ecfg: expand::ExpansionConfig<'a>,
     pub num_standard_library_imports: usize,
     pub reduced_recursion_limit: Option<Limit>,
+    pub reduced_expansion_growth_limit: Option<Limit>,
     pub root_path: PathBuf,
     pub resolver: &'a mut dyn ResolverExpand,
     pub current_expansion: ExpansionData,
@@ -1037,6 +1038,7 @@ impl<'a> ExtCtxt<'a> {
             ecfg,
             num_standard_library_imports: 0,
             reduced_recursion_limit: None,
+            reduced_expansion_growth_limit: None,
             resolver,
             lint_store,
             root_path: PathBuf::new(),

compiler/rustc_expand/src/expand.rs

@@ -621,6 +621,30 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
         self.cx.trace_macros_diag();
     }
 
+    fn error_expansion_growth_limit(&mut self) {
+        let expn_data = self.cx.current_expansion.id.expn_data();
+        let suggested_limit = match self.cx.ecfg.expansion_growth_limit {
+            Limit(0) => Limit(2),
+            limit => limit * 2,
+        };
+
+        self.cx
+            .struct_span_err(
+                expn_data.call_site,
+                &format!(
+                    "expansion grow limit reached while expanding `{}`",
+                    expn_data.kind.descr()
+                ),
+            )
+            .help(&format!(
+                "consider increasing the expansion grow limit by adding a \
+                 `#![expansion_growth_limit = \"{}\"]` attribute to your crate (`{}`)",
+                suggested_limit, self.cx.ecfg.crate_name,
+            ))
+            .emit();
+        self.cx.trace_macros_diag();
+    }
+
     /// A macro's expansion does not fit in this fragment kind.
     /// For example, a non-type macro in a type position.
     fn error_wrong_fragment_kind(&mut self, kind: AstFragmentKind, mac: &ast::MacCall, span: Span) {
@@ -629,6 +653,22 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
         self.cx.trace_macros_diag();
     }
 
+    fn reduce_expansion_growth_limit(&mut self, tokens_solved_size: usize) -> Result<(), ()> {
+        let expansion_limit =
+            self.cx.reduced_expansion_growth_limit.unwrap_or(self.cx.ecfg.expansion_growth_limit);
+        if !expansion_limit.value_within_limit(tokens_solved_size) {
+            if self.cx.reduced_expansion_growth_limit.is_none() {
+                self.error_expansion_growth_limit();
+            }
+
+            // Reduce the recursion limit by half each time it triggers.
+            self.cx.reduced_recursion_limit = Some(expansion_limit / 2);
+
+            return Err(());
+        }
+        Ok(())
+    }
+
     fn expand_invoc(
         &mut self,
         invoc: Invocation,
@@ -657,6 +697,11 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
                     self.parse_ast_fragment(tok_result, fragment_kind, &mac.path, span)
                 }
                 SyntaxExtensionKind::LegacyBang(expander) => {
+                    if self.reduce_expansion_growth_limit(mac.args.inner_tokens().len()).is_err() {
+                        return ExpandResult::Ready(fragment_kind.dummy(span));
+                    }
+                    let prev = self.cx.current_expansion.prior_type_ascription;
+                    self.cx.current_expansion.prior_type_ascription = mac.prior_type_ascription;
                     let tok_result = expander.expand(self.cx, span, mac.args.tokens.clone());
                     let result = if let Some(result) = fragment_kind.make_from(tok_result) {
                         result
@@ -1966,6 +2011,7 @@ pub struct ExpansionConfig<'feat> {
     pub crate_name: String,
     pub features: Option<&'feat Features>,
     pub recursion_limit: Limit,
+    pub expansion_growth_limit: Limit,
     pub trace_mac: bool,
     /// If false, strip `#[test]` nodes
     pub should_test: bool,
@@ -1981,6 +2027,7 @@ impl<'feat> ExpansionConfig<'feat> {
             crate_name,
             features: None,
             recursion_limit: Limit::new(1024),
+            expansion_growth_limit: Limit::new(1500),
             trace_mac: false,
             should_test: false,
             span_debug: false,

compiler/rustc_feature/src/builtin_attrs.rs

@@ -354,6 +354,7 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
 
     // Limits:
     ungated!(recursion_limit, CrateLevel, template!(NameValueStr: "N"), FutureWarnFollowing),
+    ungated!(expansion_growth_limit, CrateLevel, template!(NameValueStr: "N"), FutureWarnFollowing),
     ungated!(type_length_limit, CrateLevel, template!(NameValueStr: "N"), FutureWarnFollowing),
     gated!(
         move_size_limit, CrateLevel, template!(NameValueStr: "N"), ErrorFollowing,

compiler/rustc_interface/src/passes.rs

@@ -226,9 +226,11 @@ fn configure_and_expand(
         // Create the config for macro expansion
         let features = sess.features_untracked();
         let recursion_limit = get_recursion_limit(pre_configured_attrs, sess);
+        let expansion_growth_limit = get_expansion_growth_limit(&krate.attrs, sess);
         let cfg = rustc_expand::expand::ExpansionConfig {
             features: Some(features),
             recursion_limit,
+            expansion_growth_limit,
             trace_mac: sess.opts.unstable_opts.trace_macros,
             should_test: sess.is_test_crate(),
             span_debug: sess.opts.unstable_opts.span_debug,
@@ -903,3 +905,17 @@ fn get_recursion_limit(krate_attrs: &[ast::Attribute], sess: &Session) -> Limit
     }
     rustc_middle::middle::limits::get_recursion_limit(krate_attrs, sess)
 }
+
+fn get_expansion_growth_limit(krate_attrs: &[ast::Attribute], sess: &Session) -> Limit {
+    if let Some(attr) = krate_attrs
+        .iter()
+        .find(|attr| attr.has_name(sym::expansion_growth_limit) && attr.value_str().is_none())
+    {
+        validate_attr::emit_fatal_malformed_builtin_attribute(
+            &sess.parse_sess,
+            attr,
+            sym::expansion_growth_limit,
+        );
+    }
+    rustc_middle::middle::limits::get_expansion_growth_limit(krate_attrs, sess)
+}

compiler/rustc_middle/src/middle/limits.rs

@@ -21,6 +21,7 @@ use std::num::IntErrorKind;
 pub fn provide(providers: &mut Providers) {
     providers.limits = |tcx, ()| Limits {
         recursion_limit: get_recursion_limit(tcx.hir().krate_attrs(), tcx.sess),
+        expansion_growth_limit: get_expansion_growth_limit(tcx.hir().krate_attrs(), tcx.sess),
         move_size_limit: get_limit(
             tcx.hir().krate_attrs(),
             tcx.sess,
@@ -40,6 +41,10 @@ pub fn get_recursion_limit(krate_attrs: &[Attribute], sess: &Session) -> Limit {
     get_limit(krate_attrs, sess, sym::recursion_limit, 128)
 }
 
+pub fn get_expansion_growth_limit(krate_attrs: &[Attribute], sess: &Session) -> Limit {
+    get_limit(krate_attrs, sess, sym::expansion_growth_limit, 1500)
+}
+
 fn get_limit(krate_attrs: &[Attribute], sess: &Session, name: Symbol, default: usize) -> Limit {
     for attr in krate_attrs {
         if !attr.has_name(name) {

compiler/rustc_middle/src/ty/context.rs

@@ -1175,6 +1175,10 @@ impl<'tcx> TyCtxt<'tcx> {
         self.limits(()).recursion_limit
     }
 
+    pub fn expansion_growth_limit(self) -> Limit {
+        self.limits(()).expansion_growth_limit
+    }
+
     pub fn move_size_limit(self) -> Limit {
         self.limits(()).move_size_limit
     }

compiler/rustc_session/src/session.rs

@@ -128,6 +128,8 @@ pub struct Limits {
     pub move_size_limit: Limit,
     /// The maximum length of types during monomorphization.
     pub type_length_limit: Limit,
+    /// FIXME(vincenzopalazzo) add docs
+    pub expansion_growth_limit: Limit,
 }
 
 pub struct CompilerIO {

compiler/rustc_span/src/symbol.rs

@@ -682,6 +682,7 @@ symbols! {
         existential_type,
         exp2f32,
         exp2f64,
+        expansion_growth_limit,
         expect,
         expected,
         expf32,