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 Metadata {
36    /// Get an iterator over hax attributes for this AST fragment.
37    pub fn hax_attributes(&self) -> impl Iterator<Item = &hax_lib_macros_types::AttrPayload> {
38        crate::attributes::hax_attributes(&self.attributes)
39    }
40}
41
42impl Pat {
43    /// Expects the pattern to be a simple binding `self`.
44    pub fn expect_self(&self) -> Option<LocalId> {
45        if let PatKind::Binding { var, .. } = self.kind()
46            && var.is_self()
47        {
48            Some(var.clone())
49        } else {
50            None
51        }
52    }
53}
54
55impl Item {
56    /// Returns a `LocalId` named `self` if the item is a standalone function
57    /// whose first argument is the keyword `self`. In other words, this
58    /// function returns a local identifier only for associated methods from
59    /// inherent `impl` blocks.
60    pub fn self_id(&self) -> Option<LocalId> {
61        if let ItemKind::Fn { params, .. } = self.kind()
62            && let [first, ..] = &params[..]
63            && let Some(self_id) = first.pat.expect_self()
64        {
65            Some(self_id.clone())
66        } else {
67            None
68        }
69    }
70}