address review

3 files changed, 54 insertions, 17 deletions

diff --git a/compiler/rustc_const_eval/src/const_eval/error.rs b/compiler/rustc_const_eval/src/const_eval/error.rs
index 3bd092263c1..3eeb0138b37 100644
--- a/compiler/rustc_const_eval/src/const_eval/error.rs
+++ b/compiler/rustc_const_eval/src/const_eval/error.rs
@@ -140,6 +140,7 @@ impl<'tcx> ConstEvalErr<'tcx> {
     ///
     /// If `lint_root.is_some()` report it as a lint, else report it as a hard error.
     /// (Except that for some errors, we ignore all that -- see `must_error` below.)
+    #[instrument(skip(self, tcx, decorate, lint_root), level = "debug")]
     fn struct_generic(
         &self,
         tcx: TyCtxtAt<'tcx>,
@@ -190,6 +191,7 @@ impl<'tcx> ConstEvalErr<'tcx> {
             decorate(err);
         };
 
+        debug!("self.error: {:?}", self.error);
         // Special handling for certain errors
         match &self.error {
             // Don't emit a new diagnostic for these errors
diff --git a/compiler/rustc_const_eval/src/const_eval/mod.rs b/compiler/rustc_const_eval/src/const_eval/mod.rs
index 51edf64de80..793f02a8ade 100644
--- a/compiler/rustc_const_eval/src/const_eval/mod.rs
+++ b/compiler/rustc_const_eval/src/const_eval/mod.rs
@@ -24,12 +24,6 @@ pub use fn_queries::*;
 pub use machine::*;
 pub(crate) use valtrees::{const_to_valtree_inner, valtree_to_const_value};
 
-pub(crate) enum ValTreeCreationError {
-    NonSupportedType,
-    Other,
-}
-pub(crate) type ValTreeCreationResult<'tcx> = Result<ty::ValTree<'tcx>, ValTreeCreationError>;
-
 pub(crate) fn const_caller_location(
     tcx: TyCtxt<'_>,
     (file, line, col): (Symbol, u32, u32),
@@ -44,6 +38,16 @@ pub(crate) fn const_caller_location(
     ConstValue::Scalar(Scalar::from_maybe_pointer(loc_place.ptr, &tcx))
 }
 
+// We forbid type-level constants that contain more than `VALTREE_MAX_NODES` nodes.
+const VALTREE_MAX_NODES: usize = 100000;
+
+pub(crate) enum ValTreeCreationError {
+    NodesOverflow,
+    NonSupportedType,
+    Other,
+}
+pub(crate) type ValTreeCreationResult<'tcx> = Result<ty::ValTree<'tcx>, ValTreeCreationError>;
+
 /// Evaluates a constant and turns it into a type-level constant value.
 pub(crate) fn eval_to_valtree<'tcx>(
     tcx: TyCtxt<'tcx>,
@@ -62,11 +66,28 @@ pub(crate) fn eval_to_valtree<'tcx>(
     let place = ecx.raw_const_to_mplace(const_alloc).unwrap();
     debug!(?place);
 
-    let valtree_result = const_to_valtree_inner(&ecx, &place);
+    let mut num_nodes = 0;
+    let valtree_result = const_to_valtree_inner(&ecx, &place, &mut num_nodes);
 
     match valtree_result {
         Ok(valtree) => Ok(Some(valtree)),
-        Err(_) => Ok(None),
+        Err(err) => {
+            let did = cid.instance.def_id();
+            let s = cid.display(tcx);
+            match err {
+                ValTreeCreationError::NodesOverflow => {
+                    let msg = format!("maximum number of nodes exceeded in constant {}", &s);
+                    let mut diag = match tcx.hir().span_if_local(did) {
+                        Some(span) => tcx.sess.struct_span_err(span, &msg),
+                        None => tcx.sess.struct_err(&msg),
+                    };
+                    diag.emit();
+
+                    Ok(None)
+                }
+                ValTreeCreationError::NonSupportedType | ValTreeCreationError::Other => Ok(None),
+            }
+        }
     }
 }
 
@@ -75,7 +96,7 @@ pub(crate) fn eval_to_valtree<'tcx>(
 pub(crate) fn try_destructure_const<'tcx>(
     tcx: TyCtxt<'tcx>,
     const_: ty::Const<'tcx>,
-) -> Option<mir::DestructuredConst<'tcx>> {
+) -> Option<ty::DestructuredConst<'tcx>> {
     if let ty::ConstKind::Value(valtree) = const_.val() {
         let branches = match valtree {
             ty::ValTree::Branch(b) => b,
@@ -141,7 +162,7 @@ pub(crate) fn try_destructure_const<'tcx>(
 
         let fields = tcx.arena.alloc_from_iter(fields.into_iter());
 
-        Some(mir::DestructuredConst { variant, fields })
+        Some(ty::DestructuredConst { variant, fields })
     } else {
         None
     }
diff --git a/compiler/rustc_const_eval/src/const_eval/valtrees.rs b/compiler/rustc_const_eval/src/const_eval/valtrees.rs
index 080133275a6..4137e351344 100644
--- a/compiler/rustc_const_eval/src/const_eval/valtrees.rs
+++ b/compiler/rustc_const_eval/src/const_eval/valtrees.rs
@@ -1,6 +1,6 @@
 use super::eval_queries::{mk_eval_cx, op_to_const};
 use super::machine::CompileTimeEvalContext;
-use super::{ValTreeCreationError, ValTreeCreationResult};
+use super::{ValTreeCreationError, ValTreeCreationResult, VALTREE_MAX_NODES};
 use crate::interpret::{
     intern_const_alloc_recursive, ConstValue, ImmTy, Immediate, InternKind, MemPlaceMeta,
     MemoryKind, PlaceTy, Scalar, ScalarMaybeUninit,
@@ -16,6 +16,7 @@ fn branches<'tcx>(
     place: &MPlaceTy<'tcx>,
     n: usize,
     variant: Option<VariantIdx>,
+    num_nodes: &mut usize,
 ) -> ValTreeCreationResult<'tcx> {
     let place = match variant {
         Some(variant) => ecx.mplace_downcast(&place, variant).unwrap(),
@@ -27,7 +28,7 @@ fn branches<'tcx>(
     let mut fields = Vec::with_capacity(n);
     for i in 0..n {
         let field = ecx.mplace_field(&place, i).unwrap();
-        let valtree = const_to_valtree_inner(ecx, &field)?;
+        let valtree = const_to_valtree_inner(ecx, &field, num_nodes)?;
         fields.push(Some(valtree));
     }
 
@@ -39,6 +40,11 @@ fn branches<'tcx>(
         .collect::<Option<Vec<_>>>()
         .expect("should have already checked for errors in ValTree creation");
 
+    // Have to account for ZSTs here
+    if branches.len() == 0 {
+        *num_nodes += 1;
+    }
+
     Ok(ty::ValTree::Branch(ecx.tcx.arena.alloc_from_iter(branches)))
 }
 
@@ -46,6 +52,7 @@ fn branches<'tcx>(
 fn slice_branches<'tcx>(
     ecx: &CompileTimeEvalContext<'tcx, 'tcx>,
     place: &MPlaceTy<'tcx>,
+    num_nodes: &mut usize,
 ) -> ValTreeCreationResult<'tcx> {
     let n = place
         .len(&ecx.tcx.tcx)
@@ -54,7 +61,7 @@ fn slice_branches<'tcx>(
     let mut elems = Vec::with_capacity(n as usize);
     for i in 0..n {
         let place_elem = ecx.mplace_index(place, i).unwrap();
-        let valtree = const_to_valtree_inner(ecx, &place_elem)?;
+        let valtree = const_to_valtree_inner(ecx, &place_elem, num_nodes)?;
         elems.push(valtree);
     }
 
@@ -65,12 +72,18 @@ fn slice_branches<'tcx>(
 pub(crate) fn const_to_valtree_inner<'tcx>(
     ecx: &CompileTimeEvalContext<'tcx, 'tcx>,
     place: &MPlaceTy<'tcx>,
+    num_nodes: &mut usize,
 ) -> ValTreeCreationResult<'tcx> {
     let ty = place.layout.ty;
     debug!("ty kind: {:?}", ty.kind());
 
