1use 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
56pub trait IntoStatementRange {
58 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
88pub trait IntoMnemonicRange: Debug {
90 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
134pub 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#[derive(Debug, Clone, PartialEq, Eq)]
159pub enum CfgNode {
160 BasicBlock(BasicBlockIndex),
162 Value(Value),
164}
165
166#[derive(Clone, Copy, PartialEq, Eq, Debug)]
167pub enum CfgTarget {
168 Value(Value),
169 Resolved(u64, usize),
170}
171
172#[derive(Debug, Clone)]
179pub struct Function {
180 pub name: Str,
182 pub names: Names,
184 pub strings: Strings,
186 pub segments: Segments,
188 uuid: UUID,
190 region: UUID,
192 statements: Statements,
194 basic_blocks: Vec<BasicBlock>,
196 mnemonics: Vec<Mnemonic>,
198 cflow_graph: Graph<CfgNode,Guard>,
199 entry_point: BasicBlockIndex,
200}
201
202impl Function {
203 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 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 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 _ => { }
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 pub fn entry_point(&self) -> BasicBlockIndex { self.entry_point }
445
446 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 pub fn basic_blocks<'a>(&'a self) -> BasicBlockIterator<'a> {
454 BasicBlockIterator::new(self,0,self.basic_blocks.len())
455 }
456
457 pub fn cflow_graph<'a>(&'a self) -> &'a Graph<CfgNode,Guard> {
459 &self.cflow_graph
460 }
461
462 pub fn basic_block<'a>(&'a self, idx: BasicBlockIndex) -> &'a BasicBlock {
464 &self.basic_blocks[idx.index()]
465 }
466
467 pub fn mnemonic<'a>(&'a self, idx: MnemonicIndex) -> &'a Mnemonic {
469 &self.mnemonics[idx.index()]
470 }
471
472 pub fn uuid<'a>(&'a self) -> &'a UUID {
474 &self.uuid
475 }
476
477 pub fn region<'a>(&'a self) -> &'a UUID {
479 &self.region
480 }
481
482 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 pub fn first_address(&self) -> u64 {
491 self.basic_blocks[0].area().start
492 }
493
494 pub fn entry_address(&self) -> u64 {
496 let e = self.entry_point().index();
497 self.basic_blocks[e].area().start
498 }
499
500 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 pub fn cflow_graph_mut<'a>(&'a mut self) -> &'a mut Graph<CfgNode,Guard> {
510 &mut self.cflow_graph
511 }
512
513 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 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 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 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 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 pub fn retain_basic_blocks<F: (FnMut(&Function, BasicBlockIndex) -> bool)>(&mut self, mut f: F) -> Result<()> {
691 use vec_map::VecMap;
692
693 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 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 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 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 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 self.entry_point = old_bb_to_new_bb[self.entry_point.index()];
757 Ok(())
758
759 }
760
761 #[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 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 '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 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(_) => { }
907 }
908 }
909
910 Ok(())
911 }
912
913 pub fn pack(&mut self) -> Result<()> {
916 self.statements.pack(&mut self.basic_blocks,&mut self.mnemonics)
917 }
918
919 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 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 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 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 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 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 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 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 let mut new_bb = a.area.end != b.area.start || a.area.offset_end != b.area.offset_start;
1142
1143 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 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 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, ®, 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, ®, 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, ®, 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, ®, 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, ®, 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>((), ®).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, ®, 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>((), ®).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, ®, 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>((), ®).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, ®, 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, ®, 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 #[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, ®, 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 #[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, ®, 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 #[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, ®, 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 #[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, ®, 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, ®, 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, ®, 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, ®, 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, ®, 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, ®, 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, ®, 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 #[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, ®, 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, ®, 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, ®, 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, ®, 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, ®, 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, ®, 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 #[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, ®, 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, ®, 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, ®, 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, ®, 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, ®, uu, "test".to_string()).is_err());
2456 }
2457
2458}