Don't allow constant eval to access generator layouts

Zoxc · Zoxc · commit 0adc84b415c8 · 2023-08-29T23:25:04.000+02:00
diff --git a/compiler/rustc_const_eval/messages.ftl b/compiler/rustc_const_eval/messages.ftl
@@ -112,6 +112,8 @@ const_eval_frame_note_inner = inside {$where_ ->
     *[other] {""}
 }
 
+const_eval_generator_layout_indirect = because of contained type `{$ty}`
+
 const_eval_in_bounds_test = out-of-bounds pointer use
 const_eval_incompatible_calling_conventions =
     calling a function with calling convention {$callee_conv} using calling convention {$caller_conv}
@@ -431,6 +433,8 @@ const_eval_validation_expected_str = expected a string
 const_eval_validation_front_matter_invalid_value = constructing invalid value
 const_eval_validation_front_matter_invalid_value_with_path = constructing invalid value at {$path}
 
+const_eval_validation_generator_layout_access = cannot compute layout of `{$ty}`
+
 const_eval_validation_invalid_bool = {$front_matter}: encountered {$value}, but expected a boolean
 const_eval_validation_invalid_box_meta = {$front_matter}: encountered invalid box metadata: total size is bigger than largest supported object
 const_eval_validation_invalid_box_slice_meta = {$front_matter}: encountered invalid box metadata: slice is bigger than largest supported object
diff --git a/compiler/rustc_const_eval/src/const_eval/machine.rs b/compiler/rustc_const_eval/src/const_eval/machine.rs
@@ -357,6 +357,9 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
 
     const PANIC_ON_ALLOC_FAIL: bool = false; // will be raised as a proper error
 
+    // We don't allow access to generator layout at compile time.
+    const ACCESS_GENERATOR_LAYOUT: bool = false;
+
     #[inline(always)]
     fn enforce_alignment(ecx: &InterpCx<'mir, 'tcx, Self>) -> CheckAlignment {
         ecx.machine.check_alignment
diff --git a/compiler/rustc_const_eval/src/errors.rs b/compiler/rustc_const_eval/src/errors.rs
@@ -674,6 +674,7 @@ impl<'tcx> ReportErrorExt for ValidationErrorInfo<'tcx> {
             InvalidBool { .. } => const_eval_validation_invalid_bool,
             InvalidChar { .. } => const_eval_validation_invalid_char,
             InvalidFnPtr { .. } => const_eval_validation_invalid_fn_ptr,
+            GeneratorLayoutAccess { .. } => const_eval_validation_generator_layout_access,
         }
     }
 
