Skip to main content

charon_lib/transform/resugar/
reconstruct_copy_for_deref.rs

1//! Reconstruct places that rustc split with `CopyForDeref` in elaborated MIR.
2//!
3//! A statement like `*tuple.1 = x` can become `_tmp = noretag copy tuple.1; *_tmp = x`.
4//! The copy only reads the pointer; we put the dereference back on `tuple.1`. This also
5//! applies to shared references and boxes, and to reads of pointer metadata.
6use std::ops::ControlFlow;
7
8use crate::ids::IndexVec;
9use crate::transform::TransformCtx;
10use crate::transform::ctx::UllbcPass;
11use crate::ullbc_ast::*;
12
13#[derive(Default)]
14struct CopyInfo {
15    assignment: Option<(StmtLoc, Place)>,
16    use_loc: Option<StmtLoc>,
17    /// Whether this local is used in any way that doesn't allow the merging we want.
18    invalid: bool,
19}
20
21#[derive(Visitor)]
22struct UseCollector<'a> {
23    body: &'a ExprBody,
24    copies: &'a mut IndexVec<LocalId, CopyInfo>,
25    has_candidate: bool,
26    current_st: Option<StmtLoc>,
27}
28
29impl VisitBody for UseCollector<'_> {
30    fn visit_place(&mut self, place: &Place) -> ControlFlow<Self::Break> {
31        if let Some((subplace, ProjectionElem::Deref | ProjectionElem::PtrMetadata)) =
32            place.as_projection()
33            && let Some(local) = subplace.as_local()
34        {
35            let info = &mut self.copies[local];
36            if let Some(loc) = self.current_st
37                && (info.use_loc.is_none() || info.use_loc == self.current_st)
38            {
39                info.use_loc = self.current_st;
40                // We conservatively skip replacing into a statement that has by-move operands, to
41                // avoid cases like:
42                //   tmp = noretag copy pair._0
43                //   out = (move pair._0, copy (*tmp))
44                self.body[loc].dyn_visit_in_body(|operand: &Operand| {
45                    info.invalid |= matches!(operand, Operand::Move(_));
46                });
47            } else {
48                info.invalid = true;
49            }
50            // Don't visit the local.
51            return ControlFlow::Continue(());
52        }
53
54        if let Some(local) = place.as_local() {
55            // If we get to here then this is an invalid use.
56            self.copies[local].invalid = true;
57        }
58        self.visit_inner(place)
59    }
60
61    fn visit_ullbc_statement(&mut self, st: &Statement) -> ControlFlow<Self::Break> {
62        if let StatementKind::Assign(dest, rvalue) = &st.kind
63            && let Some(local) = dest.as_local()
64            && let Some(loc) = self.current_st
65        {
66            let info = &mut self.copies[local];
67            if info.assignment.is_some() {
68                info.invalid = true;
69            } else if let Rvalue::Use(Operand::Copy(source), WithRetag::No) = rvalue
70                && (matches!(source.ty().kind(), TyKind::Ref(..)) || source.ty().is_box())
71                && !info.invalid
72            {
73                info.assignment = Some((loc, source.clone()));
74                self.has_candidate = true;
75            } else {
76                info.invalid = true;
77            }
78            self.visit(rvalue)
79        } else {
80            self.visit_inner(st)
81        }
82    }
83}
84
85pub struct Transform;
86
87impl UllbcPass for Transform {
88    fn transform_body(&self, _ctx: &mut TransformCtx, body: &mut ExprBody) {
89        let mut copies: IndexVec<LocalId, CopyInfo> =
90            body.locals.locals.map_ref(|_| CopyInfo::default());
91
92        // Look for locals that are assigned exactly once with a non-retagging copy. A local is
93        // eligible only if every use is beneath a dereference or a metadata projection.
94        let mut collector = UseCollector {
95            body,
96            copies: &mut copies,
97            has_candidate: false,
98            current_st: None,
99        };
100        for (block, data) in body.body.iter_enumerated() {
101            for (statement, st) in data.statements.iter().enumerate() {
102                collector.current_st = Some(StmtLoc::new(block, statement));
103                collector.visit(st);
104            }
105            collector.current_st = None;
106            collector.visit(&data.terminator);
107        }
108        if !collector.has_candidate {
109            return;
110        }
111
112        // Inlining the source is safe only if nothing can change it between the
113        // pointer read and its use. Nops do not affect that ordering.
114        let uninterrupted = |body: &ExprBody, before: StmtLoc, after: StmtLoc| {
115            before.block == after.block
116                && before.statement < after.statement
117                && body.body[before.block].statements[before.statement + 1..after.statement]
118                    .iter()
119                    .all(|st| st.kind.is_nop())
120        };
121        let mut replacements: IndexMap<LocalId, Place> = IndexMap::new();
122        for (local, info) in copies.iter_enumerated() {
123            if !info.invalid
124                && let Some((assignment, source)) = &info.assignment
125                && let Some(use_loc) = info.use_loc
126                && uninterrupted(body, *assignment, use_loc)
127            {
128                replacements.insert(local, source.clone());
129                body[*assignment].kind = StatementKind::Nop;
130            }
131        }
132        if replacements.is_empty() {
133            return;
134        }
135
136        body.body.dyn_visit_in_body_mut(|place: &mut Place| {
137            // Keep the projection and replace only its base; the visitor reaches
138            // inner projections on its own.
139            if let PlaceKind::Projection(
140                subplace,
141                ProjectionElem::Deref | ProjectionElem::PtrMetadata,
142            ) = &mut place.kind
143                && let Some(local) = subplace.as_local()
144                && let Some(source) = replacements.get(local)
145            {
146                **subplace = source.clone();
147            }
148        });
149    }
150}