charon_lib/ast/
ullbc_ast_utils.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
//! Implementations for [crate::ullbc_ast]
use crate::ids::Vector;
use crate::meta::Span;
use crate::ullbc_ast::*;
use derive_visitor::{visitor_enter_fn_mut, visitor_fn_mut, DriveMut, Event};
use take_mut::take;

impl SwitchTargets {
    pub fn get_targets(&self) -> Vec<BlockId> {
        match self {
            SwitchTargets::If(then_tgt, else_tgt) => {
                vec![*then_tgt, *else_tgt]
            }
            SwitchTargets::SwitchInt(_, targets, otherwise) => {
                let mut all_targets = vec![];
                for (_, target) in targets {
                    all_targets.push(*target);
                }
                all_targets.push(*otherwise);
                all_targets
            }
        }
    }
}

impl Statement {
    pub fn new(span: Span, content: RawStatement) -> Self {
        Statement { span, content }
    }
}

impl Terminator {
    pub fn new(span: Span, content: RawTerminator) -> Self {
        Terminator { span, content }
    }
}

impl BlockData {
    pub fn targets(&self) -> Vec<BlockId> {
        match &self.terminator.content {
            RawTerminator::Goto { target } => {
                vec![*target]
            }
            RawTerminator::Switch { targets, .. } => targets.get_targets(),
            RawTerminator::Abort(..) | RawTerminator::Return => {
                vec![]
            }
        }
    }

    /// See [body_transform_operands]
    pub fn transform_operands<F: FnMut(&Span, &mut Vec<Statement>, &mut Operand)>(
        mut self,
        f: &mut F,
    ) -> Self {
        // The new vector of statements
        let mut nst = vec![];

        // Explore the operands in the statements
        for mut st in self.statements {
            st.content
                .drive_mut(&mut visitor_enter_fn_mut(|op: &mut Operand| {
                    f(&st.span, &mut nst, op)
                }));
            // Add the statement to the vector of statements
            nst.push(st)
        }

        // Explore the terminator
        self.terminator
            .content
            .drive_mut(&mut visitor_enter_fn_mut(|op: &mut Operand| {
                f(&self.terminator.span, &mut nst, op)
            }));

        // Update the vector of statements
        self.statements = nst;

        // Return
        self
    }

    /// Apply a transformer to all the statements, in a bottom-up manner.
    ///
    /// The transformer should:
    /// - mutate the current statements in place
    /// - return a list of `(i, statements)` where `statements` will be inserted before index `i`.
    pub fn transform_sequences<F>(&mut self, f: &mut F)
    where
        F: FnMut(&mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
    {
        for i in (0..self.statements.len()).rev() {
            let mut to_insert = f(&mut self.statements[i..]);
            if !to_insert.is_empty() {
                to_insert.sort_by_key(|(i, _)| *i);
                for (j, statements) in to_insert.into_iter().rev() {
                    // Insert the new elements at index `j`. This only modifies `statements[j..]`
                    // so we can keep iterating `j` (and `i`) down as if nothing happened.
                    self.statements.splice(i + j..i + j, statements);
                }
            }
        }
    }
}

impl ExprBody {
    pub fn transform_sequences<F>(&mut self, f: &mut F)
    where
        F: FnMut(&mut Locals, &mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
    {
        for block in &mut self.body {
            block.transform_sequences(&mut |seq| f(&mut self.locals, seq));
        }
    }

    /// Apply a function to all the statements, in a bottom-up manner.
    pub fn visit_statements<F: FnMut(&mut Statement)>(&mut self, f: &mut F) {
        self.drive_mut(&mut visitor_fn_mut(|st: &mut Statement, e: Event| {
            if matches!(e, Event::Exit) {
                f(st)
            }
        }))
    }
}

/// Transform a body by applying a function to its operands, and
/// inserting the statements generated by the operands at the end of the
/// block.
/// Useful to implement a pass on operands (see e.g., [crate::extract_global_assignments]).
///
/// The span argument given to `f` is the span argument of the [Terminator]
/// containing the operand. `f` should explore the operand it receives, and
/// push statements to the vector it receives as input.
pub fn body_transform_operands<F: FnMut(&Span, &mut Vec<Statement>, &mut Operand)>(
    blocks: &mut Vector<BlockId, BlockData>,
    f: &mut F,
) {
    for block in blocks.iter_mut() {
        take(block, |b| b.transform_operands(f));
    }
}