[Rust] Flow graph API improvements

- Fixed some misc bugs
- Added real FlowGraphEdge type
- Added accessors for node edges
- Added interaction handler to the flow graph example to dump the flow graph to stdin
- Split out the flow graph API into smaller files

4 files changed, 362 insertions, 172 deletions

diff --git a/rust/examples/flowgraph.rs b/rust/examples/flowgraph.rs
index 6666b7fb..515d52fe 100644
--- a/rust/examples/flowgraph.rs
+++ b/rust/examples/flowgraph.rs
@@ -1,9 +1,82 @@
+use binaryninja::flowgraph::edge::EdgeStyle;
+use binaryninja::flowgraph::FlowGraphNode;
+use binaryninja::function::FunctionViewType;
+use binaryninja::interaction::form::Form;
+use binaryninja::interaction::handler::{
+    register_interaction_handler, InteractionHandler, InteractionHandlerTask,
+};
+use binaryninja::interaction::{MessageBoxButtonResult, MessageBoxButtonSet, MessageBoxIcon};
 use binaryninja::{
     binary_view::{BinaryView, BinaryViewExt},
     disassembly::{DisassemblyTextLine, InstructionTextToken, InstructionTextTokenKind},
-    flowgraph::{BranchType, EdgePenStyle, EdgeStyle, FlowGraph, FlowGraphNode, ThemeColor},
+    flowgraph::{BranchType, EdgePenStyle, FlowGraph, ThemeColor},
 };
 
