Re-implement a type-size based limit

8 files changed, 128 insertions, 95 deletions

diff --git a/compiler/rustc_middle/messages.ftl b/compiler/rustc_middle/messages.ftl
index f4d619329eb..2b9d9a07a98 100644
--- a/compiler/rustc_middle/messages.ftl
+++ b/compiler/rustc_middle/messages.ftl
@@ -41,6 +41,9 @@ middle_cannot_be_normalized =
 middle_conflict_types =
     this expression supplies two conflicting concrete types for the same opaque type
 
+middle_consider_type_length_limit =
+    consider adding a `#![type_length_limit="{$type_length}"]` attribute to your crate
+
 middle_const_eval_non_int =
     constant evaluation of enum discriminant resulted in non-integer
 
@@ -94,8 +97,11 @@ middle_strict_coherence_needs_negative_coherence =
     to use `strict_coherence` on this trait, the `with_negative_coherence` feature must be enabled
     .label = due to this attribute
 
+middle_type_length_limit = reached the type-length limit while instantiating `{$shrunk}`
+
 middle_unknown_layout =
     the type `{$ty}` has an unknown layout
 
 middle_values_too_big =
     values of the type `{$ty}` are too big for the current architecture
+middle_written_to_path = the full type name has been written to '{$path}'
diff --git a/compiler/rustc_middle/src/error.rs b/compiler/rustc_middle/src/error.rs
index 6e622b0405f..711db4e0a6b 100644
--- a/compiler/rustc_middle/src/error.rs
+++ b/compiler/rustc_middle/src/error.rs
@@ -1,4 +1,5 @@
 use std::fmt;
+use std::path::PathBuf;
 
 use rustc_errors::{codes::*, DiagArgName, DiagArgValue, DiagMessage};
 use rustc_macros::{Diagnostic, Subdiagnostic};
@@ -149,3 +150,16 @@ pub struct ErroneousConstant {
 
 /// Used by `rustc_const_eval`
 pub use crate::fluent_generated::middle_adjust_for_foreign_abi_error;
+
+#[derive(Diagnostic)]
+#[diag(middle_type_length_limit)]
+#[help(middle_consider_type_length_limit)]
+pub struct TypeLengthLimit {
+    #[primary_span]
+    pub span: Span,
+    pub shrunk: String,
+    #[note(middle_written_to_path)]
+    pub was_written: Option<()>,
+    pub path: PathBuf,
+    pub type_length: usize,
+}
diff --git a/compiler/rustc_middle/src/middle/limits.rs b/compiler/rustc_middle/src/middle/limits.rs
index 4d698012749..d0b4f36a426 100644
--- a/compiler/rustc_middle/src/middle/limits.rs
+++ b/compiler/rustc_middle/src/middle/limits.rs
@@ -30,7 +30,7 @@ pub fn provide(providers: &mut Providers) {
             tcx.hir().krate_attrs(),
             tcx.sess,
             sym::type_length_limit,
-            1048576,
+            2usize.pow(24),
         ),
     }
 }
diff --git a/compiler/rustc_middle/src/ty/instance.rs b/compiler/rustc_middle/src/ty/instance.rs
index 6d7c04fca4b..65039d8b52c 100644
--- a/compiler/rustc_middle/src/ty/instance.rs
+++ b/compiler/rustc_middle/src/ty/instance.rs
@@ -1,23 +1,25 @@
+use crate::error;
 use crate::middle::codegen_fn_attrs::CodegenFnAttrFlags;
-use crate::ty::print::{FmtPrinter, Printer};
-use crate::ty::{self, Ty, TyCtxt, TypeFoldable, TypeSuperFoldable};
-use crate::ty::{EarlyBinder, GenericArgs, GenericArgsRef, TypeVisitableExt};
+use crate::ty::print::{shrunk_instance_name, FmtPrinter, Printer};
+use crate::ty::{
+    self, EarlyBinder, GenericArgs, GenericArgsRef, Ty, TyCtxt, TypeFoldable, TypeSuperFoldable,
+    TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor,
+};
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir as hir;
 use rustc_hir::def::Namespace;
 use rustc_hir::def_id::{CrateNum, DefId};
 use rustc_hir::lang_items::LangItem;
 use rustc_index::bit_set::FiniteBitSet;
