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

rustc_mir_transform/
check_enums.rs
