Skip to main content

hax_rust_engine/
attributes.rs

1//! Work with hax attributes.
2
3use std::collections::HashMap;
4
5use hax_lib_macros_types::{AssociationRole, AttrPayload, ItemUid, ProofMethod};
6
7use crate::ast::diagnostics::{Context, DiagnosticInfo, DiagnosticInfoKind};
8
9use super::ast::*;
10use visitors::AstVisitorMut;
11
12/// A graph of items connected via the hax attribute [`AttrPayload::AssociatedItem`] and UUIDs.
13#[derive(Clone)]
14pub struct LinkedItemGraph {
15    items: HashMap<ItemUid, Item>,
16    context: Context,
17}
18
19impl Default for LinkedItemGraph {
20    fn default() -> Self {
21        Self {
22            items: Default::default(),
23            context: Context::Unknown,
24        }
25    }
26}
27
28/// Get an iterator over hax attributes contained in the given attributes.
29pub fn hax_attributes(attrs: &Attributes) -> impl Iterator<Item = &AttrPayload> {
30    attrs.iter().flat_map(|attr| match &attr.kind {
31        AttributeKind::Hax(attr_payload) => Some(attr_payload),
32        _ => None,
33    })
34}
35
36/// Get proof attributes attached to the item
37pub fn hax_proof_attributes(item: &Item) -> Result<ProofAttributes, String> {
38    let mut proofs = hax_attributes(&item.meta.attributes).flat_map(|attr| match attr {
39        AttrPayload::Proof(proof) => Some(proof.clone()),
40        _ => None,
41    });
42    let proof = proofs.next();
43    if proofs.next().is_some() {
44        return Err("At most one `proof` attribute per item is allowed.".into());
45    }
46    let mut pure_requires_proofs =
47        hax_attributes(&item.meta.attributes).flat_map(|attr| match attr {
48            AttrPayload::PureRequiresProof(proof) => Some(proof.clone()),
49            _ => None,
50        });
51    let pure_requires_proof = pure_requires_proofs.next();
52    if pure_requires_proofs.next().is_some() {
53        return Err("At most one `pure_requires_proof` attribute per item is allowed.".into());
54    }
55    let mut pure_ensures_proofs =
56        hax_attributes(&item.meta.attributes).flat_map(|attr| match attr {
57            AttrPayload::PureEnsuresProof(proof) => Some(proof.clone()),
58            _ => None,
59        });
60    let pure_ensures_proof = pure_ensures_proofs.next();
61    if pure_ensures_proofs.next().is_some() {
62        return Err("At most one `pure_ensures_proof` attribute per item is allowed.".into());
63    }
64    let mut proof_methods = hax_attributes(&item.meta.attributes).flat_map(|attr| match attr {
65        AttrPayload::ProofMethod(method) => Some(*method),
66        _ => None,
67    });
68    let proof_method = proof_methods.next();
69    if proof_methods.next().is_some() {
70        return Err("At most one `proof_method` attribute per item is allowed.".into());
71    }
72    Ok(ProofAttributes {
73        proof,
74        pure_requires_proof,
75        pure_ensures_proof,
76        proof_method,
77    })
78}
79
80fn uuid(context: Context, item: &Item) -> Option<ItemUid> {
81    let mut uuids = hax_attributes(&item.meta.attributes).flat_map(|attr| match attr {
82        AttrPayload::Uid(item_uid) => Some(item_uid),
83        _ => None,
84    });
85    let uuid = uuids.next()?;
86    if let Some(other) = uuids.next() {
87        emit_assertion_failure(
88            context,
89            item.span(),
90            format!(
91                "Found more than one UUID hax attribute on this item. The two first UUIDs are {uuid} and {other}."
92            ),
93        );
94        None
95    } else {
96        Some(uuid.clone())
97    }
98}
99
100fn emit_assertion_failure(context: Context, span: span::Span, message: impl Into<String>) {
101    DiagnosticInfo {
102        context,
103        span,
104        kind: DiagnosticInfoKind::AssertionFailure {
105            details: message.into(),
106        },
107    }
108    .emit();
109}
110
111impl std::fmt::Debug for LinkedItemGraph {
112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113        f.debug_struct("LinkedItemGraph")
114            .field(
115                "items",
116                &self
117                    .items
118                    .iter()
119                    .map(|(id, item)| (id.to_string(), item.ident.to_debug_string()))
120                    .collect::<Vec<_>>(),
121            )
122            .field("context", &self.context)
123            .finish()
124    }
125}
126
127impl LinkedItemGraph {
128    /// Clone items marked with UUIDs attributes to build a graph of linked items.
129    /// This graph clones the items that represent linked items: e.g. pre and post conditions.
130    pub fn new(items: &[Item], context: Context) -> Self {
131        Self {
132            items: HashMap::from_iter(
133                items
134                    .iter()
135                    .filter_map(|item| Some((uuid(context.clone(), item)?, item.clone()))),
136            ),
137            context,
138        }
139    }
140
141    fn emit_assertion_failure(&self, span: span::Span, message: impl Into<String>) {
142        emit_assertion_failure(self.context.clone(), span, message)
143    }
144
145    fn emit_unimplemented(&self, span: span::Span, issue_id: u32, message: impl Into<String>) {
146        DiagnosticInfo {
147            context: self.context.clone(),
148            span,
149            kind: DiagnosticInfoKind::Unimplemented {
150                issue_id: Some(issue_id),
151                details: Some(message.into()),
152            },
153        }
154        .emit();
155    }
156
157    /// Given a graph and an item `item`, returns an iterator of the various items that are linked with `item`.
158    pub fn linked_items_iter(
159        &self,
160        item: &impl HasMetadata,
161    ) -> impl Iterator<Item = (AssociationRole, Result<&Item, DiagnosticInfo>)> {
162        let item_attributes = &item.metadata().attributes;
163        hax_attributes(item_attributes).flat_map(move |attr| match attr {
164            AttrPayload::AssociatedItem { role, item: target } => {
165                let target = self.items.get(target).map(Ok).unwrap_or_else(|| {
166                    Err(DiagnosticInfo {
167                        context: self.context.clone(),
168                        span: item.span(),
169                        kind: DiagnosticInfoKind::AssertionFailure {
170                            details: format!("An item linked via hax attributes could not be found. The UUID is {target:?}. The graph is {:#?}.", self),
171                        },
172                    })
173                });
174                Some((*role, target))
175            }
176            _ => None,
177        })
178    }
179
180    /// Returns the items linked to a given item.
181    pub fn linked_items(
182        &self,
183        item: &impl HasMetadata,
184    ) -> HashMap<AssociationRole, Vec<Result<&Item, DiagnosticInfo>>> {
185        let mut map: HashMap<AssociationRole, Vec<_>> = HashMap::new();
186        for (role, item) in self.linked_items_iter(item) {
187            map.entry(role).or_default().push(item);
188        }
189        map
190    }
191
192    /// Returns the precondition, postcondition and decreases clause, if any, for a given item.
193    /// When operating on a linked function, `self_id` is the local identifier of `self`.
194    pub fn fn_like_linked_expressions(
195        &self,
196        item: &impl HasMetadata,
197        self_id: Option<LocalId>,
198    ) -> FnLikeAssocatedExpressions {
199        let assoc_items = self.linked_items(item);
200        let get = |role| {
201            assoc_items
202                .get(&role)
203                .iter()
204                .flat_map(|vec| vec.iter())
205                .flat_map(|item| match item {
206                    Ok(item) => Some(item),
207                    Err(err) => {
208                        err.emit();
209                        None
210                    }
211                })
212                .map(|item| extract_expr(&self.context, item, self_id.clone()))
213                .collect::<Vec<_>>()
214        };
215        let precondition = {
216            let mut preconditions = get(AssociationRole::Requires).into_iter();
217            preconditions.next().map(|(e, _)| {
218                for extra in preconditions {
219                    self.emit_unimplemented(extra.0.span(), 1270, "multiple pre-conditions");
220                }
221                e
222            })
223        };
224        let decreases = {
225            let mut decreases = get(AssociationRole::Decreases).into_iter();
226            decreases.next().map(|(e, _)| {
227                for extra in decreases {
228                    self.emit_unimplemented(extra.0.span(), 1270, "multiple decreases");
229                }
230                e
231            })
232        };
233        let postcondition = {
234            let mut postconditions = get(AssociationRole::Ensures).into_iter();
235            postconditions.next().and_then(|(e, params)| {
236                for extra in postconditions {
237                    self.emit_unimplemented(extra.0.span(), 1270, "multiple post-conditions");
238                }
239                if let Some(last_param) = params.last() {
240                    Some(Postcondition {
241                        result_binder: last_param.pat.clone(),
242                        body: e.clone(),
243                    })
244                } else {
245                    self.emit_assertion_failure(
246                        e.span(),
247                        "hax ensures attribute: could not find output binder",
248                    );
249                    None
250                }
251            })
252        };
253        FnLikeAssocatedExpressions {
254            decreases,
255            precondition,
256            postcondition,
257        }
258    }
259
260    /// Is there a specification that we should prove for this item?
261    pub fn has_spec(&self, item: &Item) -> bool {
262        let spec = self.fn_like_linked_expressions(item, item.self_id());
263        spec.precondition.is_some() || spec.postcondition.is_some()
264    }
265}
266
267fn extract_expr<'a>(
268    context: &Context,
269    item: &'a Item,
270    self_id: Option<LocalId>,
271) -> (Expr, Vec<&'a Param>) {
272    let ItemKind::Fn { body, params, .. } = item.kind() else {
273        return (
274            ExprKind::Error(ErrorNode::assertion_failure(
275                item.clone(),
276                context.clone(),
277                "Expected an function",
278            ))
279            .into_expr(item.span(), Ty::prop(), vec![]),
280            vec![],
281        );
282    };
283    let mut body = body.clone();
284    if let Some(self_id) = self_id
285        && let [maybe_self, ..] = params.as_slice()
286        && let PatKind::Binding {
287            var, sub_pat: None, ..
288        } = &*maybe_self.pat.kind
289    {
290        // Here, we expect `self_id` is `self`, thus we cannot have any shadowing.
291        utils::mappers::SubstLocalIds::one(var.clone(), self_id.clone()).visit(&mut body)
292    }
293    (body, params.iter().collect())
294}
295
296/// A postcondition.
297///
298/// ## Example
299/// The expression `result != x` in the following is a postcondition.
300/// Note that `result` is an extra binder that represent the result of `f`, whose type is `u8` in this case: the return type of `f`.
301///
302/// ```rust
303/// #[hax_lib::ensures(|result| result != x)]
304/// fn f(x: u8) -> u8 { x.wrapping_add(1) }
305/// ```
306pub struct Postcondition {
307    /// In the example, this is `|result|`.
308    pub result_binder: Pat,
309    /// The formula of the postcondition, `result != x` in the example.
310    pub body: Expr,
311}
312
313/// The various linked expressions one can usually find on a (linked or not) function.
314pub struct FnLikeAssocatedExpressions {
315    /// A decreases clause, see [`hax_lib::decreases`]
316    pub decreases: Option<Expr>,
317    /// A precondition, see [`hax_lib::requires`]
318    pub precondition: Option<Expr>,
319    /// A postcondition, see [`hax_lib::ensures`]
320    pub postcondition: Option<Postcondition>,
321}
322
323/// The various linked expressions one can usually find on a (linked or not) function.
324pub struct ProofAttributes {
325    /// A custom proof, see [`hax_lib::lean::proof`]
326    pub proof: Option<String>,
327    /// A proof that the precondition is pure, see [`hax_lib::lean::pure_requires_proof`]
328    pub pure_requires_proof: Option<String>,
329    /// A proof that the postcondition is pure, see [`hax_lib::lean::pure_ensures_proof`]
330    pub pure_ensures_proof: Option<String>,
331    /// A proof method, see [`hax_lib::lean::proof_method`]
332    pub proof_method: Option<ProofMethod>,
333}