-use rustc_macros::{
-    Decodable, Encodable, HashStable, Lift, TyDecodable, TyEncodable, TypeVisitable,
-};
+use rustc_macros::{Decodable, Encodable, HashStable, Lift, TyDecodable, TyEncodable};
 use rustc_middle::ty::normalize_erasing_regions::NormalizationError;
 use rustc_span::def_id::LOCAL_CRATE;
-use rustc_span::{Span, Symbol};
+use rustc_span::{Span, Symbol, DUMMY_SP};
 use tracing::{debug, instrument};
 
 use std::assert_matches::assert_matches;
 use std::fmt;
+use std::path::PathBuf;
 
 /// An `InstanceKind` along with the args that are needed to substitute the instance.
 ///
@@ -385,7 +387,28 @@ impl<'tcx> InstanceKind<'tcx> {
     }
 }
 
-fn fmt_instance(
+fn type_length<'tcx>(item: impl TypeVisitable<TyCtxt<'tcx>>) -> usize {
+    struct Visitor {
+        type_length: usize,
+    }
+    impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for Visitor {
+        fn visit_ty(&mut self, t: Ty<'tcx>) {
+            self.type_length += 1;
+            t.super_visit_with(self);
+        }
+
+        fn visit_const(&mut self, ct: ty::Const<'tcx>) {
+            self.type_length += 1;
+            ct.super_visit_with(self);
+        }
+    }
+    let mut visitor = Visitor { type_length: 0 };
+    item.visit_with(&mut visitor);
+
+    visitor.type_length
+}
+
+pub fn fmt_instance(
     f: &mut fmt::Formatter<'_>,
     instance: Instance<'_>,
     type_length: Option<rustc_session::Limit>,
@@ -485,7 +508,8 @@ impl<'tcx> Instance<'tcx> {
     ///
     /// Presuming that coherence and type-check have succeeded, if this method is invoked
     /// in a monomorphic context (i.e., like during codegen), then it is guaranteed to return
-    /// `Ok(Some(instance))`.
+    /// `Ok(Some(instance))`, **except** for when the instance's inputs hit the type size limit,
+    /// in which case it may bail out and return `Ok(None)`.
     ///
     /// Returns `Err(ErrorGuaranteed)` when the `Instance` resolution process
     /// couldn't complete due to errors elsewhere - this is distinct
@@ -498,6 +522,16 @@ impl<'tcx> Instance<'tcx> {
         def_id: DefId,
         args: GenericArgsRef<'tcx>,
     ) -> Result<Option<Instance<'tcx>>, ErrorGuaranteed> {
+        // Rust code can easily create exponentially-long types using only a
+        // polynomial recursion depth. Even with the default recursion
+        // depth, you can easily get cases that take >2^60 steps to run,
+        // which means that rustc basically hangs.
+        //
+        // Bail out in these cases to avoid that bad user experience.
+        if !tcx.type_length_limit().value_within_limit(type_length(args)) {
+            return Ok(None);
+        }
+
         // All regions in the result of this query are erased, so it's
         // fine to erase all of the input regions.
 
@@ -517,6 +551,36 @@ impl<'tcx> Instance<'tcx> {
     ) -> Instance<'tcx> {
         match ty::Instance::resolve(tcx, param_env, def_id, args) {
             Ok(Some(instance)) => instance,
+            Ok(None) => {
+                let type_length = type_length(args);
+                if !tcx.type_length_limit().value_within_limit(type_length) {
+                    let (shrunk, written_to_path) =
+                        shrunk_instance_name(tcx, Instance::new(def_id, args));
+                    let mut path = PathBuf::new();
+                    let was_written = if let Some(path2) = written_to_path {
+                        path = path2;
+                        Some(())
+                    } else {
+                        None
+                    };
+                    tcx.dcx().emit_fatal(error::TypeLengthLimit {
+                        // We don't use `def_span(def_id)` so that diagnostics point
+                        // to the crate root during mono instead of to foreign items.
+                        // This is arguably better.
+                        span: span.unwrap_or(DUMMY_SP),
+                        shrunk,
+                        was_written,
+                        path,
+                        type_length,
+                    });
+                } else {
+                    span_bug!(
+                        span.unwrap_or(tcx.def_span(def_id)),
+                        "failed to resolve instance for {}",
+                        tcx.def_path_str_with_args(def_id, args)
+                    )
+                }
+            }
             instance => span_bug!(
                 span.unwrap_or(tcx.def_span(def_id)),
                 "failed to resolve instance for {}: {instance:#?}",
diff --git a/compiler/rustc_middle/src/ty/print/mod.rs b/compiler/rustc_middle/src/ty/print/mod.rs
index 3c27df9529a..c165790548d 100644
--- a/compiler/rustc_middle/src/ty/print/mod.rs
+++ b/compiler/rustc_middle/src/ty/print/mod.rs
@@ -1,5 +1,7 @@
+use std::path::PathBuf;
+
 use crate::ty::GenericArg;
-use crate::ty::{self, Ty, TyCtxt};
+use crate::ty::{self, ShortInstance, Ty, TyCtxt};
 
 use hir::def::Namespace;
 use rustc_data_structures::fx::FxHashSet;
@@ -356,3 +358,31 @@ where
         with_no_trimmed_paths!(Self::print(t, fmt))
     }
 }
+
+/// Format instance name that is already known to be too long for rustc.
+/// Show only the first 2 types if it is longer than 32 characters to avoid blasting
+/// the user's terminal with thousands of lines of type-name.
+///
+/// If the type name is longer than before+after, it will be written to a file.
+pub fn shrunk_instance_name<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: ty::Instance<'tcx>,
+) -> (String, Option<PathBuf>) {
+    let s = instance.to_string();
+
+    // Only use the shrunk version if it's really shorter.
+    // This also avoids the case where before and after slices overlap.
+    if s.chars().nth(33).is_some() {
+        let shrunk = format!("{}", ShortInstance(instance, 4));
+        if shrunk == s {
+            return (s, None);
+        }
+
+        let path = tcx.output_filenames(()).temp_path_ext("long-type.txt", None);
+        let written_to_path = std::fs::write(&path, s).ok().map(|_| path);
+
+        (shrunk, written_to_path)
+    } else {
+        (s, None)
+    }
+}
diff --git a/compiler/rustc_monomorphize/messages.ftl b/compiler/rustc_monomorphize/messages.ftl
index 94b553a07a7..7210701d482 100644
--- a/compiler/rustc_monomorphize/messages.ftl
+++ b/compiler/rustc_monomorphize/messages.ftl
@@ -1,6 +1,3 @@
-monomorphize_consider_type_length_limit =
-    consider adding a `#![type_length_limit="{$type_length}"]` attribute to your crate
-
 monomorphize_couldnt_dump_mono_stats =
     unexpected error occurred while dumping monomorphization stats: {$error}
 
@@ -25,8 +22,6 @@ monomorphize_start_not_found = using `fn main` requires the standard library
 
 monomorphize_symbol_already_defined = symbol `{$symbol}` is already defined
 
-monomorphize_type_length_limit = reached the type-length limit while instantiating `{$shrunk}`
-
 monomorphize_unknown_cgu_collection_mode =
     unknown codegen-item collection mode '{$mode}', falling back to 'lazy' mode
 
diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index 89158762265..f98d4700528 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/compiler/rustc_monomorphize/src/collector.rs
@@ -222,12 +222,12 @@ use rustc_middle::mir::{self, Location, MentionedItem};
 use rustc_middle::query::TyCtxtAt;
 use rustc_middle::ty::adjustment::{CustomCoerceUnsized, PointerCoercion};
 use rustc_middle::ty::layout::ValidityRequirement;
-use rustc_middle::ty::print::with_no_trimmed_paths;
+use rustc_middle::ty::print::{shrunk_instance_name, with_no_trimmed_paths};
+use rustc_middle::ty::GenericArgs;
 use rustc_middle::ty::{
     self, AssocKind, GenericParamDefKind, Instance, InstanceKind, Ty, TyCtxt, TypeFoldable,
     TypeVisitableExt, VtblEntry,
 };
-use rustc_middle::ty::{GenericArgKind, GenericArgs};
 use rustc_middle::{bug, span_bug};
 use rustc_session::config::EntryFnType;
 use rustc_session::Limit;
@@ -238,9 +238,7 @@ use rustc_target::abi::Size;
 use std::path::PathBuf;
 use tracing::{debug, instrument, trace};
 
-use crate::errors::{
-    self, EncounteredErrorWhileInstantiating, NoOptimizedMir, RecursionLimit, TypeLengthLimit,
-};
+use crate::errors::{self, EncounteredErrorWhileInstantiating, NoOptimizedMir, RecursionLimit};
 use move_check::MoveCheckState;
 
 #[derive(PartialEq)]
@@ -443,7 +441,6 @@ fn collect_items_rec<'tcx>(
                 recursion_depths,
                 recursion_limit,
             ));
