Skip to main content

charon_lib/transform/resugar/
reconstruct_box_derefs.rs

1//! Resugar the box derefs that got desugared in elaborated MIR.
2//!
3//! In elaborated MIR, box derefs become actual accesses to the contained raw pointer. Under
4//! `--treat-box-as-builtin`, we convert these back to derefs on the box.
5use std::ops::ControlFlow;
6
7use derive_generic_visitor::*;
8
9use crate::ast::ullbc_ast_utils::StmtLoc;
10use crate::ast::*;
11use crate::ids::IndexVec;
12use crate::transform::TransformCtx;
13use crate::transform::ctx::UllbcPass;
14use crate::ullbc_ast::{ExprBody, StatementKind};
15
16pub struct Transform;
17
18/// Look for
19/// ```ignore
20/// transmute::<NonNull<T>, *const T>(copy (*b).0)
21/// ```
22/// Returns `*b`
23fn box_pointee_pointer_assignment(rvalue: &Rvalue) -> Option<Place> {
24    let Rvalue::UnaryOp(
25        UnOp::Cast(CastKind::Transmute(_, raw_ptr_ty)),
26        Operand::Copy(hidden_pointer),
27    ) = &rvalue
28    else {
29        return None;
30    };
31    let (field_base, ProjectionElem::Field(_, FieldId::ZERO)) = hidden_pointer.as_projection()?
32    else {
33        return None;
34    };
35    let (box_place, ProjectionElem::Deref) = field_base.as_projection()? else {
36        return None;
37    };
38    let TyKind::Adt(TypeDeclRef {
39        id: TypeId::Builtin(BuiltinTy::Box),
40        generics: box_generics,
41    }) = box_place.ty().kind()
42    else {
43        return None;
44    };
45    if &box_generics.types[0] != raw_ptr_ty.as_raw_ptr()?.0 {
46        return None;
47    }
48    Some(box_place.clone().deref())
49}
50
51#[derive(Default)]
52struct LocalStatus {
53    /// Whether that local is the target of the special `transmute` statement we're looking for.
54    /// Stores the `*b` where `b` is the Box.
55    box_pointee_assignment: Option<(StmtLoc, Place)>,
56    /// Whether that local is the target of more than one such special `transmute`.
57    ambiguous: bool,
58    /// Whether that local is ever used outside of a deref projection (apart from the initial
59    /// assignment).
60    used_outside_derefs: bool,
61}
62
63impl LocalStatus {
64    fn rewritable_box_pointee(&self) -> Option<&(StmtLoc, Place)> {
65        if self.ambiguous || self.used_outside_derefs {
66            return None;
67        }
68        self.box_pointee_assignment.as_ref()
69    }
70}
71
72struct LocalStatusCollector {
73    local_status: IndexVec<LocalId, LocalStatus>,
74    current_statement: Option<StmtLoc>,
75}
76
77impl Visitor for LocalStatusCollector {
78    type Break = std::convert::Infallible;
79}
80
81impl VisitBody for LocalStatusCollector {
82    fn visit_place(&mut self, place: &Place) -> ControlFlow<Self::Break> {
83        match &place.kind {
84            // Skip uses of a local where the last projection is a deref.
85            PlaceKind::Projection(subplace, pj) if subplace.is_local() && pj.is_deref() => {
86                return ControlFlow::Continue(());
87            }
88            // This didn't get caught by the branch above, so it's an invalid use.
89            PlaceKind::Local(local) => self.local_status[*local].used_outside_derefs = true,
90            _ => {}
91        }
92        self.visit_inner(place)
93    }
94
95    fn visit_ullbc_statement(&mut self, st: &ullbc_ast::Statement) -> ControlFlow<Self::Break> {
96        if let StatementKind::Assign(dst, rvalue) = &st.kind
97            && let Some(local) = dst.as_local()
98            && let Some(box_pointee) = box_pointee_pointer_assignment(rvalue)
99        {
100            let status = &mut self.local_status[local];
101            if status.box_pointee_assignment.is_some() {
102                status.ambiguous = true;
103            } else {
104                let loc = self.current_statement.unwrap();
105                status.box_pointee_assignment = Some((loc, box_pointee));
106            }
107            // Ignore the assignment destination: this initialization is the one non-deref
108            // interaction allowed for the temporary pointer.
109            self.visit(rvalue)
110        } else {
111            self.visit_inner(st)
112        }
113    }
114}
115
116fn collect_local_status(body: &ExprBody) -> IndexVec<LocalId, LocalStatus> {
117    let mut collector = LocalStatusCollector {
118        local_status: body.locals.locals.map_ref(|_| Default::default()),
119        current_statement: None,
120    };
121    for (block, block_data) in body.body.iter_enumerated() {
122        for (statement, st) in block_data.statements.iter().enumerate() {
123            collector.current_statement = Some(StmtLoc::new(block, statement));
124            let _ = collector.visit(st);
125        }
126        collector.current_statement = None;
127        let _ = collector.visit(&block_data.terminator);
128    }
129    collector.local_status
130}
131
132fn rewrite_place(place: &mut Place, local_status: &IndexVec<LocalId, LocalStatus>) {
133    let Some((subplace, projection)) = place.as_projection() else {
134        return;
135    };
136    if projection == &ProjectionElem::Deref
137        && let Some(local) = subplace.as_local()
138        && let Some((_, box_pointee)) = local_status[local].rewritable_box_pointee()
139    {
140        *place = box_pointee.clone();
141        return;
142    }
143
144    let PlaceKind::Projection(subplace, _) = &mut place.kind else {
145        unreachable!();
146    };
147    rewrite_place(subplace, local_status);
148}
149
150impl UllbcPass for Transform {
151    fn should_run(&self, options: &crate::options::TranslateOptions) -> bool {
152        options.treat_box_as_builtin
153    }
154
155    fn transform_body(&self, _ctx: &mut TransformCtx, body: &mut ExprBody) {
156        let local_status = collect_local_status(body);
157
158        body.body.dyn_visit_in_body_mut(|place: &mut Place| {
159            rewrite_place(place, &local_status);
160        });
161
162        for status in local_status {
163            if let Some((assign_loc, _)) = status.rewritable_box_pointee() {
164                body[*assign_loc].kind = StatementKind::Nop;
165            }
166        }
167    }
168}