pub struct CallGraph {
    g: petgraph::graph::Graph<Node, Node>,
}

fn func_of_node(hugr: &impl HugrView, node: Node) -> Option<Node> {
    let mut n = node;
    while let Some(parent) = hugr.get_parent(n) {
        if hugr.get_optype(parent).is_func_defn() {
            return Some(parent);
        }
        n = parent;
    }
    None
}

impl CallGraph {
    pub fn new(hugr: &impl HugrView) -> Self {
        let mut g: petgraph::graph::Graph<Node, Node> = Default::default();

        let node_to_cg: HashMap<_, _> = hugr
            .nodes()
            .filter(|&n| (hugr.get_optype(n).is_func_decl() || hugr.get_optype(n).is_func_defn()))
            .map(|n| (n, g.add_node(n)))
            .collect();

        for n in hugr.nodes() {
            if let Some(call) = hugr.get_optype(n).as_call() {
                if let Some(caller_func) = func_of_node(hugr, n) {
                    if let Some((callee_func, _)) =
                        hugr.single_linked_output(n, call.called_function_port())
                    {
                        g.add_edge(node_to_cg[&caller_func], node_to_cg[&callee_func], n);
                    }
                }
            }
        }

        Self { g }
    }

    pub fn iter_nonrecursive(&self) -> Option<impl Iterator<Item = (Node, Node)> + '_> {
        let funcs = petgraph::algo::toposort(&self.g, None).ok()?;

        Some(funcs.into_iter().flat_map(move |f| {
            self.g
                .edges(f)
                .map(move |e| (*e.weight(), self.g[e.target()]))
        }))
    }
}