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::Len(place) => {
416                let place_ty = place.ty(self.local_decls, self.tcx);
417                if let ty::Array(_, len) = place_ty.ty.kind() {
418                    Const::Ty(self.tcx.types.usize, *len)
419                        .try_eval_scalar(self.tcx, self.typing_env)
420                        .map_or(FlatSet::Top, FlatSet::Elem)
421                } else if let [ProjectionElem::Deref] = place.projection[..] {
422                    state.get_len(place.local.into(), &self.map)
423                } else {
424                    FlatSet::Top
425                }
426            }
427            Rvalue::Cast(CastKind::IntToInt | CastKind::IntToFloat, operand, ty) => {
428                let Ok(layout) = self.tcx.layout_of(self.typing_env.as_query_input(*ty)) else {
429                    return ValueOrPlace::Value(FlatSet::Top);
430                };
431                match self.eval_operand(operand, state) {
432                    FlatSet::Elem(op) => self
433                        .ecx
434                        .int_to_int_or_float(&op, layout)
435                        .discard_err()
436                        .map_or(FlatSet::Top, |result| self.wrap_immediate(*result)),
437                    FlatSet::Bottom => FlatSet::Bottom,
438                    FlatSet::Top => FlatSet::Top,
439                }
440            }
441            Rvalue::Cast(CastKind::FloatToInt | CastKind::FloatToFloat, operand, ty) => {
442                let Ok(layout) = self.tcx.layout_of(self.typing_env.as_query_input(*ty)) else {
443                    return ValueOrPlace::Value(FlatSet::Top);
444                };
445                match self.eval_operand(operand, state) {
446                    FlatSet::Elem(op) => self
447                        .ecx
448                        .float_to_float_or_int(&op, layout)
449                        .discard_err()
450                        .map_or(FlatSet::Top, |result| self.wrap_immediate(*result)),
451                    FlatSet::Bottom => FlatSet::Bottom,
452                    FlatSet::Top => FlatSet::Top,
453                }
454            }
455            Rvalue::Cast(CastKind::Transmute, operand, _) => {
456                match self.eval_operand(operand, state) {
457                    FlatSet::Elem(op) => self.wrap_immediate(*op),
458                    FlatSet::Bottom => FlatSet::Bottom,
459                    FlatSet::Top => FlatSet::Top,
460                }
461            }
462            Rvalue::BinaryOp(op, box (left, right)) if !op.is_overflowing() => {
463                // Overflows must be ignored here.
464                // The overflowing operators are handled in `handle_assign`.
465                let (val, _overflow) = self.binary_op(state, *op, left, right);
466                val
467            }
468            Rvalue::UnaryOp(op, operand) => match self.eval_operand(operand, state) {
469                FlatSet::Elem(value) => self
470                    .ecx
471                    .unary_op(*op, &value)
472                    .discard_err()
473                    .map_or(FlatSet::Top, |val| self.wrap_immediate(*val)),
474                FlatSet::Bottom => FlatSet::Bottom,
475                FlatSet::Top => FlatSet::Top,
476            },
477            Rvalue::NullaryOp(null_op, ty) => {
478                let Ok(layout) = self.tcx.layout_of(self.typing_env.as_query_input(*ty)) else {
479                    return ValueOrPlace::Value(FlatSet::Top);
480                };
481                let val = match null_op {
482                    NullOp::SizeOf if layout.is_sized() => layout.size.bytes(),
483                    NullOp::AlignOf if layout.is_sized() => layout.align.abi.bytes(),
484                    NullOp::OffsetOf(fields) => self
485                        .ecx
486                        .tcx
487                        .offset_of_subfield(self.typing_env, layout, fields.iter())
488                        .bytes(),
489                    _ => return ValueOrPlace::Value(FlatSet::Top),
490                };
491                FlatSet::Elem(Scalar::from_target_usize(val, &self.tcx))
492            }
493            Rvalue::Discriminant(place) => state.get_discr(place.as_ref(), &self.map),
494            Rvalue::Use(operand) => return self.handle_operand(operand, state),
495            Rvalue::CopyForDeref(place) => {
496                return self.handle_operand(&Operand::Copy(*place), state);
497            }
498            Rvalue::Ref(..) | Rvalue::RawPtr(..) => {
499                // We don't track such places.
500                return ValueOrPlace::TOP;
501            }
502            Rvalue::Repeat(..)
503            | Rvalue::ThreadLocalRef(..)
504            | Rvalue::Cast(..)
505            | Rvalue::BinaryOp(..)
506            | Rvalue::Aggregate(..)
507            | Rvalue::ShallowInitBox(..)
508            | Rvalue::WrapUnsafeBinder(..) => {
509                // No modification is possible through these r-values.
510                return ValueOrPlace::TOP;
511            }
512        };
513        ValueOrPlace::Value(val)
514    }
515
516    fn handle_constant(
517        &self,
518        constant: &ConstOperand<'tcx>,
519        _state: &mut State<FlatSet<Scalar>>,
520    ) -> FlatSet<Scalar> {
521        constant
522            .const_
523            .try_eval_scalar(self.tcx, self.typing_env)
524            .map_or(FlatSet::Top, FlatSet::Elem)
525    }
526
527    fn handle_switch_int<'mir>(
528        &self,
529        discr: &'mir Operand<'tcx>,
530        targets: &'mir SwitchTargets,
531        state: &mut State<FlatSet<Scalar>>,
532    ) -> TerminatorEdges<'mir, 'tcx> {
533        let value = match self.handle_operand(discr, state) {
534            ValueOrPlace::Value(value) => value,
535            ValueOrPlace::Place(place) => state.get_idx(place, &self.map),
536        };
537        match value {
538            // We are branching on uninitialized data, this is UB, treat it as unreachable.
539            // This allows the set of visited edges to grow monotonically with the lattice.
540            FlatSet::Bottom => TerminatorEdges::None,
541            FlatSet::Elem(scalar) => {
542                if let Ok(scalar_int) = scalar.try_to_scalar_int() {
543                    TerminatorEdges::Single(
544                        targets.target_for_value(scalar_int.to_bits_unchecked()),
545                    )
546                } else {
547                    TerminatorEdges::SwitchInt { discr, targets }
548                }
549            }
550            FlatSet::Top => TerminatorEdges::SwitchInt { discr, targets },
551        }
552    }
553
554    /// The caller must have flooded `place`.
555    fn assign_operand(
556        &self,
557        state: &mut State<FlatSet<Scalar>>,
558        place: PlaceIndex,
559        operand: &Operand<'tcx>,
560    ) {
561        match operand {
562            Operand::Copy(rhs) | Operand::Move(rhs) => {
563                if let Some(rhs) = self.map.find(rhs.as_ref()) {
564                    state.insert_place_idx(place, rhs, &self.map);
565                } else if rhs.projection.first() == Some(&PlaceElem::Deref)
566                    && let FlatSet::Elem(pointer) = state.get(rhs.local.into(), &self.map)
567                    && let rhs_ty = self.local_decls[rhs.local].ty
568                    && let Ok(rhs_layout) =
569                        self.tcx.layout_of(self.typing_env.as_query_input(rhs_ty))
570                {
571                    let op = ImmTy::from_scalar(pointer, rhs_layout).into();
572                    self.assign_constant(state, place, op, rhs.projection);
573                }
574            }
575            Operand::Constant(box constant) => {
576                if let Some(constant) =
577                    self.ecx.eval_mir_constant(&constant.const_, constant.span, None).discard_err()
578                {
579                    self.assign_constant(state, place, constant, &[]);
580                }
581            }
582        }
583    }
584
585    /// The caller must have flooded `place`.
586    ///
587    /// Perform: `place = operand.projection`.
588    #[instrument(level = "trace", skip(self, state))]
589    fn assign_constant(
590        &self,
591        state: &mut State<FlatSet<Scalar>>,
592        place: PlaceIndex,
593        mut operand: OpTy<'tcx>,
594        projection: &[PlaceElem<'tcx>],
595    ) {
596        for &(mut proj_elem) in projection {
597            if let PlaceElem::Index(index) = proj_elem {
598                if let FlatSet::Elem(index) = state.get(index.into(), &self.map)
599                    && let Some(offset) = index.to_target_usize(&self.tcx).discard_err()
600                    && let Some(min_length) = offset.checked_add(1)
601                {
602                    proj_elem = PlaceElem::ConstantIndex { offset, min_length, from_end: false };
603                } else {
604                    return;
605                }
606            }
607            operand = if let Some(operand) = self.ecx.project(&operand, proj_elem).discard_err() {
608                operand
609            } else {
610                return;
611            }
612        }
613
614        self.map.for_each_projection_value(
615            place,
616            operand,
617            &mut |elem, op| match elem {
618                TrackElem::Field(idx) => self.ecx.project_field(op, idx).discard_err(),
619                TrackElem::Variant(idx) => self.ecx.project_downcast(op, idx).discard_err(),
620                TrackElem::Discriminant => {
621                    let variant = self.ecx.read_discriminant(op).discard_err()?;
622                    let discr_value =
623                        self.ecx.discriminant_for_variant(op.layout.ty, variant).discard_err()?;
624                    Some(discr_value.into())
625                }
626                TrackElem::DerefLen => {
627                    let op: OpTy<'_> = self.ecx.deref_pointer(op).discard_err()?.into();
628                    let len_usize = op.len(&self.ecx).discard_err()?;
629                    let layout = self
630                        .tcx
631                        .layout_of(self.typing_env.as_query_input(self.tcx.types.usize))
632                        .unwrap();
633                    Some(ImmTy::from_uint(len_usize, layout).into())
634                }
635            },
636            &mut |place, op| {
637                if let Some(imm) = self.ecx.read_immediate_raw(op).discard_err()
638                    && let Some(imm) = imm.right()
639                {
640                    let elem = self.wrap_immediate(*imm);
641                    state.insert_value_idx(place, elem, &self.map);
642                }
643            },
644        );
645    }
646
647    fn binary_op(
648        &self,
649        state: &mut State<FlatSet<Scalar>>,
650        op: BinOp,
651        left: &Operand<'tcx>,
652        right: &Operand<'tcx>,
653    ) -> (FlatSet<Scalar>, FlatSet<Scalar>) {
654        let left = self.eval_operand(left, state);
655        let right = self.eval_operand(right, state);
656
657        match (left, right) {
658            (FlatSet::Bottom, _) | (_, FlatSet::Bottom) => (FlatSet::Bottom, FlatSet::Bottom),
659            // Both sides are known, do the actual computation.
660            (FlatSet::Elem(left), FlatSet::Elem(right)) => {
661                match self.ecx.binary_op(op, &left, &right).discard_err() {
662                    // Ideally this would return an Immediate, since it's sometimes
663                    // a pair and sometimes not. But as a hack we always return a pair
664                    // and just make the 2nd component `Bottom` when it does not exist.
665                    Some(val) => {
666                        if matches!(val.layout.backend_repr, BackendRepr::ScalarPair(..)) {
667                            let (val, overflow) = val.to_scalar_pair();
668                            (FlatSet::Elem(val), FlatSet::Elem(overflow))
669                        } else {
670                            (FlatSet::Elem(val.to_scalar()), FlatSet::Bottom)
671                        }
672                    }
673                    _ => (FlatSet::Top, FlatSet::Top),
674                }
675            }
676            // Exactly one side is known, attempt some algebraic simplifications.
677            (FlatSet::Elem(const_arg), _) | (_, FlatSet::Elem(const_arg)) => {
678                let layout = const_arg.layout;
679                if !matches!(layout.backend_repr, rustc_abi::BackendRepr::Scalar(..)) {
680                    return (FlatSet::Top, FlatSet::Top);
681                }
682
683                let arg_scalar = const_arg.to_scalar();
684                let Some(arg_value) = arg_scalar.to_bits(layout.size).discard_err() else {
685                    return (FlatSet::Top, FlatSet::Top);
686                };
687
688                match op {
689                    BinOp::BitAnd if arg_value == 0 => (FlatSet::Elem(arg_scalar), FlatSet::Bottom),
690                    BinOp::BitOr
691                        if arg_value == layout.size.truncate(u128::MAX)
692                            || (layout.ty.is_bool() && arg_value == 1) =>
693                    {
694                        (FlatSet::Elem(arg_scalar), FlatSet::Bottom)
695                    }
696                    BinOp::Mul if layout.ty.is_integral() && arg_value == 0 => {
697                        (FlatSet::Elem(arg_scalar), FlatSet::Elem(Scalar::from_bool(false)))
698                    }
699                    _ => (FlatSet::Top, FlatSet::Top),
700                }
701            }
702            (FlatSet::Top, FlatSet::Top) => (FlatSet::Top, FlatSet::Top),
703        }
704    }
705
706    fn eval_operand(
707        &self,
708        op: &Operand<'tcx>,
709        state: &mut State<FlatSet<Scalar>>,
710    ) -> FlatSet<ImmTy<'tcx>> {
711        let value = match self.handle_operand(op, state) {
712            ValueOrPlace::Value(value) => value,
713            ValueOrPlace::Place(place) => state.get_idx(place, &self.map),
714        };
715        match value {
716            FlatSet::Top => FlatSet::Top,
717            FlatSet::Elem(scalar) => {
718                let ty = op.ty(self.local_decls, self.tcx);
719                self.tcx
720                    .layout_of(self.typing_env.as_query_input(ty))
721                    .map_or(FlatSet::Top, |layout| {
722                        FlatSet::Elem(ImmTy::from_scalar(scalar, layout))
723                    })
724            }
725            FlatSet::Bottom => FlatSet::Bottom,
726        }
727    }
728
729    fn eval_discriminant(&self, enum_ty: Ty<'tcx>, variant_index: VariantIdx) -> Option<Scalar> {
730        if !enum_ty.is_enum() {
731            return None;
732        }
733        let enum_ty_layout = self.tcx.layout_of(self.typing_env.as_query_input(enum_ty)).ok()?;
734        let discr_value =
735            self.ecx.discriminant_for_variant(enum_ty_layout.ty, variant_index).discard_err()?;
736        Some(discr_value.to_scalar())
737    }
738
739    fn wrap_immediate(&self, imm: Immediate) -> FlatSet<Scalar> {
740        match imm {
741            Immediate::Scalar(scalar) => FlatSet::Elem(scalar),
742            Immediate::Uninit => FlatSet::Bottom,
743            _ => FlatSet::Top,
744        }
745    }
746}
747
748/// This is used to visualize the dataflow analysis.
749impl<'tcx> DebugWithContext<ConstAnalysis<'_, 'tcx>> for State<FlatSet<Scalar>> {
750    fn fmt_with(&self, ctxt: &ConstAnalysis<'_, 'tcx>, f: &mut Formatter<'_>) -> std::fmt::Result {
751        match self {
752            State::Reachable(values) => debug_with_context(values, None, &ctxt.map, f),
753            State::Unreachable => write!(f, "unreachable"),
754        }
755    }
756
757    fn fmt_diff_with(
758        &self,
759        old: &Self,
760        ctxt: &ConstAnalysis<'_, 'tcx>,
761        f: &mut Formatter<'_>,
762    ) -> std::fmt::Result {
763        match (self, old) {
764            (State::Reachable(this), State::Reachable(old)) => {
765                debug_with_context(this, Some(old), &ctxt.map, f)
766            }
767            _ => Ok(()), // Consider printing something here.
768        }
769    }
770}
771
772struct Patch<'tcx> {
773    tcx: TyCtxt<'tcx>,
774
775    /// For a given MIR location, this stores the values of the operands used by that location. In
776    /// particular, this is before the effect, such that the operands of `_1 = _1 + _2` are
777    /// properly captured. (This may become UB soon, but it is currently emitted even by safe code.)
778    before_effect: FxHashMap<(Location, Place<'tcx>), Const<'tcx>>,
779
780    /// Stores the assigned values for assignments where the Rvalue is constant.
781    assignments: FxHashMap<Location, Const<'tcx>>,
782}
783
784impl<'tcx> Patch<'tcx> {
785    pub(crate) fn new(tcx: TyCtxt<'tcx>) -> Self {
786        Self { tcx, before_effect: FxHashMap::default(), assignments: FxHashMap::default() }
787    }
788
789    fn make_operand(&self, const_: Const<'tcx>) -> Operand<'tcx> {
790        Operand::Constant(Box::new(ConstOperand { span: DUMMY_SP, user_ty: None, const_ }))
791    }
792}
793
794struct Collector<'a, 'tcx> {
795    patch: Patch<'tcx>,
796    local_decls: &'a LocalDecls<'tcx>,
797}
798
799impl<'a, 'tcx> Collector<'a, 'tcx> {
800    pub(crate) fn new(tcx: TyCtxt<'tcx>, local_decls: &'a LocalDecls<'tcx>) -> Self {
801        Self { patch: Patch::new(tcx), local_decls }
802    }
803
804    #[instrument(level = "trace", skip(self, ecx, map), ret)]
805    fn try_make_constant(
806        &self,
807        ecx: &mut InterpCx<'tcx, DummyMachine>,
808        place: Place<'tcx>,
809        state: &State<FlatSet<Scalar>>,
810        map: &Map<'tcx>,
811    ) -> Option<Const<'tcx>> {
812        let ty = place.ty(self.local_decls, self.patch.tcx).ty;
813        let layout = ecx.layout_of(ty).ok()?;
814
815        if layout.is_zst() {
816            return Some(Const::zero_sized(ty));
817        }
818
819        if layout.is_unsized() {
820            return None;
821        }
822
823        let place = map.find(place.as_ref())?;
824        if layout.backend_repr.is_scalar()
825            && let Some(value) = propagatable_scalar(place, state, map)
826        {
827            return Some(Const::Val(ConstValue::Scalar(value), ty));
828        }
829
830        if matches!(layout.backend_repr, BackendRepr::Scalar(..) | BackendRepr::ScalarPair(..)) {
831            let alloc_id = ecx
832                .intern_with_temp_alloc(layout, |ecx, dest| {
833                    try_write_constant(ecx, dest, place, ty, state, map)
834                })
835                .discard_err()?;
836            return Some(Const::Val(ConstValue::Indirect { alloc_id, offset: Size::ZERO }, ty));
837        }
838
839        None
840    }
841}
842
843#[instrument(level = "trace", skip(map), ret)]
844fn propagatable_scalar(
845    place: PlaceIndex,
846    state: &State<FlatSet<Scalar>>,
847    map: &Map<'_>,
848) -> Option<Scalar> {
849    if let FlatSet::Elem(value) = state.get_idx(place, map)
850        && value.try_to_scalar_int().is_ok()
851    {
852        // Do not attempt to propagate pointers, as we may fail to preserve their identity.
853        Some(value)
854    } else {
855        None
856    }
857}
858
859#[instrument(level = "trace", skip(ecx, state, map), ret)]
860fn try_write_constant<'tcx>(
861    ecx: &mut InterpCx<'tcx, DummyMachine>,
862    dest: &PlaceTy<'tcx>,
863    place: PlaceIndex,
864    ty: Ty<'tcx>,
865    state: &State<FlatSet<Scalar>>,
866    map: &Map<'tcx>,
867) -> InterpResult<'tcx> {
868    let layout = ecx.layout_of(ty)?;
869
870    // Fast path for ZSTs.
871    if layout.is_zst() {
872        return interp_ok(());
873    }
874
875    // Fast path for scalars.
876    if layout.backend_repr.is_scalar()
877        && let Some(value) = propagatable_scalar(place, state, map)
878    {
879        return ecx.write_immediate(Immediate::Scalar(value), dest);
880    }
881
882    match ty.kind() {
883        // ZSTs. Nothing to do.
884        ty::FnDef(..) => {}
885
886        // Those are scalars, must be handled above.
887        ty::Bool | ty::Int(_) | ty::Uint(_) | ty::Float(_) | ty::Char =>
888            throw_machine_stop_str!("primitive type with provenance"),
889
890        ty::Tuple(elem_tys) => {
891            for (i, elem) in elem_tys.iter().enumerate() {
892                let i = FieldIdx::from_usize(i);
893                let Some(field) = map.apply(place, TrackElem::Field(i)) else {
894                    throw_machine_stop_str!("missing field in tuple")
895                };
896                let field_dest = ecx.project_field(dest, i)?;
897                try_write_constant(ecx, &field_dest, field, elem, state, map)?;
898            }
899        }
900
901        ty::Adt(def, args) => {
902            if def.is_union() {
903                throw_machine_stop_str!("cannot propagate unions")
904            }
905
906            let (variant_idx, variant_def, variant_place, variant_dest) = if def.is_enum() {
907                let Some(discr) = map.apply(place, TrackElem::Discriminant) else {
908                    throw_machine_stop_str!("missing discriminant for enum")
909                };
910                let FlatSet::Elem(Scalar::Int(discr)) = state.get_idx(discr, map) else {
911                    throw_machine_stop_str!("discriminant with provenance")
912                };
913                let discr_bits = discr.to_bits(discr.size());
914                let Some((variant, _)) = def.discriminants(*ecx.tcx).find(|(_, var)| discr_bits == var.val) else {
915                    throw_machine_stop_str!("illegal discriminant for enum")
916                };
917                let Some(variant_place) = map.apply(place, TrackElem::Variant(variant)) else {
918                    throw_machine_stop_str!("missing variant for enum")
919                };
920                let variant_dest = ecx.project_downcast(dest, variant)?;
921                (variant, def.variant(variant), variant_place, variant_dest)
922            } else {
923                (FIRST_VARIANT, def.non_enum_variant(), place, dest.clone())
924            };
925
926            for (i, field) in variant_def.fields.iter_enumerated() {
927                let ty = field.ty(*ecx.tcx, args);
928                let Some(field) = map.apply(variant_place, TrackElem::Field(i)) else {
929                    throw_machine_stop_str!("missing field in ADT")
930                };
931                let field_dest = ecx.project_field(&variant_dest, i)?;
932                try_write_constant(ecx, &field_dest, field, ty, state, map)?;
933            }
934            ecx.write_discriminant(variant_idx, dest)?;
935        }
936
937        // Unsupported for now.
938        ty::Array(_, _)
939        | ty::Pat(_, _)
940
941        // Do not attempt to support indirection in constants.
942        | ty::Ref(..) | ty::RawPtr(..) | ty::FnPtr(..) | ty::Str | ty::Slice(_)
943
944        | ty::Never
945        | ty::Foreign(..)
946        | ty::Alias(..)
947        | ty::Param(_)
948        | ty::Bound(..)
949        | ty::Placeholder(..)
950        | ty::Closure(..)
951        | ty::CoroutineClosure(..)
952        | ty::Coroutine(..)
953        | ty::Dynamic(..)
954        | ty::UnsafeBinder(_) => throw_machine_stop_str!("unsupported type"),
955
956        ty::Error(_) | ty::Infer(..) | ty::CoroutineWitness(..) => bug!(),
957    }
958
959    interp_ok(())
960}
961
962impl<'tcx> ResultsVisitor<'tcx, ConstAnalysis<'_, 'tcx>> for Collector<'_, 'tcx> {
963    #[instrument(level = "trace", skip(self, analysis, statement))]
964    fn visit_after_early_statement_effect(
965        &mut self,
966        analysis: &mut ConstAnalysis<'_, 'tcx>,
967        state: &State<FlatSet<Scalar>>,
968        statement: &Statement<'tcx>,
969        location: Location,
970    ) {
971        match &statement.kind {
972            StatementKind::Assign(box (_, rvalue)) => {
973                OperandCollector {
974                    state,
975                    visitor: self,
976                    ecx: &mut analysis.ecx,
977                    map: &analysis.map,
978                }
979                .visit_rvalue(rvalue, location);
980            }
981            _ => (),
982        }
983    }
984
985    #[instrument(level = "trace", skip(self, analysis, statement))]
986    fn visit_after_primary_statement_effect(
987        &mut self,
988        analysis: &mut ConstAnalysis<'_, 'tcx>,
989        state: &State<FlatSet<Scalar>>,
990        statement: &Statement<'tcx>,
991        location: Location,
992    ) {
993        match statement.kind {
994            StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(_)))) => {
995                // Don't overwrite the assignment if it already uses a constant (to keep the span).
996            }
997            StatementKind::Assign(box (place, _)) => {
998                if let Some(value) =
999                    self.try_make_constant(&mut analysis.ecx, place, state, &analysis.map)
1000                {
1001                    self.patch.assignments.insert(location, value);
1002                }
1003            }
1004            _ => (),
1005        }
1006    }
1007
1008    fn visit_after_early_terminator_effect(
1009        &mut self,
1010        analysis: &mut ConstAnalysis<'_, 'tcx>,
1011        state: &State<FlatSet<Scalar>>,
1012        terminator: &Terminator<'tcx>,
1013        location: Location,
1014    ) {
1015        OperandCollector { state, visitor: self, ecx: &mut analysis.ecx, map: &analysis.map }
1016            .visit_terminator(terminator, location);
1017    }
1018}
1019
1020impl<'tcx> MutVisitor<'tcx> for Patch<'tcx> {
1021    fn tcx(&self) -> TyCtxt<'tcx> {
1022        self.tcx
1023    }
1024
1025    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
1026        if let Some(value) = self.assignments.get(&location) {
1027            match &mut statement.kind {
1028                StatementKind::Assign(box (_, rvalue)) => {
1029                    *rvalue = Rvalue::Use(self.make_operand(*value));
1030                }
1031                _ => bug!("found assignment info for non-assign statement"),
1032            }
1033        } else {
1034            self.super_statement(statement, location);
1035        }
1036    }
1037
1038    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
1039        match operand {
1040            Operand::Copy(place) | Operand::Move(place) => {
1041                if let Some(value) = self.before_effect.get(&(location, *place)) {
1042                    *operand = self.make_operand(*value);
1043                } else if !place.projection.is_empty() {
1044                    self.super_operand(operand, location)
1045                }
1046            }
1047            Operand::Constant(_) => {}
1048        }
1049    }
1050
1051    fn process_projection_elem(
1052        &mut self,
1053        elem: PlaceElem<'tcx>,
1054        location: Location,
1055    ) -> Option<PlaceElem<'tcx>> {
1056        if let PlaceElem::Index(local) = elem {
1057            let offset = self.before_effect.get(&(location, local.into()))?;
1058            let offset = offset.try_to_scalar()?;
1059            let offset = offset.to_target_usize(&self.tcx).discard_err()?;
1060            let min_length = offset.checked_add(1)?;
1061            Some(PlaceElem::ConstantIndex { offset, min_length, from_end: false })
1062        } else {
1063            None
1064        }
1065    }
1066}
1067
1068struct OperandCollector<'a, 'b, 'tcx> {
1069    state: &'a State<FlatSet<Scalar>>,
1070    visitor: &'a mut Collector<'b, 'tcx>,
1071    ecx: &'a mut InterpCx<'tcx, DummyMachine>,
1072    map: &'a Map<'tcx>,
1073}
1074
1075impl<'tcx> Visitor<'tcx> for OperandCollector<'_, '_, 'tcx> {
1076    fn visit_projection_elem(
1077        &mut self,
1078        _: PlaceRef<'tcx>,
1079        elem: PlaceElem<'tcx>,
1080        _: PlaceContext,
1081        location: Location,
1082    ) {
1083        if let PlaceElem::Index(local) = elem
1084            && let Some(value) =
1085                self.visitor.try_make_constant(self.ecx, local.into(), self.state, self.map)
1086        {
1087            self.visitor.patch.before_effect.insert((location, local.into()), value);
1088        }
1089    }
1090
1091    fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
1092        if let Some(place) = operand.place() {
1093            if let Some(value) =
1094                self.visitor.try_make_constant(self.ecx, place, self.state, self.map)
1095            {
1096                self.visitor.patch.before_effect.insert((location, place), value);
1097            } else if !place.projection.is_empty() {
1098                // Try to propagate into `Index` projections.
1099                self.super_operand(operand, location)
1100            }
1101        }
1102    }
1103}