+pub struct GraphPrinter;
+
+impl GraphPrinter {
+    pub fn print_graph(&self, _view: &BinaryView, graph: &FlowGraph) {
+        println!("Printing flow graph:");
+        for node in &graph.nodes() {
+            // Print all disassembly lines in the node
+            println!("Node @ {:?}:", node.position());
+            println!("------------------");
+            println!("Disassembly lines:");
+            for line in &node.lines() {
+                println!("  {}", line);
+            }
+
+            // Print outgoing edges
+            println!("Outgoing edges:");
+            for edge in &node.outgoing_edges() {
+                println!("  {:?} => {:?}", edge.branch_type, edge.target.position());
+            }
+            println!("------------------");
+        }
+    }
+}
+
+impl InteractionHandler for GraphPrinter {
+    fn show_message_box(
+        &mut self,
+        _title: &str,
+        _text: &str,
+        _buttons: MessageBoxButtonSet,
+        _icon: MessageBoxIcon,
+    ) -> MessageBoxButtonResult {
+        MessageBoxButtonResult::CancelButton
+    }
+
+    fn open_url(&mut self, _url: &str) -> bool {
+        false
+    }
+
+    fn run_progress_dialog(
+        &mut self,
+        _title: &str,
+        _can_cancel: bool,
+        _task: &InteractionHandlerTask,
+    ) -> bool {
+        false
+    }
+
+    fn show_plain_text_report(&mut self, _view: &BinaryView, title: &str, contents: &str) {
+        println!("Plain text report");
+        println!("Title: {}", title);
+        println!("Contents: {}", contents);
+    }
+
+    fn show_graph_report(&mut self, view: &BinaryView, title: &str, graph: &FlowGraph) {
+        println!("Graph report");
+        println!("Title: {}", title);
+        self.print_graph(view, graph);
+    }
+
+    fn get_form_input(&mut self, _form: &mut Form) -> bool {
+        false
+    }
+}
+
 fn test_graph(view: &BinaryView) {
     let graph = FlowGraph::new();
 
@@ -15,8 +88,10 @@ fn test_graph(view: &BinaryView) {
 
     let node_a = FlowGraphNode::new(&graph);
     node_a.set_lines(disassembly_lines_a);
+    node_a.set_position(1337, 7331);
 
     let node_b = FlowGraphNode::new(&graph);
+    node_b.set_position(100, 200);
     let disassembly_lines_b = vec![DisassemblyTextLine::new(vec![
         InstructionTextToken::new("Li", InstructionTextTokenKind::Text),
         InstructionTextToken::new("ne", InstructionTextTokenKind::Text),
@@ -49,8 +124,16 @@ fn main() {
         .load("/bin/cat")
         .expect("Couldn't open `/bin/cat`");
 
-    // TODO: Register BNInteractionHandlerCallbacks with showGraphReport pointing at our function
-    // TODO: Idea: register showGraphReport that dumps a dotgraph to stdin
+    // Register the interaction handler so we can see the graph report headlessly.
+    register_interaction_handler(GraphPrinter);
 
     test_graph(&bv);
+
+    for func in bv.functions().iter().take(5) {
+        let graph = func.create_graph(FunctionViewType::MediumLevelIL, None);
+        assert!(!graph.nodes().is_empty());
+        let func_name = func.symbol().short_name();
+        let title = func_name.to_string_lossy();
+        bv.show_graph_report(&title, &graph);
+    }
 }
diff --git a/rust/src/flowgraph.rs b/rust/src/flowgraph.rs
index 976764eb..0faab0e4 100644
--- a/rust/src/flowgraph.rs
+++ b/rust/src/flowgraph.rs
@@ -14,17 +14,21 @@
 
 //! Interfaces for creating and displaying pretty CFGs in Binary Ninja.
 
-use crate::disassembly::DisassemblyTextLine;
-use crate::rc::*;
 use binaryninjacore_sys::*;
 
-use crate::basic_block::{BasicBlock, BlockContext};
-use crate::function::HighlightColor;
 use crate::high_level_il::HighLevelILFunction;
 use crate::low_level_il::RegularLowLevelILFunction;
 use crate::medium_level_il::MediumLevelILFunction;
+use crate::rc::*;
 use crate::render_layer::CoreRenderLayer;
 
+pub mod edge;
+pub mod node;
+
+pub use edge::EdgeStyle;
+pub use edge::FlowGraphEdge;
+pub use node::FlowGraphNode;
+
 pub type BranchType = BNBranchType;
 pub type EdgePenStyle = BNEdgePenStyle;
 pub type ThemeColor = BNThemeColor;
@@ -162,168 +166,3 @@ impl ToOwned for FlowGraph {
         unsafe { RefCountable::inc_ref(self) }
     }
 }
-
-#[derive(PartialEq, Eq, Hash)]
-pub struct FlowGraphNode {
-    pub(crate) handle: *mut BNFlowGraphNode,
-}
-
-impl FlowGraphNode {
-    pub(crate) unsafe fn from_raw(raw: *mut BNFlowGraphNode) -> Self {
-        Self { handle: raw }
-    }
-
-    pub(crate) unsafe fn ref_from_raw(raw: *mut BNFlowGraphNode) -> Ref<Self> {
-        Ref::new(Self { handle: raw })
-    }
-
-    pub fn new(graph: &FlowGraph) -> Ref<Self> {
-        unsafe { FlowGraphNode::ref_from_raw(BNCreateFlowGraphNode(graph.handle)) }
-    }
-
-    pub fn basic_block<C: BlockContext>(&self, context: C) -> Option<Ref<BasicBlock<C>>> {
-        let block_ptr = unsafe { BNGetFlowGraphBasicBlock(self.handle) };
-        if block_ptr.is_null() {
-            return None;
-        }
-        Some(unsafe { BasicBlock::ref_from_raw(block_ptr, context) })
-    }
-
-    pub fn set_basic_block<C: BlockContext>(&self, block: Option<&BasicBlock<C>>) {
-        match block {
-            Some(block) => unsafe { BNSetFlowGraphBasicBlock(self.handle, block.handle) },
-            None => unsafe { BNSetFlowGraphBasicBlock(self.handle, std::ptr::null_mut()) },
-        }
-    }
-
-    pub fn lines(&self) -> Array<DisassemblyTextLine> {
-        let mut count = 0;
-        let result = unsafe { BNGetFlowGraphNodeLines(self.handle, &mut count) };
-        assert!(!result.is_null());
-        unsafe { Array::new(result, count, ()) }
-    }
-
-    pub fn set_lines(&self, lines: impl IntoIterator<Item = DisassemblyTextLine>) {
-        // NOTE: This will create allocations and increment tag refs, we must call DisassemblyTextLine::free_raw
-        let mut raw_lines: Vec<BNDisassemblyTextLine> = lines
-            .into_iter()
-            .map(DisassemblyTextLine::into_raw)
-            .collect();
-        unsafe {
-            BNSetFlowGraphNodeLines(self.handle, raw_lines.as_mut_ptr(), raw_lines.len());
-            for raw_line in raw_lines {
-                DisassemblyTextLine::free_raw(raw_line);
-            }
-        }
-    }
-
-    /// Returns the graph position of the node in X, Y form.
-    pub fn position(&self) -> (i32, i32) {
-        let pos_x = unsafe { BNGetFlowGraphNodeX(self.handle) };
-        let pos_y = unsafe { BNGetFlowGraphNodeY(self.handle) };
-        (pos_x, pos_y)
-    }
-
-    /// Sets the graph position of the node.
-    pub fn set_position(&self, x: i32, y: i32) {
-        unsafe { BNFlowGraphNodeSetX(self.handle, x) };
-        unsafe { BNFlowGraphNodeSetX(self.handle, y) };
-    }
-
-    pub fn highlight_color(&self) -> HighlightColor {
-        let raw = unsafe { BNGetFlowGraphNodeHighlight(self.handle) };
-        HighlightColor::from(raw)
-    }
-
-    pub fn set_highlight_color(&self, highlight: HighlightColor) {
-        unsafe { BNSetFlowGraphNodeHighlight(self.handle, highlight.into()) };
-    }
-
-    // TODO: Add getters and setters for edges
-
-    pub fn add_outgoing_edge(
-        &self,
-        type_: BranchType,
-        target: &FlowGraphNode,
-        edge_style: EdgeStyle,
-    ) {
-        unsafe {
-            BNAddFlowGraphNodeOutgoingEdge(self.handle, type_, target.handle, edge_style.into())
-        }
-    }
-}
-
-unsafe impl RefCountable for FlowGraphNode {
-    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
-        Ref::new(Self {
-            handle: BNNewFlowGraphNodeReference(handle.handle),
-        })
-    }
-
-    unsafe fn dec_ref(handle: &Self) {
-        BNFreeFlowGraphNode(handle.handle);
-    }
-}
-
-impl ToOwned for FlowGraphNode {
-    type Owned = Ref<Self>;
-
-    fn to_owned(&self) -> Self::Owned {
-        unsafe { RefCountable::inc_ref(self) }
-    }
-}
-
-impl CoreArrayProvider for FlowGraphNode {
-    type Raw = *mut BNFlowGraphNode;
-    type Context = ();
-    type Wrapped<'a> = Guard<'a, FlowGraphNode>;
-}
-
-unsafe impl CoreArrayProviderInner for FlowGraphNode {
-    unsafe fn free(raw: *mut Self::Raw, count: usize, _: &Self::Context) {
-        BNFreeFlowGraphNodeList(raw, count);
-    }
-
-    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> {
-        Guard::new(Self::from_raw(*raw), context)
-    }
-}
-
-#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
-pub struct EdgeStyle {
-    style: EdgePenStyle,
-    width: usize,
-    color: ThemeColor,
-}
-
-impl EdgeStyle {
-    pub fn new(style: EdgePenStyle, width: usize, color: ThemeColor) -> Self {
-        Self {
-            style,
-            width,
-            color,
-        }
-    }
-}
-
-impl Default for EdgeStyle {
-    fn default() -> Self {
-        Self::new(EdgePenStyle::SolidLine, 0, ThemeColor::AddressColor)
-    }
-}
-
-impl From<BNEdgeStyle> for EdgeStyle {
-    fn from(style: BNEdgeStyle) -> Self {
-        Self::new(style.style, style.width, style.color)
-    }
-}
-
-impl From<EdgeStyle> for BNEdgeStyle {
-    fn from(style: EdgeStyle) -> Self {
-        Self {
-            style: style.style,
-            width: style.width,
-            color: style.color,
-        }
-    }
-}
diff --git a/rust/src/flowgraph/edge.rs b/rust/src/flowgraph/edge.rs
new file mode 100644
index 00000000..7862f156
--- /dev/null
+++ b/rust/src/flowgraph/edge.rs
@@ -0,0 +1,113 @@
+use binaryninjacore_sys::*;
+
+use crate::flowgraph::node::FlowGraphNode;
+use crate::flowgraph::{BranchType, EdgePenStyle, ThemeColor};
+use crate::rc::{CoreArrayProvider, CoreArrayProviderInner, Ref};
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct FlowGraphEdge {
+    pub branch_type: BranchType,
+    pub target: Ref<FlowGraphNode>,
+    pub points: Vec<Point>,
+    pub back_edge: bool,
+    pub style: EdgeStyle,
+}
+
+impl FlowGraphEdge {
+    pub fn from_raw(value: &BNFlowGraphEdge) -> Self {
+        let raw_points = unsafe { std::slice::from_raw_parts(value.points, value.pointCount) };
+        let points: Vec<_> = raw_points.iter().copied().map(Point::from).collect();
+        Self {
+            branch_type: value.type_,
+            target: unsafe { FlowGraphNode::from_raw(value.target) }.to_owned(),
+            points,
+            back_edge: value.backEdge,
+            style: value.style.into(),
+        }
+    }
+}
+
+impl CoreArrayProvider for FlowGraphEdge {
+    type Raw = BNFlowGraphEdge;
+    type Context = ();
+    type Wrapped<'a> = FlowGraphEdge;
+}
+
+unsafe impl CoreArrayProviderInner for FlowGraphEdge {
+    unsafe fn free(raw: *mut Self::Raw, count: usize, _: &Self::Context) {
+        BNFreeFlowGraphNodeEdgeList(raw, count);
+    }
+
+    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
+        Self::from_raw(raw)
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Default)]
+pub struct Point {
+    pub x: f32,
+    pub y: f32,
+}
+
+impl Point {
+    pub fn new(x: f32, y: f32) -> Self {
+        Self { x, y }
+    }
+}
+
+impl From<BNPoint> for Point {
+    fn from(value: BNPoint) -> Self {
+        Self {
+            x: value.x,
+            y: value.y,
+        }
+    }
+}
+
+impl From<Point> for BNPoint {
+    fn from(value: Point) -> Self {
+        Self {
+            x: value.x,
+            y: value.y,
+        }
+    }
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
+pub struct EdgeStyle {
+    style: EdgePenStyle,
+    width: usize,
+    color: ThemeColor,
+}
+
+impl EdgeStyle {
+    pub fn new(style: EdgePenStyle, width: usize, color: ThemeColor) -> Self {
+        Self {
+            style,
+            width,
+            color,
+        }
+    }
+}
+
+impl Default for EdgeStyle {
+    fn default() -> Self {
+        Self::new(EdgePenStyle::SolidLine, 0, ThemeColor::AddressColor)
+    }
+}
+
+impl From<BNEdgeStyle> for EdgeStyle {
+    fn from(style: BNEdgeStyle) -> Self {
+        Self::new(style.style, style.width, style.color)
+    }
+}
+
+impl From<EdgeStyle> for BNEdgeStyle {
+    fn from(style: EdgeStyle) -> Self {
+        Self {
+            style: style.style,
+            width: style.width,
+            color: style.color,
+        }
+    }
+}
diff --git a/rust/src/flowgraph/node.rs b/rust/src/flowgraph/node.rs
new file mode 100644
index 00000000..5c2b9248
--- /dev/null
+++ b/rust/src/flowgraph/node.rs
@@ -0,0 +1,155 @@
+use crate::basic_block::{BasicBlock, BlockContext};
+use crate::disassembly::DisassemblyTextLine;
+use crate::flowgraph::edge::{EdgeStyle, FlowGraphEdge};
+use crate::flowgraph::{BranchType, FlowGraph};
+use crate::function::HighlightColor;
+use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable};
+use binaryninjacore_sys::*;
+use std::fmt::{Debug, Formatter};
+
+#[derive(PartialEq, Eq, Hash)]
+pub struct FlowGraphNode {
+    pub(crate) handle: *mut BNFlowGraphNode,
+}
+
+impl FlowGraphNode {
+    pub(crate) unsafe fn from_raw(raw: *mut BNFlowGraphNode) -> Self {
+        Self { handle: raw }
+    }
+
+    pub(crate) unsafe fn ref_from_raw(raw: *mut BNFlowGraphNode) -> Ref<Self> {
+        Ref::new(Self { handle: raw })
+    }
+
+    pub fn new(graph: &FlowGraph) -> Ref<Self> {
+        unsafe { FlowGraphNode::ref_from_raw(BNCreateFlowGraphNode(graph.handle)) }
+    }
+
+    pub fn basic_block<C: BlockContext>(&self, context: C) -> Option<Ref<BasicBlock<C>>> {
+        let block_ptr = unsafe { BNGetFlowGraphBasicBlock(self.handle) };
+        if block_ptr.is_null() {
+            return None;
+        }
+        Some(unsafe { BasicBlock::ref_from_raw(block_ptr, context) })
+    }
+
+    pub fn set_basic_block<C: BlockContext>(&self, block: Option<&BasicBlock<C>>) {
+        match block {
+            Some(block) => unsafe { BNSetFlowGraphBasicBlock(self.handle, block.handle) },
+            None => unsafe { BNSetFlowGraphBasicBlock(self.handle, std::ptr::null_mut()) },
+        }
+    }
+
+    pub fn lines(&self) -> Array<DisassemblyTextLine> {
+        let mut count = 0;
+        let result = unsafe { BNGetFlowGraphNodeLines(self.handle, &mut count) };
+        assert!(!result.is_null());
+        unsafe { Array::new(result, count, ()) }
+    }
+
+    pub fn set_lines(&self, lines: impl IntoIterator<Item = DisassemblyTextLine>) {
+        // NOTE: This will create allocations and increment tag refs, we must call DisassemblyTextLine::free_raw
+        let mut raw_lines: Vec<BNDisassemblyTextLine> = lines
+            .into_iter()
+            .map(DisassemblyTextLine::into_raw)
+            .collect();
+        unsafe {
+            BNSetFlowGraphNodeLines(self.handle, raw_lines.as_mut_ptr(), raw_lines.len());
+            for raw_line in raw_lines {
+                DisassemblyTextLine::free_raw(raw_line);
+            }
+        }
+    }
+
+    /// Returns the graph position of the node in X, Y form.
+    pub fn position(&self) -> (i32, i32) {
+        let pos_x = unsafe { BNGetFlowGraphNodeX(self.handle) };
+        let pos_y = unsafe { BNGetFlowGraphNodeY(self.handle) };
+        (pos_x, pos_y)
+    }
+
+    /// Sets the graph position of the node.
+    pub fn set_position(&self, x: i32, y: i32) {
+        unsafe { BNFlowGraphNodeSetX(self.handle, x) };
+        unsafe { BNFlowGraphNodeSetY(self.handle, y) };
+    }
+
+    pub fn highlight_color(&self) -> HighlightColor {
+        let raw = unsafe { BNGetFlowGraphNodeHighlight(self.handle) };
+        HighlightColor::from(raw)
+    }
+
+    pub fn set_highlight_color(&self, highlight: HighlightColor) {
+        unsafe { BNSetFlowGraphNodeHighlight(self.handle, highlight.into()) };
+    }
+
+    pub fn incoming_edges(&self) -> Array<FlowGraphEdge> {
+        let mut count = 0;
+        let result = unsafe { BNGetFlowGraphNodeIncomingEdges(self.handle, &mut count) };
+        assert!(!result.is_null());
+        unsafe { Array::new(result, count, ()) }
+    }
+
+    pub fn outgoing_edges(&self) -> Array<FlowGraphEdge> {
+        let mut count = 0;
+        let result = unsafe { BNGetFlowGraphNodeOutgoingEdges(self.handle, &mut count) };
+        assert!(!result.is_null());
+        unsafe { Array::new(result, count, ()) }
+    }
+
+    /// Connects two flow graph nodes with an edge.
+    pub fn add_outgoing_edge(
+        &self,
+        type_: BranchType,
+        target: &FlowGraphNode,
+        edge_style: EdgeStyle,
+    ) {
+        unsafe {
+            BNAddFlowGraphNodeOutgoingEdge(self.handle, type_, target.handle, edge_style.into())
+        }
+    }
+}
+
+impl Debug for FlowGraphNode {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("FlowGraphNode")
+            .field("lines", &self.lines().to_vec())
+            .finish()
+    }
+}
+
+unsafe impl RefCountable for FlowGraphNode {
+    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
+        Ref::new(Self {
+            handle: BNNewFlowGraphNodeReference(handle.handle),
+        })
+    }
+
+    unsafe fn dec_ref(handle: &Self) {
+        BNFreeFlowGraphNode(handle.handle);
+    }
+}
+
+impl ToOwned for FlowGraphNode {
+    type Owned = Ref<Self>;
+
+    fn to_owned(&self) -> Self::Owned {
+        unsafe { RefCountable::inc_ref(self) }
+    }
+}
+
+impl CoreArrayProvider for FlowGraphNode {
+    type Raw = *mut BNFlowGraphNode;
+    type Context = ();
+    type Wrapped<'a> = Guard<'a, FlowGraphNode>;
+}
+
+unsafe impl CoreArrayProviderInner for FlowGraphNode {
+    unsafe fn free(raw: *mut Self::Raw, count: usize, _: &Self::Context) {
+        BNFreeFlowGraphNodeList(raw, count);
+    }
+
+    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> {
+        Guard::new(Self::from_raw(*raw), context)
+    }
+}