-            check_type_length_limit(tcx, instance);
 
             rustc_data_structures::stack::ensure_sufficient_stack(|| {
                 collect_items_of_instance(
@@ -554,34 +551,6 @@ fn collect_items_rec<'tcx>(
     }
 }
 
-/// Format instance name that is already known to be too long for rustc.
-/// Show only the first 2 types if it is longer than 32 characters to avoid blasting
-/// the user's terminal with thousands of lines of type-name.
-///
-/// If the type name is longer than before+after, it will be written to a file.
-fn shrunk_instance_name<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    instance: Instance<'tcx>,
-) -> (String, Option<PathBuf>) {
-    let s = instance.to_string();
-
-    // Only use the shrunk version if it's really shorter.
-    // This also avoids the case where before and after slices overlap.
-    if s.chars().nth(33).is_some() {
-        let shrunk = format!("{}", ty::ShortInstance(instance, 4));
-        if shrunk == s {
-            return (s, None);
-        }
-
-        let path = tcx.output_filenames(()).temp_path_ext("long-type.txt", None);
-        let written_to_path = std::fs::write(&path, s).ok().map(|_| path);
-
-        (shrunk, written_to_path)
-    } else {
-        (s, None)
-    }
-}
-
 fn check_recursion_limit<'tcx>(
     tcx: TyCtxt<'tcx>,
     instance: Instance<'tcx>,
