p8n_types/
function.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
18use std::ops::{RangeFull,Range};
19use std::iter::FromIterator;
20use std::cmp;
21use std::usize;
22use std::collections::{HashSet,HashMap};
23use std::fmt::Debug;
24use std::borrow::Cow;
25
26use petgraph::Graph;
27use petgraph::graph::{NodeIndices, NodeIndex};
28use petgraph::visit::{Walker, DfsPostOrder};
29use smallvec::SmallVec;
30use ron_uuid::UUID;
31
32use {
33    Architecture,
34    Region,
35    Guard,
36    Names,
37    Strings,
38    Str,
39    StrRef,
40    Result,
41    Statement,
42    Value,
43    Constant,
44    Variable,
45    Segments,
46    BasicBlock,
47    BasicBlockIndex,
48    Mnemonic,
49    RewriteControl,
50    Area,
51};
52use statements::{StatementsIter,Statements};
53use mnemonic::{MnemonicIterator,MnemonicIndex};
54use basic_block::BasicBlockIterator;
55
56/// Convertable into a range of statements.
57pub trait IntoStatementRange {
58    /// Converts Self into a range of statements for function `func`.
59    fn into_statement_range(self, func: &Function) -> Range<usize>;
60}
61
62impl IntoStatementRange for RangeFull {
63    fn into_statement_range(self, func: &Function) -> Range<usize> {
64        0..func.statements.len()
65    }
66}
67
68impl IntoStatementRange for Range<usize> {
69    fn into_statement_range(self, _: &Function) -> Range<usize> {
70        self
71    }
72}
73
74impl IntoStatementRange for BasicBlockIndex {
75    fn into_statement_range(self, func: &Function) -> Range<usize> {
76        debug!("into rgn: {}",self.index());
77        let bb = &func.basic_blocks[self.index()];
78        bb.into_statement_range(func)
79    }
80}
81
82impl<'a> IntoStatementRange for &'a BasicBlock {
83    fn into_statement_range(self, _: &Function) -> Range<usize> {
84        self.statements.clone()
85    }
86}
87
88/// Convertable into a range of mnemonics.
89pub trait IntoMnemonicRange: Debug {
90    /// Converts Self into a range of mnemonics for function `func`.
91    fn into_mnemonic_range(self, func: &Function) -> Range<usize>;
92}
93
94impl IntoMnemonicRange for RangeFull {
95    fn into_mnemonic_range(self, func: &Function) -> Range<usize> {
96        0..func.mnemonics.len()
97    }
98}
99
100impl IntoMnemonicRange for Range<usize> {
101    fn into_mnemonic_range(self, _: &Function) -> Range<usize> {
102        self
103    }
104}
105
106impl IntoMnemonicRange for Range<MnemonicIndex> {
107    fn into_mnemonic_range(self, _: &Function) -> Range<usize> {
108        self.start.index()..self.end.index()
109    }
110}
111
112impl IntoMnemonicRange for BasicBlockIndex {
113    fn into_mnemonic_range(self, func: &Function) -> Range<usize> {
114        let bb = &func.basic_blocks[self.index()];
115        bb.into_mnemonic_range(func)
116    }
117}
118
119impl<'a> IntoMnemonicRange for &'a BasicBlock {
120    fn into_mnemonic_range(self, _: &Function) -> Range<usize> {
121        let start = self.mnemonics.start.index();
122        let end = self.mnemonics.end.index();
123
124        start..end
125    }
126}
127
128impl IntoMnemonicRange for MnemonicIndex {
129    fn into_mnemonic_range(self, _: &Function) -> Range<usize> {
130        self.index()..(self.index() + 1)
131    }
132}
133
134/// Iterator over all indirect, unresolved jumps/branches.
135pub struct IndirectJumps<'a> {
136    pub graph: &'a Graph<CfgNode,Guard>,
137    pub iterator: NodeIndices<u32>,
138}
139
140impl<'a> Iterator for IndirectJumps<'a> {
141    type Item = Variable;
142
143    fn next(&mut self) -> Option<Variable> {
144        while let Some(idx) = self.iterator.next() {
145            match self.graph.node_weight(idx) {
146                Some(&CfgNode::Value(Value::Variable(ref v))) => {
147                    return Some(v.clone());
148                }
149                _ => {}
150            }
151        }
152
153        None
154    }
155}
156
157/// Node in the control flow graph.
158#[derive(Debug, Clone, PartialEq, Eq)]
159pub enum CfgNode {
160    /// Basic block
161    BasicBlock(BasicBlockIndex),
162    /// Indirect jump to this value.
163    Value(Value),
164}
165
166#[derive(Clone, Copy, PartialEq, Eq, Debug)]
167pub enum CfgTarget {
168    Value(Value),
169    Resolved(u64, usize),
170}
171
172/// A single function in the binary.
173///
174/// A function is a set of basic blocks, with one serving as an entry point. Basic blocks form a
175/// control flow graph.
176///
177/// Each function also has a user-changeable `name` and a unchangeable `uuid`.
178#[derive(Debug, Clone)]
179pub struct Function {
180    /// Canonical name of the function. Can be anything.
181    pub name: Str,
182    /// Table mapping SSA variable names to integers.
183    pub names: Names,
184    /// Table mapping strings to integers.
185    pub strings: Strings,
186    /// Table mapping segments identifiers to integers.
187    pub segments: Segments,
188    /// Fixed UUID of this functions. Never changes.
189    uuid: UUID,
190    /// Region this function is in.
191    region: UUID,
192    /// Sort by rev. post order
193    statements: Statements,
194    /// Sort by rev. post order
195    basic_blocks: Vec<BasicBlock>,
196    /// Sort by area.start
197    mnemonics: Vec<Mnemonic>,
198    cflow_graph: Graph<CfgNode,Guard>,
199    entry_point: BasicBlockIndex,
200}
201
202impl Function {
203    /// Creates a new function by disassembling from `region` starting at `start`.
204    pub fn new<A: Architecture>(init: A::Configuration, start: u64, region: &Region, uuid: UUID) -> Result<Function> {
205        let mut mnemonics = Vec::new();
206        let mut by_source = HashMap::new();
207        let mut by_destination = HashMap::new();
208        let mut func = Function{
209            name: format!("fn_{:x}", start).into(),
210            names: Names::default(),
211            strings: Strings::default(),
212            segments: Segments::default(),
213            uuid: uuid,
214            region: region.uuid().clone(),
215            statements: Statements::default(),
216            basic_blocks: Vec::new(),
217            mnemonics: Vec::new(),
218            cflow_graph: Graph::new(),
219            entry_point: BasicBlockIndex::new(0),
220        };
221
222        disassemble::<A>(init,vec![start],region, &mut func.names, &mut func.strings,
223                         &mut func.segments, &mut mnemonics, &mut by_source, &mut by_destination)?;
224        assemble_function(&mut func, start, mnemonics, by_source, by_destination)?;
225
226        Ok(func)
227    }
228
229    /// Disassemble a function with known control flow graph.
230    pub fn known_cflow_graph<A>(init: A::Configuration, mut names: Names,
231                                cfg: Vec<(Range<u64>, SmallVec<[(Value, Guard); 2]>)>,
232                                entry: u64, region: &Region, uuid: UUID,
233                                name: String)
234        -> Result<Function> where A: Architecture
235    {
236        let mut strings = Strings::default();
237        let mut segs = Segments::default();
238        let mut basic_blocks = Vec::with_capacity(cfg.len());
239        let mut cflow_graph = Graph::<CfgNode, Guard>::new();
240        let mut start_to_vx = HashMap::<u64, NodeIndex>::new();
241        let mut entry_vx = None;
242
243        for &(ref rgn, _) in cfg.iter() {
244            let mut bb = BasicBlock{
245                uuid: UUID::now(),
246                mnemonics: MnemonicIndex::new(0)..MnemonicIndex::new(0),
247                node: NodeIndex::new(0),
248                area: rgn.clone().into(),
249                statements: 0..0,
250            };
251            let nx = cflow_graph.add_node(
252                CfgNode::BasicBlock(BasicBlockIndex::new(basic_blocks.len())));
253
254            bb.node = nx;
255            basic_blocks.push(bb);
256            start_to_vx.insert(rgn.start, nx);
257
258            if rgn.start == entry {
259                entry_vx = Some(nx);
260            }
261        }
262
263        for &(ref rgn, ref jmps) in cfg.iter() {
264            let from = start_to_vx[&rgn.start];
265
266            for &(ref tgt, ref g) in jmps.iter() {
267                match tgt {
268                    Value::Constant(Constant{ value,.. }) => {
269                        match start_to_vx.get(&value) {
270                            Some(&to) => { cflow_graph.add_edge(from, to, g.clone()); }
271                            None => { return Err("unknown branch target".into()); }
272                        }
273                    }
274                    val => {
275                        let to = cflow_graph.add_node(CfgNode::Value(val.clone()));
276                        cflow_graph.add_edge(from, to, g.clone());
277                    }
278                }
279            }
280        }
281
282        if entry_vx.is_none() {
283            return Err("entrypoint not found".into());
284        }
285
286        let mut postorder = DfsPostOrder::new(&cflow_graph,entry_vx.unwrap())
287            .iter(&cflow_graph).collect::<Vec<_>>();
288        let mut postorder_rev = vec![0; basic_blocks.len()];
289        let mut po_idx = 0;
290        let mut mnemonics = Vec::default();
291        let mut statements = Vec::default();
292
293        postorder.reverse();
294
295        for &n in postorder.iter() {
296            if let Some(&CfgNode::BasicBlock(bb_idx)) = cflow_graph.node_weight(n) {
297                postorder_rev[bb_idx.index()] = po_idx;
298                po_idx += 1;
299            }
300        }
301
302        basic_blocks.sort_by_key(|x| {
303            if let Some(&CfgNode::BasicBlock(bb_idx)) = cflow_graph.node_weight(x.node) {
304                postorder_rev[bb_idx.index()]
305            } else {
306                unreachable!()
307            }
308        });
309
310        for (idx,bb) in basic_blocks.iter().enumerate() {
311            *cflow_graph.node_weight_mut(bb.node).unwrap() = CfgNode::BasicBlock(BasicBlockIndex::new(idx));
312        }
313
314        for bb in basic_blocks.iter_mut() {
315            let mut pos = bb.area.start;
316            let mut matches = Vec::default();
317            let mne_start = MnemonicIndex::new(mnemonics.len());
318            let stmt_cnt = statements.len();
319
320            while pos < bb.area.end {
321                A::decode(region, pos, &init, &mut names, &mut segs,
322                          &mut strings, &mut matches)?;
323
324                for mut m in matches.drain(..) {
325                    pos = cmp::max(pos, m.area.end);
326                    mnemonics.push(Mnemonic{
327                        area: m.area.into(),
328                        opcode: strings.insert(&m.opcode),
329                        operands: m.operands,
330                        statement_count: m.instructions.len(),
331                    });
332                    statements.append(&mut m.instructions);
333                }
334            }
335
336            bb.mnemonics = mne_start..MnemonicIndex::new(mnemonics.len());
337            bb.statements = stmt_cnt..statements.len();
338        }
339
340        let entry = match cflow_graph.node_weight(entry_vx.unwrap()) {
341            Some(CfgNode::BasicBlock(bb)) => bb.clone(),
342            _ => { return Err("entry point is not a basic block".into()); }
343        };
344        let f = Function{
345            name: name.into(),
346            uuid: uuid,
347            region: region.uuid().clone(),
348            basic_blocks: basic_blocks,
349            mnemonics: mnemonics,
350            statements: Statements::Vector(statements),
351            strings: strings,
352            names: names,
353            segments: segs,
354            cflow_graph: cflow_graph,
355            entry_point: entry,
356        };
357
358        Ok(f)
359    }
360
361    /// Continues disassembling of `region`. The function looks for resolved, indirect control flow
362    /// edges.
363    pub fn extend<A: Architecture>(&mut self, init: A::Configuration, region: &Region) -> Result<()> {
364        use petgraph::visit::EdgeRef;
365
366        let mut by_source = HashMap::new();
367        let mut by_destination = HashMap::new();
368        let mut mnemonics = self.basic_blocks.iter().flat_map(|bb| {
369            let mut stmt_idx = bb.statements.start;
370
371            (bb.mnemonics.start.index()..bb.mnemonics.end.index()).map(|mne_idx| {
372                let mne = &self.mnemonics[mne_idx];
373                let stmt_rgn = stmt_idx..(stmt_idx + mne.statement_count);
374                let stmts = self.statements.iter_range(stmt_rgn).map(|x| {
375                    match x {
376                        Cow::Borrowed(s) => s.clone(),
377                        Cow::Owned(s) => s,
378                    }
379                }).collect::<Vec<_>>();
380
381                if mne_idx != bb.mnemonics.end.index() - 1 {
382                    by_source.entry((mne.area.start, mne.area.offset_start))
383                        .or_insert_with(|| Vec::new())
384                        .push((CfgTarget::Resolved(mne.area.end, mne.area.offset_end), Guard::True));
385                    by_destination.entry((mne.area.end, mne.area.offset_end))
386                        .or_insert_with(|| Vec::new())
387                        .push((CfgTarget::Resolved(mne.area.start, mne.area.offset_start), Guard::True));
388                }
389                stmt_idx += self.mnemonics[mne_idx].statement_count;
390                (mne.clone(),stmts)
391            }).collect::<Vec<_>>()
392        }).collect::<Vec<_>>();
393        let mut starts = Vec::new();
394
395        for e in self.cflow_graph.edge_references() {
396            let src = match self.cflow_graph.node_weight(e.source()) {
397                Some(&CfgNode::BasicBlock(bb_idx)) => {
398                    let bb = &self.basic_blocks[bb_idx.index()];
399                    let mne = &self.mnemonics[bb.mnemonics.end.index() - 1];
400                    (mne.area.start, mne.area.offset_start)
401                }
402                _ => unreachable!()
403            };
404            let dst = match self.cflow_graph.node_weight(e.target()) {
405                Some(&CfgNode::BasicBlock(bb_idx)) => {
406                    let bb = &self.basic_blocks[bb_idx.index()];
407                    let mne = &self.mnemonics[bb.mnemonics.start.index()];
408                    CfgTarget::Resolved(mne.area.start, mne.area.offset_start)
409                }
410                Some(&CfgNode::Value(ref val)) => {
411                    CfgTarget::Value(val.clone())
412                }
413                None => unreachable!()
414            };
415
416            by_source.entry(src).or_insert_with(|| Vec::new())
417                .push((dst.clone(),e.weight().clone()));
418
419            match dst {
420                CfgTarget::Value(Value::Constant(Constant{ value,.. })) => {
421                    by_destination.entry((value, 0)).or_insert_with(|| Vec::new())
422                        .push((CfgTarget::Resolved(src.0, src.1),e.weight().clone()));
423                    starts.push(value);
424                }
425                CfgTarget::Resolved(real, virt) => {
426                    by_destination.entry((real, virt)).or_insert_with(|| Vec::new())
427                        .push((CfgTarget::Resolved(src.0, src.1),e.weight().clone()));
428                }
429                _ => { /* skip */ }
430            }
431        }
432
433        mnemonics.sort_by_key(|&(ref mne, _)| (mne.area.start, mne.area.offset_start));
434
435        let entry = self.entry_address();
436        disassemble::<A>(init, starts, region, &mut self.names,
437                         &mut self.strings, &mut self.segments,
438                         &mut mnemonics, &mut by_source, &mut by_destination)?;
439
440        assemble_function(self,entry,mnemonics,by_source,by_destination)
441    }
442
443    /// Function entry point.
444    pub fn entry_point(&self) -> BasicBlockIndex { self.entry_point }
445
446    /// Iterator over all mnemonics in basic block `idx`.
447    pub fn mnemonics<'a,Idx: IntoMnemonicRange + Sized>(&'a self, idx: Idx) -> MnemonicIterator<'a> {
448        let idx = idx.into_mnemonic_range(self);
449        MnemonicIterator::new(self,idx.start,idx.end - 1)
450    }
451
452    /// Iterator over all basic blocks in reverse post order.
453    pub fn basic_blocks<'a>(&'a self) -> BasicBlockIterator<'a> {
454        BasicBlockIterator::new(self,0,self.basic_blocks.len())
455    }
456
457    /// The Functions control flow graph.
458    pub fn cflow_graph<'a>(&'a self) -> &'a Graph<CfgNode,Guard> {
459        &self.cflow_graph
460    }
461
462    /// Returns a reference to the basic block `idx`.
463    pub fn basic_block<'a>(&'a self, idx: BasicBlockIndex) -> &'a BasicBlock {
464        &self.basic_blocks[idx.index()]
465    }
466
467    /// Returns a reference to the mnemonic `idx`.
468    pub fn mnemonic<'a>(&'a self, idx: MnemonicIndex) -> &'a Mnemonic {
469        &self.mnemonics[idx.index()]
470    }
471
472    /// The functions uuid.
473    pub fn uuid<'a>(&'a self) -> &'a UUID {
474        &self.uuid
475    }
476
477    /// The functions region.
478    pub fn region<'a>(&'a self) -> &'a UUID {
479        &self.region
480    }
481
482    /// Iterator over all IL statements in `rgn`.
483    pub fn statements<'a, Idx: IntoStatementRange + Sized>(&'a self, rgn: Idx) -> StatementsIter<'a> {
484        let rgn = rgn.into_statement_range(self);
485        debug!("read statements {:?}",rgn);
486        self.statements.iter_range(rgn)
487    }
488
489    /// Lowest address occupied by a basic block. Not neccecarly the entry point.
490    pub fn first_address(&self) -> u64 {
491        self.basic_blocks[0].area().start
492    }
493
494    /// First address of the functions entry basic block.
495    pub fn entry_address(&self) -> u64 {
496        let e = self.entry_point().index();
497        self.basic_blocks[e].area().start
498    }
499
500    /// Iterator over all indirect, unresolved jumps.
501    pub fn indirect_jumps<'a>(&'a self) -> IndirectJumps<'a> {
502        IndirectJumps{
503            graph: &self.cflow_graph,
504            iterator: self.cflow_graph.node_indices()
505        }
506    }
507
508    /// Mutable reference to the control flow graph.
509    pub fn cflow_graph_mut<'a>(&'a mut self) -> &'a mut Graph<CfgNode,Guard> {
510        &mut self.cflow_graph
511    }
512
513    /// Calls `func` on each indirect jump and call instruction and rewrites it to jump/call the
514    /// returned concrete addresses.
515    pub fn resolve_indirect_jumps<F: FnMut(&Variable) -> SmallVec<[Constant; 1]>>(&mut self, mut func: F) -> bool {
516        use petgraph::Direction;
517        use petgraph::visit::EdgeRef;
518
519        let mut retval = false;
520
521        'outer: loop {
522            let nodes = self.cflow_graph.node_indices().collect::<Vec<_>>();
523
524            for n in nodes {
525                let vals = match self.cflow_graph.node_weight(n) {
526                    Some(&CfgNode::Value(Value::Variable(ref var))) => {
527                        func(var)
528                    }
529                    _ => SmallVec::default(),
530                };
531
532                if !vals.is_empty() {
533                    let nodes = vals.into_iter().map(|c| {
534                        self.cflow_graph.add_node(CfgNode::Value(Value::Constant(c)))
535                    }).collect::<Vec<_>>();
536                    let edges = self.cflow_graph.edges_directed(n, Direction::Incoming)
537                        .map(|e| (e.source(), e.weight().clone())).collect::<Vec<_>>();
538
539                    for (w, g) in edges {
540                        for &m in nodes.iter() {
541                            self.cflow_graph.add_edge(w, m, g.clone());
542                        }
543                    }
544
545                    self.cflow_graph.remove_node(n);
546                    retval = true;
547                    continue 'outer;
548                }
549            }
550
551            self.reindex_basic_blocks().unwrap();
552            return retval;
553        }
554    }
555
556    /// Iterates thru all statements in `basic_block` calling `func` one each. The function is
557    /// allowed to modify the IL.
558    pub fn rewrite_basic_block<F: FnMut(&mut Statement,&mut Names,&mut Strings,&mut Segments) -> Result<RewriteControl> + Sized>(&mut self, basic_block: BasicBlockIndex, mut func: F) -> Result<()> {
559        debug!("start func rewrite of {:?}",basic_block);
560
561        let mut bb_offset = 0;
562        let mut stmt_offset = 0isize;
563
564        {
565            let bb = &self.basic_blocks[basic_block.index()];
566
567            for mne_idx in bb.mnemonics.start.index()..bb.mnemonics.end.index() {
568                let mne = &mut self.mnemonics[mne_idx];
569
570                debug!("mne {} at {:#x}",self.strings.value(mne.opcode)?,mne.area.start);
571                if mne.statement_count == 0 {
572                    debug!("skip this mnemonic (empty)");
573                    continue;
574                }
575
576                let stmt_idx = (bb.statements.start as isize + stmt_offset) as usize;
577                debug!("from {:?}",stmt_idx..(stmt_idx + mne.statement_count));
578                let new_rgn = self.statements.rewrite(stmt_idx..(stmt_idx + mne.statement_count),&mut self.names,&mut self.strings,&mut self.segments,&mut func)?;
579
580                let new_stmt_num = new_rgn.end - new_rgn.start;
581                let offset = new_stmt_num as isize - mne.statement_count as isize;
582
583                debug!("...to {:?}",new_rgn);
584                bb_offset += offset;
585                mne.statement_count = new_stmt_num;
586                stmt_offset += new_stmt_num as isize;
587            }
588        }
589
590        if bb_offset != 0 {
591            for bb_idx in (basic_block.index() + 1)..self.basic_blocks.len() {
592                let bb = &mut self.basic_blocks[bb_idx];
593
594                bb.statements.start = (bb.statements.start as isize + bb_offset) as usize;
595                bb.statements.end = (bb.statements.end as isize + bb_offset) as usize;
596            }
597
598            let bb_stmt = &mut self.basic_blocks[basic_block.index()].statements;
599            bb_stmt.end = (bb_stmt.end as isize + bb_offset) as usize;
600        }
601
602        Ok(())
603    }
604
605    /// Inserts a new mnemonic at position `pos` inside `basic_block` with `opcode` and semantics
606    /// described by `stmts`.
607    pub fn insert_mnemonic(&mut self, basic_block: BasicBlockIndex, pos: usize, opcode: StrRef, args: SmallVec<[Value; 3]>, stmts: Vec<Statement>) -> Result<()> {
608        debug!("insert mne {} in {:?} as {}",self.strings.value(opcode)?,basic_block,pos);
609
610        let mne_idx = self.basic_blocks[basic_block.index()].mnemonics.start.index() + pos;
611        let stmt_rgn = {
612            let bb = &self.basic_blocks[basic_block.index()];
613
614            if bb.mnemonics.end == bb.mnemonics.start {
615                return Err("Internal error: empty basic block".into());
616            }
617
618            if mne_idx > bb.mnemonics.end.index() {
619                return Err(format!("Internal error: mnemonic position out of range: {} > {:?}",mne_idx,bb.mnemonics).into());
620            }
621
622            let stmts_pos = bb.statements.start +
623                self.mnemonics[bb.mnemonics.start.index()..mne_idx].iter().map(|x| x.statement_count).sum::<usize>();
624
625            debug!("prepend mne at {} in {:?} {}: {:?}",stmts_pos,basic_block,self.strings.value(opcode)?,stmts);
626            self.statements.insert(stmts_pos,stmts)?
627        };
628
629        self.shift_areas(basic_block, pos, 1);
630        self.shift_mnemonics(MnemonicIndex::new(mne_idx),1);
631        self.shift_statements(BasicBlockIndex::new(basic_block.index() + 1),(stmt_rgn.end - stmt_rgn.start) as isize);
632
633        let bb = &mut self.basic_blocks[basic_block.index()];
634        let addr = if pos == 0 { (bb.area.start, 0) }
635                   else { let a = &self.mnemonics[mne_idx - 1].area; (a.end, a.offset_end) };
636        let mne = Mnemonic{
637            opcode: opcode,
638            area: Area{
639                start: addr.0,
640                end: addr.0,
641                offset_start: addr.1,
642                offset_end: addr.1 + 1,
643            },
644            operands: args,
645            statement_count: stmt_rgn.end - stmt_rgn.start,
646        };
647
648        self.mnemonics.insert(mne_idx,mne);
649
650        bb.mnemonics.end = MnemonicIndex::new(bb.mnemonics.end.index() + 1);
651        bb.statements.end += stmt_rgn.end - stmt_rgn.start;
652
653        Ok(())
654    }
655
656    /// Removes `mnemonic`.
657    pub fn remove_mnemonic(&mut self, mnemonic: MnemonicIndex) -> Result<()> {
658        let bb: Result<usize> = self.basic_blocks.iter().position(|bb| {
659            bb.mnemonics.start <= mnemonic && bb.mnemonics.end > mnemonic
660        }).ok_or("no such mnemonic".into());
661        let bb = bb?;
662        let move_by = {
663            let bb = &self.basic_blocks[bb];
664            let stmt_cnt = self.mnemonics[mnemonic.index()].statement_count;
665            let stmt_off: usize = self.mnemonics[bb.mnemonics.start.index()..mnemonic.index()].iter().map(|mne| mne.statement_count).sum();
666            let stmt_rgn = (bb.statements.start + stmt_off)..(bb.statements.start + stmt_off + stmt_cnt);
667
668            self.statements.remove(stmt_rgn.clone());
669            self.mnemonics.remove(mnemonic.index());
670            stmt_rgn.end - stmt_rgn.start
671        };
672
673        self.shift_mnemonics(mnemonic,-1);
674        self.shift_statements(BasicBlockIndex::new(bb + 1),-1 * move_by as isize);
675
676        let bb = &mut self.basic_blocks[bb];
677        bb.mnemonics.end = MnemonicIndex::new(bb.mnemonics.end.index() - 1);
678        bb.statements.end -= move_by;
679
680        Ok(())
681    }
682
683    /// Removes the first mnemonic of `basic_block`. Fails of `basic_block` has no mnemonics.
684    pub fn drop_first_mnemonic(&mut self, basic_block: BasicBlockIndex) -> Result<()> {
685        let mne = self.basic_blocks[basic_block.index()].mnemonics.start;
686        self.remove_mnemonic(mne)
687    }
688
689    /// Deletes all basic blocks for which `f` returns `false`.
690    pub fn retain_basic_blocks<F: (FnMut(&Function, BasicBlockIndex) -> bool)>(&mut self, mut f: F) -> Result<()> {
691        use vec_map::VecMap;
692
693        // map from original block indices to start addresses
694        let addr_to_old_bb = HashMap::<(u64, usize), BasicBlockIndex>::from_iter(
695            self.basic_blocks().map(|x| {
696                let a = x.1.area.clone();
697                ((a.start, a.offset_start), x.0)
698            }));
699        // blocks to delete
700        let mut to_del = self.basic_blocks()
701            .filter(|x| self.entry_point != x.0)
702            .filter(|x| !f(&*self, x.0)).map(|x| x.0.index()).collect::<Vec<_>>();
703
704        to_del.sort();
705
706        // delete mnemonics
707        for &bb in to_del.iter() {
708            loop {
709                let mnes = self.basic_blocks[bb].mnemonics.clone();
710
711                if mnes.start == mnes.end {
712                    break;
713                } else {
714                    self.drop_first_mnemonic(BasicBlockIndex::new(bb))?;
715                }
716            }
717        }
718
719        // delete basic blocks
720        for bb in to_del.into_iter().rev() {
721            self.basic_blocks.remove(bb);
722        }
723
724        let old_bb_to_new_bb = VecMap::<BasicBlockIndex>::from_iter(
725            self.basic_blocks().map(|x| {
726                let a = x.1.area.clone();
727                (addr_to_old_bb[&(a.start, a.offset_start)].index(), x.0)
728            }));
729
730        // delete cfg nodes
731        self.cflow_graph.retain_nodes(|mut g,vx| {
732            match &mut g[vx] {
733                &mut CfgNode::BasicBlock(ref mut bb) => {
734                    match old_bb_to_new_bb.get(bb.index()) {
735                        Some(new_bb) => {
736                            *bb = *new_bb;
737                            true
738                        }
739                        None => false,
740                    }
741                }
742                &mut CfgNode::Value(_) => true
743            }
744        });
745
746        for vx in self.cflow_graph.node_indices() {
747            match self.cflow_graph.node_weight(vx) {
748                Some(&CfgNode::BasicBlock(bb)) => {
749                    self.basic_blocks[bb.index()].node = vx;
750                }
751                _ => {}
752            }
753        }
754
755        // update entry point
756        self.entry_point = old_bb_to_new_bb[self.entry_point.index()];
757        Ok(())
758
759    }
760
761    /// Create a dummy function with no basic blocks and an invalid entry point.
762    #[cfg(test)]
763    pub fn facade() -> Function {
764        use Name;
765
766        Function{
767            name: "(none)".into(),
768            names: Names::facade(&Name::new("name-facade".into(), 0)),
769            strings: Strings::facade(&"string-facade".into()),
770            segments: Segments::facade(&Name::new("segment-facade".into(), 0)),
771            uuid: UUID::now(),
772            region: UUID::now(),
773            statements: Statements::default(),
774            basic_blocks: Vec::new(),
775            mnemonics: Vec::new(),
776            cflow_graph: Graph::new(),
777            entry_point: BasicBlockIndex::new(0),
778        }
779    }
780
781    /// Assembles a new function from a list of mnemonics and control flow edges. Used for
782    /// deserializing functions.
783    pub fn assemble(name: Str, names: Names, segments: Segments, region: UUID,
784                    strings: Strings, uuid: UUID, start: u64,
785                    mnemonics: Vec<(Mnemonic,Vec<Statement>)>,
786                    by_source: HashMap<(u64, usize), Vec<(CfgTarget, Guard)>>,
787                    by_destination: HashMap<(u64, usize), Vec<(CfgTarget, Guard)>>)
788        -> Result<Function> {
789
790        let mut func = Function{
791            name: name,
792            names: names,
793            strings: strings,
794            segments: segments,
795            uuid: uuid,
796            region: region,
797            statements: Statements::default(),
798            basic_blocks: Vec::new(),
799            mnemonics: Vec::new(),
800            cflow_graph: Graph::new(),
801            entry_point: BasicBlockIndex::new(0),
802        };
803
804        assemble_function(&mut func, start, mnemonics, by_source, by_destination)?;
805
806        Ok(func)
807    }
808
809    fn shift_mnemonics(&mut self, start: MnemonicIndex, change: isize) {
810        if start.index() >= self.mnemonics.len() { return; }
811        for bb in self.basic_blocks.iter_mut() {
812            if bb.mnemonics.start.index() > start.index() {
813                bb.mnemonics.start = MnemonicIndex::new((bb.mnemonics.start.index() as isize + change) as usize);
814                bb.mnemonics.end = MnemonicIndex::new((bb.mnemonics.end.index() as isize + change) as usize);
815            }
816        }
817    }
818
819    fn shift_statements(&mut self, start: BasicBlockIndex, change: isize) {
820        if start.index() >= self.basic_blocks.len() { return; }
821
822        let start_index = self.basic_blocks[start.index()].statements.start;
823
824        for bb_idx in start.index()..self.basic_blocks.len() {
825            let bb = &mut self.basic_blocks[bb_idx];
826            let rgn = bb.statements.clone();
827            let after_modification = rgn.start >= start_index;
828            let no_underflow = change >= 0 || rgn.start as isize >= -change;
829
830            if after_modification && no_underflow {
831                bb.statements.start = (bb.statements.start as isize + change) as usize;
832                bb.statements.end = (bb.statements.end as isize + change) as usize;
833            }
834        }
835    }
836
837    fn shift_areas(&mut self, mut bb: BasicBlockIndex, mnemonic: usize, change: usize) {
838        let mut idx = self.basic_blocks[bb.index()].mnemonics.start.index() + mnemonic;
839        let start_bb = bb;
840        let start = if self.mnemonics.len() > idx {
841            self.mnemonics[idx].area.start
842        } else {
843            self.mnemonics.last().unwrap().area.end
844        };
845
846        // adjust mnemonics
847        'outer: loop {
848            let idx_end = self.basic_blocks[bb.index()].mnemonics.end.index();
849
850            for i in idx..idx_end {
851                let mne = &mut self.mnemonics[i];
852
853                if mne.area.start != start {
854                   break 'outer;
855                } else {
856                    mne.area.offset_start += change;
857
858                    if mne.area.start == mne.area.end {
859                        mne.area.offset_end += change;
860                    }
861                }
862            }
863
864            if self.basic_blocks[bb.index()].area.end == start {
865                let area_end = self.basic_blocks[bb.index()].area.end;
866                let area_offset_end = self.basic_blocks[bb.index()].area.offset_end;
867
868                for (next_idx, next_bb) in self.basic_blocks.iter().enumerate() {
869                    if next_bb.area.start == area_end &&
870                       next_bb.area.offset_start == area_offset_end
871                    {
872                        bb = BasicBlockIndex::new(next_idx);
873                        idx = self.basic_blocks[bb.index()].mnemonics.start.index();
874                        continue 'outer;
875                    }
876                }
877            }
878
879            break;
880        }
881
882        // adjust basic blocks
883        for (idx, bb) in self.basic_blocks.iter_mut().enumerate() {
884            if !(idx == start_bb.index() && mnemonic == 0) {
885                bb.area.start = self.mnemonics[bb.mnemonics.start.index()].area.start;
886                bb.area.offset_start = self.mnemonics[bb.mnemonics.start.index()].area.offset_start;
887            }
888
889            if bb.mnemonics.start.index() < bb.mnemonics.end.index() {
890                bb.area.end = self.mnemonics[bb.mnemonics.end.index() - 1].area.end;
891                bb.area.offset_end = self.mnemonics[bb.mnemonics.end.index() - 1].area.offset_end;
892            }
893        }
894    }
895
896    fn reindex_basic_blocks(&mut self) -> Result<()> {
897        for n in self.cflow_graph.node_indices() {
898            match self.cflow_graph.node_weight(n).unwrap() {
899                &CfgNode::BasicBlock(bb_idx) => {
900                    let bb: Result<_> = self.basic_blocks.get_mut(bb_idx.index()).ok_or(
901                        "Internal error: re-indexing of basic blocks failed, missing block".into());
902                    let bb = bb?;
903
904                    bb.node = n;
905                }
906                &CfgNode::Value(_) => { /* skip */ }
907            }
908        }
909
910        Ok(())
911    }
912
913    /// Compresses the internal representation of the function. Call this if you're running low on
914    /// memory.
915    pub fn pack(&mut self) -> Result<()> {
916        self.statements.pack(&mut self.basic_blocks,&mut self.mnemonics)
917    }
918
919    /// Uncompresses the function into a faster to process representation.
920    pub fn unpack(&mut self) -> Result<()> {
921        self.statements.unpack(&mut self.basic_blocks,&mut self.mnemonics)
922    }
923}
924
925fn disassemble<A: Architecture>(init: A::Configuration, starts: Vec<u64>, region: &Region,
926                                names: &mut Names, strings: &mut Strings, segments: &mut Segments,
927                                mnemonics: &mut Vec<(Mnemonic,Vec<Statement>)>,
928                                by_source: &mut HashMap<(u64, usize),Vec<(CfgTarget, Guard)>>,
929                                by_destination: &mut HashMap<(u64, usize),Vec<(CfgTarget, Guard)>>) -> Result<()> {
930    let mut todo = HashSet::<u64>::from_iter(starts.into_iter());
931
932    while let Some(addr) = todo.iter().next().cloned() {
933        assert!(todo.remove(&addr));
934
935        match mnemonics.binary_search_by_key(&addr,|&(ref x,_)| x.area.start) {
936            // Already disassembled here
937            Ok(pos) => {
938                let mne = &mnemonics[pos].0;
939
940                if mne.area.start != addr {
941                    error!("{:#x}: Jump inside mnemonic {} at {:#x}",addr,strings.value(mne.opcode)?,mne.area.start);
942                }
943            }
944            // New mnemonic
945            Err(pos) => {
946                let mut matches = Vec::with_capacity(2);
947
948                match A::decode(region,addr,&init,names,segments,strings,&mut matches) {
949                    Ok(()) => {
950                        // Matched mnemonics
951                        if matches.is_empty() {
952                            error!("{:#x}: Unrecognized instruction",addr);
953                        } else {
954                            for m in matches {
955                                trace!("{:x}: {}",m.area.start,m.opcode);
956                                let this_mne = Mnemonic{
957                                    area: m.area.into(),
958                                    opcode: strings.insert(&m.opcode),
959                                    operands: m.operands,
960                                    statement_count: 0,
961                                };
962                                mnemonics.insert(pos,(this_mne,m.instructions));
963
964                                // New control transfers
965                                for (origin, tgt, gu) in m.jumps {
966                                    trace!("jump to {:?}", tgt);
967                                    match tgt {
968                                        Value::Constant(Constant{ value,.. }) => {
969                                            by_source.entry((origin, 0))
970                                                .or_insert(Vec::new())
971                                                .push((CfgTarget::Resolved(value, 0), gu.clone()));
972                                            by_destination.entry((value, 0))
973                                                .or_insert(Vec::new())
974                                                .push((CfgTarget::Resolved(origin, 0), gu));
975                                            todo.insert(value);
976                                        }
977                                        Value::Variable(Variable{ name, bits }) => {
978                                            by_source.entry((origin, 0))
979                                                .or_insert(Vec::new())
980                                                .push((CfgTarget::Value(Value::var(name,bits)?), gu));
981                                        }
982                                        Value::Undefined => {
983                                            by_source.entry((origin, 0))
984                                                .or_insert(Vec::new())
985                                                .push((CfgTarget::Value(Value::undef()), gu));
986                                        }
987                                    }
988                                }
989                            }
990                        }
991                    }
992                    Err(e) => {
993                        error!("{:#x} Failed to disassemble: {}",addr, e);
994                    }
995                }
996            }
997        }
998    }
999
1000    Ok(())
1001}
1002
1003fn assemble_function(function: &mut Function, entry: u64,
1004                     mut mnemonics: Vec<(Mnemonic,Vec<Statement>)>,
1005                     by_source: HashMap<(u64, usize),Vec<(CfgTarget,Guard)>>,
1006                     by_destination: HashMap<(u64, usize),Vec<(CfgTarget,Guard)>>) -> Result<()> {
1007
1008    let mut basic_blocks = Vec::<BasicBlock>::new();
1009    let mut idx = 0;
1010
1011    // Partition mnemonics into basic blocks
1012    while idx < mnemonics.len() {
1013        if mnemonics.len() - idx > 1 {
1014            let next_bb = mnemonics
1015                .as_slice()[idx..].windows(2)
1016                .position(|x| is_basic_block_boundary(&x[0].0,&x[1].0,entry,&by_source,&by_destination))
1017                .map(|x| x + 1 + idx)
1018                .unwrap_or(mnemonics.len());
1019            let bb = BasicBlock{
1020                uuid: UUID::now(),
1021                mnemonics: MnemonicIndex::new(idx)..MnemonicIndex::new(next_bb),
1022                area: Area{
1023                    start: mnemonics[idx].0.area.start,
1024                    offset_start: mnemonics[idx].0.area.offset_start,
1025                    end: mnemonics[next_bb - 1].0.area.end,
1026                    offset_end: mnemonics[next_bb - 1].0.area.offset_end,
1027                },
1028                node: NodeIndex::new(0),
1029                statements: 0..0,
1030            };
1031
1032            basic_blocks.push(bb);
1033            idx = next_bb;
1034        } else {
1035            let bb = BasicBlock{
1036                uuid: UUID::now(),
1037                mnemonics: MnemonicIndex::new(idx)..MnemonicIndex::new(mnemonics.len()),
1038                area: mnemonics[idx].0.area.clone(),
1039                node: NodeIndex::new(0),
1040                statements: 0..0,
1041            };
1042
1043            basic_blocks.push(bb);
1044            break;
1045        }
1046    }
1047
1048    // Build control flow graph
1049    let mut cfg = Graph::<CfgNode,Guard>::with_capacity(basic_blocks.len(),3*basic_blocks.len() / 2);
1050
1051    for (i,bb) in basic_blocks.iter_mut().enumerate() {
1052        bb.node = cfg.add_node(CfgNode::BasicBlock(BasicBlockIndex::new(i)));
1053    }
1054
1055    for bb in basic_blocks.iter() {
1056        let last_mne = &mnemonics[bb.mnemonics.end.index() - 1].0;
1057        if let Some(ct) = by_source.get(&(last_mne.area.start, last_mne.area.offset_start)) {
1058            for &(ref val,ref guard) in ct.iter() {
1059                match val {
1060                    &CfgTarget::Value(Value::Constant(_)) | &CfgTarget::Resolved(_, _) => {
1061                        let (real, virt) = match val {
1062                            &CfgTarget::Value(Value::Constant(Constant{ value,.. })) => (value, 0),
1063                            &CfgTarget::Resolved(r, v) => (r, v),
1064                            _ => unreachable!(),
1065                        };
1066                        if let Ok(pos) = basic_blocks.binary_search_by_key(&(real, virt),|bb| (bb.area.start, bb.area.offset_start)) {
1067                            cfg.update_edge(bb.node,basic_blocks[pos].node,guard.clone());
1068                        } else if virt == 0 {
1069                            let n = cfg.add_node(CfgNode::Value(Value::val(real, 64)?));
1070                            cfg.update_edge(bb.node,n,guard.clone());
1071                        } else {
1072                            error!("Internal error: drop jump to {}.{}", real, virt);
1073                        }
1074                    }
1075                    &CfgTarget::Value(ref val) => {
1076                        let n = cfg.add_node(CfgNode::Value(val.clone()));
1077                        cfg.update_edge(bb.node,n,guard.clone());
1078                    }
1079                }
1080            }
1081        }
1082    }
1083
1084    let entry_idx = basic_blocks
1085        .iter().position(|x| x.area.start == entry)
1086        .ok_or("Internal error: no basic block at the entry point")?;
1087
1088    // Generate bitcode
1089    let mut postorder = DfsPostOrder::new(&cfg,basic_blocks[entry_idx].node).iter(&cfg).collect::<Vec<_>>();
1090    let mut bitcode = Statements::default();
1091    let mut postorder_rev = vec![0; basic_blocks.len()];
1092    let mut po_idx = 0;
1093
1094    postorder.reverse();
1095
1096    for &n in postorder.iter() {
1097        if let Some(&CfgNode::BasicBlock(bb_idx)) = cfg.node_weight(n) {
1098            let bb = &mut basic_blocks[bb_idx.index()];
1099            let sl = bb.mnemonics.start.index()..bb.mnemonics.end.index();
1100
1101            bb.statements = usize::MAX..usize::MIN;
1102
1103            for &mut (ref mut mne,ref mut instr) in mnemonics.as_mut_slice()[sl].iter_mut() {
1104                let rgn = bitcode.append(instr.drain(..))?;
1105
1106                mne.statement_count = rgn.end - rgn.start;
1107                bb.statements.start = cmp::min(bb.statements.start,rgn.start);
1108                bb.statements.end = cmp::max(bb.statements.end,rgn.end);
1109            }
1110
1111            postorder_rev[bb_idx.index()] = po_idx;
1112            po_idx += 1;
1113        }
1114    }
1115
1116    basic_blocks.sort_by_key(|x| {
1117        if let Some(&CfgNode::BasicBlock(bb_idx)) = cfg.node_weight(x.node) {
1118            postorder_rev[bb_idx.index()]
1119        } else {
1120            unreachable!()
1121        }
1122    });
1123
1124    for (idx,bb) in basic_blocks.iter().enumerate() {
1125        *cfg.node_weight_mut(bb.node).unwrap() = CfgNode::BasicBlock(BasicBlockIndex::new(idx));
1126    }
1127
1128    function.statements = bitcode;
1129    function.basic_blocks = basic_blocks;
1130    function.mnemonics = mnemonics.into_iter().map(|(mne,_)| mne).collect();
1131    function.cflow_graph = cfg;
1132    function.entry_point = BasicBlockIndex::new(0);
1133
1134    Ok(())
1135}
1136
1137fn is_basic_block_boundary(a: &Mnemonic, b: &Mnemonic, entry: u64,
1138                           by_source: &HashMap<(u64, usize),Vec<(CfgTarget,Guard)>>,
1139                           by_destination: &HashMap<(u64, usize), Vec<(CfgTarget,Guard)>>) -> bool {
1140    // if next mnemonics aren't adjacent
1141    let mut new_bb = a.area.end != b.area.start || a.area.offset_end != b.area.offset_start;
1142
1143    // or any following jumps aren't to adjacent mnemonics
1144    new_bb |= by_source
1145        .get(&(a.area.start, a.area.offset_start))
1146        .unwrap_or(&Vec::new())
1147        .iter()
1148        .any(|&(ref opt_dest, _)| {
1149            match opt_dest {
1150                &CfgTarget::Resolved(r, v) => b.area.start != r || b.area.offset_start != v,
1151                &CfgTarget::Value(Value::Constant(Constant{ value, .. })) => {
1152                    value != b.area.start || b.area.offset_start != 0
1153                }
1154                _ => false,
1155            }
1156        });
1157
1158    // or any jumps pointing to the next that aren't from here
1159    new_bb |= by_destination
1160        .get(&(b.area.start, b.area.offset_start))
1161        .unwrap_or(&Vec::new())
1162        .iter()
1163        .any(|&(ref opt_src, _)| {
1164            match opt_src {
1165                &CfgTarget::Resolved(r, v) => a.area.start != r || a.area.offset_start != v,
1166                &CfgTarget::Value(Value::Constant(Constant{ value, .. })) => {
1167                    value != a.area.start || a.area.offset_start != 0
1168                }
1169                _ => false,
1170            }
1171        });
1172
1173    // or the entry point does not point here
1174    new_bb |= b.area.start == entry && b.area.offset_start == 0;
1175
1176    new_bb
1177}
1178
1179#[cfg(test)]
1180mod tests {
1181    use super::*;
1182    use {Region, TestArch, Name, Operation};
1183    use simple_logger;
1184
1185    #[test]
1186    fn single_instruction() {
1187        let _ = simple_logger::init();
1188        let reg = Region::from_buf("", 16, b"Ma0R".to_vec(), 0, None);
1189        let func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1190
1191        assert_eq!(func.cflow_graph().node_count(), 1);
1192        assert_eq!(func.cflow_graph().edge_count(), 0);
1193
1194        let node = func.cflow_graph().node_indices().next().unwrap();
1195        assert!(if let Some(&CfgNode::BasicBlock(_)) = func.cflow_graph().node_weight(node) { true } else { false });
1196
1197        assert_eq!(func.entry_address(), 0);
1198        assert_eq!(func.basic_blocks().len(), 1);
1199
1200        let (bb_idx,bb) = func.basic_blocks().next().unwrap();
1201        assert_eq!(bb.area(), &(0..4).into());
1202        assert_eq!(func.mnemonics(bb_idx).len(), 2);
1203
1204        let (_,mne) = func.mnemonics(bb_idx).next().unwrap();
1205        assert_eq!(func.strings.value(mne.opcode).unwrap(), "mov");
1206        let (_,mne) = func.mnemonics(bb_idx).skip(1).next().unwrap();
1207        assert_eq!(func.strings.value(mne.opcode).unwrap(), "ret");
1208
1209    }
1210
1211    #[test]
1212    fn branch() {
1213        let _ = simple_logger::init();
1214        let reg = Region::from_buf("", 16, b"Bf4RR".to_vec(), 0, None);
1215        let func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1216
1217        assert_eq!(func.cflow_graph.node_count(), 3);
1218        assert_eq!(func.cflow_graph.edge_count(), 2);
1219
1220        let mut bb0_vx = None;
1221        let mut bb1_vx = None;
1222        let mut bb2_vx = None;
1223
1224        for n in func.cflow_graph.node_indices() {
1225            match func.cflow_graph().node_weight(n) {
1226                Some(&CfgNode::BasicBlock(bb_idx)) => {
1227                    let bb = func.basic_block(bb_idx);
1228                    let mnes = func.mnemonics(bb_idx).collect::<Vec<_>>();
1229
1230                    assert_eq!(bb.node, n);
1231
1232                    if bb.area.start == 0 {
1233                        assert_eq!(mnes.len(), 1);
1234                        assert_eq!(func.strings.value(mnes[0].1.opcode).unwrap(), "br");
1235                        assert_eq!(mnes[0].1.area, (0..3).into());
1236                        assert_eq!(bb.area, (0..3).into());
1237                        bb0_vx = Some(n);
1238                    } else if bb.area.start == 3 {
1239                        assert_eq!(mnes.len(), 1);
1240                        assert_eq!(func.strings.value(mnes[0].1.opcode).unwrap(), "ret");
1241                        assert_eq!(mnes[0].1.area, (3..4).into());
1242                        assert_eq!(bb.area, (3..4).into());
1243                        bb1_vx = Some(n);
1244                    } else if bb.area.start == 4 {
1245                        assert_eq!(mnes.len(), 1);
1246                        assert_eq!(func.strings.value(mnes[0].1.opcode).unwrap(), "ret");
1247                        assert_eq!(mnes[0].1.area, (4..5).into());
1248                        assert_eq!(bb.area, (4..5).into());
1249                        bb2_vx = Some(n);
1250                    } else {
1251                        unreachable!();
1252                    }
1253                }
1254                _ => { unreachable!(); }
1255            }
1256        }
1257
1258        assert!(bb0_vx.is_some() && bb1_vx.is_some() && bb2_vx.is_some());
1259
1260        debug!("{:?}, {:?}, {:?}",bb0_vx, bb1_vx, bb2_vx);
1261
1262        let entry_bb = func.entry_point();
1263        assert_eq!(func.basic_block(entry_bb).node, bb0_vx.unwrap());
1264        assert!(func.cflow_graph().find_edge(bb0_vx.unwrap(), bb1_vx.unwrap()).is_some());
1265        assert!(func.cflow_graph().find_edge(bb0_vx.unwrap(), bb2_vx.unwrap()).is_some());
1266    }
1267
1268    #[test]
1269    fn single_loop() {
1270        let _ = simple_logger::init();
1271        let reg = Region::from_buf("", 16, b"AaaaAxxxJ0".to_vec(), 0, None);
1272        let func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1273
1274        assert_eq!(func.cflow_graph.node_count(), 1);
1275        assert_eq!(func.cflow_graph.edge_count(), 1);
1276
1277        let vx = func.cflow_graph.node_indices().next().unwrap();
1278        if let Some(&CfgNode::BasicBlock(bb_idx)) = func.cflow_graph().node_weight(vx) {
1279            let bb = func.basic_block(bb_idx);
1280            let mnes = func.mnemonics(bb_idx).collect::<Vec<_>>();
1281
1282            if bb.area.start == 0 {
1283                assert_eq!(mnes.len(), 3);
1284                assert_eq!(func.strings.value(mnes[0].1.opcode).unwrap(), "add");
1285                assert_eq!(mnes[0].1.area, (0..4).into());
1286                assert_eq!(func.strings.value(mnes[1].1.opcode).unwrap(), "add");
1287                assert_eq!(mnes[1].1.area, (4..8).into());
1288                assert_eq!(func.strings.value(mnes[2].1.opcode).unwrap(), "jmp");
1289                assert_eq!(mnes[2].1.area, (8..10).into());
1290                assert_eq!(bb.area, (0..10).into());
1291            } else {
1292                unreachable!();
1293            }
1294        }
1295
1296        let entry_idx = func.entry_point();
1297        assert_eq!(func.basic_block(entry_idx).node, vx);
1298        assert!(func.cflow_graph.find_edge(vx, vx).is_some());
1299    }
1300
1301    #[test]
1302    fn empty_function() {
1303        let reg = Region::from_buf("", 16, vec![], 0, None);
1304        assert!(Function::new::<TestArch>((), 0, &reg, UUID::now()).is_err());
1305    }
1306
1307    #[test]
1308    fn resolve_indirect() {
1309        let _ = simple_logger::init();
1310        let reg = Region::from_buf("", 16, b"AaaaJxAxxxR".to_vec(), 0, None);
1311        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1312        let a = func.names.index(&Name::new("x".into(),None)).unwrap();
1313
1314        assert_eq!(func.cflow_graph().node_count(), 2);
1315        assert_eq!(func.cflow_graph().edge_count(), 1);
1316
1317        for n in func.cflow_graph().node_indices() {
1318            match func.cflow_graph.node_weight(n) {
1319                Some(&CfgNode::BasicBlock(bb)) => {
1320                    assert_eq!(func.basic_block(bb).area, (0..6).into());
1321                }
1322                Some(&CfgNode::Value(Value::Variable(Variable{ ref name, bits: 32 }))) if *name == a => {}
1323                a => unreachable!("got: {:?}",a)
1324            }
1325        }
1326
1327        let unres = func.indirect_jumps().collect::<Vec<_>>();
1328        assert_eq!(unres.len(), 1);
1329        assert_eq!(unres[0], Variable{ name: a, bits: 32 });
1330
1331        assert!(func.resolve_indirect_jumps(|x| { if *x == Variable::new(a, 32).unwrap() { vec![Constant::new(6,32).unwrap()].into() } else { Default::default() } }));
1332        assert!(func.extend::<TestArch>((), &reg).is_ok());
1333
1334        assert_eq!(func.cflow_graph().node_count(), 1);
1335        assert_eq!(func.cflow_graph().edge_count(), 0);
1336
1337        for n in func.cflow_graph().node_indices() {
1338            match func.cflow_graph.node_weight(n) {
1339                Some(&CfgNode::BasicBlock(bb)) => {
1340                    assert_eq!(func.basic_block(bb).area, (0..11).into());
1341                }
1342                _ => unreachable!()
1343            }
1344        }
1345    }
1346
1347    #[test]
1348    fn entry_split() {
1349        use petgraph::dot::Dot;
1350
1351        let _ = simple_logger::init();
1352        let reg = Region::from_buf("", 16, b"AaaaAaaaJx".to_vec(), 0, None);
1353        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1354        let unres = func.indirect_jumps().collect::<Vec<_>>();
1355        let a = func.names.index(&Name::new("x".into(),None)).unwrap();
1356        assert_eq!(unres.len(), 1);
1357        assert_eq!(unres[0], Variable{ name: a, bits: 32 });
1358
1359        assert!(func.resolve_indirect_jumps(|x| { if *x == Variable::new(a, 32).unwrap() { vec![Constant::new(4,32).unwrap()].into() } else { Default::default() } }));
1360        assert!(func.extend::<TestArch>((), &reg).is_ok());
1361
1362        println!("{:?}", Dot::new(func.cflow_graph()));
1363
1364        assert_eq!(func.cflow_graph().node_count(), 2);
1365        assert_eq!(func.cflow_graph().edge_count(), 2);
1366
1367        let mut bb0_vx = None;
1368        let mut bb1_vx = None;
1369
1370        for n in func.cflow_graph().node_indices() {
1371            match func.cflow_graph.node_weight(n) {
1372                Some(&CfgNode::BasicBlock(bb)) => {
1373                    if func.basic_block(bb).area == (4..10).into() {
1374                        bb1_vx = Some(n);
1375                    } else if func.basic_block(bb).area == (0..4).into() {
1376                        bb0_vx = Some(n);
1377                    } else {
1378                        unreachable!();
1379                    }
1380                }
1381                _ => unreachable!()
1382            }
1383        }
1384
1385        assert!(bb0_vx.is_some() && bb1_vx.is_some());
1386        let entry_idx = func.entry_point();
1387        assert_eq!(func.basic_block(entry_idx).node, bb0_vx.unwrap());
1388    }
1389
1390    #[test]
1391    fn resolve_indirect_to_multiple() {
1392        let _ = simple_logger::init();
1393        let reg = Region::from_buf("", 16, b"AaaaJxAxxxRRRR".to_vec(), 0, None);
1394        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1395        let a = func.names.index(&Name::new("x".into(),None)).unwrap();
1396
1397        assert_eq!(func.cflow_graph().node_count(), 2);
1398        assert_eq!(func.cflow_graph().edge_count(), 1);
1399
1400        for n in func.cflow_graph().node_indices() {
1401            match func.cflow_graph.node_weight(n) {
1402                Some(&CfgNode::BasicBlock(bb)) => {
1403                    assert_eq!(func.basic_block(bb).area, (0..6).into());
1404                }
1405                Some(&CfgNode::Value(Value::Variable(Variable{ ref name, bits: 32 }))) if *name == a => {}
1406                a => unreachable!("got: {:?}",a)
1407            }
1408        }
1409
1410        let unres = func.indirect_jumps().collect::<Vec<_>>();
1411        assert_eq!(unres.len(), 1);
1412        assert_eq!(unres[0], Variable{ name: a, bits: 32 });
1413
1414        assert!(func.resolve_indirect_jumps(|x| {
1415            if *x == Variable::new(a, 32).unwrap() {
1416                vec![
1417                    Constant::new(6,32).unwrap(),
1418                    Constant::new(11,32).unwrap(),
1419                    Constant::new(12,32).unwrap(),
1420                    Constant::new(13,32).unwrap(),
1421                ].into()
1422            } else {
1423                Default::default()
1424            }
1425        }));
1426        assert!(func.extend::<TestArch>((), &reg).is_ok());
1427
1428        assert_eq!(func.cflow_graph().node_count(), 5);
1429        assert_eq!(func.cflow_graph().edge_count(), 4);
1430
1431        for n in func.cflow_graph().node_indices() {
1432            match func.cflow_graph.node_weight(n) {
1433                Some(&CfgNode::BasicBlock(bb)) => {
1434                    assert!(
1435                        func.basic_block(bb).area == (0..6).into() ||
1436                        func.basic_block(bb).area == (11..12).into() ||
1437                        func.basic_block(bb).area == (12..13).into() ||
1438                        func.basic_block(bb).area == (13..14).into() ||
1439                        func.basic_block(bb).area == (6..11).into());
1440                }
1441                _ => unreachable!()
1442            }
1443        }
1444    }
1445
1446    #[test]
1447    fn issue_51_treat_entry_point_as_incoming_edge() {
1448        let _ = simple_logger::init();
1449        let reg = Region::from_buf("", 16, b"AaaaAxxxJ0".to_vec(), 0, None);
1450        let func = Function::new::<TestArch>((), 4, &reg, UUID::now()).unwrap();
1451
1452        assert_eq!(func.cflow_graph.node_count(), 2);
1453        assert_eq!(func.cflow_graph.edge_count(), 2);
1454
1455        let mut bb0_vx = None;
1456        let mut bb1_vx = None;
1457
1458        for vx in func.cflow_graph.node_indices() {
1459            if let Some(&CfgNode::BasicBlock(bb_idx)) = func.cflow_graph().node_weight(vx) {
1460                let bb = func.basic_block(bb_idx);
1461                let mnes = func.mnemonics(bb_idx).collect::<Vec<_>>();
1462
1463                if bb.area.start == 0 {
1464                    assert_eq!(mnes.len(), 1);
1465                    assert_eq!(bb.area, (0..4).into());
1466                    bb0_vx = Some(vx);
1467                } else if bb.area.start == 4 {
1468                    assert_eq!(mnes.len(), 2);
1469                    assert_eq!(bb.area, (4..10).into());
1470                    bb1_vx = Some(vx);
1471                } else {
1472                    unreachable!();
1473                }
1474            } else {
1475                unreachable!();
1476            }
1477        }
1478
1479        assert!(bb0_vx.is_some() && bb1_vx.is_some());
1480        let entry_idx = func.entry_point();
1481        assert_eq!(func.basic_block(entry_idx).node, bb1_vx.unwrap());
1482        assert!(func.cflow_graph.find_edge(bb0_vx.unwrap(), bb1_vx.unwrap()).is_some());
1483        assert!(func.cflow_graph.find_edge(bb1_vx.unwrap(), bb0_vx.unwrap()).is_some());
1484    }
1485
1486    #[test]
1487    fn iter_range() {
1488        let data = b"AabcAabcAabcAabcAabcAabcAabcAabcR".to_vec();
1489        let reg = Region::from_buf("", 16, data, 0, None);
1490        let func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1491
1492        let bb_idx = func.basic_blocks().map(|x| x.0).collect::<Vec<_>>();
1493        assert_eq!(bb_idx.len(), 1);
1494        let stmts = func.statements(bb_idx[0]).collect::<Vec<_>>();
1495        assert_eq!(stmts.len(), 9);
1496
1497        let bb = func.basic_blocks().map(|x| x.1).collect::<Vec<_>>();
1498        assert_eq!(bb.len(), 1);
1499        let stmts = func.statements(bb[0]).collect::<Vec<_>>();
1500        assert_eq!(stmts.len(), 9);
1501
1502        let stmts = func.statements(..).collect::<Vec<_>>();
1503        assert_eq!(stmts.len(), 9);
1504
1505        let mne_idx = func.mnemonics(..).map(|x| x.0).collect::<Vec<_>>();
1506        assert_eq!(mne_idx.len(), 9);
1507    }
1508
1509    /*
1510     * (B0)
1511     * 0:  Mi1  ; mov i 1
1512     * 3:  Cfi0 ; cmp f i 0
1513     * 7:  Bf18 ; br f (B2)
1514     *
1515     * (B1)
1516     * 11: Aii3 ; add i i 3
1517     * 15: J22  ; jmp (B3)
1518     *
1519     * (B2)
1520     * 18: Aii3 ; add i i 3
1521     *
1522     * (B3)
1523     * 22: Ms3  ; mov s 3
1524     * 25: R    ; ret
1525     */
1526    #[test]
1527    fn iter_basic_blocks() {
1528        let _ = simple_logger::init();
1529        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
1530        let reg = Region::from_buf("", 16, data, 0, None);
1531        let func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1532        let mut fwd = func.basic_blocks().map(|(idx,_)| idx).collect::<Vec<_>>();
1533        let mut bwd = func.basic_blocks().rev().map(|(idx,_)| idx).collect::<Vec<_>>();
1534
1535        assert_eq!(fwd.len(), 4);
1536        assert_eq!(fwd[0], BasicBlockIndex::new(0));
1537        assert!(fwd[1] == BasicBlockIndex::new(1) || fwd[1] == BasicBlockIndex::new(2));
1538        assert!(fwd[2] == BasicBlockIndex::new(1) || fwd[2] == BasicBlockIndex::new(2));
1539        assert!(fwd[1] != fwd[2]);
1540        assert_eq!(fwd[3], BasicBlockIndex::new(3));
1541
1542        fwd.sort();
1543        fwd.dedup();
1544        assert_eq!(fwd.len(), 4);
1545
1546        assert_eq!(bwd.len(), 4);
1547        assert_eq!(bwd[0], BasicBlockIndex::new(3));
1548        assert!(bwd[1] == BasicBlockIndex::new(1) || bwd[1] == BasicBlockIndex::new(2));
1549        assert!(bwd[2] == BasicBlockIndex::new(1) || bwd[2] == BasicBlockIndex::new(2));
1550        assert!(bwd[1] != bwd[2]);
1551        assert_eq!(bwd[3], BasicBlockIndex::new(0));
1552
1553        bwd.sort();
1554        bwd.dedup();
1555        assert_eq!(bwd.len(), 4);
1556    }
1557
1558    /*
1559     * (B0)
1560     * 0:  Mi1  ; mov i 1
1561     * 3:  Cfi0 ; cmp f i 0
1562     * 7:  Bf18 ; br f (B2)
1563     *
1564     * (B1)
1565     * 11: Aii3 ; add i i 3
1566     * 15: J22  ; jmp (B3)
1567     *
1568     * (B2)
1569     * 18: Aii3 ; add i i 3
1570     *
1571     * (B3)
1572     * 22: Ms3  ; mov s 3
1573     * 25: R    ; ret
1574     */
1575    #[test]
1576    fn rewrite_increase() {
1577        let _ = simple_logger::init();
1578        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
1579        let reg = Region::from_buf("", 16, data, 0, None);
1580        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1581        let mut b0_idx = None;
1582        let mut b1_idx = None;
1583        let mut b2_idx = None;
1584        let mut b3_idx = None;
1585
1586        func.pack().unwrap();
1587
1588        for (idx,bb) in func.basic_blocks() {
1589            if bb.area.start == 0 {
1590                b0_idx = Some(idx);
1591            } else if bb.area.start == 11 {
1592                b1_idx = Some(idx);
1593            } else if bb.area.start == 18 {
1594                b2_idx = Some(idx);
1595            } else if bb.area.start == 22 {
1596                b3_idx = Some(idx);
1597            } else {
1598                unreachable!()
1599            }
1600        }
1601
1602        assert!(b0_idx.is_some() && b1_idx.is_some() && b2_idx.is_some() && b3_idx.is_some());
1603
1604        let _ = func.rewrite_basic_block(b2_idx.unwrap(),|stmt,_,_,_| {
1605            match stmt {
1606                &mut Statement::Expression{ op: Operation::Add(Value::Constant(ref mut a),Value::Constant(ref mut b)),.. } => {
1607                    *a = Constant::new(0xffffffff,32).unwrap();
1608                    *b = Constant::new(0x11111111,32).unwrap();
1609                }
1610                _ => {}
1611            }
1612
1613            Ok(RewriteControl::Continue)
1614        }).unwrap();
1615
1616        let b0 = func.statements(b0_idx.unwrap()).collect::<Vec<_>>();
1617        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b0[0] {} else { unreachable!() }
1618        if let &Statement::Expression{ op: Operation::LessOrEqualUnsigned(Value::Variable(_),Value::Constant(_)),.. } = &*b0[1] {} else { unreachable!() }
1619        assert_eq!(b0.len(), 2);
1620        let mne0 = func.mnemonics(b0_idx.unwrap()).collect::<Vec<_>>();
1621        assert_eq!(mne0.len(), 3);
1622
1623        let b1 = func.statements(b1_idx.unwrap()).collect::<Vec<_>>();
1624        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b1[0] {} else { unreachable!() }
1625        assert_eq!(b1.len(), 1);
1626        let mne1 = func.mnemonics(b1_idx.unwrap()).collect::<Vec<_>>();
1627        assert_eq!(mne1.len(), 2);
1628
1629        let b2 = func.statements(b2_idx.unwrap()).collect::<Vec<_>>();
1630        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b2[0] {} else { unreachable!() }
1631        assert_eq!(b2.len(), 1);
1632        let mne2 = func.mnemonics(b2_idx.unwrap()).collect::<Vec<_>>();
1633        assert_eq!(mne2.len(), 1);
1634
1635        let b3 = func.statements(b3_idx.unwrap()).collect::<Vec<_>>();
1636        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b3[0] {} else { unreachable!() }
1637        assert_eq!(b3.len(), 2);
1638        let mne3 = func.mnemonics(b3_idx.unwrap()).collect::<Vec<_>>();
1639        assert_eq!(mne3.len(), 2);
1640    }
1641
1642    /*
1643     * (B0)
1644     * 0:  Mi1  ; mov i 1
1645     * 3:  Cfi0 ; cmp f i 0
1646     * 7:  Bf18 ; br f (B2)
1647     *
1648     * (B1)
1649     * 11: Aii3 ; add i i 3
1650     * 15: J22  ; jmp (B3)
1651     *
1652     * (B2)
1653     * 18: Aii3 ; add i i 3
1654     *
1655     * (B3)
1656     * 22: Ms3  ; mov s 3
1657     * 25: R    ; ret
1658     */
1659    #[test]
1660    fn rewrite_rename() {
1661        let _ = simple_logger::init();
1662        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
1663        let reg = Region::from_buf("", 16, data, 0, None);
1664        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1665        let mut b0_idx = None;
1666        let mut b1_idx = None;
1667        let mut b2_idx = None;
1668        let mut b3_idx = None;
1669
1670        for (idx,bb) in func.basic_blocks() {
1671            if bb.area.start == 0 {
1672                b0_idx = Some(idx);
1673            } else if bb.area.start == 11 {
1674                b1_idx = Some(idx);
1675            } else if bb.area.start == 18 {
1676                b2_idx = Some(idx);
1677            } else if bb.area.start == 22 {
1678                b3_idx = Some(idx);
1679            } else {
1680                unreachable!()
1681            }
1682        }
1683
1684        assert!(b0_idx.is_some() && b1_idx.is_some() && b2_idx.is_some() && b3_idx.is_some());
1685        fn f(stmt: &mut Statement, names: &mut Names,_: &mut Strings, _: &mut Segments) -> ::Result<::RewriteControl> {
1686            match stmt {
1687                &mut Statement::Expression{ result: Variable{ ref mut name,.. },.. } => {
1688                    let new_name = names.value(*name)?.clone();
1689                    let new_name = Name::new(new_name.base().to_string().to_uppercase().into(),new_name.subscript());
1690                    *name = names.insert(&new_name);
1691                }
1692                _ => {}
1693            }
1694
1695            Ok(RewriteControl::Continue)
1696        }
1697
1698        let _ = func.rewrite_basic_block(b0_idx.unwrap(),&f).unwrap();
1699        let _ = func.rewrite_basic_block(b1_idx.unwrap(),&f).unwrap();
1700        let _ = func.rewrite_basic_block(b2_idx.unwrap(),&f).unwrap();
1701        let _ = func.rewrite_basic_block(b3_idx.unwrap(),&f).unwrap();
1702
1703        let b0 = func.statements(b0_idx.unwrap()).collect::<Vec<_>>();
1704        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b0[0] {} else { unreachable!() }
1705        if let &Statement::Expression{ op: Operation::LessOrEqualUnsigned(Value::Variable(_),Value::Constant(_)),.. } = &*b0[1] {} else { unreachable!() }
1706        assert_eq!(b0.len(), 2);
1707
1708        let b1 = func.statements(b1_idx.unwrap()).collect::<Vec<_>>();
1709        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b1[0] {} else { unreachable!() }
1710        assert_eq!(b1.len(), 1);
1711
1712        let b2 = func.statements(b2_idx.unwrap()).collect::<Vec<_>>();
1713        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b2[0] {} else { unreachable!() }
1714        assert_eq!(b2.len(), 1);
1715
1716        let b3 = func.statements(b3_idx.unwrap()).collect::<Vec<_>>();
1717        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b3[0] {} else { unreachable!() }
1718        assert_eq!(b3.len(), 2);
1719
1720        for stmt in func.statements(..) {
1721            match &*stmt {
1722                &Statement::Expression{ result: Variable{ ref name,.. },.. } => {
1723                    assert!(func.names.value(*name).unwrap().base().chars().all(|x| x.is_uppercase()));
1724                }
1725                _ => {}
1726            }
1727        }
1728    }
1729
1730    /*
1731     * (B0)
1732     * 0:  Mi1  ; mov i 1
1733     * 3:  Cfi0 ; cmp f i 0
1734     * 7:  Bf18 ; br f (B2)
1735     *
1736     * (B1)
1737     *          ; test
1738     * 11: Aii3 ; add i i 3
1739     * 15: J22  ; jmp (B3)
1740     *
1741     * (B2)
1742     * 18: Ai2i ; add i 2 i
1743     *
1744     * (B3)
1745     * 22: Ms3  ; mov s 3
1746     * 25: R    ; ret
1747     */
1748    #[test]
1749    fn rewrite_prepend_mnemonic_unpacked() {
1750        let _ = simple_logger::init();
1751        let data = b"Mi1Cfi0Bf18Aii3J22Ai2iMs3R".to_vec();
1752        let reg = Region::from_buf("", 16, data, 0, None);
1753        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1754        let mut b0_idx = None;
1755        let mut b1_idx = None;
1756        let mut b2_idx = None;
1757        let mut b3_idx = None;
1758
1759        for (idx,bb) in func.basic_blocks() {
1760            if bb.area.start == 0 {
1761                b0_idx = Some(idx);
1762            } else if bb.area.start == 11 {
1763                b1_idx = Some(idx);
1764            } else if bb.area.start == 18 {
1765                b2_idx = Some(idx);
1766            } else if bb.area.start == 22 {
1767                b3_idx = Some(idx);
1768            } else {
1769                unreachable!()
1770            }
1771        }
1772
1773        assert!(b0_idx.is_some() && b1_idx.is_some() && b2_idx.is_some() && b3_idx.is_some());
1774
1775        let x = func.names.insert(&Name::new("x".into(),None));
1776        let stmts = vec![
1777            Statement::Expression{
1778                op: Operation::And(Value::val(42,32).unwrap(),Value::var(x,32).unwrap()),
1779                result: Variable::new(x,32).unwrap()
1780            },
1781            Statement::Expression{
1782                op: Operation::Subtract(Value::val(42,32).unwrap(),Value::var(x,32).unwrap()),
1783                result: Variable::new(x,32).unwrap()
1784            },
1785        ];
1786        debug!("{:?}",func);
1787
1788        {
1789            debug!("pre-read bb 0");
1790            let b0 = func.basic_block(b0_idx.unwrap());
1791            debug!("b0: {:?}",b0);
1792
1793            debug!("pre-read bb 1");
1794            let b1 = func.basic_block(b1_idx.unwrap());
1795            debug!("b1: {:?}",b1);
1796        }
1797
1798        let test = func.strings.insert(&"test".into());
1799        let _ = func.insert_mnemonic(b1_idx.unwrap(),1,test,SmallVec::default(),stmts).unwrap();
1800        debug!("{:?}",func);
1801
1802        debug!("read bb 0");
1803        let b0 = func.statements(b0_idx.unwrap()).collect::<Vec<_>>();
1804        assert_eq!(b0.len(), 2);
1805        debug!("b0: {:?}",b0);
1806        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b0[0] {} else { unreachable!() }
1807        if let &Statement::Expression{ op: Operation::LessOrEqualUnsigned(Value::Variable(_),Value::Constant(_)),.. } = &*b0[1] {} else { unreachable!() }
1808        let mne0 = func.mnemonics(b0_idx.unwrap()).collect::<Vec<_>>();
1809        assert_eq!(mne0.len(), 3);
1810        let b0 = func.basic_block(b0_idx.unwrap());
1811        debug!("b0: {:?}",b0);
1812        assert_eq!(func.mnemonic(b0.mnemonics.start).area.start, 0);
1813        assert_eq!(func.mnemonic(MnemonicIndex::new(b0.mnemonics.start.index() + 1)).area.start, 3);
1814        assert_eq!(func.mnemonic(MnemonicIndex::new(b0.mnemonics.start.index() + 2)).area.start, 7);
1815
1816        debug!("read bb 1");
1817        let b1 = func.statements(b1_idx.unwrap()).collect::<Vec<_>>();
1818        debug!("b1: {:?}",b1);
1819        assert_eq!(b1.len(), 3);
1820        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b1[0] {} else { unreachable!() }
1821        if let &Statement::Expression{ op: Operation::And(Value::Constant(_),Value::Variable(_)),.. } = &*b1[1] {} else { unreachable!() }
1822        if let &Statement::Expression{ op: Operation::Subtract(Value::Constant(_),Value::Variable(_)),.. } = &*b1[2] {} else { unreachable!() }
1823        let mne1 = func.mnemonics(b1_idx.unwrap()).collect::<Vec<_>>();
1824        assert_eq!(mne1.len(), 3);
1825        let b1 = func.basic_block(b1_idx.unwrap());
1826        debug!("b1: {:?}",b1);
1827        assert_eq!(func.mnemonic(b1.mnemonics.start).area.start, 11);
1828        assert_eq!(func.mnemonic(MnemonicIndex::new(b1.mnemonics.start.index() + 1)).area.start, 15);
1829        assert_eq!(func.mnemonic(MnemonicIndex::new(b1.mnemonics.start.index() + 2)).area.start, 15);
1830
1831        debug!("read bb 2");
1832        let b2 = func.statements(b2_idx.unwrap()).collect::<Vec<_>>();
1833        assert_eq!(b2.len(), 1);
1834        if let &Statement::Expression{ op: Operation::Add(Value::Constant(_),Value::Variable(_)),.. } = &*b2[0] {} else { unreachable!() }
1835        let mne2 = func.mnemonics(b2_idx.unwrap()).collect::<Vec<_>>();
1836        assert_eq!(mne2.len(), 1);
1837        let b2 = func.basic_block(b2_idx.unwrap());
1838        debug!("b2: {:?}",b2);
1839
1840        debug!("read bb 3");
1841        let b3 = func.statements(b3_idx.unwrap()).collect::<Vec<_>>();
1842        assert_eq!(b3.len(), 2);
1843        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b3[0] {} else { unreachable!() }
1844        let mne3 = func.mnemonics(b3_idx.unwrap()).collect::<Vec<_>>();
1845        assert_eq!(mne3.len(), 2);
1846        let b3 = func.basic_block(b3_idx.unwrap());
1847        debug!("b3: {:?}",b3);
1848    }
1849
1850    #[test]
1851    fn rewrite_remove_mnemonic() {
1852        let _ = simple_logger::init();
1853        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
1854        let reg = Region::from_buf("", 16, data, 0, None);
1855        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1856        let mut b0_idx = None;
1857        let mut b1_idx = None;
1858        let mut b2_idx = None;
1859        let mut b3_idx = None;
1860
1861        for (idx,bb) in func.basic_blocks() {
1862            if bb.area.start == 0 {
1863                b0_idx = Some(idx);
1864            } else if bb.area.start == 11 {
1865                b1_idx = Some(idx);
1866            } else if bb.area.start == 18 {
1867                b2_idx = Some(idx);
1868            } else if bb.area.start == 22 {
1869                b3_idx = Some(idx);
1870            } else {
1871                unreachable!()
1872            }
1873        }
1874
1875        assert!(b0_idx.is_some() && b1_idx.is_some() && b2_idx.is_some() && b3_idx.is_some());
1876
1877        debug!("{:?}",func);
1878        let _ = func.drop_first_mnemonic(b1_idx.unwrap()).unwrap();
1879        debug!("{:?}",func);
1880
1881        let b0 = func.statements(b0_idx.unwrap()).collect::<Vec<_>>();
1882        assert_eq!(b0.len(), 2);
1883        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b0[0] {} else { unreachable!() }
1884        if let &Statement::Expression{ op: Operation::LessOrEqualUnsigned(Value::Variable(_),Value::Constant(_)),.. } = &*b0[1] {} else { unreachable!() }
1885
1886        let b1 = func.statements(b1_idx.unwrap()).collect::<Vec<_>>();
1887        assert_eq!(b1.len(), 0);
1888
1889        let b2 = func.statements(b2_idx.unwrap()).collect::<Vec<_>>();
1890        assert_eq!(b2.len(), 1);
1891        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b2[0] {} else { unreachable!() }
1892
1893        let b3 = func.statements(b3_idx.unwrap()).collect::<Vec<_>>();
1894        assert_eq!(b3.len(), 2);
1895        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b3[0] {} else { unreachable!() }
1896    }
1897
1898    #[test]
1899    fn remove_mnemonic_mid() {
1900        let _ = simple_logger::init();
1901        let data = b"Mi1Mi2Cfi0Bf21Aii3J25Aii3Ms3R".to_vec();
1902        let reg = Region::from_buf("", 16, data, 0, None);
1903        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1904        let mut b0_idx = None;
1905        let mut b1_idx = None;
1906        let mut b2_idx = None;
1907        let mut b3_idx = None;
1908
1909        for (idx,bb) in func.basic_blocks() {
1910            if bb.area.start == 0 {
1911                b0_idx = Some(idx);
1912            } else if bb.area.start == 14 {
1913                b1_idx = Some(idx);
1914            } else if bb.area.start == 21 {
1915                b2_idx = Some(idx);
1916            } else if bb.area.start == 25 {
1917                b3_idx = Some(idx);
1918            } else {
1919                unreachable!()
1920            }
1921        }
1922
1923        assert!(b0_idx.is_some() && b1_idx.is_some() && b2_idx.is_some() && b3_idx.is_some());
1924
1925        debug!("{:?}",func);
1926        let _ = func.remove_mnemonic(MnemonicIndex::new(1)).unwrap();
1927        debug!("{:?}",func);
1928
1929        let b0 = func.statements(b0_idx.unwrap()).collect::<Vec<_>>();
1930        assert_eq!(func.mnemonics(b0_idx.unwrap()).count(), 3);
1931        assert_eq!(b0.len(), 2);
1932        if let &Statement::Expression{ op: Operation::Move(Value::Constant(Constant{ value: 1,.. })),.. } = &*b0[0] {} else { unreachable!() }
1933        if let &Statement::Expression{ op: Operation::LessOrEqualUnsigned(Value::Variable(_),Value::Constant(_)),.. } = &*b0[1] {} else { unreachable!() }
1934
1935        let b1 = func.statements(b1_idx.unwrap()).collect::<Vec<_>>();
1936        assert_eq!(b1.len(), 1);
1937        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b1[0] {} else { unreachable!() }
1938
1939        let b2 = func.statements(b2_idx.unwrap()).collect::<Vec<_>>();
1940        assert_eq!(b2.len(), 1);
1941        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b2[0] {} else { unreachable!() }
1942
1943        let b3 = func.statements(b3_idx.unwrap()).collect::<Vec<_>>();
1944        assert_eq!(b3.len(), 2);
1945        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b3[0] {} else { unreachable!() }
1946    }
1947
1948    #[test]
1949    fn packed_iter_range() {
1950        let data = b"AabcAabcAabcAabcAabcAabcAabcAabcR".to_vec();
1951        let reg = Region::from_buf("", 16, data, 0, None);
1952        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1953
1954        func.pack().unwrap();
1955
1956        let bb_idx = func.basic_blocks().map(|x| x.0).collect::<Vec<_>>();
1957        assert_eq!(bb_idx.len(), 1);
1958        let stmts = func.statements(bb_idx[0]).collect::<Vec<_>>();
1959        assert_eq!(stmts.len(), 9);
1960
1961        let bb = func.basic_blocks().map(|x| x.1).collect::<Vec<_>>();
1962        assert_eq!(bb.len(), 1);
1963        let stmts = func.statements(bb[0]).collect::<Vec<_>>();
1964        assert_eq!(stmts.len(), 9);
1965
1966        let stmts = func.statements(..).collect::<Vec<_>>();
1967        assert_eq!(stmts.len(), 9);
1968
1969        let mne_idx = func.mnemonics(..).map(|x| x.0).collect::<Vec<_>>();
1970        assert_eq!(mne_idx.len(), 9);
1971    }
1972
1973    #[test]
1974    fn rewrite_prepend_mnemonic_packed() {
1975        let _ = simple_logger::init();
1976        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
1977        let reg = Region::from_buf("", 16, data, 0, None);
1978        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
1979        let mut b0_idx = None;
1980        let mut b1_idx = None;
1981        let mut b2_idx = None;
1982        let mut b3_idx = None;
1983
1984        func.pack().unwrap();
1985
1986        for (idx,bb) in func.basic_blocks() {
1987            if bb.area.start == 0 {
1988                b0_idx = Some(idx);
1989            } else if bb.area.start == 11 {
1990                b1_idx = Some(idx);
1991            } else if bb.area.start == 18 {
1992                b2_idx = Some(idx);
1993            } else if bb.area.start == 22 {
1994                b3_idx = Some(idx);
1995            } else {
1996                unreachable!()
1997            }
1998        }
1999
2000        assert!(b0_idx.is_some() && b1_idx.is_some() && b2_idx.is_some() && b3_idx.is_some());
2001
2002        let x = func.names.insert(&Name::new("x".into(),None));
2003        let stmts = vec![
2004            Statement::Expression{
2005                op: Operation::And(Value::val(42,32).unwrap(),Value::var(x,32).unwrap()),
2006                result: Variable::new(x,32).unwrap()
2007            },
2008            Statement::Expression{
2009                op: Operation::Subtract(Value::val(42,32).unwrap(),Value::var(x,32).unwrap()),
2010                result: Variable::new(x,32).unwrap()
2011            },
2012        ];
2013        debug!("{:?}",func);
2014        let test = func.strings.insert(&"test".into());
2015        let _ = func.insert_mnemonic(b1_idx.unwrap(),0,test,SmallVec::default(),stmts).unwrap();
2016        debug!("{:?}",func);
2017
2018        debug!("read bb 0");
2019        let b0 = func.statements(b0_idx.unwrap()).collect::<Vec<_>>();
2020        assert_eq!(b0.len(), 2);
2021        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b0[0] {} else { unreachable!() }
2022        if let &Statement::Expression{ op: Operation::LessOrEqualUnsigned(Value::Variable(_),Value::Constant(_)),.. } = &*b0[1] {} else { unreachable!() }
2023
2024        debug!("read bb 1");
2025        let b1 = func.statements(b1_idx.unwrap()).collect::<Vec<_>>();
2026        debug!("{:?}",b1);
2027        assert_eq!(b1.len(), 3);
2028        if let &Statement::Expression{ op: Operation::And(Value::Constant(_),Value::Variable(_)),.. } = &*b1[0] {} else { unreachable!() }
2029        if let &Statement::Expression{ op: Operation::Subtract(Value::Constant(_),Value::Variable(_)),.. } = &*b1[1] {} else { unreachable!() }
2030        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b1[2] {} else { unreachable!() }
2031
2032        debug!("read bb 2");
2033        let b2 = func.statements(b2_idx.unwrap()).collect::<Vec<_>>();
2034        assert_eq!(b2.len(), 1);
2035        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b2[0] {} else { unreachable!() }
2036
2037        debug!("read bb 3");
2038        let b3 = func.statements(b3_idx.unwrap()).collect::<Vec<_>>();
2039        assert_eq!(b3.len(), 2);
2040        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b3[0] {} else { unreachable!() }
2041    }
2042
2043    #[test]
2044    fn rewrite_remove_mnemonic_packed() {
2045        let _ = simple_logger::init();
2046        let reg = Region::from_buf("", 16, b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec(), None, None);
2047        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
2048        let mut b0_idx = None;
2049        let mut b1_idx = None;
2050        let mut b2_idx = None;
2051        let mut b3_idx = None;
2052
2053        func.pack().unwrap();
2054
2055        for (idx,bb) in func.basic_blocks() {
2056            if bb.area.start == 0 {
2057                b0_idx = Some(idx);
2058            } else if bb.area.start == 11 {
2059                b1_idx = Some(idx);
2060            } else if bb.area.start == 18 {
2061                b2_idx = Some(idx);
2062            } else if bb.area.start == 22 {
2063                b3_idx = Some(idx);
2064            } else {
2065                unreachable!()
2066            }
2067        }
2068
2069        assert!(b0_idx.is_some() && b1_idx.is_some() && b2_idx.is_some() && b3_idx.is_some());
2070
2071        debug!("{:?}",func);
2072        let _ = func.drop_first_mnemonic(b1_idx.unwrap()).unwrap();
2073        debug!("{:?}",func);
2074
2075        let b0 = func.statements(b0_idx.unwrap()).collect::<Vec<_>>();
2076        assert_eq!(b0.len(), 2);
2077        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b0[0] {} else { unreachable!() }
2078        if let &Statement::Expression{ op: Operation::LessOrEqualUnsigned(Value::Variable(_),Value::Constant(_)),.. } = &*b0[1] {} else { unreachable!() }
2079
2080        let b1 = func.statements(b1_idx.unwrap()).collect::<Vec<_>>();
2081        assert_eq!(b1.len(), 0);
2082
2083        let b2 = func.statements(b2_idx.unwrap()).collect::<Vec<_>>();
2084        assert_eq!(b2.len(), 1);
2085        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b2[0] {} else { unreachable!() }
2086
2087        let b3 = func.statements(b3_idx.unwrap()).collect::<Vec<_>>();
2088        assert_eq!(b3.len(), 2);
2089        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b3[0] {} else { unreachable!() }
2090    }
2091
2092    #[test]
2093    fn remove_mnemonic_mid_packed() {
2094        let _ = simple_logger::init();
2095        let data = b"Mi1Mi2Cfi0Bf21Aii3J25Aii3Ms3R".to_vec();
2096        let reg = Region::from_buf("", 16, data, 0, None);
2097        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
2098        let mut b0_idx = None;
2099        let mut b1_idx = None;
2100        let mut b2_idx = None;
2101        let mut b3_idx = None;
2102
2103        func.pack().unwrap();
2104
2105        for (idx,bb) in func.basic_blocks() {
2106            if bb.area.start == 0 {
2107                b0_idx = Some(idx);
2108            } else if bb.area.start == 14 {
2109                b1_idx = Some(idx);
2110            } else if bb.area.start == 21 {
2111                b2_idx = Some(idx);
2112            } else if bb.area.start == 25 {
2113                b3_idx = Some(idx);
2114            } else {
2115                unreachable!()
2116            }
2117        }
2118
2119        assert!(b0_idx.is_some() && b1_idx.is_some() && b2_idx.is_some() && b3_idx.is_some());
2120
2121        debug!("{:?}",func);
2122        let _ = func.remove_mnemonic(MnemonicIndex::new(1)).unwrap();
2123        debug!("{:?}",func);
2124
2125        let b0 = func.statements(b0_idx.unwrap()).collect::<Vec<_>>();
2126        assert_eq!(func.mnemonics(b0_idx.unwrap()).count(), 3);
2127        assert_eq!(b0.len(), 2);
2128        if let &Statement::Expression{ op: Operation::Move(Value::Constant(Constant{ value: 1,.. })),.. } = &*b0[0] {} else { unreachable!() }
2129        if let &Statement::Expression{ op: Operation::LessOrEqualUnsigned(Value::Variable(_),Value::Constant(_)),.. } = &*b0[1] {} else { unreachable!() }
2130
2131        let b1 = func.statements(b1_idx.unwrap()).collect::<Vec<_>>();
2132        assert_eq!(b1.len(), 1);
2133        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b1[0] {} else { unreachable!() }
2134
2135        let b2 = func.statements(b2_idx.unwrap()).collect::<Vec<_>>();
2136        assert_eq!(b2.len(), 1);
2137        if let &Statement::Expression{ op: Operation::Add(Value::Variable(_),Value::Constant(_)),.. } = &*b2[0] {} else { unreachable!() }
2138
2139        let b3 = func.statements(b3_idx.unwrap()).collect::<Vec<_>>();
2140        assert_eq!(b3.len(), 2);
2141        if let &Statement::Expression{ op: Operation::Move(Value::Constant(_)),.. } = &*b3[0] {} else { unreachable!() }
2142    }
2143
2144    #[test]
2145    fn shift_areas() {
2146        let mut func = Function::facade();
2147        let uu0 = UUID::now();
2148        let uu1 = UUID::now();
2149        let uu2 = UUID::now();
2150
2151        let mne0 = Mnemonic{ area: (0..1).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2152        let mne1 = Mnemonic{ area: (1..2).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2153        let mne2 = Mnemonic{ area: (2..3).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2154        let mne22 = Mnemonic{ area: Area{ start: 3, end: 3, offset_start: 0, offset_end: 1 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2155        let mne3 = Mnemonic{ area: Area{ start: 3, end: 4, offset_start: 1, offset_end: 0 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2156        let mne4 = Mnemonic{ area: (4..5).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2157        let mne5 = Mnemonic{ area: (5..6).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2158
2159        let bb0 = BasicBlock{ uuid: uu0.clone(), area: (0..2).into(), node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(0)..MnemonicIndex::new(2), statements: 0..0 };
2160        let bb1 = BasicBlock{ uuid: uu1.clone(), area: (2..4).into(), node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(2)..MnemonicIndex::new(5), statements: 0..0 };
2161        let bb2 = BasicBlock{ uuid: uu2.clone(), area: (4..6).into(), node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(5)..MnemonicIndex::new(7), statements: 0..0 };
2162
2163        func.mnemonics = vec![mne0,mne1,mne2,mne22,mne3,mne4,mne5];
2164        func.basic_blocks = vec![bb0,bb1,bb2];
2165
2166        func.shift_areas(BasicBlockIndex::new(1), 0, 2);
2167
2168        let mne0 = Mnemonic{ area: (0..1).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2169        let mne1 = Mnemonic{ area: (1..2).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2170        let mne2 = Mnemonic{ area: Area{ start: 2, end: 3, offset_start: 2, offset_end: 0 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2171        let mne22 = Mnemonic{ area: Area{ start: 3, end: 3, offset_start: 0, offset_end: 1 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2172        let mne3 = Mnemonic{ area: Area{ start: 3, end: 4, offset_start: 1, offset_end: 0 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2173        let mne4 = Mnemonic{ area: (4..5).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2174        let mne5 = Mnemonic{ area: (5..6).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2175
2176        let bb0 = BasicBlock{ uuid: uu0.clone(), area: (0..2).into(), node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(0)..MnemonicIndex::new(2), statements: 0..0 };
2177        let bb1 = BasicBlock{ uuid: uu1.clone(), area: (2..4).into(), node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(2)..MnemonicIndex::new(5), statements: 0..0 };
2178        let bb2 = BasicBlock{ uuid: uu2.clone(), area: (4..6).into(), node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(5)..MnemonicIndex::new(7), statements: 0..0 };
2179
2180
2181        assert_eq!(func.mnemonics, vec![mne0,mne1,mne2,mne22,mne3,mne4,mne5]);
2182        assert_eq!(func.basic_blocks, vec![bb0,bb1,bb2]);
2183    }
2184
2185    #[test]
2186    fn shift_areas_multiple_block() {
2187        let mut func = Function::facade();
2188        let uu0 = UUID::now();
2189        let uu1 = UUID::now();
2190        let uu2 = UUID::now();
2191
2192        let mne0 = Mnemonic{ area: (0..2).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2193        let mne1 = Mnemonic{ area: Area{ start: 2, end: 2, offset_start: 0, offset_end: 1 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2194        let mne2 = Mnemonic{ area: Area{ start: 2, end: 2, offset_start: 1, offset_end: 2 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2195        let mne3 = Mnemonic{ area: Area{ start: 2, end: 2, offset_start: 2, offset_end: 3 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2196        let mne4 = Mnemonic{ area: Area{ start: 2, end: 5, offset_start: 3, offset_end: 0 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2197        let mne5 = Mnemonic{ area: (5..6).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2198
2199        let bb0 = BasicBlock{ uuid: uu0.clone(), area:  Area{ start: 0, end: 2, offset_start: 0, offset_end: 1 }, node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(0)..MnemonicIndex::new(2), statements: 0..0 };
2200        let bb1 = BasicBlock{ uuid: uu1.clone(), area:  Area{ start: 2, end: 2, offset_start: 1, offset_end: 3 }, node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(2)..MnemonicIndex::new(4), statements: 0..0 };
2201        let bb2 = BasicBlock{ uuid: uu2.clone(), area:  Area{ start: 2, end: 6, offset_start: 3, offset_end: 0 }, node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(4)..MnemonicIndex::new(6), statements: 0..0 };
2202
2203        func.mnemonics = vec![mne0,mne1,mne2,mne3,mne4,mne5];
2204        func.basic_blocks = vec![bb0,bb1,bb2];
2205        func.shift_areas(BasicBlockIndex::new(0), 1, 2);
2206
2207        let mne0 = Mnemonic{ area: (0..2).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2208        let mne1 = Mnemonic{ area: Area{ start: 2, end: 2, offset_start: 2, offset_end: 3 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2209        let mne2 = Mnemonic{ area: Area{ start: 2, end: 2, offset_start: 3, offset_end: 4 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2210        let mne3 = Mnemonic{ area: Area{ start: 2, end: 2, offset_start: 4, offset_end: 5 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2211        let mne4 = Mnemonic{ area: Area{ start: 2, end: 5, offset_start: 5, offset_end: 0 }, opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2212        let mne5 = Mnemonic{ area: (5..6).into(), opcode: StrRef::new(0), operands: SmallVec::default(), statement_count: 0 };
2213
2214        let bb0 = BasicBlock{ uuid: uu0.clone(), area:  Area{ start: 0, end: 2, offset_start: 0, offset_end: 3 }, node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(0)..MnemonicIndex::new(2), statements: 0..0 };
2215        let bb1 = BasicBlock{ uuid: uu1.clone(), area:  Area{ start: 2, end: 2, offset_start: 3, offset_end: 5 }, node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(2)..MnemonicIndex::new(4), statements: 0..0 };
2216        let bb2 = BasicBlock{ uuid: uu2.clone(), area:  Area{ start: 2, end: 6, offset_start: 5, offset_end: 0 }, node: NodeIndex::new(0), mnemonics: MnemonicIndex::new(4)..MnemonicIndex::new(6), statements: 0..0 };
2217
2218        assert_eq!(func.mnemonics, vec![mne0,mne1,mne2,mne3,mne4,mne5]);
2219        assert_eq!(func.basic_blocks, vec![bb0,bb1,bb2]);
2220    }
2221
2222    /*
2223     * (B0)
2224     * 0:  Mi1  ; mov i 1
2225     * 3:  Cfi0 ; cmp f i 0
2226     * 7:  Bf18 ; br f (B2)
2227     *
2228     * (B1)
2229     * 11: Aii3 ; add i i 3
2230     * 15: J22  ; jmp (B3)
2231     *
2232     * (B2)
2233     * 18: Aii3 ; add i i 3
2234     *
2235     * (B3)
2236     * 22: Ms3  ; mov s 3
2237     * 25: R    ; ret
2238     */
2239    #[test]
2240    fn delete_single_dead_end() {
2241        let _ = simple_logger::init();
2242        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2243        let reg = Region::from_buf("", 16, data, 0, None);
2244        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
2245
2246        func.retain_basic_blocks(|func, bb| {
2247            let bb = func.basic_block(bb);
2248            bb.area.start != 22
2249        }).unwrap();
2250
2251        assert_eq!(func.cflow_graph().node_count(), 3);
2252    }
2253
2254    #[test]
2255    fn delete_middle() {
2256        let _ = simple_logger::init();
2257        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2258        let reg = Region::from_buf("", 16, data, 0, None);
2259        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
2260
2261        func.retain_basic_blocks(|func, bb| {
2262            let bb = func.basic_block(bb);
2263            bb.area.start != 11
2264        }).unwrap();
2265
2266        assert_eq!(func.cflow_graph().node_count(), 3);
2267    }
2268
2269    #[test]
2270    fn delete_multiple() {
2271        let _ = simple_logger::init();
2272        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2273        let reg = Region::from_buf("", 16, data, 0, None);
2274        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
2275
2276        func.retain_basic_blocks(|func, bb| {
2277            let bb = func.basic_block(bb);
2278            bb.area.start == 11
2279        }).unwrap();
2280
2281        assert_eq!(func.cflow_graph().node_count(), 2);
2282    }
2283
2284    #[test]
2285    fn delete_all() {
2286        let _ = simple_logger::init();
2287        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2288        let reg = Region::from_buf("", 16, data, 0, None);
2289        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
2290
2291        func.retain_basic_blocks(|_, _| {
2292            false
2293        }).unwrap();
2294
2295        assert_eq!(func.cflow_graph().node_count(), 1);
2296    }
2297
2298    #[test]
2299    fn delete_none() {
2300        let _ = simple_logger::init();
2301        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2302        let reg = Region::from_buf("", 16, data, 0, None);
2303        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
2304
2305        func.retain_basic_blocks(|_, _| {
2306            true
2307        }).unwrap();
2308
2309        assert_eq!(func.cflow_graph().node_count(), 4);
2310    }
2311
2312    #[test]
2313    fn delete_twice() {
2314        let _ = simple_logger::init();
2315        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2316        let reg = Region::from_buf("", 16, data, 0, None);
2317        let mut func = Function::new::<TestArch>((), 0, &reg, UUID::now()).unwrap();
2318
2319        func.retain_basic_blocks(|func, bb| {
2320            let bb = func.basic_block(bb);
2321            bb.area.start != 11
2322        }).unwrap();
2323
2324        func.retain_basic_blocks(|func, bb| {
2325            let bb = func.basic_block(bb);
2326            bb.area.start != 11
2327        }).unwrap();
2328
2329        assert_eq!(func.cflow_graph().node_count(), 3);
2330    }
2331
2332    /*
2333     * (B0)
2334     * 0:  Mi1  ; mov i 1
2335     * 3:  Cfi0 ; cmp f i 0
2336     * 7:  Bf18 ; br f (B2)
2337     *
2338     * (B1)
2339     * 11: Aii3 ; add i i 3
2340     * 15: J22  ; jmp (B3)
2341     *
2342     * (B2)
2343     * 18: Aii3 ; add i i 3
2344     *
2345     * (B3)
2346     * 22: Ms3  ; mov s 3
2347     * 25: R    ; ret
2348     */
2349    #[test]
2350    fn recover_cfg() {
2351        let _ = simple_logger::init();
2352        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2353        let reg = Region::from_buf("", 16, data, 0, None);
2354        let cfg = vec![
2355            (0..11, SmallVec::from_vec(vec![
2356                (Value::val(11,32).unwrap(), Guard::True),
2357                (Value::val(18,32).unwrap(), Guard::True)
2358            ])),
2359            (11..18, vec![
2360                (Value::val(22,32).unwrap(), Guard::True)
2361            ].into()),
2362            (18..22, vec![
2363                (Value::val(22,32).unwrap(), Guard::True)
2364            ].into()),
2365            (22..26, vec![].into()),
2366        ];
2367        let uu = UUID::now();
2368
2369        let func1 = Function::new::<TestArch>((), 0, &reg, uu).unwrap();
2370        let bb1a = func1.basic_blocks().map(|(_,bb)| bb.area.clone()).collect::<Vec<_>>();
2371        let bb1m = func1.basic_blocks().map(|(_,bb)| bb.mnemonics.clone()).collect::<Vec<_>>();
2372        let mne1 = func1.mnemonics(..).map(|(_,bb)| bb.area.clone()).collect::<Vec<_>>();
2373
2374        let func2 = Function::known_cflow_graph::<TestArch>((), Names::default(), cfg, 0, &reg, uu, "test".to_string()).unwrap();
2375        let bb2a = func2.basic_blocks().map(|(_,bb)| bb.area.clone()).collect::<Vec<_>>();
2376        let bb2m = func2.basic_blocks().map(|(_,bb)| bb.mnemonics.clone()).collect::<Vec<_>>();
2377        let mne2 = func2.mnemonics(..).map(|(_,bb)| bb.area.clone()).collect::<Vec<_>>();
2378
2379        assert_eq!(bb1a.len(), 4);
2380        assert_eq!(bb2a.len(), 4);
2381        assert_eq!(mne1.len(), 8);
2382        assert_eq!(mne2.len(), 8);
2383        assert_eq!(bb1a[0], bb2a[0]);
2384        assert_eq!(mne1[0..3], mne2[0..3]);
2385        assert_eq!(bb1a[3], bb2a[3]);
2386        assert_eq!(mne1[6..8], mne2[6..8]);
2387
2388        let swapped = bb1a[1] == bb2a[2] && bb1a[2] == bb2a[1];
2389        let bb11 = func1.mnemonics(bb1m[1].clone()).map(|(_,m)| m.area.clone()).collect::<Vec<_>>();
2390        let bb12 = func1.mnemonics(bb1m[2].clone()).map(|(_,m)| m.area.clone()).collect::<Vec<_>>();
2391        let bb21 = func2.mnemonics(bb2m[1].clone()).map(|(_,m)| m.area.clone()).collect::<Vec<_>>();
2392        let bb22 = func2.mnemonics(bb2m[2].clone()).map(|(_,m)| m.area.clone()).collect::<Vec<_>>();
2393
2394        if swapped {
2395            assert!(bb1a[2] == bb2a[1] && bb1a[2] == bb2a[1]);
2396            assert_eq!(bb11,bb22);
2397            assert_eq!(bb12,bb21);
2398        } else {
2399            assert!(bb1a[1] == bb2a[1] && bb1a[2] == bb2a[2]);
2400            assert_eq!(bb11,bb21);
2401            assert_eq!(bb12,bb22);
2402        }
2403    }
2404
2405    #[test]
2406    fn recover_cfg_bad_entry() {
2407        let _ = simple_logger::init();
2408        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2409        let reg = Region::from_buf("", 16, data, 0, None);
2410        let cfg = vec![
2411            (0..11, SmallVec::from_vec(vec![
2412                (Value::val(11,32).unwrap(), Guard::True),
2413                (Value::val(18,32).unwrap(), Guard::True)
2414            ])),
2415            (11..18, vec![
2416                (Value::val(22,32).unwrap(), Guard::True)
2417            ].into()),
2418            (18..22, vec![
2419                (Value::val(22,32).unwrap(), Guard::True)
2420            ].into()),
2421            (22..26, vec![].into()),
2422        ];
2423        let uu = UUID::now();
2424        assert!(Function::known_cflow_graph::<TestArch>((), Names::default(), cfg, 1, &reg, uu, "test".to_string()).is_err());
2425    }
2426
2427    #[test]
2428    fn recover_cfg_indirect_jmp() {
2429        let _ = simple_logger::init();
2430        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2431        let reg = Region::from_buf("", 16, data, 0, None);
2432        let mut names = Names::default();
2433        let a = names.var("a", None, 32).unwrap();
2434        let cfg = vec![
2435            (0..1, SmallVec::from_vec(vec![
2436                (Value::Variable(a), Guard::True),
2437            ])),
2438        ];
2439        let uu = UUID::now();
2440        assert!(Function::known_cflow_graph::<TestArch>((), Names::default(), cfg, 0, &reg, uu, "test".to_string()).is_ok());
2441    }
2442
2443    #[test]
2444    fn recover_cfg_bad_blocks() {
2445        let _ = simple_logger::init();
2446        let data = b"Mi1Cfi0Bf18Aii3J22Aii3Ms3R".to_vec();
2447        let reg = Region::from_buf("", 16, data, 0, None);
2448        let cfg = vec![
2449            (0..12, SmallVec::from_vec(vec![
2450                (Value::val(11,32).unwrap(), Guard::True),
2451                (Value::val(18,32).unwrap(), Guard::True)
2452            ])),
2453        ];
2454        let uu = UUID::now();
2455        assert!(Function::known_cflow_graph::<TestArch>((), Names::default(), cfg, 0, &reg, uu, "test".to_string()).is_err());
2456    }
2457
2458}