1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

mod hazard;
use hazard::load_graph;

mod callgraph;
use callgraph::Callgraph;

#[macro_use]
extern crate lazy_static;

#[macro_use]
extern crate cpython;

use cpython::{PyResult, PyErr};
use cpython::exc;
use petgraph::graph::NodeIndex;
use std::cell;
use std::collections::HashSet;
use std::iter::FromIterator;

// impl cpython::ToPyObject for NodeIndex {
//     type ObjectType = PyInt;
//     fn to_py_object(&self, py: Python) -> Self::ObjectType {
//         PyInt::new(py, self.index())
//     }
// }

py_class!(class HazGraph |py| {
    data callgraph: cell::RefCell<Callgraph>;

    def __new__(_cls, filename: &str) -> PyResult<HazGraph> {
        let callgraph = load_graph(filename, 0);
        let callgraph = callgraph.unwrap();
        HazGraph::create_instance(py, cell::RefCell::new(callgraph))
    }

    def resolve(&self, query: &str) -> PyResult<Vec<usize>> {
        let cg = self.callgraph(py).borrow();
        if &query[0..1] == "#" {
            match query[1..].parse::<usize>() {
                Ok(n) => {
                    Ok(vec![n])
                },
                Err(_) =>
                    Err(PyErr::new::<exc::ValueError, _>(py, "invalid node id"))
            }
        } else {
            match cg.resolve(query) {
                None => Ok(vec![]),
                Some(matches) => Ok(matches.iter().map(|&x| x.index()).collect())
            }
        }
    }

    def callees(&self, func: usize) -> PyResult<Vec<usize>> {
        let cg = self.callgraph(py).borrow();
        let callees = cg.callees(NodeIndex::new(func));
        Ok(callees.iter().map(|&x| x.index()).collect())
    }

    def callers(&self, func: usize) -> PyResult<Vec<usize>> {
        let cg = self.callgraph(py).borrow();
        let callers = cg.callers(NodeIndex::new(func));
        Ok(callers.iter().map(|&x| x.index()).collect())
    }

    def route(&self, src: usize, goal: Vec<usize>, avoid: Vec<usize>) -> PyResult<Vec<usize>> {
        let cg = self.callgraph(py).borrow();
        let src = NodeIndex::new(src);
        let goal : Vec<NodeIndex> = goal.iter().map(|&x| NodeIndex::new(x)).collect();
        let goal = HashSet::from_iter(goal);
        let avoid : Vec<NodeIndex> = avoid.iter().map(|&x| NodeIndex::new(x)).collect();
        let avoid = HashSet::from_iter(avoid);

        match cg.any_route(src, goal, avoid) {
            None => Ok(vec![]),
            Some(route) => Ok(route.iter().map(|&x| x.index()).collect())
        }
    }

    def names(&self, func: usize) -> PyResult<Vec<String>> {
        let cg = self.callgraph(py).borrow();
        let names = cg.names(NodeIndex::new(func));
        Ok(names.iter().map(|&x| x.to_string()).collect())
    }

    // Err(PyErr::new::<exc::TypeError, _>(py, "unimplemented"))
});

py_module_initializer!(hazgraph, inithazgraph, PyInit_hazgraph, |py, m| {
    m.add(py, "__doc__", "Python wrapper for Callgraph.")?;
    m.add_class::<HazGraph>(py)?;
    Ok(())
});