@@ -630,38 +599,6 @@ fn check_recursion_limit<'tcx>(
     (def_id, recursion_depth)
 }
 
-fn check_type_length_limit<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
-    let type_length = instance
-        .args
-        .iter()
-        .flat_map(|arg| arg.walk())
-        .filter(|arg| match arg.unpack() {
-            GenericArgKind::Type(_) | GenericArgKind::Const(_) => true,
-            GenericArgKind::Lifetime(_) => false,
-        })
-        .count();
-    debug!(" => type length={}", type_length);
-
-    // Rust code can easily create exponentially-long types using only a
-    // polynomial recursion depth. Even with the default recursion
-    // depth, you can easily get cases that take >2^60 steps to run,
-    // which means that rustc basically hangs.
-    //
-    // Bail out in these cases to avoid that bad user experience.
-    if !tcx.type_length_limit().value_within_limit(type_length) {
-        let (shrunk, written_to_path) = shrunk_instance_name(tcx, instance);
-        let span = tcx.def_span(instance.def_id());
-        let mut path = PathBuf::new();
-        let was_written = if let Some(path2) = written_to_path {
-            path = path2;
-            Some(())
-        } else {
-            None
-        };
-        tcx.dcx().emit_fatal(TypeLengthLimit { span, shrunk, was_written, path, type_length });
-    }
-}
-
 struct MirUsedCollector<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     body: &'a mir::Body<'tcx>,
diff --git a/compiler/rustc_monomorphize/src/errors.rs b/compiler/rustc_monomorphize/src/errors.rs
index c0d1efd96c5..9548c46e6fa 100644
--- a/compiler/rustc_monomorphize/src/errors.rs
+++ b/compiler/rustc_monomorphize/src/errors.rs
@@ -20,19 +20,6 @@ pub struct RecursionLimit {
 }
 
 #[derive(Diagnostic)]
-#[diag(monomorphize_type_length_limit)]
-#[help(monomorphize_consider_type_length_limit)]
-pub struct TypeLengthLimit {
-    #[primary_span]
-    pub span: Span,
-    pub shrunk: String,
-    #[note(monomorphize_written_to_path)]
-    pub was_written: Option<()>,
-    pub path: PathBuf,
-    pub type_length: usize,
-}
-
-#[derive(Diagnostic)]
 #[diag(monomorphize_no_optimized_mir)]
 pub struct NoOptimizedMir {
     #[note]