charon_lib/ast/
llbc_ast_utils.rs

1//! Implementations for [crate::llbc_ast]
2
3use crate::llbc_ast::*;
4use crate::meta;
5use crate::meta::Span;
6use derive_generic_visitor::*;
7use std::mem;
8
9/// Combine the span information from a [Switch]
10pub fn combine_switch_targets_span(targets: &Switch) -> Span {
11    match targets {
12        Switch::If(_, st1, st2) => meta::combine_span(&st1.span, &st2.span),
13        Switch::SwitchInt(_, _, branches, otherwise) => {
14            let branches = branches.iter().map(|b| &b.1.span);
15            let mbranches = meta::combine_span_iter(branches);
16            meta::combine_span(&mbranches, &otherwise.span)
17        }
18        Switch::Match(_, branches, otherwise) => {
19            let branches = branches.iter().map(|b| &b.1.span);
20            let mbranches = meta::combine_span_iter(branches);
21            if let Some(otherwise) = otherwise {
22                meta::combine_span(&mbranches, &otherwise.span)
23            } else {
24                mbranches
25            }
26        }
27    }
28}
29
30impl Switch {
31    pub fn iter_targets(&self) -> impl Iterator<Item = &Block> {
32        use itertools::Either;
33        match self {
34            Switch::If(_, exp1, exp2) => Either::Left([exp1, exp2].into_iter()),
35            Switch::SwitchInt(_, _, targets, otherwise) => Either::Right(Either::Left(
36                targets.iter().map(|(_, tgt)| tgt).chain([otherwise]),
37            )),
38            Switch::Match(_, targets, otherwise) => Either::Right(Either::Right(
39                targets.iter().map(|(_, tgt)| tgt).chain(otherwise.as_ref()),
40            )),
41        }
42    }
43
44    pub fn iter_targets_mut(&mut self) -> impl Iterator<Item = &mut Block> {
45        use itertools::Either;
46        match self {
47            Switch::If(_, exp1, exp2) => Either::Left([exp1, exp2].into_iter()),
48            Switch::SwitchInt(_, _, targets, otherwise) => Either::Right(Either::Left(
49                targets.iter_mut().map(|(_, tgt)| tgt).chain([otherwise]),
50            )),
51            Switch::Match(_, targets, otherwise) => Either::Right(Either::Right(
52                targets
53                    .iter_mut()
54                    .map(|(_, tgt)| tgt)
55                    .chain(otherwise.as_mut()),
56            )),
57        }
58    }
59}
60
61impl StatementId {
62    pub fn fresh() -> StatementId {
63        use std::sync::atomic::AtomicUsize;
64        use std::sync::atomic::Ordering;
65        static COUNTER: AtomicUsize = AtomicUsize::new(0);
66        let id = COUNTER.fetch_add(1, Ordering::Relaxed);
67        StatementId::new(id)
68    }
69}
70
71impl Statement {
72    pub fn new(span: Span, content: RawStatement) -> Self {
73        Statement {
74            span,
75            id: StatementId::fresh(),
76            content,
77            comments_before: vec![],
78        }
79    }
80
81    pub fn into_box(self) -> Box<Self> {
82        Box::new(self)
83    }
84
85    pub fn into_block(self) -> Block {
86        Block {
87            span: self.span,
88            statements: vec![self],
89        }
90    }
91}
92
93impl Block {
94    pub fn from_seq(seq: Vec<Statement>) -> Option<Self> {
95        if seq.is_empty() {
96            None
97        } else {
98            let span = seq
99                .iter()
100                .map(|st| st.span)
101                .reduce(|a, b| meta::combine_span(&a, &b))
102                .unwrap();
103            Some(Block {
104                span,
105                statements: seq,
106            })
107        }
108    }
109
110    pub fn merge(mut self, mut other: Self) -> Self {
111        self.span = meta::combine_span(&self.span, &other.span);
112        self.statements.append(&mut other.statements);
113        self
114    }
115
116    pub fn then(mut self, r: Statement) -> Self {
117        self.span = meta::combine_span(&self.span, &r.span);
118        self.statements.push(r);
119        self
120    }
121
122    pub fn then_opt(self, other: Option<Statement>) -> Self {
123        if let Some(other) = other {
124            self.then(other)
125        } else {
126            self
127        }
128    }
129
130    /// Apply a function to all the statements, in a top-down manner.
131    pub fn visit_statements<F: FnMut(&mut Statement)>(&mut self, f: F) {
132        let _ = BlockVisitor::new(|_| {}, f).visit(self);
133    }
134
135    /// Apply a transformer to all the statements, in a bottom-up manner.
136    ///
137    /// The transformer should:
138    /// - mutate the current statement in place
139    /// - return the sequence of statements to introduce before the current statement
140    pub fn transform<F: FnMut(&mut Statement) -> Vec<Statement>>(&mut self, mut f: F) {
141        self.transform_sequences(|slice| f(&mut slice[0]));
142    }
143
144    /// Apply a transformer to all the statements, in a bottom-up manner. Compared to `transform`,
145    /// this also gives access to the following statements if any. Statements that are not part of
146    /// a sequence will be traversed as `[st]`. Statements that are will be traversed twice: once
147    /// as `[st]`, and then as `[st, ..]` with the following statements if any.
148    ///
149    /// The transformer should:
150    /// - mutate the current statements in place
151    /// - return the sequence of statements to introduce before the current statements
152    pub fn transform_sequences<F: FnMut(&mut [Statement]) -> Vec<Statement>>(&mut self, mut f: F) {
153        self.visit_blocks_bwd(|blk: &mut Block| {
154            let mut final_len = blk.statements.len();
155            let mut to_insert = vec![];
156            for i in (0..blk.statements.len()).rev() {
157                let new_to_insert = f(&mut blk.statements[i..]);
158                final_len += new_to_insert.len();
159                to_insert.push((i, new_to_insert));
160            }
161            if !to_insert.is_empty() {
162                to_insert.sort_by_key(|(i, _)| *i);
163                // Make it so the first element is always at the end so we can pop it.
164                to_insert.reverse();
165                // Construct the merged list of statements.
166                let old_statements =
167                    mem::replace(&mut blk.statements, Vec::with_capacity(final_len));
168                for (i, stmt) in old_statements.into_iter().enumerate() {
169                    while let Some((j, _)) = to_insert.last()
170                        && *j == i
171                    {
172                        let (_, mut stmts) = to_insert.pop().unwrap();
173                        blk.statements.append(&mut stmts);
174                    }
175                    blk.statements.push(stmt);
176                }
177            }
178        })
179    }
180
181    /// Visit `self` and its sub-blocks in a bottom-up (post-order) traversal.
182    pub fn visit_blocks_bwd<F: FnMut(&mut Block)>(&mut self, f: F) {
183        let _ = BlockVisitor::new(f, |_| {}).visit(self);
184    }
185}
186
187/// Small visitor to visit statements and blocks.
188#[derive(Visitor)]
189pub struct BlockVisitor<F: FnMut(&mut Block), G: FnMut(&mut Statement)> {
190    exit_blk: F,
191    enter_stmt: G,
192}
193
194impl<F: FnMut(&mut Block), G: FnMut(&mut Statement)> BlockVisitor<F, G> {
195    pub fn new(exit_blk: F, enter_stmt: G) -> Self {
196        Self {
197            exit_blk,
198            enter_stmt,
199        }
200    }
201}
202
203impl<F: FnMut(&mut Block), G: FnMut(&mut Statement)> VisitBodyMut for BlockVisitor<F, G> {
204    fn exit_llbc_block(&mut self, x: &mut Block) {
205        (self.exit_blk)(x)
206    }
207    fn enter_llbc_statement(&mut self, x: &mut Statement) {
208        (self.enter_stmt)(x)
209    }
210}
charon_lib/ast/llbc_ast_utils.rs

charon_lib/ast/
llbc_ast_utils.rs
