hax_rust_engine/
resugarings.rs

1//! The "resugaring" phases used by printers.
2
3//! This module defines resugarings instances (see
4//! [`hax_rust_engine::ast::Resugaring`] for the definition of a
5//! resugaring). Each backend defines its own set of resugaring phases.
6
7use crate::ast::identifiers::GlobalId;
8use crate::ast::resugared::*;
9use crate::ast::visitors::*;
10use crate::ast::*;
11use crate::printer::*;
12use std::collections::HashSet;
13
14/// Transforms [`ItemKind::Fn`] of arity zero into [`ResugaredItemKind::Constant`].
15/// Rust `const` items are encoded by the `ImportThir` phase of the hax engine as function of arity zero.
16/// Functions of arity zero themselves are encoded as functions operating on one argument of type `()`.
17#[derive(Copy, Clone, Default)]
18pub struct FunctionsToConstants;
19
20impl AstVisitorMut for FunctionsToConstants {
21    fn enter_item_kind(&mut self, item_kind: &mut ItemKind) {
22        let ItemKind::Fn {
23            name,
24            generics,
25            body,
26            params,
27            safety: SafetyKind::Safe,
28        } = item_kind
29        else {
30            return;
31        };
32        if !params.is_empty() {
33            return;
34        }
35        *item_kind = ItemKind::Resugared(ResugaredItemKind::Constant {
36            name: *name,
37            body: body.clone(),
38            generics: generics.clone(),
39        });
40    }
41}
42
43impl Resugaring for FunctionsToConstants {
44    fn name(&self) -> String {
45        "functions-to-constants".to_string()
46    }
47}
48
49/// Binop resugaring. Used to identify expressions of the form `(f e1 e2)` where
50/// `f` is a known identifier.
51pub struct BinOp {
52    /// Stores a set of identifiers that should be resugared as binary
53    /// operations. Usually, those identifiers come from the hax encoding. Each
54    /// backend can select its own set of identifiers Typically, if the backend
55    /// has a special support for addition, `known_ops` will contain
56    /// `hax::machine::int::add`
57    pub known_ops: HashSet<GlobalId>,
58}
59
60impl BinOp {
61    /// Adds a new binary operation from a list of (hax-introduced) names
62    pub fn new(known_ops: &[GlobalId]) -> Self {
63        Self {
64            known_ops: HashSet::from_iter(known_ops.iter().cloned()),
65        }
66    }
67}
68
69impl AstVisitorMut for BinOp {
70    fn enter_expr_kind(&mut self, x: &mut ExprKind) {
71        let ExprKind::App {
72            head,
73            args,
74            generic_args,
75            bounds_impls,
76            trait_,
77        }: &mut ExprKind = x
78        else {
79            return;
80        };
81        let ExprKind::GlobalId(id) = &*head.kind else {
82            return;
83        };
84        let [lhs, rhs] = &args[..] else { return };
85        if self.known_ops.iter().any(|defid| id == defid) {
86            *x = ExprKind::Resugared(ResugaredExprKind::BinOp {
87                op: *id,
88                lhs: lhs.clone(),
89                rhs: rhs.clone(),
90                generic_args: generic_args.clone(),
91                bounds_impls: bounds_impls.clone(),
92                trait_: trait_.clone(),
93            });
94        }
95    }
96}
97
98impl Resugaring for BinOp {
99    fn name(&self) -> String {
100        "binop".to_string()
101    }
102}
103
104/// Tuples resugaring. Resugars tuple constructors to the dedicated expression variant [`ResugaredExprKind::Tuple`],
105/// and tuple types to the dedicated type variant [`ResugaredTyKind::Tuple`].
106pub struct Tuples;
107
108impl AstVisitorMut for Tuples {
109    fn enter_expr_kind(&mut self, x: &mut ExprKind) {
110        let (constructor, fields) = match x {
111            ExprKind::Construct {
112                constructor,
113                is_record: false,
114                is_struct: true,
115                base: None,
116                fields,
117            } => (constructor, &fields[..]),
118            ExprKind::GlobalId(constructor) => (constructor, &[][..]),
119            _ => return,
120        };
121        if constructor.expect_tuple().is_some() {
122            let args = fields.iter().map(|(_, e)| e).cloned().collect();
123            *x = ExprKind::Resugared(ResugaredExprKind::Tuple(args))
124        }
125    }
126    fn enter_ty_kind(&mut self, x: &mut TyKind) {
127        let TyKind::App { head, args } = x else {
128            return;
129        };
130        if head.expect_tuple().is_some() {
131            let Some(args) = args
132                .iter()
133                .map(GenericValue::expect_ty)
134                .collect::<Option<Vec<_>>>()
135            else {
136                return;
137            };
138            *x = TyKind::Resugared(ResugaredTyKind::Tuple(args.into_iter().cloned().collect()))
139        }
140    }
141}
142
143impl Resugaring for Tuples {
144    fn name(&self) -> String {
145        "tuples".to_string()
146    }
147}
148
149/// Let-pure resugaring. Use to identify expressions of the form `let x ← pure ..`, where the arrow
150/// can be turned into a normal assignment `:=`
151pub struct LetPure;
152
153impl AstVisitorMut for LetPure {
154    fn enter_expr_kind(&mut self, expr: &mut ExprKind) {
155        const PURE: GlobalId = crate::names::rust_primitives::hax::explicit_monadic::pure;
156        if let ExprKind::Let { lhs, rhs, body } = expr
157            && let ExprKind::App {
158                head,
159                args,
160                generic_args,
161                bounds_impls,
162                trait_: None,
163            } = rhs.kind()
164            && *head.kind() == ExprKind::GlobalId(PURE)
165            && let ([pure_rhs], [], []) = (&args[..], &generic_args[..], &bounds_impls[..])
166        {
167            *expr = ExprKind::Resugared(ResugaredExprKind::LetPure {
168                lhs: lhs.clone(),
169                rhs: pure_rhs.clone(),
170                body: body.clone(),
171            })
172        }
173    }
174}
175
176impl Resugaring for LetPure {
177    fn name(&self) -> String {
178        "let_pure".to_string()
179    }
180}