hax_rust_engine/ast/
utils.rs

1//! This module provides a collection of utilities to work on AST.
2
3use super::visitors::*;
4use super::*;
5use identifiers::*;
6use std::collections::HashMap;
7
8/// Useful visitor to map AST fragments.
9pub mod mappers {
10    use super::*;
11
12    /// Visitor that substitutes local identifiers in ASTs.
13    pub struct SubstLocalIds(HashMap<LocalId, LocalId>);
14
15    impl SubstLocalIds {
16        /// Create a substituer given one replacement couple.
17        pub fn one(from: LocalId, to: LocalId) -> Self {
18            Self::many([(from, to)])
19        }
20        /// Create a substituer given a bunch of replacement couples.
21        pub fn many(replacements: impl IntoIterator<Item = (LocalId, LocalId)>) -> Self {
22            Self(replacements.into_iter().collect())
23        }
24    }
25
26    impl AstVisitorMut for SubstLocalIds {
27        fn visit_local_id(&mut self, local_id: &mut LocalId) {
28            if let Some(replacement) = self.0.get(local_id) {
29                *local_id = replacement.clone();
30            }
31        }
32    }
33}
34
35impl Expr {
36    /// Create a tuple expression out of components.
37    pub fn tuple(components: Vec<Expr>, span: Span) -> Self {
38        let ty = TyKind::tuple(
39            components
40                .iter()
41                .map(Typed::ty)
42                .cloned()
43                .map(GenericValue::Ty)
44                .collect(),
45        )
46        .promote();
47        ExprKind::tuple(components).promote(ty, span)
48    }
49
50    /// Create a unit (tuple of size 0) expression.
51    pub fn unit(span: Span) -> Self {
52        ExprKind::GlobalId(global_id::TupleId::Constructor { length: 0 }.into())
53            .promote(TyKind::unit().promote(), span)
54    }
55
56    /// Creates a `App` node for a standalone function.
57    pub fn standalone_fn_app(
58        head: impl Into<FnAppHead>,
59        generic_args: Vec<GenericValue>,
60        args: Vec<Expr>,
61        output_type: Ty,
62        span: Span,
63    ) -> Self {
64        ExprKind::standalone_fn_app(head, generic_args, args, output_type.clone(), span)
65            .promote(output_type, span)
66    }
67
68    /// Creates a `App` node.
69    pub fn fn_app(
70        head: impl Into<FnAppHead>,
71        generic_args: Vec<GenericValue>,
72        args: Vec<Expr>,
73        output_type: Ty,
74        bounds_impls: Vec<ImplExpr>,
75        trait_: Option<(ImplExpr, Vec<GenericValue>)>,
76        span: Span,
77    ) -> Self {
78        ExprKind::fn_app(
79            head,
80            generic_args,
81            args,
82            output_type.clone(),
83            bounds_impls,
84            trait_,
85            span,
86        )
87        .promote(output_type, span)
88    }
89
90    /// Removes a box
91    pub fn unbox_once(&self) -> Option<&Expr> {
92        if let ExprKind::App { head, args, .. } = self.kind()
93            && let [arg] = &**args
94            && let ExprKind::GlobalId(head) = head.kind()
95            && let crate::names::alloc::boxed::Impl::new
96            | crate::names::rust_primitives::hax::box_new = *head
97        {
98            Some(arg)
99        } else {
100            None
101        }
102    }
103
104    /// Removes a deref
105    pub fn underef_once(&self) -> Option<&Expr> {
106        if let ExprKind::App { head, args, .. } = self.kind()
107            && let [arg] = &**args
108            && let ExprKind::GlobalId(head) = head.kind()
109            && let crate::names::rust_primitives::hax::deref_op = *head
110        {
111            Some(arg)
112        } else {
113            None
114        }
115    }
116
117    /// Removes all boxes and derefs wrapping the expression
118    pub fn unbox_underef(&self) -> &Expr {
119        let mut current = self;
120        while let Some(e) = current.unbox_once().or_else(|| current.underef_once()) {
121            current = e
122        }
123        current
124    }
125}
126
127impl ExprKind {
128    /// Creates a `App` node for a standalone function.
129    pub fn standalone_fn_app(
130        head: impl Into<FnAppHead>,
131        generic_args: Vec<GenericValue>,
132        args: Vec<Expr>,
133        output_type: Ty,
134        span: Span,
135    ) -> Self {
136        Self::fn_app(head, generic_args, args, output_type, vec![], None, span)
137    }
138
139    /// Creates a `App` node.
140    pub fn fn_app(
141        head: impl Into<FnAppHead>,
142        generic_args: Vec<GenericValue>,
143        args: Vec<Expr>,
144        output_type: Ty,
145        bounds_impls: Vec<ImplExpr>,
146        trait_: Option<(ImplExpr, Vec<GenericValue>)>,
147        span: Span,
148    ) -> Self {
149        let head = 'head: {
150            let kind = match head.into() {
151                FnAppHead::GlobalId(global_id) => ExprKind::GlobalId(global_id),
152                FnAppHead::ExprKind(expr_kind) => expr_kind,
153                FnAppHead::Expr(expr) => break 'head expr,
154            };
155            let head_ty = TyKind::Arrow {
156                inputs: args.iter().map(Typed::ty).cloned().collect(),
157                output: output_type.clone(),
158            }
159            .promote();
160            kind.promote(head_ty, span)
161        };
162
163        Self::App {
164            head,
165            args,
166            generic_args,
167            bounds_impls,
168            trait_,
169        }
170    }
171
172    /// Creates a tuple out of a vector of components.
173    pub fn tuple(components: Vec<Expr>) -> Self {
174        let length = components.len();
175        ExprKind::Construct {
176            constructor: global_id::TupleId::Constructor { length }.into(),
177            is_record: false,
178            is_struct: true,
179            fields: components
180                .into_iter()
181                .enumerate()
182                .map(|(field, expr)| (global_id::TupleId::Field { length, field }.into(), expr))
183                .collect(),
184            base: None,
185        }
186    }
187
188    /// Promote to an `Expr`
189    pub fn promote(self, ty: Ty, span: Span) -> Expr {
190        Expr {
191            kind: Box::new(self),
192            ty,
193            meta: Metadata {
194                span,
195                attributes: Vec::new(),
196            },
197        }
198    }
199}
200
201impl Metadata {
202    /// Get an iterator over hax attributes for this AST fragment.
203    pub fn hax_attributes(&self) -> impl Iterator<Item = &hax_lib_macros_types::AttrPayload> {
204        crate::attributes::hax_attributes(&self.attributes)
205    }
206}
207
208impl Pat {
209    /// Expects the pattern to be a simple binding `self`.
210    pub fn expect_self(&self) -> Option<LocalId> {
211        if let PatKind::Binding { var, .. } = self.kind()
212            && var.is_self()
213        {
214            Some(var.clone())
215        } else {
216            None
217        }
218    }
219}
220
221/// Helper enum that describes what can serve as function application heads.
222/// This is an helper that is useful for [`ExprKind::fn_application`].
223pub enum FnAppHead {
224    /// A global identifier
225    GlobalId(GlobalId),
226    /// An expression kind
227    ExprKind(ExprKind),
228    /// A full blown expression
229    Expr(Expr),
230}
231
232impl From<GlobalId> for FnAppHead {
233    fn from(value: GlobalId) -> Self {
234        Self::GlobalId(value)
235    }
236}
237impl From<ExprKind> for FnAppHead {
238    fn from(value: ExprKind) -> Self {
239        Self::ExprKind(value)
240    }
241}
242impl From<Expr> for FnAppHead {
243    fn from(value: Expr) -> Self {
244        Self::Expr(value)
245    }
246}
247
248impl Generics {
249    /// Concatenate two generics
250    pub fn concat(mut self, other: Self) -> Self {
251        self.constraints.extend(other.constraints);
252        self.params.extend(other.params);
253        use std::cmp::Ordering;
254        self.params.sort_by(|a, b| match (a.kind(), b.kind()) {
255            (GenericParamKind::Lifetime, GenericParamKind::Lifetime) => Ordering::Equal,
256            (GenericParamKind::Lifetime, _) => Ordering::Less,
257            (_, GenericParamKind::Lifetime) => Ordering::Greater,
258            _ => Ordering::Equal,
259        });
260        self
261    }
262    /// Empty generics
263    pub fn empty() -> Self {
264        Self {
265            params: Vec::new(),
266            constraints: Vec::new(),
267        }
268    }
269}
270
271impl Item {
272    /// Returns a `LocalId` named `self` if the item is a standalone function
273    /// whose first argument is the keyword `self`. In other words, this
274    /// function returns a local identifier only for associated methods from
275    /// inherent `impl` blocks.
276    pub fn self_id(&self) -> Option<LocalId> {
277        if let ItemKind::Fn { params, .. } = self.kind()
278            && let [first, ..] = &params[..]
279            && let Some(self_id) = first.pat.expect_self()
280        {
281            Some(self_id.clone())
282        } else {
283            None
284        }
285    }
286}
287
288impl ItemKind {
289    /// Promote to an item
290    pub fn promote(self, ident: GlobalId, span: Span) -> Item {
291        Item {
292            ident,
293            kind: self,
294            meta: Metadata {
295                span,
296                attributes: Vec::new(),
297            },
298        }
299    }
300}
301
302impl GenericValue {
303    /// Tries to extract a [`Ty`] out of a [`GenericValue`].
304    pub fn expect_ty(&self) -> Option<&Ty> {
305        let Self::Ty(ty) = self else { return None };
306        Some(ty)
307    }
308}
309
310impl TyKind {
311    /// Tuple type
312    pub fn tuple(args: Vec<GenericValue>) -> Self {
313        let head = global_id::TupleId::Type { length: args.len() }.into();
314        Self::App { head, args }
315    }
316    /// Unit type
317    pub fn unit() -> Self {
318        Self::tuple(Vec::new())
319    }
320    /// Promote to a Ty
321    pub fn promote(self) -> Ty {
322        Ty(Box::new(self))
323    }
324}
325
326impl Arm {
327    /// Create a non-guarded arm
328    pub fn non_guarded(pat: Pat, body: Expr, span: Span) -> Self {
329        Self {
330            pat,
331            body,
332            guard: None,
333            meta: Metadata {
334                span,
335                attributes: Vec::new(),
336            },
337        }
338    }
339}
340
341impl PatKind {
342    /// Pattern for binding to a single variable
343    pub fn var_pat(var: LocalId) -> Self {
344        Self::Binding {
345            mutable: false,
346            var,
347            mode: BindingMode::ByValue,
348            sub_pat: None,
349        }
350    }
351    /// Promote to a `Pat`
352    pub fn promote(self, ty: Ty, span: Span) -> Pat {
353        Pat {
354            kind: Box::new(self),
355            ty,
356            meta: Metadata {
357                span,
358                attributes: Vec::new(),
359            },
360        }
361    }
362}
363
364impl Variant {
365    /// Whether a variant has fields or not.
366    /// See https://doc.rust-lang.org/reference/items/enumerations.html#field-less-enum.
367    pub fn is_fieldless(&self) -> bool {
368        self.arguments.is_empty()
369    }
370}
hax_rust_engine/ast/utils.rs

hax_rust_engine/ast/
utils.rs
