1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

use std::collections::HashSet;
use std::collections::HashMap;
use bb_disassembler::BBDisassembler;
use basic_block::BasicBlock;

#[derive(Debug, Eq, PartialEq)]
pub struct Disassembly{
    pub offset: u64,
    pub data: Vec<u8>,
    pub bbs: HashMap<u64, BasicBlock>,
}

pub struct RecursiveDisassembler<T: BBDisassembler> {
    pub offset: u64,
    pub data: Vec<u8>,
    pub roots: HashSet<u64>,
    pub bbs: HashMap<u64, BasicBlock>,
    dis: T,
    unprocessed_roots: HashSet<u64>,
}

impl<T: BBDisassembler> RecursiveDisassembler<T> {
    pub fn new(data: Vec<u8>, offset: u64, dis: T) -> Self {
        let unprocessed_roots = HashSet::new();
        let roots = HashSet::new();
        let bbs = HashMap::new();
        return RecursiveDisassembler { offset, data, dis, roots, unprocessed_roots, bbs };
    }

    pub fn add_root(&mut self, addr: u64) {
        if !self.roots.contains(&addr) {
            self.unprocessed_roots.insert(addr);
        }
    }

    pub fn disassemble(&mut self) {
        while let Some(&root) = self.unprocessed_roots.iter().next() {
            self.unprocessed_roots.remove(&root);
            if !self.in_range(root) { continue; }
            self.roots.insert(root);
            let bb = self.dis.get_basic_block(root, self.offset, &self.data);

            for addr in bb.successors() {
                self.add_root(addr);
            }
            self.bbs.insert(root, bb);
        }
    }

    fn in_range(&self, addr: u64) -> bool {
        return (self.offset <= addr) && (self.offset + self.data.len() as u64 > addr);
    }

    pub fn into_disassembly(self) -> Disassembly{
        let RecursiveDisassembler{ offset, data, bbs, ..} = self;
        return Disassembly{offset, data, bbs}
    }
}