p8n_types/
statements.rs

1// Panopticon - A libre program analysis library for machine code
2// Copyright (C) 2014-2018  The Panopticon Developers
3//
4// This library is free software; you can redistribute it and/or
5// modify it under the terms of the GNU Lesser General Public
6// License as published by the Free Software Foundation; either
7// version 2.1 of the License, or (at your option) any later version.
8//
9// This library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
12// Lesser General Public License for more details.
13//
14// You should have received a copy of the GNU Lesser General Public
15// License along with this library; if not, write to the Free Software
16// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
17
18//! IL Statement encoding.
19//!
20//! To save space IL statements inside a function can be "compressed" by enconding the into a more
21//! space conserving, binary representation.
22//!
23//! This module handles convertion between those two.
24
25use std::ops::Range;
26use std::slice::Iter;
27use std::usize;
28use std::borrow::Cow;
29
30use {
31    Mnemonic,BasicBlock,Names,Strings,
32    Result,
33    Statement,Bitcode,BitcodeIter,Segments
34};
35
36/// Iterator over IL statements
37pub enum StatementsIter<'a> {
38    Bitcode(BitcodeIter<'a>),
39    Vector(Iter<'a,Statement>),
40}
41
42impl<'a> Iterator for StatementsIter<'a> {
43    type Item = Cow<'a,Statement>;
44
45    fn next(&mut self) -> Option<Cow<'a,Statement>> {
46        match self {
47            &mut StatementsIter::Bitcode(ref mut bt) => bt.next().map(|x| Cow::Owned(x)),
48            &mut StatementsIter::Vector(ref mut v) => v.next().map(|x| Cow::Borrowed(x)),
49        }
50    }
51}
52
53/// Signaling for `Statement::rewrite`.
54pub enum RewriteControl {
55    /// Continue with the next IL statement.
56    Continue,
57    /// Stop iteration.
58    Break,
59}
60
61/// A possibly compressed IL sequence.
62#[derive(Clone, Debug)]
63pub enum Statements {
64    /// Compressed binary encoding for IL statements.
65    Bitcode(Bitcode),
66    /// Uncompressed statement vector for processing.
67    Vector(Vec<Statement>),
68}
69
70impl Default for Statements {
71    fn default() -> Statements {
72        Statements::Vector(Vec::new())
73    }
74}
75
76impl Statements {
77    /// Creates an iterator over IL statements inside `rgn`.
78    pub fn iter_range<'a>(&'a self, rgn: Range<usize>) -> StatementsIter<'a> {
79        match self {
80            &Statements::Bitcode(ref bt) => StatementsIter::Bitcode(bt.iter_range(rgn)),
81            &Statements::Vector(ref v) => StatementsIter::Vector(v[rgn].iter()),
82        }
83    }
84
85    /// Mutate all IL statements in `rgn` using `f`.
86    pub fn rewrite<F: FnMut(&mut Statement, &mut Names, &mut Strings, &mut Segments) -> Result<RewriteControl> + Sized>(&mut self, rgn: Range<usize>, names: &mut Names, strings: &mut Strings, segments: &mut Segments, mut f: F) -> Result<Range<usize>> {
87        match self {
88            &mut Statements::Bitcode(ref mut bt) => bt.rewrite(rgn,names,strings,segments,f),
89            &mut Statements::Vector(ref mut v) => {
90                for stmt in v[rgn.clone()].iter_mut() {
91                    match f(stmt,names,strings,segments)? {
92                        RewriteControl::Break => { break; }
93                        RewriteControl::Continue => { continue; }
94                    }
95                }
96                Ok(rgn)
97            }
98        }
99    }
100
101    /// Inserts IL statements `stmts` at position `pos`.
102    pub fn insert(&mut self, pos: usize, stmts: Vec<Statement>) -> Result<Range<usize>> {
103        match self {
104            &mut Statements::Bitcode(ref mut bt) => bt.insert(pos,stmts),
105            &mut Statements::Vector(ref mut v) => {
106                let rgn = pos..(pos + stmts.len());
107                for (i,stmt) in stmts.into_iter().enumerate() {
108                    v.insert(i + pos,stmt);
109                }
110                Ok(rgn)
111            }
112        }
113    }
114
115    /// Removes IL statements in `rgn`
116    pub fn remove(&mut self, rgn: Range<usize>) {
117        match self {
118            &mut Statements::Bitcode(ref mut bt) => bt.remove(rgn),
119            &mut Statements::Vector(ref mut v) => { v.drain(rgn); }
120        }
121    }
122
123    /// Appends all IL statements in `i` to the instance.
124    pub fn append<I: IntoIterator<Item=Statement> + Sized>(&mut self, i: I) -> Result<Range<usize>> {
125        match self {
126            &mut Statements::Bitcode(ref mut bt) => bt.append(i),
127            &mut Statements::Vector(ref mut v) => {
128                let fst = v.len();
129                v.extend(i);
130                Ok(fst..v.len())
131            }
132        }
133    }
134
135    /// Upper end of the allowable IL statement range.
136    pub fn len(&self) -> usize {
137        match self {
138            &Statements::Bitcode(ref bt) => bt.num_bytes(),
139            &Statements::Vector(ref v) => v.len(),
140        }
141    }
142
143    /// Compress IL statements into bitcode, updates `basic_blocks` and `mnemonics` statement
144    /// ranges.
145    pub fn pack(&mut self, basic_blocks: &mut Vec<BasicBlock>, mnemonics: &mut Vec<Mnemonic>) -> Result<()> {
146        *self = match self {
147            &mut Statements::Bitcode(_) => { return Ok(()); },
148            &mut Statements::Vector(ref mut vec) => {
149                let mut bc = Bitcode::with_capacity(vec.len() * 15);
150
151                for bb in basic_blocks.iter_mut() {
152                    let mut old_idx = bb.statements.start;
153                    let mut new_idx = bc.num_bytes();
154                    let new_start = new_idx;
155                    let mne_rgn = bb.mnemonics.start.index()..bb.mnemonics.end.index();
156
157                    for mne in mnemonics[mne_rgn].iter_mut() {
158                        let old_rgn = old_idx..(old_idx + mne.statement_count);
159                        let new_rgn = bc.append(vec[old_rgn].iter().cloned())?;
160
161                        old_idx += mne.statement_count;
162                        mne.statement_count = new_rgn.end - new_rgn.start;
163                        new_idx += mne.statement_count;
164                    }
165
166                    assert_eq!(bb.statements.end, old_idx);
167                    bb.statements = new_start..new_idx;
168                }
169
170                Statements::Bitcode(bc)
171            }
172        };
173
174        Ok(())
175    }
176
177    /// Uncompress IL statements into bitcode, updates `basic_blocks` and `mnemonics` statement
178    /// ranges.
179     pub fn unpack(&mut self, basic_blocks: &mut Vec<BasicBlock>, mnemonics: &mut Vec<Mnemonic>) -> Result<()> {
180        *self = match self {
181            &mut Statements::Vector(_) => { return Ok(()); }
182            &mut Statements::Bitcode(ref bc) => {
183                let mut vec = Vec::with_capacity(bc.num_bytes() / 15);
184
185                for bb in basic_blocks.iter_mut() {
186                    let mut old_idx = bb.statements.start;
187                    let mut new_idx = vec.len();
188                    let new_start = new_idx;
189                    let mne_rgn = bb.mnemonics.start.index()..bb.mnemonics.end.index();
190
191                    for mne in mnemonics[mne_rgn].iter_mut() {
192                        let old_rgn = old_idx..(old_idx + mne.statement_count);
193
194                        vec.extend(bc.iter_range(old_rgn));
195                        old_idx += mne.statement_count;
196                        mne.statement_count = vec.len() - new_idx;
197                        new_idx += mne.statement_count;
198                    }
199
200                    assert_eq!(bb.statements.end, old_idx);
201                    bb.statements = new_start..new_idx;
202                }
203
204                Statements::Vector(vec)
205            }
206        };
207
208        Ok(())
209    }
210}
211