charon_lib/ast/
llbc_ast_utils.rs

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

charon_lib/ast/
llbc_ast_utils.rs
