rustc_mir_transform/
check_enums.rs

1use rustc_abi::{Scalar, Size, TagEncoding, Variants, WrappingRange};
2use rustc_data_structures::thin_vec::ThinVec;
3use rustc_hir::attrs::lang_items::LangItem;
4use rustc_index::IndexVec;
5use rustc_middle::bug;
6use rustc_middle::mir::visit::Visitor;
7use rustc_middle::mir::*;
8use rustc_middle::ty::layout::PrimitiveExt;
9use rustc_middle::ty::{self, Ty, TyCtxt, TypingEnv};
10use rustc_session::Session;
11use tracing::debug;
12
13use crate::PassPolicy;
14
15/// This pass inserts checks for a valid enum discriminant where they are most
16/// likely to find UB, because checking everywhere like Miri would generate too
17/// much MIR.
18pub(super) struct CheckEnums;
19
20impl<'tcx> crate::MirPass<'tcx> for CheckEnums {
21    fn policy(&self, sess: &Session) -> PassPolicy {
22        // When UB checks are enabled this is part of their semantics, not an optimization.
23        PassPolicy::optional_non_optimization(sess.ub_checks())
24    }
25
26    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
27        // This pass emits new panics. If for whatever reason we do not have a panic
28        // implementation, running this pass may cause otherwise-valid code to not compile.
29        if tcx.lang_items().get(LangItem::PanicImpl).is_none() {
30            return;
31        }
32
33        let typing_env = body.typing_env(tcx);
34        let basic_blocks = body.basic_blocks.as_mut();
35        let local_decls = &mut body.local_decls;
36
37        // This operation inserts new blocks. Each insertion changes the Location for all
38        // statements/blocks after. Iterating or visiting the MIR in order would require updating
39        // our current location after every insertion. By iterating backwards, we dodge this issue:
40        // The only Locations that an insertion changes have already been handled.
41        for block in basic_blocks.indices().rev() {
42            for statement_index in (0..basic_blocks[block].statements.len()).rev() {
43                let location = Location { block, statement_index };
44                let statement = &basic_blocks[block].statements[statement_index];
45                let source_info = statement.source_info;
46
47                let mut finder = EnumFinder::new(tcx, local_decls, typing_env);
48                finder.visit_statement(statement, location);
49
50                for check in finder.into_found_enums() {
51                    debug!("Inserting enum check");
52                    let new_block = split_block(basic_blocks, location);
53
54                    match check {
55                        EnumCheckType::Direct { op_size, .. }
56                        | EnumCheckType::WithNiche { op_size, .. }
57                            if op_size.bytes() == 0 =>
58                        {
59                            // It is never valid to use a ZST as a discriminant for an inhabited enum, but that will
60                            // have been caught by the type checker. Do nothing but ensure that a bug has been signaled.
61                            tcx.dcx().span_delayed_bug(
62                                source_info.span,
63                                "cannot build enum discriminant from zero-sized type",
64                            );
65                            basic_blocks[block].terminator = Some(Terminator {
66                                source_info,
67                                kind: TerminatorKind::Goto { target: new_block },
68                                attributes: ThinVec::new(),
69                            });
70                        }
71                        EnumCheckType::Direct { source_op, discr, op_size, valid_discrs } => {
72                            insert_direct_enum_check(
73                                tcx,
74                                local_decls,
75                                basic_blocks,
76                                block,
77                                source_op,
78                                discr,
79                                op_size,
80                                valid_discrs,
81                                source_info,
82                                new_block,
83                            )
84                        }
85                        EnumCheckType::Uninhabited => insert_uninhabited_enum_check(
86                            tcx,
87                            local_decls,
88                            &mut basic_blocks[block],
89                            source_info,
90                            new_block,
91                        ),
92                        EnumCheckType::WithNiche {
93                            source_op,
94                            discr,
95                            op_size,
96                            offset,
97                            valid_range,
98                        } => insert_niche_check(
99                            tcx,
100                            local_decls,
101                            &mut basic_blocks[block],
102                            source_op,
103                            valid_range,
104                            discr,
105                            op_size,
106                            offset,
107                            source_info,
108                            new_block,
109                        ),
110                    }
111                }
112            }
113        }
114    }
115}
116
117/// Represent the different kind of enum checks we can insert.
118enum EnumCheckType<'tcx> {
119    /// We know we try to create an uninhabited enum from an inhabited variant.
120    Uninhabited,
121    /// We know the enum does no niche optimizations and can thus easily compute
122    /// the valid discriminants.
123    Direct {
124        source_op: Operand<'tcx>,
125        discr: TyAndSize<'tcx>,
126        op_size: Size,
127        valid_discrs: Vec<u128>,
128    },
129    /// We try to construct an enum that has a niche.
130    WithNiche {
131        source_op: Operand<'tcx>,
132        discr: TyAndSize<'tcx>,
133        op_size: Size,
134        offset: Size,
135        valid_range: WrappingRange,
136    },
137}
138
139#[derive(Debug, Copy, Clone)]
140struct TyAndSize<'tcx> {
141    pub ty: Ty<'tcx>,
142    pub size: Size,
143}
144
145/// A [Visitor] that finds the construction of enums and evaluates which checks
146/// we should apply.
147struct EnumFinder<'a, 'tcx> {
148    tcx: TyCtxt<'tcx>,
149    local_decls: &'a mut LocalDecls<'tcx>,
150    typing_env: TypingEnv<'tcx>,
151    enums: Vec<EnumCheckType<'tcx>>,
152}
153
154impl<'a, 'tcx> EnumFinder<'a, 'tcx> {
155    fn new(
156        tcx: TyCtxt<'tcx>,
157        local_decls: &'a mut LocalDecls<'tcx>,
158        typing_env: TypingEnv<'tcx>,
159    ) -> Self {
160        EnumFinder { tcx, local_decls, typing_env, enums: Vec::new() }
161    }
162
163    /// Returns the found enum creations and which checks should be inserted.
164    fn into_found_enums(self) -> Vec<EnumCheckType<'tcx>> {
165        self.enums
166    }
167}
168
169impl<'a, 'tcx> Visitor<'tcx> for EnumFinder<'a, 'tcx> {
170    fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) {
171        if let Rvalue::Cast(CastKind::Transmute, op, ty) = rvalue {
172            let ty::Adt(adt_def, _) = ty.kind() else {
173                return;
174            };
175            if !adt_def.is_enum() {
176                return;
177            }
178
179            let Ok(enum_layout) = self.tcx.layout_of(self.typing_env.as_query_input(*ty)) else {
180                return;
181            };
182            let Ok(op_layout) = self
183                .tcx
184                .layout_of(self.typing_env.as_query_input(op.ty(self.local_decls, self.tcx)))
185            else {
186                return;
187            };
188
189            match enum_layout.variants {
190                Variants::Empty if op_layout.is_uninhabited() => return,
191                // An empty enum that tries to be constructed from an inhabited value, this
192                // is never correct.
193                Variants::Empty => {
194                    // The enum layout is uninhabited but we construct it from sth inhabited.
195                    // This is always UB.
196                    self.enums.push(EnumCheckType::Uninhabited);
197                }
198                // Construction of Single value enums is always fine.
199                Variants::Single { .. } => {}
200                // Construction of an enum with multiple variants but no niche optimizations.
201                Variants::Multiple {
202                    tag_encoding: TagEncoding::Direct,
203                    tag: Scalar::Initialized { value, .. },
204                    ..
205                } => {
206                    let valid_discrs =
207                        adt_def.discriminants(self.tcx).map(|(_, discr)| discr.val).collect();
208
209                    let discr =
210                        TyAndSize { ty: value.to_int_ty(self.tcx), size: value.size(&self.tcx) };
211                    self.enums.push(EnumCheckType::Direct {
212                        source_op: op.to_copy(),
213                        discr,
214                        op_size: op_layout.size,
215                        valid_discrs,
216                    });
217                }
218                // Construction of an enum with multiple variants and niche optimizations.
219                Variants::Multiple {
220                    tag_encoding: TagEncoding::Niche { .. },
221                    tag: Scalar::Initialized { value, valid_range, .. },
222                    tag_field,
223                    ..
224                } => {
225                    let discr =
226                        TyAndSize { ty: value.to_int_ty(self.tcx), size: value.size(&self.tcx) };
227                    self.enums.push(EnumCheckType::WithNiche {
228                        source_op: op.to_copy(),
229                        discr,
230                        op_size: op_layout.size,
231                        offset: enum_layout.fields.offset(tag_field.as_usize()),
232                        valid_range,
233                    });
234                }
235                _ => return,
236            }
237
238            self.super_rvalue(rvalue, location);
239        }
240    }
241}
242
243fn split_block(
244    basic_blocks: &mut IndexVec<BasicBlock, BasicBlockData<'_>>,
245    location: Location,
246) -> BasicBlock {
247    let block_data = &mut basic_blocks[location.block];
248
249    // Drain every statement after this one and move the current terminator to a new basic block.
250    let new_block = BasicBlockData::new_stmts(
251        block_data.statements.split_off(location.statement_index),
252        block_data.terminator.take(),
253        block_data.is_cleanup,
254    );
255
256    basic_blocks.push(new_block)
257}
258
259/// Inserts the cast of an operand (any type) to a u128 value that holds the discriminant value.
260fn insert_discr_cast_to_u128<'tcx>(
261    tcx: TyCtxt<'tcx>,
262    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
263    block_data: &mut BasicBlockData<'tcx>,
264    source_op: Operand<'tcx>,
265    discr: TyAndSize<'tcx>,
266    op_size: Size,
267    offset: Option<Size>,
268    source_info: SourceInfo,
269) -> Place<'tcx> {
270    let get_ty_for_size = |tcx: TyCtxt<'tcx>, size: Size| -> Ty<'tcx> {
271        match size.bytes() {
272            1 => tcx.types.u8,
273            2 => tcx.types.u16,
274            4 => tcx.types.u32,
275            8 => tcx.types.u64,
276            16 => tcx.types.u128,
277            invalid => bug!("Found discriminant with invalid size, has {} bytes", invalid),
278        }
279    };
280
281    let (cast_kind, discr_ty_bits) = if discr.size.bytes() < op_size.bytes() {
282        // The discriminant is less wide than the operand, cast the operand into
283        // [MaybeUninit; N] and then index into it.
284        let mu = Ty::new_maybe_uninit(tcx, tcx.types.u8);
285        let array_len = op_size.bytes();
286        let mu_array_ty = Ty::new_array(tcx, mu, array_len);
287        let mu_array =
288            local_decls.push(LocalDecl::with_source_info(mu_array_ty, source_info)).into();
289        let rvalue = Rvalue::Cast(CastKind::Transmute, source_op, mu_array_ty);
290        block_data
291            .statements
292            .push(Statement::new(source_info, StatementKind::Assign(Box::new((mu_array, rvalue)))));
293
294        // Index into the array of MaybeUninit to get something that is actually
295        // as wide as the discriminant.
296        let offset = offset.unwrap_or(Size::ZERO);
297        let smaller_mu_array = mu_array.project_deeper(
298            &[ProjectionElem::Subslice {
299                from: offset.bytes(),
300                to: offset.bytes() + discr.size.bytes(),
301                from_end: false,
302            }],
303            tcx,
304        );
305
306        (CastKind::Transmute, Operand::Copy(smaller_mu_array))
307    } else {
308        let operand_int_ty = get_ty_for_size(tcx, op_size);
309
310        let op_as_int =
311            local_decls.push(LocalDecl::with_source_info(operand_int_ty, source_info)).into();
312        let rvalue = Rvalue::Cast(CastKind::Transmute, source_op, operand_int_ty);
313        block_data.statements.push(Statement::new(
314            source_info,
315            StatementKind::Assign(Box::new((op_as_int, rvalue))),
316        ));
317
318        (CastKind::IntToInt, Operand::Copy(op_as_int))
319    };
320
321    // Cast the resulting value to the actual discriminant integer type.
322    let rvalue = Rvalue::Cast(cast_kind, discr_ty_bits, discr.ty);
323    let discr_in_discr_ty =
324        local_decls.push(LocalDecl::with_source_info(discr.ty, source_info)).into();
325    block_data.statements.push(Statement::new(
326        source_info,
327        StatementKind::Assign(Box::new((discr_in_discr_ty, rvalue))),
328    ));
329
330    // Cast the discriminant to a u128 (base for comparisons of enum discriminants).
331    let const_u128 = Ty::new_uint(tcx, ty::UintTy::U128);
332    let rvalue = Rvalue::Cast(CastKind::IntToInt, Operand::Copy(discr_in_discr_ty), const_u128);
333    let discr = local_decls.push(LocalDecl::with_source_info(const_u128, source_info)).into();
334    block_data
335        .statements
336        .push(Statement::new(source_info, StatementKind::Assign(Box::new((discr, rvalue)))));
337
338    discr
339}
340
341fn insert_direct_enum_check<'tcx>(
342    tcx: TyCtxt<'tcx>,
343    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
344    basic_blocks: &mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
345    current_block: BasicBlock,
346    source_op: Operand<'tcx>,
347    discr: TyAndSize<'tcx>,
348    op_size: Size,
349    discriminants: Vec<u128>,
350    source_info: SourceInfo,
351    new_block: BasicBlock,
352) {
353    // Insert a new target block that is branched to in case of an invalid discriminant.
354    let invalid_discr_block_data = BasicBlockData::new(None, false);
355    let invalid_discr_block = basic_blocks.push(invalid_discr_block_data);
356    let block_data = &mut basic_blocks[current_block];
357    let discr_place = insert_discr_cast_to_u128(
358        tcx,
359        local_decls,
360        block_data,
361        source_op,
362        discr,
363        op_size,
364        None,
365        source_info,
366    );
367
368    // Mask out the bits of the discriminant type.
369    let mask = discr.size.unsigned_int_max();
370    let discr_masked =
371        local_decls.push(LocalDecl::with_source_info(tcx.types.u128, source_info)).into();
372    let rvalue = Rvalue::BinaryOp(
373        BinOp::BitAnd,
374        Box::new((
375            Operand::Copy(discr_place),
376            Operand::Constant(Box::new(ConstOperand {
377                span: source_info.span,
378                user_ty: None,
379                const_: Const::Val(ConstValue::from_u128(mask), tcx.types.u128),
380            })),
381        )),
382    );
383    block_data
384        .statements
385        .push(Statement::new(source_info, StatementKind::Assign(Box::new((discr_masked, rvalue)))));
386
387    // Branch based on the discriminant value.
388    block_data.terminator = Some(Terminator {
389        source_info,
390        kind: TerminatorKind::SwitchInt {
391            discr: Operand::Copy(discr_masked),
392            targets: SwitchTargets::new(
393                discriminants
394                    .into_iter()
395                    .map(|discr_val| (discr.size.truncate(discr_val), new_block)),
396                invalid_discr_block,
397            ),
398        },
399        attributes: ThinVec::new(),
400    });
401
402    // Abort in case of an invalid enum discriminant.
403    basic_blocks[invalid_discr_block].terminator = Some(Terminator {
404        source_info,
405        kind: TerminatorKind::Assert {
406            cond: Operand::Constant(Box::new(ConstOperand {
407                span: source_info.span,
408                user_ty: None,
409                const_: Const::Val(ConstValue::from_bool(false), tcx.types.bool),
410            })),
411            expected: true,
412            target: new_block,
413            msg: Box::new(AssertKind::InvalidEnumConstruction(Operand::Copy(discr_masked))),
414            // This calls panic_invalid_enum_construction, which is #[rustc_nounwind].
415            // We never want to insert an unwind into unsafe code, because unwinding could
416            // make a failing UB check turn into much worse UB when we start unwinding.
417            unwind: UnwindAction::Unreachable,
418        },
419        attributes: ThinVec::new(),
420    });
421}
422
423fn insert_uninhabited_enum_check<'tcx>(
424    tcx: TyCtxt<'tcx>,
425    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
426    block_data: &mut BasicBlockData<'tcx>,
427    source_info: SourceInfo,
428    new_block: BasicBlock,
429) {
430    let is_ok: Place<'_> =
431        local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into();
432    block_data.statements.push(Statement::new(
433        source_info,
434        StatementKind::Assign(Box::new((
435            is_ok,
436            Rvalue::Use(
437                Operand::Constant(Box::new(ConstOperand {
438                    span: source_info.span,
439                    user_ty: None,
440                    const_: Const::Val(ConstValue::from_bool(false), tcx.types.bool),
441                })),
442                WithRetag::Yes, // it's a bool, retag doesn't matter
443            ),
444        ))),
445    ));
446
447    block_data.terminator = Some(Terminator {
448        source_info,
449        kind: TerminatorKind::Assert {
450            cond: Operand::Copy(is_ok),
451            expected: true,
452            target: new_block,
453            msg: Box::new(AssertKind::InvalidEnumConstruction(Operand::Constant(Box::new(
454                ConstOperand {
455                    span: source_info.span,
456                    user_ty: None,
457                    const_: Const::Val(ConstValue::from_u128(0), tcx.types.u128),
458                },
459            )))),
460            // This calls panic_invalid_enum_construction, which is #[rustc_nounwind].
461            // We never want to insert an unwind into unsafe code, because unwinding could
462            // make a failing UB check turn into much worse UB when we start unwinding.
463            unwind: UnwindAction::Unreachable,
464        },
465        attributes: ThinVec::new(),
466    });
467}
468
469fn insert_niche_check<'tcx>(
470    tcx: TyCtxt<'tcx>,
471    local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
472    block_data: &mut BasicBlockData<'tcx>,
473    source_op: Operand<'tcx>,
474    valid_range: WrappingRange,
475    discr: TyAndSize<'tcx>,
476    op_size: Size,
477    offset: Size,
478    source_info: SourceInfo,
479    new_block: BasicBlock,
480) {
481    let discr = insert_discr_cast_to_u128(
482        tcx,
483        local_decls,
484        block_data,
485        source_op,
486        discr,
487        op_size,
488        Some(offset),
489        source_info,
490    );
491
492    // Compare the discriminant against the valid_range.
493    let start_const = Operand::Constant(Box::new(ConstOperand {
494        span: source_info.span,
495        user_ty: None,
496        const_: Const::Val(ConstValue::from_u128(valid_range.start), tcx.types.u128),
497    }));
498    let end_start_diff_const = Operand::Constant(Box::new(ConstOperand {
499        span: source_info.span,
500        user_ty: None,
501        const_: Const::Val(
502            ConstValue::from_u128(u128::wrapping_sub(valid_range.end, valid_range.start)),
503            tcx.types.u128,
504        ),
505    }));
506
507    let discr_diff: Place<'_> =
508        local_decls.push(LocalDecl::with_source_info(tcx.types.u128, source_info)).into();
509    block_data.statements.push(Statement::new(
510        source_info,
511        StatementKind::Assign(Box::new((
512            discr_diff,
513            Rvalue::BinaryOp(BinOp::Sub, Box::new((Operand::Copy(discr), start_const))),
514        ))),
515    ));
516
517    let is_ok: Place<'_> =
518        local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into();
519    block_data.statements.push(Statement::new(
520        source_info,
521        StatementKind::Assign(Box::new((
522            is_ok,
523            Rvalue::BinaryOp(
524                // This is a `WrappingRange`, so make sure to get the wrapping right.
525                BinOp::Le,
526                Box::new((Operand::Copy(discr_diff), end_start_diff_const)),
527            ),
528        ))),
529    ));
530
531    block_data.terminator = Some(Terminator {
532        source_info,
533        kind: TerminatorKind::Assert {
534            cond: Operand::Copy(is_ok),
535            expected: true,
536            target: new_block,
537            msg: Box::new(AssertKind::InvalidEnumConstruction(Operand::Copy(discr))),
538            // This calls panic_invalid_enum_construction, which is #[rustc_nounwind].
539            // We never want to insert an unwind into unsafe code, because unwinding could
540            // make a failing UB check turn into much worse UB when we start unwinding.
541            unwind: UnwindAction::Unreachable,
542        },
543        attributes: ThinVec::new(),
544    });
545}
rustc_mir_transform/check_enums.rs

rustc_mir_transform/
check_enums.rs
