1 files changed, 81 insertions, 64 deletions

diff --git a/rust/src/flowgraph.rs b/rust/src/flowgraph.rs
index c8505187..ead984f9 100644
--- a/rust/src/flowgraph.rs
+++ b/rust/src/flowgraph.rs
@@ -27,27 +27,51 @@ pub type EdgePenStyle = BNEdgePenStyle;
 pub type ThemeColor = BNThemeColor;
 pub type FlowGraphOption = BNFlowGraphOption;
 
-#[repr(transparent)]
-pub struct EdgeStyle(pub(crate) BNEdgeStyle);
+#[derive(PartialEq, Eq, Hash)]
+pub struct FlowGraph {
+    pub(crate) handle: *mut BNFlowGraph,
+}
 
-impl EdgeStyle {
-    pub fn new(style: EdgePenStyle, width: usize, color: ThemeColor) -> Self {
-        EdgeStyle(BNEdgeStyle {
-            style,
-            width,
-            color,
-        })
+impl FlowGraph {
+    pub(crate) unsafe fn from_raw(raw: *mut BNFlowGraph) -> Self {
+        Self { handle: raw }
+    }
+
+    pub fn new() -> Ref<Self> {
+        unsafe { Ref::new(FlowGraph::from_raw(BNCreateFlowGraph())) }
+    }
+
+    pub fn append(&self, node: &FlowGraphNode) -> usize {
+        unsafe { BNAddFlowGraphNode(self.handle, node.handle) }
+    }
+
+    pub fn set_option(&self, option: FlowGraphOption, value: bool) {
+        unsafe { BNSetFlowGraphOption(self.handle, option, value) }
+    }
+
+    pub fn is_option_set(&self, option: FlowGraphOption) -> bool {
+        unsafe { BNIsFlowGraphOptionSet(self.handle, option) }
     }
 }
 
-impl Default for EdgeStyle {
-    fn default() -> Self {
-        EdgeStyle(BNEdgeStyle {
-            style: EdgePenStyle::SolidLine,
-            width: 0,
-            color: ThemeColor::AddressColor,
+unsafe impl RefCountable for FlowGraph {
+    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
+        Ref::new(Self {
+            handle: BNNewFlowGraphReference(handle.handle),
         })
     }
+
+    unsafe fn dec_ref(handle: &Self) {
+        BNFreeFlowGraph(handle.handle);
+    }
+}
+
+impl ToOwned for FlowGraph {
+    type Owned = Ref<Self>;
+
+    fn to_owned(&self) -> Self::Owned {
+        unsafe { RefCountable::inc_ref(self) }
+    }
 }
 
 #[derive(PartialEq, Eq, Hash)]
@@ -68,32 +92,33 @@ impl<'a> FlowGraphNode<'a> {
         unsafe { FlowGraphNode::from_raw(BNCreateFlowGraphNode(graph.handle)) }
     }
 
-    pub fn set_disassembly_lines(&self, lines: &'a [DisassemblyTextLine]) {
+    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, lines.as_ptr() as *mut _, lines.len());
-            // BNFreeDisassemblyTextLines(lines.as_ptr() as *mut _, lines.len());  // Shouldn't need...would be a double free?
+            BNSetFlowGraphNodeLines(self.handle, raw_lines.as_mut_ptr(), raw_lines.len());
+            for raw_line in raw_lines {
+                DisassemblyTextLine::free_raw(raw_line);
+            }
         }
     }
 
-    pub fn set_lines(&self, lines: Vec<&str>) {
-        let lines = lines
-            .iter()
-            .map(|&line| DisassemblyTextLine::from(&vec![line]))
-            .collect::<Vec<_>>();
-        self.set_disassembly_lines(&lines);
-    }
-
     pub fn add_outgoing_edge(
         &self,
         type_: BranchType,
         target: &'a FlowGraphNode,
-        edge_style: &'a EdgeStyle,
+        edge_style: EdgeStyle,
     ) {
-        unsafe { BNAddFlowGraphNodeOutgoingEdge(self.handle, type_, target.handle, edge_style.0) }
+        unsafe {
+            BNAddFlowGraphNodeOutgoingEdge(self.handle, type_, target.handle, edge_style.into())
+        }
     }
 }
 
-unsafe impl<'a> RefCountable for FlowGraphNode<'a> {
+unsafe impl RefCountable for FlowGraphNode<'_> {
     unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
         Ref::new(Self {
             handle: BNNewFlowGraphNodeReference(handle.handle),
@@ -106,7 +131,7 @@ unsafe impl<'a> RefCountable for FlowGraphNode<'a> {
     }
 }
 
-impl<'a> ToOwned for FlowGraphNode<'a> {
+impl ToOwned for FlowGraphNode<'_> {
     type Owned = Ref<Self>;
 
     fn to_owned(&self) -> Self::Owned {
@@ -114,49 +139,41 @@ impl<'a> ToOwned for FlowGraphNode<'a> {
     }
 }
 
-#[derive(PartialEq, Eq, Hash)]
-pub struct FlowGraph {
-    pub(crate) handle: *mut BNFlowGraph,
+#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
+pub struct EdgeStyle {
+    style: EdgePenStyle,
+    width: usize,
+    color: ThemeColor,
 }
 
-impl FlowGraph {
-    pub(crate) unsafe fn from_raw(raw: *mut BNFlowGraph) -> Self {
-        Self { handle: raw }
-    }
-
-    pub fn new() -> Ref<Self> {
-        unsafe { Ref::new(FlowGraph::from_raw(BNCreateFlowGraph())) }
-    }
-
-    pub fn append(&self, node: &FlowGraphNode) -> usize {
-        unsafe { BNAddFlowGraphNode(self.handle, node.handle) }
-    }
-
-    pub fn set_option(&self, option: FlowGraphOption, value: bool) {
-        unsafe { BNSetFlowGraphOption(self.handle, option, value) }
-    }
-
-    pub fn is_option_set(&self, option: FlowGraphOption) -> bool {
-        unsafe { BNIsFlowGraphOptionSet(self.handle, option) }
+impl EdgeStyle {
+    pub fn new(style: EdgePenStyle, width: usize, color: ThemeColor) -> Self {
+        Self {
+            style,
+            width,
+            color,
+        }
     }
 }
 
-unsafe impl RefCountable for FlowGraph {
-    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
-        Ref::new(Self {
-            handle: BNNewFlowGraphReference(handle.handle),
-        })
+impl Default for EdgeStyle {
+    fn default() -> Self {
+        Self::new(EdgePenStyle::SolidLine, 0, ThemeColor::AddressColor)
     }
+}
 
-    unsafe fn dec_ref(handle: &Self) {
-        BNFreeFlowGraph(handle.handle);
+impl From<BNEdgeStyle> for EdgeStyle {
+    fn from(style: BNEdgeStyle) -> Self {
+        Self::new(style.style, style.width, style.color)
     }
 }
 
-impl ToOwned for FlowGraph {
-    type Owned = Ref<Self>;
-
-    fn to_owned(&self) -> Self::Owned {
-        unsafe { RefCountable::inc_ref(self) }
+impl From<EdgeStyle> for BNEdgeStyle {
+    fn from(style: EdgeStyle) -> Self {
+        Self {
+            style: style.style,
+            width: style.width,
+            color: style.color,
+        }
     }
 }