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 Vector<VarId, Var>, &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));
}
}