@@ -773,6 +774,20 @@ impl<'tcx> ReportErrorExt for ValidationErrorInfo<'tcx> {
             DanglingPtrNoProvenance { pointer, .. } => {
                 err.set_arg("pointer", pointer);
             }
+            GeneratorLayoutAccess { ty, generators } => {
+                err.set_arg("ty", ty);
+                for generator in generators {
+                    if generator != ty {
+                        let message = handler.eagerly_translate_to_string(
+                            fluent::const_eval_generator_layout_indirect,
+                            [("ty".into(), DiagnosticArgValue::Str(generator.to_string().into()))]
+                                .iter()
+                                .map(|(a, b)| (a, b)),
+                        );
+                        err.help(message);
+                    }
+                }
+            }
             NullPtr { .. }
             | PtrToStatic { .. }
             | PtrToMut { .. }
diff --git a/compiler/rustc_const_eval/src/interpret/machine.rs b/compiler/rustc_const_eval/src/interpret/machine.rs
@@ -130,6 +130,9 @@ pub trait Machine<'mir, 'tcx: 'mir>: Sized {
     /// Should the machine panic on allocation failures?
     const PANIC_ON_ALLOC_FAIL: bool;
 
+    /// Should the machine allow access to generator layout.
+    const ACCESS_GENERATOR_LAYOUT: bool;
+
     /// Whether memory accesses should be alignment-checked.
     fn enforce_alignment(ecx: &InterpCx<'mir, 'tcx, Self>) -> CheckAlignment;
 
diff --git a/compiler/rustc_const_eval/src/interpret/step.rs b/compiler/rustc_const_eval/src/interpret/step.rs
@@ -267,6 +267,11 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
 
             NullaryOp(ref null_op, ty) => {
                 let ty = self.subst_from_current_frame_and_normalize_erasing_regions(ty)?;
+
+                // Ensure we don't need the layout of any generators if applicable.
+                // This prevents typeck from depending on MIR optimizations.
+                self.validate_generator_layout_access(ty)?;
+
                 let layout = self.layout_of(ty)?;
                 if let mir::NullOp::SizeOf | mir::NullOp::AlignOf = null_op && layout.is_unsized() {
                     span_bug!(
diff --git a/compiler/rustc_const_eval/src/interpret/validity.rs b/compiler/rustc_const_eval/src/interpret/validity.rs
@@ -16,8 +16,8 @@ use rustc_middle::mir::interpret::{
     ExpectedKind, InterpError, InvalidMetaKind, PointerKind, ValidationErrorInfo,
     ValidationErrorKind, ValidationErrorKind::*,
 };
-use rustc_middle::ty;
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
+use rustc_middle::ty::{self, Ty};
 use rustc_span::symbol::{sym, Symbol};
 use rustc_target::abi::{
     Abi, FieldIdx, Scalar as ScalarAbi, Size, VariantIdx, Variants, WrappingRange,
@@ -950,4 +950,16 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         // recurse through references which, for now, we don't want here, either.
         self.validate_operand_internal(op, vec![], None, None)
     }
+
+    /// This function checks if the given `ty`'s layout depends on generators.
+    #[inline(always)]
+    pub fn validate_generator_layout_access(&self, ty: Ty<'tcx>) -> InterpResult<'tcx> {
+        if !M::ACCESS_GENERATOR_LAYOUT {
+            let generators = self.tcx.layout_generators(ty);
+            if !generators.is_empty() {
+                throw_validation_failure!(Vec::new(), GeneratorLayoutAccess { ty, generators });
+            }
+        }
+        Ok(())
+    }
 }
diff --git a/compiler/rustc_middle/src/mir/interpret/error.rs b/compiler/rustc_middle/src/mir/interpret/error.rs
@@ -2,7 +2,7 @@ use super::{AllocId, AllocRange, ConstAlloc, Pointer, Scalar};
 
 use crate::mir::interpret::ConstValue;
 use crate::query::TyCtxtAt;
-use crate::ty::{layout, tls, Ty, ValTree};
+use crate::ty::{layout, tls, List, Ty, ValTree};
 
 use rustc_data_structures::sync::Lock;
 use rustc_errors::{
@@ -404,6 +404,7 @@ pub enum ValidationErrorKind<'tcx> {
     InvalidBool { value: String },
     InvalidChar { value: String },
     InvalidFnPtr { value: String },
+    GeneratorLayoutAccess { ty: Ty<'tcx>, generators: &'tcx List<Ty<'tcx>> },
 }
 
 /// Error information for when the program did something that might (or might not) be correct
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -1405,6 +1405,15 @@ rustc_queries! {
         desc { "computing layout of `{}`", key.value }
     }
 
+    /// Computes the generators present in the layout of a type.
+    /// This expects a normalized input type with regions erased.
+    query layout_generators(
+        key: Ty<'tcx>
+    ) -> &'tcx ty::List<Ty<'tcx>> {
+        depth_limit
+        desc { "calculating generators in layout of `{}`", key }
+    }
+
     /// Compute a `FnAbi` suitable for indirect calls, i.e. to `fn` pointers.
     ///
     /// NB: this doesn't handle virtual calls - those should use `fn_abi_of_instance`
diff --git a/compiler/rustc_mir_transform/src/const_prop.rs b/compiler/rustc_mir_transform/src/const_prop.rs
@@ -147,6 +147,7 @@ impl ConstPropMachine<'_, '_> {
 impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for ConstPropMachine<'mir, 'tcx> {
     compile_time_machine!(<'mir, 'tcx>);
     const PANIC_ON_ALLOC_FAIL: bool = true; // all allocations are small (see `MAX_ALLOC_LIMIT`)
+    const ACCESS_GENERATOR_LAYOUT: bool = true;
 
     type MemoryKind = !;
 
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -512,6 +512,7 @@ impl<'mir, 'tcx: 'mir> rustc_const_eval::interpret::Machine<'mir, 'tcx> for Dumm
     rustc_const_eval::interpret::compile_time_machine!(<'mir, 'tcx>);
     type MemoryKind = !;
     const PANIC_ON_ALLOC_FAIL: bool = true;
+    const ACCESS_GENERATOR_LAYOUT: bool = true;
 
     fn enforce_alignment(_ecx: &InterpCx<'mir, 'tcx, Self>) -> CheckAlignment {
         unimplemented!()
diff --git a/compiler/rustc_ty_utils/src/layout_generators.rs b/compiler/rustc_ty_utils/src/layout_generators.rs
@@ -0,0 +1,78 @@
+use rustc_middle::query::Providers;
+use rustc_middle::ty::{self, Ty, TyCtxt};
+
+pub fn provide(providers: &mut Providers) {
+    *providers = Providers { layout_generators, ..*providers };
+}
+
+/// Computes the generators present in the layout of a type.
+/// This expects a normalized input type with regions erased.
+fn layout_generators<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> &'tcx ty::List<Ty<'tcx>> {
+    let mut generators = Vec::new();
+
+    let inner = |generators: &mut Vec<_>, ty: Ty<'tcx>| {
+        let list = tcx.layout_generators(ty);
+        for generator in list.iter() {
+            generators.push(generator);
+        }
+    };
+
+    match *ty.kind() {
+        // These can't contain generators in their layout
+        ty::Bool
+        | ty::Char
+        | ty::Int(_)
+        | ty::Uint(_)
+        | ty::Float(_)
+        | ty::FnPtr(_)
+        | ty::FnDef(..)
+        | ty::Never
+        | ty::Ref(..)
+        | ty::RawPtr(..)
+        | ty::Str => {}
+
+        ty::Array(element, _) => {
+            inner(&mut generators, element);
+        }
+
+        ty::Generator(..) => {
+            generators.push(ty);
+        }
+
+        ty::Closure(_, ref args) => {
+            let tys = args.as_closure().upvar_tys();
+            tys.iter().for_each(|ty| inner(&mut generators, ty));
+        }
+
+        ty::Tuple(tys) => {
+            tys.iter().for_each(|ty| inner(&mut generators, ty));
+        }
+
+        ty::Adt(def, args) => {
+            def.variants().iter().for_each(|v| {
+                v.fields.iter().for_each(|field| {
+                    let ty = field.ty(tcx, args);
+                    let ty = tcx.normalize_erasing_regions(ty::ParamEnv::reveal_all(), ty);
+                    inner(&mut generators, ty)
+                })
+            });
+        }
+
+        ty::Slice(..) | ty::Dynamic(..) | ty::Foreign(..) => {
+            bug!("these are unsized")
+        }
+
+        ty::Alias(..)
+        | ty::Bound(..)
+        | ty::GeneratorWitness(..)
+        | ty::GeneratorWitnessMIR(..)
+        | ty::Infer(_)
+        | ty::Error(_)
+        | ty::Placeholder(..)
+        | ty::Param(_) => {
+            bug!("unexpected type")
+        }
+    }
+
+    tcx.mk_type_list(&generators)
+}
diff --git a/compiler/rustc_ty_utils/src/lib.rs b/compiler/rustc_ty_utils/src/lib.rs
@@ -32,6 +32,7 @@ mod errors;
 mod implied_bounds;
 pub mod instance;
 mod layout;
+mod layout_generators;
 mod layout_sanity_check;
 mod needs_drop;
 mod opaque_types;
@@ -48,6 +49,7 @@ pub fn provide(providers: &mut Providers) {
     consts::provide(providers);
     implied_bounds::provide(providers);
     layout::provide(providers);
+    layout_generators::provide(providers);
     needs_drop::provide(providers);
     opaque_types::provide(providers);
     representability::provide(providers);
diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs
@@ -884,6 +884,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
     const GLOBAL_KIND: Option<MiriMemoryKind> = Some(MiriMemoryKind::Global);
 
     const PANIC_ON_ALLOC_FAIL: bool = false;
+    const ACCESS_GENERATOR_LAYOUT: bool = true;
 
     #[inline(always)]
     fn enforce_alignment(ecx: &MiriInterpCx<'mir, 'tcx>) -> CheckAlignment {
@@ -1410,17 +1411,14 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         ecx: &mut InterpCx<'mir, 'tcx, Self>,
         frame: usize,
         local: mir::Local,
-        mplace: &MPlaceTy<'tcx, Provenance>
+        mplace: &MPlaceTy<'tcx, Provenance>,
     ) -> InterpResult<'tcx> {
         let Some(Provenance::Concrete { alloc_id, .. }) = mplace.ptr.provenance else {
             panic!("after_local_allocated should only be called on fresh allocations");
         };
         let local_decl = &ecx.active_thread_stack()[frame].body.local_decls[local];
         let span = local_decl.source_info.span;
-        ecx.machine
-            .allocation_spans
-            .borrow_mut()
-            .insert(alloc_id, (span, None));
+        ecx.machine.allocation_spans.borrow_mut().insert(alloc_id, (span, None));
         Ok(())
     }
 }
diff --git a/tests/ui/generator/const-eval-size.rs b/tests/ui/generator/const-eval-size.rs
@@ -0,0 +1,20 @@
+// error-pattern: evaluation of constant value failed
+
+#![feature(generators)]
+#![feature(generator_trait)]
+#![feature(type_alias_impl_trait)]
+#![allow(dead_code)]
+
+type Gen = impl std::ops::Generator;
+
+const A: usize = std::mem::size_of::<Gen>();
+
+const B: usize = std::mem::size_of::<Option<Gen>>();
+
+fn gen() -> Gen {
+    move || {
+        yield;
+    }
+}
+
+fn main() {}
diff --git a/tests/ui/generator/const-eval-size.stderr b/tests/ui/generator/const-eval-size.stderr
@@ -0,0 +1,30 @@
+error[E0080]: evaluation of constant value failed
+  --> $SRC_DIR/core/src/mem/mod.rs:LL:COL
+   |
+   = note: cannot compute layout of `[generator@$DIR/const-eval-size.rs:15:5: 15:12]`
+   |
+note: inside `std::mem::size_of::<[generator@$DIR/const-eval-size.rs:15:5: 15:12]>`
+  --> $SRC_DIR/core/src/mem/mod.rs:LL:COL
+note: inside `A`
+  --> $DIR/const-eval-size.rs:10:18
+   |
+LL | const A: usize = std::mem::size_of::<Gen>();
+   |                  ^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error[E0080]: evaluation of constant value failed
+  --> $SRC_DIR/core/src/mem/mod.rs:LL:COL
+   |
+   = note: cannot compute layout of `Option<[generator@$DIR/const-eval-size.rs:15:5: 15:12]>`
+   |
+note: inside `std::mem::size_of::<Option<[generator@$DIR/const-eval-size.rs:15:5: 15:12]>>`
+  --> $SRC_DIR/core/src/mem/mod.rs:LL:COL
+note: inside `B`
+  --> $DIR/const-eval-size.rs:12:18
+   |
+LL | const B: usize = std::mem::size_of::<Option<Gen>>();
+   |                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   = help: because of contained type `[generator@$DIR/const-eval-size.rs:15:5: 15:12]`
+
+error: aborting due to 2 previous errors
+
+For more information about this error, try `rustc --explain E0080`.
