rustc_mir_transform/
dataflow_const_prop.rs

1//! A constant propagation optimization pass based on dataflow analysis.
2//!
3//! Currently, this pass only propagates scalar values.
4
5use std::assert_matches::assert_matches;
6use std::fmt::Formatter;
7
8use rustc_abi::{BackendRepr, FIRST_VARIANT, FieldIdx, Size, VariantIdx};
9use rustc_const_eval::const_eval::{DummyMachine, throw_machine_stop_str};
10use rustc_const_eval::interpret::{
11    ImmTy, Immediate, InterpCx, OpTy, PlaceTy, Projectable, interp_ok,
12};
13use rustc_data_structures::fx::FxHashMap;
14use rustc_hir::def::DefKind;
15use rustc_middle::bug;
16use rustc_middle::mir::interpret::{InterpResult, Scalar};
17use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
18use rustc_middle::mir::*;
19use rustc_middle::ty::{self, Ty, TyCtxt};
20use rustc_mir_dataflow::fmt::DebugWithContext;
21use rustc_mir_dataflow::lattice::{FlatSet, HasBottom};
22use rustc_mir_dataflow::value_analysis::{
23    Map, PlaceIndex, State, TrackElem, ValueOrPlace, debug_with_context,
24};
25use rustc_mir_dataflow::{Analysis, ResultsVisitor, visit_reachable_results};
26use rustc_span::DUMMY_SP;
27use tracing::{debug, debug_span, instrument};
28
29// These constants are somewhat random guesses and have not been optimized.
30// If `tcx.sess.mir_opt_level() >= 4`, we ignore the limits (this can become very expensive).
31const BLOCK_LIMIT: usize = 100;
32const PLACE_LIMIT: usize = 100;
33
34pub(super) struct DataflowConstProp;
35
36impl<'tcx> crate::MirPass<'tcx> for DataflowConstProp {
37    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
38        sess.mir_opt_level() >= 3
39    }
40
41    #[instrument(skip_all level = "debug")]
42    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
43        debug!(def_id = ?body.source.def_id());
44        if tcx.sess.mir_opt_level() < 4 && body.basic_blocks.len() > BLOCK_LIMIT {
45            debug!("aborted dataflow const prop due too many basic blocks");
46            return;
47        }
48
49        // We want to have a somewhat linear runtime w.r.t. the number of statements/terminators.
50        // Let's call this number `n`. Dataflow analysis has `O(h*n)` transfer function
51        // applications, where `h` is the height of the lattice. Because the height of our lattice
52        // is linear w.r.t. the number of tracked places, this is `O(tracked_places * n)`. However,
53        // because every transfer function application could traverse the whole map, this becomes
54        // `O(num_nodes * tracked_places * n)` in terms of time complexity. Since the number of
55        // map nodes is strongly correlated to the number of tracked places, this becomes more or
56        // less `O(n)` if we place a constant limit on the number of tracked places.
57        let place_limit = if tcx.sess.mir_opt_level() < 4 { Some(PLACE_LIMIT) } else { None };
58
59        // Decide which places to track during the analysis.
60        let map = Map::new(tcx, body, place_limit);
61
62        // Perform the actual dataflow analysis.
63        let mut const_ = debug_span!("analyze")
64            .in_scope(|| ConstAnalysis::new(tcx, body, map).iterate_to_fixpoint(tcx, body, None));
65
66        // Collect results and patch the body afterwards.
67        let mut visitor = Collector::new(tcx, &body.local_decls);
68        debug_span!("collect").in_scope(|| {
69            visit_reachable_results(body, &mut const_.analysis, &const_.results, &mut visitor)
70        });
71        let mut patch = visitor.patch;
72        debug_span!("patch").in_scope(|| patch.visit_body_preserves_cfg(body));
73    }
74
75    fn is_required(&self) -> bool {
76        false
77    }
78}
79
80// Note: Currently, places that have their reference taken cannot be tracked. Although this would
81// be possible, it has to rely on some aliasing model, which we are not ready to commit to yet.
82// Because of that, we can assume that the only way to change the value behind a tracked place is
83// by direct assignment.
84struct ConstAnalysis<'a, 'tcx> {
85    map: Map<'tcx>,
86    tcx: TyCtxt<'tcx>,
87    local_decls: &'a LocalDecls<'tcx>,
88    ecx: InterpCx<'tcx, DummyMachine>,
89    typing_env: ty::TypingEnv<'tcx>,
90}
91
92impl<'tcx> Analysis<'tcx> for ConstAnalysis<'_, 'tcx> {
93    type Domain = State<FlatSet<Scalar>>;
94
95    const NAME: &'static str = "ConstAnalysis";
96
97    // The bottom state denotes uninitialized memory. Because we are only doing a sound
98    // approximation of the actual execution, we can also use this state for places where access
99    // would be UB.
100    fn bottom_value(&self, _body: &Body<'tcx>) -> Self::Domain {
101        State::Unreachable
102    }
103
104    fn initialize_start_block(&self, body: &Body<'tcx>, state: &mut Self::Domain) {
105        // The initial state maps all tracked places of argument projections to ⊤ and the rest to ⊥.
106        assert_matches!(state, State::Unreachable);
107        *state = State::new_reachable();
108        for arg in body.args_iter() {
109            state.flood(PlaceRef { local: arg, projection: &[] }, &self.map);
110        }
111    }
112
113    fn apply_primary_statement_effect(
114        &mut self,
115        state: &mut Self::Domain,
116        statement: &Statement<'tcx>,
117        _location: Location,
118    ) {
119        if state.is_reachable() {
120            self.handle_statement(statement, state);
121        }
122    }
123
124    fn apply_primary_terminator_effect<'mir>(
125        &mut self,
126        state: &mut Self::Domain,
127        terminator: &'mir Terminator<'tcx>,
128        _location: Location,
129    ) -> TerminatorEdges<'mir, 'tcx> {
130        if state.is_reachable() {
131            self.handle_terminator(terminator, state)
132        } else {
133            TerminatorEdges::None
134        }
135    }
136
137    fn apply_call_return_effect(
138        &mut self,
139        state: &mut Self::Domain,
140        _block: BasicBlock,
141        return_places: CallReturnPlaces<'_, 'tcx>,
142    ) {
143        if state.is_reachable() {
144            self.handle_call_return(return_places, state)
145        }
146    }
147}
148
149impl<'a, 'tcx> ConstAnalysis<'a, 'tcx> {
150    fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, map: Map<'tcx>) -> Self {
151        let typing_env = body.typing_env(tcx);
152        Self {
153            map,
154            tcx,
155            local_decls: &body.local_decls,
156            ecx: InterpCx::new(tcx, DUMMY_SP, typing_env, DummyMachine),
157            typing_env,
158        }
159    }
160
161    fn handle_statement(&self, statement: &Statement<'tcx>, state: &mut State<FlatSet<Scalar>>) {
162        match &statement.kind {
163            StatementKind::Assign(box (place, rvalue)) => {
164                self.handle_assign(*place, rvalue, state);
165            }
166            StatementKind::SetDiscriminant { box place, variant_index } => {
167                self.handle_set_discriminant(*place, *variant_index, state);
168            }
169            StatementKind::Intrinsic(box intrinsic) => {
170                self.handle_intrinsic(intrinsic);
171            }
172            StatementKind::StorageLive(local) | StatementKind::StorageDead(local) => {
173                // StorageLive leaves the local in an uninitialized state.
174                // StorageDead makes it UB to access the local afterwards.
175                state.flood_with(
176                    Place::from(*local).as_ref(),
177                    &self.map,
178                    FlatSet::<Scalar>::BOTTOM,
179                );
180            }
181            StatementKind::Deinit(box place) => {
182                // Deinit makes the place uninitialized.
183                state.flood_with(place.as_ref(), &self.map, FlatSet::<Scalar>::BOTTOM);
184            }
185            StatementKind::Retag(..) => {
186                // We don't track references.
187            }
188            StatementKind::ConstEvalCounter
189            | StatementKind::Nop
190            | StatementKind::FakeRead(..)
191            | StatementKind::PlaceMention(..)
192            | StatementKind::Coverage(..)
193            | StatementKind::BackwardIncompatibleDropHint { .. }
194            | StatementKind::AscribeUserType(..) => {}
195        }
196    }
197
198    fn handle_intrinsic(&self, intrinsic: &NonDivergingIntrinsic<'tcx>) {
199        match intrinsic {
200            NonDivergingIntrinsic::Assume(..) => {
201                // Could use this, but ignoring it is sound.
202            }
203            NonDivergingIntrinsic::CopyNonOverlapping(CopyNonOverlapping {
204                dst: _,
205                src: _,
206                count: _,
207            }) => {
208                // This statement represents `*dst = *src`, `count` times.
209            }
210        }
211    }
212
213    fn handle_operand(
214        &self,
215        operand: &Operand<'tcx>,
216        state: &mut State<FlatSet<Scalar>>,
217    ) -> ValueOrPlace<FlatSet<Scalar>> {
218        match operand {
219            Operand::Constant(box constant) => {
220                ValueOrPlace::Value(self.handle_constant(constant, state))
221            }
222            Operand::Copy(place) | Operand::Move(place) => {
223                // On move, we would ideally flood the place with bottom. But with the current
224                // framework this is not possible (similar to `InterpCx::eval_operand`).
225                self.map.find(place.as_ref()).map(ValueOrPlace::Place).unwrap_or(ValueOrPlace::TOP)
226            }
227        }
228    }
229
230    /// The effect of a successful function call return should not be
231    /// applied here, see [`Analysis::apply_primary_terminator_effect`].
232    fn handle_terminator<'mir>(
233        &self,
234        terminator: &'mir Terminator<'tcx>,
235        state: &mut State<FlatSet<Scalar>>,
236    ) -> TerminatorEdges<'mir, 'tcx> {
237        match &terminator.kind {
238            TerminatorKind::Call { .. } | TerminatorKind::InlineAsm { .. } => {
239                // Effect is applied by `handle_call_return`.
240            }
241            TerminatorKind::Drop { place, .. } => {
242                state.flood_with(place.as_ref(), &self.map, FlatSet::<Scalar>::BOTTOM);
243            }
244            TerminatorKind::Yield { .. } => {
245                // They would have an effect, but are not allowed in this phase.
246                bug!("encountered disallowed terminator");
247            }
248            TerminatorKind::SwitchInt { discr, targets } => {
249                return self.handle_switch_int(discr, targets, state);
250            }
251            TerminatorKind::TailCall { .. } => {
252                // FIXME(explicit_tail_calls): determine if we need to do something here (probably
253                // not)
254            }
255            TerminatorKind::Goto { .. }
256            | TerminatorKind::UnwindResume
257            | TerminatorKind::UnwindTerminate(_)
258            | TerminatorKind::Return
259            | TerminatorKind::Unreachable
260            | TerminatorKind::Assert { .. }
261            | TerminatorKind::CoroutineDrop
262            | TerminatorKind::FalseEdge { .. }
263            | TerminatorKind::FalseUnwind { .. } => {
264                // These terminators have no effect on the analysis.
265            }
266        }
267        terminator.edges()
268    }
269
270    fn handle_call_return(
271        &self,
272        return_places: CallReturnPlaces<'_, 'tcx>,
273        state: &mut State<FlatSet<Scalar>>,
274    ) {
275        return_places.for_each(|place| {
276            state.flood(place.as_ref(), &self.map);
277        })
278    }
279
280    fn handle_set_discriminant(
281        &self,
282        place: Place<'tcx>,
283        variant_index: VariantIdx,
284        state: &mut State<FlatSet<Scalar>>,
285    ) {
286        state.flood_discr(place.as_ref(), &self.map);
287        if self.map.find_discr(place.as_ref()).is_some() {
288            let enum_ty = place.ty(self.local_decls, self.tcx).ty;
289            if let Some(discr) = self.eval_discriminant(enum_ty, variant_index) {
290                state.assign_discr(
291                    place.as_ref(),
292                    ValueOrPlace::Value(FlatSet::Elem(discr)),
293                    &self.map,
294                );
295            }
296        }
297    }
298
299    fn handle_assign(
300        &self,
301        target: Place<'tcx>,
302        rvalue: &Rvalue<'tcx>,
303        state: &mut State<FlatSet<Scalar>>,
304    ) {
305        match rvalue {
306            Rvalue::Use(operand) => {
307                state.flood(target.as_ref(), &self.map);
308                if let Some(target) = self.map.find(target.as_ref()) {
309                    self.assign_operand(state, target, operand);
310                }
311            }
312            Rvalue::CopyForDeref(rhs) => {
313                state.flood(target.as_ref(), &self.map);
314                if let Some(target) = self.map.find(target.as_ref()) {
315                    self.assign_operand(state, target, &Operand::Copy(*rhs));
316                }
317            }
318            Rvalue::Aggregate(kind, operands) => {
319                // If we assign `target = Enum::Variant#0(operand)`,
320                // we must make sure that all `target as Variant#i` are `Top`.
321                state.flood(target.as_ref(), &self.map);
322
323                let Some(target_idx) = self.map.find(target.as_ref()) else { return };
324
325                let (variant_target, variant_index) = match **kind {
326                    AggregateKind::Tuple | AggregateKind::Closure(..) => (Some(target_idx), None),
327                    AggregateKind::Adt(def_id, variant_index, ..) => {
328                        match self.tcx.def_kind(def_id) {
329                            DefKind::Struct => (Some(target_idx), None),
330                            DefKind::Enum => (
331                                self.map.apply(target_idx, TrackElem::Variant(variant_index)),
332                                Some(variant_index),
333                            ),
334                            _ => return,
335                        }
336                    }
337                    _ => return,
338                };
339                if let Some(variant_target_idx) = variant_target {
340                    for (field_index, operand) in operands.iter_enumerated() {
341                        if let Some(field) =
342                            self.map.apply(variant_target_idx, TrackElem::Field(field_index))
343                        {
344                            self.assign_operand(state, field, operand);
345                        }
346                    }
347                }
348                if let Some(variant_index) = variant_index
349                    && let Some(discr_idx) = self.map.apply(target_idx, TrackElem::Discriminant)
350                {
351                    // We are assigning the discriminant as part of an aggregate.
352                    // This discriminant can only alias a variant field's value if the operand
353                    // had an invalid value for that type.
354                    // Using invalid values is UB, so we are allowed to perform the assignment
355                    // without extra flooding.
356                    let enum_ty = target.ty(self.local_decls, self.tcx).ty;
357                    if let Some(discr_val) = self.eval_discriminant(enum_ty, variant_index) {
358                        state.insert_value_idx(discr_idx, FlatSet::Elem(discr_val), &self.map);
359                    }
360                }
361            }
362            Rvalue::BinaryOp(op, box (left, right)) if op.is_overflowing() => {
363                // Flood everything now, so we can use `insert_value_idx` directly later.
364                state.flood(target.as_ref(), &self.map);
365
366                let Some(target) = self.map.find(target.as_ref()) else { return };
367
368                let value_target = self.map.apply(target, TrackElem::Field(0_u32.into()));
369                let overflow_target = self.map.apply(target, TrackElem::Field(1_u32.into()));
370
371                if value_target.is_some() || overflow_target.is_some() {
372                    let (val, overflow) = self.binary_op(state, *op, left, right);
373
374                    if let Some(value_target) = value_target {
375                        // We have flooded `target` earlier.
376                        state.insert_value_idx(value_target, val, &self.map);
377                    }
378                    if let Some(overflow_target) = overflow_target {
379                        // We have flooded `target` earlier.
380                        state.insert_value_idx(overflow_target, overflow, &self.map);
381                    }
382                }
383            }
384            Rvalue::Cast(
385                CastKind::PointerCoercion(ty::adjustment::PointerCoercion::Unsize, _),
386                operand,
387                _,
388            ) => {
389                let pointer = self.handle_operand(operand, state);
390                state.assign(target.as_ref(), pointer, &self.map);
391
392                if let Some(target_len) = self.map.find_len(target.as_ref())
393                    && let operand_ty = operand.ty(self.local_decls, self.tcx)
394                    && let Some(operand_ty) = operand_ty.builtin_deref(true)
395                    && let ty::Array(_, len) = operand_ty.kind()
396                    && let Some(len) = Const::Ty(self.tcx.types.usize, *len)
397                        .try_eval_scalar_int(self.tcx, self.typing_env)
398                {
399                    state.insert_value_idx(target_len, FlatSet::Elem(len.into()), &self.map);
400                }
401            }
402            _ => {
403                let result = self.handle_rvalue(rvalue, state);
404                state.assign(target.as_ref(), result, &self.map);
405            }
406        }
407    }
408
409    fn handle_rvalue(
410        &self,
411        rvalue: &Rvalue<'tcx>,
412        state: &mut State<FlatSet<Scalar>>,
413    ) -> ValueOrPlace<FlatSet<Scalar>> {
414        let val = match rvalue {
415            Rvalue::Cast(CastKind::IntToInt | CastKind::IntToFloat, operand, ty) => {
416                let Ok(layout) = self.tcx.layout_of(self.typing_env.as_query_input(*ty)) else {
417                    return ValueOrPlace::Value(FlatSet::Top);
418                };
419                match self.eval_operand(operand, state) {
420                    FlatSet::Elem(op) => self
421                        .ecx
422                        .int_to_int_or_float(&op, layout)
423                        .discard_err()
424                        .map_or(FlatSet::Top, |result| self.wrap_immediate(*result)),
425                    FlatSet::Bottom => FlatSet::Bottom,
426                    FlatSet::Top => FlatSet::Top,
427                }
428            }
429            Rvalue::Cast(CastKind::FloatToInt | CastKind::FloatToFloat, operand, ty) => {
430                let Ok(layout) = self.tcx.layout_of(self.typing_env.as_query_input(*ty)) else {
431                    return ValueOrPlace::Value(FlatSet::Top);
432                };
433                match self.eval_operand(operand, state) {
434                    FlatSet::Elem(op) => self
435                        .ecx
436                        .float_to_float_or_int(&op, layout)
437                        .discard_err()
438                        .map_or(FlatSet::Top, |result| self.wrap_immediate(*result)),
439                    FlatSet::Bottom => FlatSet::Bottom,
440                    FlatSet::Top => FlatSet::Top,
441                }
442            }
443            Rvalue::Cast(CastKind::Transmute, operand, _) => {
444                match self.eval_operand(operand, state) {
445                    FlatSet::Elem(op) => self.wrap_immediate(*op),
446                    FlatSet::Bottom => FlatSet::Bottom,
447                    FlatSet::Top => FlatSet::Top,
448                }
449            }
450            Rvalue::BinaryOp(op, box (left, right)) if !op.is_overflowing() => {
451                // Overflows must be ignored here.
452                // The overflowing operators are handled in `handle_assign`.
453                let (val, _overflow) = self.binary_op(state, *op, left, right);
454                val
455            }
456            Rvalue::UnaryOp(op, operand) => {
457                if let UnOp::PtrMetadata = op
458                    && let Some(place) = operand.place()
459                    && let Some(len) = self.map.find_len(place.as_ref())
460                {
461                    return ValueOrPlace::Place(len);
462                }
463                match self.eval_operand(operand, state) {
464                    FlatSet::Elem(value) => self
465                        .ecx
466                        .unary_op(*op, &value)
467                        .discard_err()
468                        .map_or(FlatSet::Top, |val| self.wrap_immediate(*val)),
469                    FlatSet::Bottom => FlatSet::Bottom,
470                    FlatSet::Top => FlatSet::Top,
471                }
472            }
473            Rvalue::NullaryOp(null_op, ty) => {
474                let Ok(layout) = self.tcx.layout_of(self.typing_env.as_query_input(*ty)) else {
475                    return ValueOrPlace::Value(FlatSet::Top);
476                };
477                let val = match null_op {
478                    NullOp::SizeOf if layout.is_sized() => layout.size.bytes(),
479                    NullOp::AlignOf if layout.is_sized() => layout.align.bytes(),
480                    NullOp::OffsetOf(fields) => self
481                        .ecx
482                        .tcx
483                        .offset_of_subfield(self.typing_env, layout, fields.iter())
484                        .bytes(),
485                    _ => return ValueOrPlace::Value(FlatSet::Top),
486                };
487                FlatSet::Elem(Scalar::from_target_usize(val, &self.tcx))
488            }
489            Rvalue::Discriminant(place) => state.get_discr(place.as_ref(), &self.map),
490            Rvalue::Use(operand) => return self.handle_operand(operand, state),
491            Rvalue::CopyForDeref(place) => {
492                return self.handle_operand(&Operand::Copy(*place), state);
493            }
494            Rvalue::Ref(..) | Rvalue::RawPtr(..) => {
495                // We don't track such places.
496                return ValueOrPlace::TOP;
497            }
498            Rvalue::Repeat(..)
499            | Rvalue::ThreadLocalRef(..)
500            | Rvalue::Cast(..)
501            | Rvalue::BinaryOp(..)
502            | Rvalue::Aggregate(..)
503            | Rvalue::ShallowInitBox(..)
504            | Rvalue::WrapUnsafeBinder(..) => {
505                // No modification is possible through these r-values.
506                return ValueOrPlace::TOP;
507            }
508        };
509        ValueOrPlace::Value(val)
510    }
511
512    fn handle_constant(
513        &self,
514        constant: &ConstOperand<'tcx>,
515        _state: &mut State<FlatSet<Scalar>>,
516    ) -> FlatSet<Scalar> {
517        constant
518            .const_
519            .try_eval_scalar(self.tcx, self.typing_env)
520            .map_or(FlatSet::Top, FlatSet::Elem)
521    }
522
523    fn handle_switch_int<'mir>(
524        &self,
525        discr: &'mir Operand<'tcx>,
526        targets: &'mir SwitchTargets,
527        state: &mut State<FlatSet<Scalar>>,
528    ) -> TerminatorEdges<'mir, 'tcx> {
529        let value = match self.handle_operand(discr, state) {
530            ValueOrPlace::Value(value) => value,
531            ValueOrPlace::Place(place) => state.get_idx(place, &self.map),
532        };
533        match value {
534            // We are branching on uninitialized data, this is UB, treat it as unreachable.
535            // This allows the set of visited edges to grow monotonically with the lattice.
536            FlatSet::Bottom => TerminatorEdges::None,
537            FlatSet::Elem(scalar) => {
538                if let Ok(scalar_int) = scalar.try_to_scalar_int() {
539                    TerminatorEdges::Single(
540                        targets.target_for_value(scalar_int.to_bits_unchecked()),
541                    )
542                } else {
543                    TerminatorEdges::SwitchInt { discr, targets }
544                }
545            }
546            FlatSet::Top => TerminatorEdges::SwitchInt { discr, targets },
547        }
548    }
549
550    /// The caller must have flooded `place`.
551    fn assign_operand(
552        &self,
553        state: &mut State<FlatSet<Scalar>>,
554        place: PlaceIndex,
555        operand: &Operand<'tcx>,
556    ) {
557        match operand {
558            Operand::Copy(rhs) | Operand::Move(rhs) => {
559                if let Some(rhs) = self.map.find(rhs.as_ref()) {
560                    state.insert_place_idx(place, rhs, &self.map);
561                } else if rhs.projection.first() == Some(&PlaceElem::Deref)
562                    && let FlatSet::Elem(pointer) = state.get(rhs.local.into(), &self.map)
563                    && let rhs_ty = self.local_decls[rhs.local].ty
564                    && let Ok(rhs_layout) =
565                        self.tcx.layout_of(self.typing_env.as_query_input(rhs_ty))
566                {
567                    let op = ImmTy::from_scalar(pointer, rhs_layout).into();
568                    self.assign_constant(state, place, op, rhs.projection);
569                }
570            }
571            Operand::Constant(box constant) => {
572                if let Some(constant) =
573                    self.ecx.eval_mir_constant(&constant.const_, constant.span, None).discard_err()
574                {
575                    self.assign_constant(state, place, constant, &[]);
576                }
577            }
578        }
579    }
580
581    /// The caller must have flooded `place`.
582    ///
583    /// Perform: `place = operand.projection`.
584    #[instrument(level = "trace", skip(self, state))]
585    fn assign_constant(
586        &self,
587        state: &mut State<FlatSet<Scalar>>,
588        place: PlaceIndex,
589        mut operand: OpTy<'tcx>,
590        projection: &[PlaceElem<'tcx>],
591    ) {
592        for &(mut proj_elem) in projection {
593            if let PlaceElem::Index(index) = proj_elem {
594                if let FlatSet::Elem(index) = state.get(index.into(), &self.map)
595                    && let Some(offset) = index.to_target_usize(&self.tcx).discard_err()
596                    && let Some(min_length) = offset.checked_add(1)
597                {
598                    proj_elem = PlaceElem::ConstantIndex { offset, min_length, from_end: false };
599                } else {
600                    return;
601                }
602            }
603            operand = if let Some(operand) = self.ecx.project(&operand, proj_elem).discard_err() {
604                operand
605            } else {
606                return;
607            }
608        }
609
610        self.map.for_each_projection_value(
611            place,
612            operand,
613            &mut |elem, op| match elem {
614                TrackElem::Field(idx) => self.ecx.project_field(op, idx).discard_err(),
615                TrackElem::Variant(idx) => self.ecx.project_downcast(op, idx).discard_err(),
616                TrackElem::Discriminant => {
617                    let variant = self.ecx.read_discriminant(op).discard_err()?;
618                    let discr_value =
619                        self.ecx.discriminant_for_variant(op.layout.ty, variant).discard_err()?;
620                    Some(discr_value.into())
621                }
622                TrackElem::DerefLen => {
623                    let op: OpTy<'_> = self.ecx.deref_pointer(op).discard_err()?.into();
624                    let len_usize = op.len(&self.ecx).discard_err()?;
625                    let layout = self
626                        .tcx
627                        .layout_of(self.typing_env.as_query_input(self.tcx.types.usize))
628                        .unwrap();
629                    Some(ImmTy::from_uint(len_usize, layout).into())
630                }
631            },
632            &mut |place, op| {
633                if let Some(imm) = self.ecx.read_immediate_raw(op).discard_err()
634                    && let Some(imm) = imm.right()
635                {
636                    let elem = self.wrap_immediate(*imm);
637                    state.insert_value_idx(place, elem, &self.map);
638                }
639            },
640        );
641    }
642
643    fn binary_op(
644        &self,
645        state: &mut State<FlatSet<Scalar>>,
646        op: BinOp,
647        left: &Operand<'tcx>,
648        right: &Operand<'tcx>,
649    ) -> (FlatSet<Scalar>, FlatSet<Scalar>) {
650        let left = self.eval_operand(left, state);
651        let right = self.eval_operand(right, state);
652
653        match (left, right) {
654            (FlatSet::Bottom, _) | (_, FlatSet::Bottom) => (FlatSet::Bottom, FlatSet::Bottom),
655            // Both sides are known, do the actual computation.
656            (FlatSet::Elem(left), FlatSet::Elem(right)) => {
657                match self.ecx.binary_op(op, &left, &right).discard_err() {
658                    // Ideally this would return an Immediate, since it's sometimes
659                    // a pair and sometimes not. But as a hack we always return a pair
660                    // and just make the 2nd component `Bottom` when it does not exist.
661                    Some(val) => {
662                        if matches!(val.layout.backend_repr, BackendRepr::ScalarPair(..)) {
663                            let (val, overflow) = val.to_scalar_pair();
664                            (FlatSet::Elem(val), FlatSet::Elem(overflow))
665                        } else {
666                            (FlatSet::Elem(val.to_scalar()), FlatSet::Bottom)
667                        }
668                    }
669                    _ => (FlatSet::Top, FlatSet::Top),
670                }
671            }
672            // Exactly one side is known, attempt some algebraic simplifications.
673            (FlatSet::Elem(const_arg), _) | (_, FlatSet::Elem(const_arg)) => {
674                let layout = const_arg.layout;
675                if !matches!(layout.backend_repr, rustc_abi::BackendRepr::Scalar(..)) {
676                    return (FlatSet::Top, FlatSet::Top);
677                }
678
679                let arg_scalar = const_arg.to_scalar();
680                let Some(arg_value) = arg_scalar.to_bits(layout.size).discard_err() else {
681                    return (FlatSet::Top, FlatSet::Top);
682                };
683
684                match op {
685                    BinOp::BitAnd if arg_value == 0 => (FlatSet::Elem(arg_scalar), FlatSet::Bottom),
686                    BinOp::BitOr
687                        if arg_value == layout.size.truncate(u128::MAX)
688                            || (layout.ty.is_bool() && arg_value == 1) =>
689                    {
690                        (FlatSet::Elem(arg_scalar), FlatSet::Bottom)
691                    }
692                    BinOp::Mul if layout.ty.is_integral() && arg_value == 0 => {
693                        (FlatSet::Elem(arg_scalar), FlatSet::Elem(Scalar::from_bool(false)))
694                    }
695                    _ => (FlatSet::Top, FlatSet::Top),
696                }
697            }
698            (FlatSet::Top, FlatSet::Top) => (FlatSet::Top, FlatSet::Top),
699        }
700    }
701
702    fn eval_operand(
703        &self,
704        op: &Operand<'tcx>,
705        state: &mut State<FlatSet<Scalar>>,
706    ) -> FlatSet<ImmTy<'tcx>> {
707        let value = match self.handle_operand(op, state) {
708            ValueOrPlace::Value(value) => value,
709            ValueOrPlace::Place(place) => state.get_idx(place, &self.map),
710        };
711        match value {
712            FlatSet::Top => FlatSet::Top,
713            FlatSet::Elem(scalar) => {
714                let ty = op.ty(self.local_decls, self.tcx);
715                self.tcx
716                    .layout_of(self.typing_env.as_query_input(ty))
717                    .map_or(FlatSet::Top, |layout| {
718                        FlatSet::Elem(ImmTy::from_scalar(scalar, layout))
719                    })
720            }
721            FlatSet::Bottom => FlatSet::Bottom,
722        }
723    }
724
725    fn eval_discriminant(&self, enum_ty: Ty<'tcx>, variant_index: VariantIdx) -> Option<Scalar> {
726        if !enum_ty.is_enum() {
727            return None;
728        }
729        let enum_ty_layout = self.tcx.layout_of(self.typing_env.as_query_input(enum_ty)).ok()?;
730        let discr_value =
731            self.ecx.discriminant_for_variant(enum_ty_layout.ty, variant_index).discard_err()?;
732        Some(discr_value.to_scalar())
733    }
734
735    fn wrap_immediate(&self, imm: Immediate) -> FlatSet<Scalar> {
736        match imm {
737            Immediate::Scalar(scalar) => FlatSet::Elem(scalar),
738            Immediate::Uninit => FlatSet::Bottom,
739            _ => FlatSet::Top,
740        }
741    }
742}
743
744/// This is used to visualize the dataflow analysis.
745impl<'tcx> DebugWithContext<ConstAnalysis<'_, 'tcx>> for State<FlatSet<Scalar>> {
746    fn fmt_with(&self, ctxt: &ConstAnalysis<'_, 'tcx>, f: &mut Formatter<'_>) -> std::fmt::Result {
747        match self {
748            State::Reachable(values) => debug_with_context(values, None, &ctxt.map, f),
749            State::Unreachable => write!(f, "unreachable"),
750        }
751    }
752
753    fn fmt_diff_with(
754        &self,
755        old: &Self,
756        ctxt: &ConstAnalysis<'_, 'tcx>,
757        f: &mut Formatter<'_>,
758    ) -> std::fmt::Result {
759        match (self, old) {
760            (State::Reachable(this), State::Reachable(old)) => {
761                debug_with_context(this, Some(old), &ctxt.map, f)
762            }
763            _ => Ok(()), // Consider printing something here.
764        }
765    }
766}
767
768struct Patch<'tcx> {
769    tcx: TyCtxt<'tcx>,
770
771    /// For a given MIR location, this stores the values of the operands used by that location. In
772    /// particular, this is before the effect, such that the operands of `_1 = _1 + _2` are
773    /// properly captured. (This may become UB soon, but it is currently emitted even by safe code.)
774    before_effect: FxHashMap<(Location, Place<'tcx>), Const<'tcx>>,
775
776    /// Stores the assigned values for assignments where the Rvalue is constant.
777    assignments: FxHashMap<Location, Const<'tcx>>,
778}
779
780impl<'tcx> Patch<'tcx> {
781    pub(crate) fn new(tcx: TyCtxt<'tcx>) -> Self {
782        Self { tcx, before_effect: FxHashMap::default(), assignments: FxHashMap::default() }
783    }
784
785    fn make_operand(&self, const_: Const<'tcx>) -> Operand<'tcx> {
786        Operand::Constant(Box::new(ConstOperand { span: DUMMY_SP, user_ty: None, const_ }))
787    }
788}
789
790struct Collector<'a, 'tcx> {
791    patch: Patch<'tcx>,
792    local_decls: &'a LocalDecls<'tcx>,
793}
794
795impl<'a, 'tcx> Collector<'a, 'tcx> {
796    pub(crate) fn new(tcx: TyCtxt<'tcx>, local_decls: &'a LocalDecls<'tcx>) -> Self {
797        Self { patch: Patch::new(tcx), local_decls }
798    }
799
800    #[instrument(level = "trace", skip(self, ecx, map), ret)]
801    fn try_make_constant(
802        &self,
803        ecx: &mut InterpCx<'tcx, DummyMachine>,
804        place: Place<'tcx>,
805        state: &State<FlatSet<Scalar>>,
806        map: &Map<'tcx>,
807    ) -> Option<Const<'tcx>> {
808        let ty = place.ty(self.local_decls, self.patch.tcx).ty;
809        let layout = ecx.layout_of(ty).ok()?;
810
811        if layout.is_zst() {
812            return Some(Const::zero_sized(ty));
813        }
814
815        if layout.is_unsized() {
816            return None;
817        }
818
819        let place = map.find(place.as_ref())?;
820        if layout.backend_repr.is_scalar()
821            && let Some(value) = propagatable_scalar(place, state, map)
822        {
823            return Some(Const::Val(ConstValue::Scalar(value), ty));
824        }
825
826        if matches!(layout.backend_repr, BackendRepr::Scalar(..) | BackendRepr::ScalarPair(..)) {
827            let alloc_id = ecx
828                .intern_with_temp_alloc(layout, |ecx, dest| {
829                    try_write_constant(ecx, dest, place, ty, state, map)
830                })
831                .discard_err()?;
832            return Some(Const::Val(ConstValue::Indirect { alloc_id, offset: Size::ZERO }, ty));
833        }
834
835        None
836    }
837}
838
839#[instrument(level = "trace", skip(map), ret)]
840fn propagatable_scalar(
841    place: PlaceIndex,
842    state: &State<FlatSet<Scalar>>,
843    map: &Map<'_>,
844) -> Option<Scalar> {
845    if let FlatSet::Elem(value) = state.get_idx(place, map)
846        && value.try_to_scalar_int().is_ok()
847    {
848        // Do not attempt to propagate pointers, as we may fail to preserve their identity.
849        Some(value)
850    } else {
851        None
852    }
853}
854
855#[instrument(level = "trace", skip(ecx, state, map), ret)]
856fn try_write_constant<'tcx>(
857    ecx: &mut InterpCx<'tcx, DummyMachine>,
858    dest: &PlaceTy<'tcx>,
859    place: PlaceIndex,
860    ty: Ty<'tcx>,
861    state: &State<FlatSet<Scalar>>,
862    map: &Map<'tcx>,
863) -> InterpResult<'tcx> {
864    let layout = ecx.layout_of(ty)?;
865
866    // Fast path for ZSTs.
867    if layout.is_zst() {
868        return interp_ok(());
869    }
870
871    // Fast path for scalars.
872    if layout.backend_repr.is_scalar()
873        && let Some(value) = propagatable_scalar(place, state, map)
874    {
875        return ecx.write_immediate(Immediate::Scalar(value), dest);
876    }
877
878    match ty.kind() {
879        // ZSTs. Nothing to do.
880        ty::FnDef(..) => {}
881
882        // Those are scalars, must be handled above.
883        ty::Bool | ty::Int(_) | ty::Uint(_) | ty::Float(_) | ty::Char =>
884            throw_machine_stop_str!("primitive type with provenance"),
885
886        ty::Tuple(elem_tys) => {
887            for (i, elem) in elem_tys.iter().enumerate() {
888                let i = FieldIdx::from_usize(i);
889                let Some(field) = map.apply(place, TrackElem::Field(i)) else {
890                    throw_machine_stop_str!("missing field in tuple")
891                };
892                let field_dest = ecx.project_field(dest, i)?;
893                try_write_constant(ecx, &field_dest, field, elem, state, map)?;
894            }
895        }
896
897        ty::Adt(def, args) => {
898            if def.is_union() {
899                throw_machine_stop_str!("cannot propagate unions")
900            }
901
902            let (variant_idx, variant_def, variant_place, variant_dest) = if def.is_enum() {
903                let Some(discr) = map.apply(place, TrackElem::Discriminant) else {
904                    throw_machine_stop_str!("missing discriminant for enum")
905                };
906                let FlatSet::Elem(Scalar::Int(discr)) = state.get_idx(discr, map) else {
907                    throw_machine_stop_str!("discriminant with provenance")
908                };
909                let discr_bits = discr.to_bits(discr.size());
910                let Some((variant, _)) = def.discriminants(*ecx.tcx).find(|(_, var)| discr_bits == var.val) else {
911                    throw_machine_stop_str!("illegal discriminant for enum")
912                };
913                let Some(variant_place) = map.apply(place, TrackElem::Variant(variant)) else {
914                    throw_machine_stop_str!("missing variant for enum")
915                };
916                let variant_dest = ecx.project_downcast(dest, variant)?;
917                (variant, def.variant(variant), variant_place, variant_dest)
918            } else {
919                (FIRST_VARIANT, def.non_enum_variant(), place, dest.clone())
920            };
921
922            for (i, field) in variant_def.fields.iter_enumerated() {
923                let ty = field.ty(*ecx.tcx, args);
924                let Some(field) = map.apply(variant_place, TrackElem::Field(i)) else {
925                    throw_machine_stop_str!("missing field in ADT")
926                };
927                let field_dest = ecx.project_field(&variant_dest, i)?;
928                try_write_constant(ecx, &field_dest, field, ty, state, map)?;
929            }
930            ecx.write_discriminant(variant_idx, dest)?;
931        }
932
933        // Unsupported for now.
934        ty::Array(_, _)
935        | ty::Pat(_, _)
936
937        // Do not attempt to support indirection in constants.
938        | ty::Ref(..) | ty::RawPtr(..) | ty::FnPtr(..) | ty::Str | ty::Slice(_)
939
940        | ty::Never
941        | ty::Foreign(..)
942        | ty::Alias(..)
943        | ty::Param(_)
944        | ty::Bound(..)
945        | ty::Placeholder(..)
946        | ty::Closure(..)
947        | ty::CoroutineClosure(..)
948        | ty::Coroutine(..)
949        | ty::Dynamic(..)
950        | ty::UnsafeBinder(_) => throw_machine_stop_str!("unsupported type"),
951
952        ty::Error(_) | ty::Infer(..) | ty::CoroutineWitness(..) => bug!(),
953    }
954
955    interp_ok(())
956}
957
958impl<'tcx> ResultsVisitor<'tcx, ConstAnalysis<'_, 'tcx>> for Collector<'_, 'tcx> {
959    #[instrument(level = "trace", skip(self, analysis, statement))]
960    fn visit_after_early_statement_effect(
961        &mut self,
962        analysis: &mut ConstAnalysis<'_, 'tcx>,
963        state: &State<FlatSet<Scalar>>,
964        statement: &Statement<'tcx>,
965        location: Location,
966    ) {
967        match &statement.kind {
968            StatementKind::Assign(box (_, rvalue)) => {
969                OperandCollector {
970                    state,
971                    visitor: self,
972                    ecx: &mut analysis.ecx,
973                    map: &analysis.map,
974                }
975                .visit_rvalue(rvalue, location);
976            }
977            _ => (),
978        }
979    }
980
981    #[instrument(level = "trace", skip(self, analysis, statement))]
982    fn visit_after_primary_statement_effect(
983        &mut self,
984        analysis: &mut ConstAnalysis<'_, 'tcx>,
985        state: &State<FlatSet<Scalar>>,
986        statement: &Statement<'tcx>,
987        location: Location,
988    ) {
989        match statement.kind {
990            StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(_)))) => {
991                // Don't overwrite the assignment if it already uses a constant (to keep the span).
992            }
993            StatementKind::Assign(box (place, _)) => {
994                if let Some(value) =
995                    self.try_make_constant(&mut analysis.ecx, place, state, &analysis.map)
996                {
997                    self.patch.assignments.insert(location, value);
998                }
999            }
1000            _ => (),
1001        }
1002    }
1003
1004    fn visit_after_early_terminator_effect(
1005        &mut self,
1006        analysis: &mut ConstAnalysis<'_, 'tcx>,
1007        state: &State<FlatSet<Scalar>>,
1008        terminator: &Terminator<'tcx>,
1009        location: Location,
1010    ) {
1011        OperandCollector { state, visitor: self, ecx: &mut analysis.ecx, map: &analysis.map }
1012            .visit_terminator(terminator, location);
1013    }
1014}
1015
1016impl<'tcx> MutVisitor<'tcx> for Patch<'tcx> {
1017    fn tcx(&self) -> TyCtxt<'tcx> {
1018        self.tcx
1019    }
1020
1021    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
1022        if let Some(value) = self.assignments.get(&location) {
1023            match &mut statement.kind {
1024                StatementKind::Assign(box (_, rvalue)) => {
1025                    *rvalue = Rvalue::Use(self.make_operand(*value));
1026                }
1027                _ => bug!("found assignment info for non-assign statement"),
1028            }
1029        } else {
1030            self.super_statement(statement, location);
1031        }
1032    }
1033
1034    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
1035        match operand {
1036            Operand::Copy(place) | Operand::Move(place) => {
1037                if let Some(value) = self.before_effect.get(&(location, *place)) {
1038                    *operand = self.make_operand(*value);
1039                } else if !place.projection.is_empty() {
1040                    self.super_operand(operand, location)
1041                }
1042            }
1043            Operand::Constant(_) => {}
1044        }
1045    }
1046
1047    fn process_projection_elem(
1048        &mut self,
1049        elem: PlaceElem<'tcx>,
1050        location: Location,
1051    ) -> Option<PlaceElem<'tcx>> {
1052        if let PlaceElem::Index(local) = elem {
1053            let offset = self.before_effect.get(&(location, local.into()))?;
1054            let offset = offset.try_to_scalar()?;
1055            let offset = offset.to_target_usize(&self.tcx).discard_err()?;
1056            let min_length = offset.checked_add(1)?;
1057            Some(PlaceElem::ConstantIndex { offset, min_length, from_end: false })
1058        } else {
1059            None
1060        }
1061    }
1062}
1063
1064struct OperandCollector<'a, 'b, 'tcx> {
1065    state: &'a State<FlatSet<Scalar>>,
1066    visitor: &'a mut Collector<'b, 'tcx>,
1067    ecx: &'a mut InterpCx<'tcx, DummyMachine>,
1068    map: &'a Map<'tcx>,
1069}
1070
1071impl<'tcx> Visitor<'tcx> for OperandCollector<'_, '_, 'tcx> {
1072    fn visit_projection_elem(
1073        &mut self,
1074        _: PlaceRef<'tcx>,
1075        elem: PlaceElem<'tcx>,
1076        _: PlaceContext,
1077        location: Location,
1078    ) {
1079        if let PlaceElem::Index(local) = elem
1080            && let Some(value) =
1081                self.visitor.try_make_constant(self.ecx, local.into(), self.state, self.map)
1082        {
1083            self.visitor.patch.before_effect.insert((location, local.into()), value);
1084        }
1085    }
1086
1087    fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
1088        if let Some(place) = operand.place() {
1089            if let Some(value) =
1090                self.visitor.try_make_constant(self.ecx, place, self.state, self.map)
1091            {
1092                self.visitor.patch.before_effect.insert((location, place), value);
1093            } else if !place.projection.is_empty() {
1094                // Try to propagate into `Index` projections.
1095                self.super_operand(operand, location)
1096            }
1097        }
1098    }
1099}