+    if *num_nodes >= VALTREE_MAX_NODES {
+        return Err(ValTreeCreationError::NodesOverflow);
+    }
+
     match ty.kind() {
         ty::FnDef(..) => {
+            *num_nodes += 1;
             Ok(ty::ValTree::zst())
         }
         ty::Bool | ty::Int(_) | ty::Uint(_) | ty::Float(_) | ty::Char => {
@@ -78,6 +91,7 @@ pub(crate) fn const_to_valtree_inner<'tcx>(
                 return Err(ValTreeCreationError::Other);
             };
             let val = val.to_scalar().unwrap();
+            *num_nodes += 1;
 
             Ok(ty::ValTree::Leaf(val.assert_int()))
         }
@@ -94,11 +108,11 @@ pub(crate) fn const_to_valtree_inner<'tcx>(
             };
             debug!(?derefd_place);
 
-            const_to_valtree_inner(ecx, &derefd_place)
+            const_to_valtree_inner(ecx, &derefd_place, num_nodes)
         }
 
         ty::Str | ty::Slice(_) | ty::Array(_, _) => {
-            slice_branches(ecx, place)
+            slice_branches(ecx, place, num_nodes)
         }
         // Trait objects are not allowed in type level constants, as we have no concept for
         // resolving their backing type, even if we can do that at const eval time. We may
@@ -107,7 +121,7 @@ pub(crate) fn const_to_valtree_inner<'tcx>(
         ty::Dynamic(..) => Err(ValTreeCreationError::NonSupportedType),
 
         ty::Tuple(elem_tys) => {
-            branches(ecx, place, elem_tys.len(), None)
+            branches(ecx, place, elem_tys.len(), None, num_nodes)
         }
 
         ty::Adt(def, _) => {
@@ -120,7 +134,7 @@ pub(crate) fn const_to_valtree_inner<'tcx>(
             let Ok((_, variant)) = ecx.read_discriminant(&place.into()) else {
                 return Err(ValTreeCreationError::Other);
             };
-            branches(ecx, place, def.variant(variant).fields.len(), def.is_enum().then_some(variant))
+            branches(ecx, place, def.variant(variant).fields.len(), def.is_enum().then_some(variant), num_nodes)
         }
 
         ty::Never