rustc_mir_transform/
check_enums.rs

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

rustc_mir_transform/
check_enums.rs
