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