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 adt = bval_ty.as_adt().expect("Non-adt slice sub-constant");
76            assert_eq!(adt.id, TypeId::Builtin(BuiltinTy::Array));
77            let len = adt.generics.const_generics[0].clone();
78
79            let array_place = ctx.rval_to_place(Rvalue::Use(bval), bval_ty.clone());
80
81            let (_, _, rk) = val.ty.as_ref().unwrap();
82            let array_ref = ctx.borrow_to_new_var(array_place, rk.clone().into(), None);
83            Rvalue::UnaryOp(
84                UnOp::Cast(CastKind::Unsize(
85                    array_ref.ty.clone(),
86                    val.ty.clone(),
87                    UnsizingMetadata::Length(len),
88                )),
89                Operand::Move(array_ref),
90            )
91        }
92        ConstantExprKind::Ref(bval) => {
93            let place = match bval.kind {
94                ConstantExprKind::Global(global_ref) => Place::new_global(global_ref, bval.ty),
95                _ => {
96                    // Recurse on the borrowed value
97                    let bval = transform_constant_expr(ctx, bval);
98
99                    // Evaluate the referenced value
100                    let bval_ty = bval.ty().clone();
101                    ctx.rval_to_place(Rvalue::Use(bval), bval_ty)
102                }
103            };
104            // Borrow the place.
105            ctx.borrow(place, BorrowKind::Shared)
106        }
107        ConstantExprKind::Ptr(rk, bval) => {
108            // As the value is originally an argument, it must be Sized, hence no metadata
109            let place = match bval.kind {
110                ConstantExprKind::Global(global_ref) => Place::new_global(global_ref, bval.ty),
111                _ => {
112                    // Recurse on the borrowed value
113                    let bval = transform_constant_expr(ctx, bval);
114
115                    // Evaluate the referenced value
116                    let bval_ty = bval.ty().clone();
117                    ctx.rval_to_place(Rvalue::Use(bval), bval_ty)
118                }
119            };
120            // Borrow the value
121            ctx.raw_borrow(place, rk)
122        }
123        ConstantExprKind::Adt(variant, fields) => {
124            let fields = fields
125                .into_iter()
126                .map(|x| transform_constant_expr(ctx, Box::new(x)))
127                .collect();
128
129            // Build an `Aggregate` rvalue.
130            let tref = val.ty.kind().as_adt().unwrap();
131            let aggregate_kind = AggregateKind::Adt(tref.clone(), variant, None);
132            Rvalue::Aggregate(aggregate_kind, fields)
133        }
134        ConstantExprKind::Array(fields) | ConstantExprKind::Slice(fields) => {
135            let fields = fields
136                .into_iter()
137                .map(|x| transform_constant_expr(ctx, Box::new(x)))
138                .collect_vec();
139
140            let len = ConstGeneric::Value(Literal::Scalar(ScalarValue::Unsigned(
141                UIntTy::Usize,
142                fields.len() as u128,
143            )));
144            let mut tref = val.ty.kind().as_adt().unwrap().clone();
145            let ty = tref.generics.types.pop().unwrap();
146            match tref.id.as_builtin().unwrap() {
147                BuiltinTy::Array => {}
148                BuiltinTy::Slice => val_ty = Ty::mk_array(ty.clone(), len.clone()),
149                _ => {
150                    unreachable!("Unpexected builtin type in array/slice constant")
151                }
152            }
153            Rvalue::Aggregate(AggregateKind::Array(ty, len), fields)
154        }
155        ConstantExprKind::FnPtr(fptr) => {
156            let TyKind::FnPtr(sig) = val.ty.kind() else {
157                unreachable!("FnPtr constant must have FnPtr type");
158            };
159            let from_ty =
160                TyKind::FnDef(sig.clone().map(|_| fptr.clone().move_under_binder())).into_ty();
161            let to_ty = TyKind::FnPtr(sig.clone()).into_ty();
162
163            Rvalue::UnaryOp(
164                UnOp::Cast(CastKind::FnPtr(from_ty.clone(), to_ty)),
165                Operand::Const(Box::new(ConstantExpr {
166                    kind: ConstantExprKind::FnDef(fptr),
167                    ty: from_ty,
168                })),
169            )
170        }
171    };
172    Operand::Move(ctx.rval_to_place(rval, val_ty))
173}
174
175fn transform_operand(ctx: &mut UllbcStatementTransformCtx<'_>, op: &mut Operand) {
176    // Transform the constant operands (otherwise do nothing)
177    take_mut::take(op, |op| {
178        if let Operand::Const(val) = op {
179            transform_constant_expr(ctx, val)
180        } else {
181            op
182        }
183    })
184}
185
186pub struct Transform;
187impl UllbcPass for Transform {
188    fn should_run(&self, options: &crate::options::TranslateOptions) -> bool {
189        !options.raw_consts
190    }
191
192    fn transform_function(&self, ctx: &mut TransformCtx, fun_decl: &mut FunDecl) {
193        fun_decl.transform_ullbc_operands(ctx, transform_operand);
194        if let Some(body) = fun_decl.body.as_unstructured_mut() {
195            for block in body.body.iter_mut() {
196                // Simplify array with repeated constants into array repeats.
197                block.dyn_visit_in_body_mut(|rvalue: &mut Rvalue| {
198                    take_mut::take(rvalue, |rvalue| match rvalue {
199                        Rvalue::Aggregate(AggregateKind::Array(ty, len), ref fields)
200                            if fields.len() >= 2
201                                && fields.iter().all(|x| x.is_const())
202                                && let Ok(op) = fields.iter().dedup().exactly_one() =>
203                        {
204                            Rvalue::Repeat(op.clone(), ty.clone(), len)
205                        }
206                        _ => rvalue,
207                    });
208                });
209            }
210        }
211    }
212}