1414
1515//! Interfaces for creating and displaying pretty CFGs in Binary Ninja.
1616
17- use binaryninjacore_sys:: * ;
18-
1917use crate :: disassembly:: DisassemblyTextLine ;
18+ use binaryninjacore_sys:: * ;
19+ use std:: slice;
2020
2121use crate :: rc:: * ;
2222
23+ use crate :: basic_block:: { BasicBlock , BlockContext } ;
24+ use crate :: function:: HighlightColor ;
2325use std:: marker:: PhantomData ;
2426
2527pub type BranchType = BNBranchType ;
@@ -37,14 +39,61 @@ impl FlowGraph {
3739 Self { handle : raw }
3840 }
3941
42+ pub ( crate ) unsafe fn ref_from_raw ( raw : * mut BNFlowGraph ) -> Ref < Self > {
43+ Ref :: new ( Self { handle : raw } )
44+ }
45+
4046 pub fn new ( ) -> Ref < Self > {
41- unsafe { Ref :: new ( FlowGraph :: from_raw ( BNCreateFlowGraph ( ) ) ) }
47+ unsafe { FlowGraph :: ref_from_raw ( BNCreateFlowGraph ( ) ) }
48+ }
49+
50+ pub fn nodes < ' a > ( & self ) -> Vec < Ref < FlowGraphNode < ' a > > > {
51+ let mut count: usize = 0 ;
52+ let nodes_ptr = unsafe { BNGetFlowGraphNodes ( self . handle , & mut count as * mut usize ) } ;
53+
54+ let nodes = unsafe { slice:: from_raw_parts_mut ( nodes_ptr, count) } ;
55+
56+ let mut result = vec ! [ ] ;
57+ result. reserve ( count) ;
58+
59+ for i in 0 ..count {
60+ result. push ( unsafe { RefCountable :: inc_ref ( & FlowGraphNode :: from_raw ( nodes[ i] ) ) } ) ;
61+ }
62+
63+ unsafe { BNFreeFlowGraphNodeList ( nodes_ptr, count) } ;
64+
65+ result
66+ }
67+
68+ pub fn get_node < ' a > ( & self , i : usize ) -> Option < Ref < FlowGraphNode < ' a > > > {
69+ let node_ptr = unsafe { BNGetFlowGraphNode ( self . handle , i) } ;
70+ if node_ptr. is_null ( ) {
71+ None
72+ } else {
73+ Some ( unsafe { Ref :: new ( FlowGraphNode :: from_raw ( node_ptr) ) } )
74+ }
75+ }
76+
77+ pub fn get_node_count ( & self ) -> usize {
78+ unsafe { BNGetFlowGraphNodeCount ( self . handle ) }
79+ }
80+
81+ pub fn has_nodes ( & self ) -> bool {
82+ unsafe { BNFlowGraphHasNodes ( self . handle ) }
4283 }
4384
4485 pub fn append ( & self , node : & FlowGraphNode ) -> usize {
4586 unsafe { BNAddFlowGraphNode ( self . handle , node. handle ) }
4687 }
4788
89+ pub fn replace ( & self , index : usize , node : & FlowGraphNode ) {
90+ unsafe { BNReplaceFlowGraphNode ( self . handle , index, node. handle ) }
91+ }
92+
93+ pub fn clear ( & self ) {
94+ unsafe { BNClearFlowGraphNodes ( self . handle ) }
95+ }
96+
4897 pub fn set_option ( & self , option : FlowGraphOption , value : bool ) {
4998 unsafe { BNSetFlowGraphOption ( self . handle , option, value) }
5099 }
@@ -88,8 +137,37 @@ impl<'a> FlowGraphNode<'a> {
88137 }
89138 }
90139
91- pub fn new ( graph : & FlowGraph ) -> Self {
92- unsafe { FlowGraphNode :: from_raw ( BNCreateFlowGraphNode ( graph. handle ) ) }
140+ pub ( crate ) unsafe fn ref_from_raw ( raw : * mut BNFlowGraphNode ) -> Ref < Self > {
141+ Ref :: new ( Self {
142+ handle : raw,
143+ _data : PhantomData ,
144+ } )
145+ }
146+
147+ pub fn new ( graph : & FlowGraph ) -> Ref < Self > {
148+ unsafe { FlowGraphNode :: ref_from_raw ( BNCreateFlowGraphNode ( graph. handle ) ) }
149+ }
150+
151+ pub fn basic_block < C : BlockContext > ( & self , context : C ) -> Option < Ref < BasicBlock < C > > > {
152+ let block_ptr = unsafe { BNGetFlowGraphBasicBlock ( self . handle ) } ;
153+ if block_ptr. is_null ( ) {
154+ return None ;
155+ }
156+ Some ( unsafe { BasicBlock :: ref_from_raw ( block_ptr, context) } )
157+ }
158+
159+ pub fn set_basic_block < C : BlockContext > ( & self , block : Option < & BasicBlock < C > > ) {
160+ match block {
161+ Some ( block) => unsafe { BNSetFlowGraphBasicBlock ( self . handle , block. handle ) } ,
162+ None => unsafe { BNSetFlowGraphBasicBlock ( self . handle , std:: ptr:: null_mut ( ) ) } ,
163+ }
164+ }
165+
166+ pub fn lines ( & self ) -> Array < DisassemblyTextLine > {
167+ let mut count = 0 ;
168+ let result = unsafe { BNGetFlowGraphNodeLines ( self . handle , & mut count) } ;
169+ assert ! ( !result. is_null( ) ) ;
170+ unsafe { Array :: new ( result, count, ( ) ) }
93171 }
94172
95173 pub fn set_lines ( & self , lines : impl IntoIterator < Item = DisassemblyTextLine > ) {
@@ -106,6 +184,30 @@ impl<'a> FlowGraphNode<'a> {
106184 }
107185 }
108186
187+ /// Returns the graph position of the node in X, Y form.
188+ pub fn position ( & self ) -> ( i32 , i32 ) {
189+ let pos_x = unsafe { BNGetFlowGraphNodeX ( self . handle ) } ;
190+ let pos_y = unsafe { BNGetFlowGraphNodeY ( self . handle ) } ;
191+ ( pos_x, pos_y)
192+ }
193+
194+ /// Sets the graph position of the node.
195+ pub fn set_position ( & self , x : i32 , y : i32 ) {
196+ unsafe { BNFlowGraphNodeSetX ( self . handle , x) } ;
197+ unsafe { BNFlowGraphNodeSetX ( self . handle , y) } ;
198+ }
199+
200+ pub fn highlight_color ( & self ) -> HighlightColor {
201+ let raw = unsafe { BNGetFlowGraphNodeHighlight ( self . handle ) } ;
202+ HighlightColor :: from ( raw)
203+ }
204+
205+ pub fn set_highlight_color ( & self , highlight : HighlightColor ) {
206+ unsafe { BNSetFlowGraphNodeHighlight ( self . handle , highlight. into ( ) ) } ;
207+ }
208+
209+ // TODO: Add getters and setters for edges
210+
109211 pub fn add_outgoing_edge (
110212 & self ,
111213 type_ : BranchType ,
0 commit comments