charon_lib/transform/simplify_output/
simplify_constants.rs

1//! The MIR constant expressions lead to a lot of duplication: there are
2//! for instance constant ADTs which duplicate the "regular" aggregated
3//! ADTs in the operands, constant references, etc. This reduces the number
4//! of cases to handle and eases the function translation in Aeneas.
5//!
6//! This pass removes all those occurrences so that only the
7//! `ConstantExpression::Literal`. It does so by introducing intermediate statements.
8//!
9//! A small remark about the intermediate statements we introduce for the globals:
10//! we do so because, when evaluating the code in "concrete" mode, it allows to
11//! handle the globals like function calls.
12
13use itertools::Itertools;
14
15use crate::transform::TransformCtx;
16use crate::transform::ctx::{BodyTransformCtx, UllbcPass, UllbcStatementTransformCtx};
17use crate::ullbc_ast::*;
18
19/// If the constant value is a constant ADT, push `Assign::Aggregate` statements
20/// to the vector of statements, that bind new variables to the ADT parts and
21/// the variable assigned to the complete ADT.
22///
23/// Goes fom e.g. `f(T::A(x, y))` to `let a = T::A(x, y); f(a)`.
24/// The function is recursively called on the aggregate fields (e.g. here x and y).
25fn transform_constant_expr(
26    ctx: &mut UllbcStatementTransformCtx<'_>,
27    val: Box<ConstantExpr>,
28) -> Operand {
29    let mut val_ty = val.ty.clone();
30    let rval = match val.kind {
31        ConstantExprKind::Literal(_)
32        | ConstantExprKind::Var(_)
33        | ConstantExprKind::RawMemory(..)
34        | ConstantExprKind::TraitConst(..)
35        | ConstantExprKind::FnDef(..)
36        | ConstantExprKind::Opaque(_) => {
37            // Nothing to do
38            // TODO: for trait const: might come from a top-level impl, so we might
39            // want to introduce an intermediate statement to be able to evaluate
40            // it as a function call, like for globals.
41            return Operand::Const(val);
42        }
43        // Here we use a copy, rather than a move -- moving a global would leave it uninitialized,
44        // which would e.g. make the following code fail:
45        //     const GLOBAL: usize = 0;
46        //     let x = GLOBAL;
47        //     let y = GLOBAL; // if moving, at this point GLOBAL would be uninitialized
48        ConstantExprKind::Global(global_ref) => {
49            return Operand::Copy(Place::new_global(global_ref, val.ty));
50        }
51        ConstantExprKind::PtrNoProvenance(ptr) => {
52            let usize_ty = TyKind::Literal(LiteralTy::UInt(UIntTy::Usize)).into_ty();
53            let ptr_usize = ConstantExprKind::Literal(Literal::Scalar(ScalarValue::Unsigned(
54                UIntTy::Usize,
55                ptr,
56            )));
57            let cast = UnOp::Cast(CastKind::RawPtr(usize_ty.clone(), val.ty.clone()));
58            Rvalue::UnaryOp(
59                cast,
60                Operand::Const(Box::new(ConstantExpr {
61                    kind: ptr_usize,
62                    ty: usize_ty,
63                })),
64            )
65        }
66        ConstantExprKind::Ref(bval) if bval.ty.is_slice() => {
67            // We need to special-case slices; we don't take a reference to the slice,
68            // but an unsizing cast
69            // We generate the following code:
70            // let array: [T; N] = [...];
71            // let array_ref: &[T; N] = &array;
72            // return unsize::<&[T; N], &[T]>(array_ref);
73            let bval = transform_constant_expr(ctx, bval);
74            let bval_ty = bval.ty().clone();
75            let (_, len) = bval_ty.as_array().expect("Non-adt slice sub-constant");
76
77            let array_place = ctx.rval_to_place(Rvalue::Use(bval), bval_ty.clone());
78
79            let (_, _, rk) = val.ty.as_ref().unwrap();
80            let array_ref = ctx.borrow_to_new_var(array_place, rk.clone().into(), None);
81            Rvalue::UnaryOp(
82                UnOp::Cast(CastKind::Unsize(
83                    array_ref.ty.clone(),
84                    val.ty.clone(),
85                    UnsizingMetadata::Length(len.clone()),
86                )),
87                Operand::Move(array_ref),
88            )
89        }
90        ConstantExprKind::Ref(bval) => {
91            let place = match bval.kind {
92                ConstantExprKind::Global(global_ref) => Place::new_global(global_ref, bval.ty),
93                _ => {
94                    // Recurse on the borrowed value
95                    let bval = transform_constant_expr(ctx, bval);
96
97                    // Evaluate the referenced value
98                    let bval_ty = bval.ty().clone();
99                    ctx.rval_to_place(Rvalue::Use(bval), bval_ty)
100                }
101            };
102            // Borrow the place.
103            ctx.borrow(place, BorrowKind::Shared)
104        }
105        ConstantExprKind::Ptr(rk, bval) => {
106            // As the value is originally an argument, it must be Sized, hence no metadata
107            let place = match bval.kind {
108                ConstantExprKind::Global(global_ref) => Place::new_global(global_ref, bval.ty),
109                _ => {
110                    // Recurse on the borrowed value
111                    let bval = transform_constant_expr(ctx, bval);
112
113                    // Evaluate the referenced value
114                    let bval_ty = bval.ty().clone();
115                    ctx.rval_to_place(Rvalue::Use(bval), bval_ty)
116                }
117            };
118            // Borrow the value
119            ctx.raw_borrow(place, rk)
120        }
121        ConstantExprKind::Adt(variant, fields) => {
122            let fields = fields
123                .into_iter()
124                .map(|x| transform_constant_expr(ctx, Box::new(x)))
125                .collect();
126
127            // Build an `Aggregate` rvalue.
128            let tref = val.ty.kind().as_adt().unwrap();
129            let aggregate_kind = AggregateKind::Adt(tref.clone(), variant, None);
130            Rvalue::Aggregate(aggregate_kind, fields)
131        }
132        ConstantExprKind::Array(fields) | ConstantExprKind::Slice(fields) => {
133            let fields = fields
134                .into_iter()
135                .map(|x| transform_constant_expr(ctx, Box::new(x)))
136                .collect_vec();
137
138            let len = ConstGeneric::Value(Literal::Scalar(ScalarValue::Unsigned(
139                UIntTy::Usize,
140                fields.len() as u128,
141            )));
142            let ty = match val.ty.kind() {
143                TyKind::Array(ty, _) => ty.clone(),
144                TyKind::Slice(ty) => {
145                    val_ty = Ty::mk_array(ty.clone(), len.clone());
146                    ty.clone()
147                }
148                _ => unreachable!("Unexpected type in array/slice constant"),
149            };
150            Rvalue::Aggregate(AggregateKind::Array(ty, len), fields)
151        }
152        ConstantExprKind::FnPtr(fptr) => {
153            let TyKind::FnPtr(sig) = val.ty.kind() else {
154                unreachable!("FnPtr constant must have FnPtr type");
155            };
156            let from_ty =
157                TyKind::FnDef(sig.clone().map(|_| fptr.clone().move_under_binder())).into_ty();
158            let to_ty = TyKind::FnPtr(sig.clone()).into_ty();
159
160            Rvalue::UnaryOp(
161                UnOp::Cast(CastKind::FnPtr(from_ty.clone(), to_ty)),
162                Operand::Const(Box::new(ConstantExpr {
163                    kind: ConstantExprKind::FnDef(fptr),
164                    ty: from_ty,
165                })),
166            )
167        }
168    };
169    Operand::Move(ctx.rval_to_place(rval, val_ty))
170}
171
172fn transform_operand(ctx: &mut UllbcStatementTransformCtx<'_>, op: &mut Operand) {
173    // Transform the constant operands (otherwise do nothing)
174    take_mut::take(op, |op| {
175        if let Operand::Const(val) = op {
176            transform_constant_expr(ctx, val)
177        } else {
178            op
179        }
180    })
181}
182
183pub struct Transform;
184impl UllbcPass for Transform {
185    fn should_run(&self, options: &crate::options::TranslateOptions) -> bool {
186        !options.raw_consts
187    }
188
189    fn transform_function(&self, ctx: &mut TransformCtx, fun_decl: &mut FunDecl) {
190        fun_decl.transform_ullbc_operands(ctx, transform_operand);
191        if let Some(body) = fun_decl.body.as_unstructured_mut() {
192            for block in body.body.iter_mut() {
193                // Simplify array with repeated constants into array repeats.
194                block.dyn_visit_in_body_mut(|rvalue: &mut Rvalue| {
195                    take_mut::take(rvalue, |rvalue| match rvalue {
196                        Rvalue::Aggregate(AggregateKind::Array(ty, len), ref fields)
197                            if fields.len() >= 2
198                                && fields.iter().all(|x| x.is_const())
199                                && let Ok(op) = fields.iter().dedup().exactly_one() =>
200                        {
201                            Rvalue::Repeat(op.clone(), ty.clone(), len)
202                        }
203                        _ => rvalue,
204                    });
205                });
206            }
207        }
208    }
209}