Add Render Layer API to Rust

Author: emesare

rust/src/basic_block.rs

@@ -19,6 +19,7 @@ use crate::BranchType;
 use binaryninjacore_sys::*;
 use std::fmt;
 use std::fmt::Debug;
+use std::hash::{Hash, Hasher};
 
 enum EdgeDirection {
     Incoming,
@@ -97,7 +98,14 @@ pub trait BlockContext: Clone + Sync + Send + Sized {
     fn iter(&self, block: &BasicBlock<Self>) -> Self::Iter;
 }
 
-#[derive(PartialEq, Eq, Hash)]
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, PartialOrd, Ord)]
+pub enum BasicBlockType {
+    Native,
+    LowLevelIL,
+    MediumLevelIL,
+    HighLevelIL,
+}
+
 pub struct BasicBlock<C: BlockContext> {
     pub(crate) handle: *mut BNBasicBlock,
     context: C,
@@ -127,6 +135,19 @@ impl<C: BlockContext> BasicBlock<C> {
         }
     }
 
+    pub fn block_type(&self) -> BasicBlockType {
+        if unsafe { !BNIsILBasicBlock(self.handle) } {
+            BasicBlockType::Native
+        } else if unsafe { BNIsLowLevelILBasicBlock(self.handle) } {
+            BasicBlockType::LowLevelIL
+        } else if unsafe { BNIsMediumLevelILBasicBlock(self.handle) } {
+            BasicBlockType::MediumLevelIL
+        } else {
+            // We checked all other IL levels, so this is safe.
+            BasicBlockType::HighLevelIL
+        }
+    }
+
     pub fn iter(&self) -> C::Iter {
         self.context.iter(self)
     }
@@ -194,7 +215,7 @@ impl<C: BlockContext> BasicBlock<C> {
             if block.is_null() {
                 return None;
             }
-            Some(Ref::new(BasicBlock::from_raw(block, self.context.clone())))
+            Some(BasicBlock::ref_from_raw(block, self.context.clone()))
         }
     }
 
@@ -237,6 +258,24 @@ impl<C: BlockContext> BasicBlock<C> {
     // TODO iterated dominance frontier
 }
 
+impl<C: BlockContext> Hash for BasicBlock<C> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.function().hash(state);
+        self.block_type().hash(state);
+        state.write_usize(self.index());
+    }
+}
+
+impl<C: BlockContext> PartialEq for BasicBlock<C> {
+    fn eq(&self, other: &Self) -> bool {
+        self.function() == other.function()
+            && self.index() == other.index()
+            && self.block_type() == other.block_type()
+    }
+}
+
+impl<C: BlockContext> Eq for BasicBlock<C> {}
+
 impl<C: BlockContext> IntoIterator for &BasicBlock<C> {
     type Item = C::Instruction;
     type IntoIter = C::Iter;

rust/src/binary_view.rs

@@ -1209,17 +1209,17 @@ pub trait BinaryViewExt: BinaryViewBase {
 
     /// Retrieves a list of the next disassembly lines.
     ///
-    /// `get_next_linear_disassembly_lines` retrieves an [Array] over [LinearDisassemblyLine] objects for the
-    /// next disassembly lines, and updates the [LinearViewCursor] passed in. This function can be called
+    /// Retrieves an [`Array`] over [`LinearDisassemblyLine`] objects for the
+    /// next disassembly lines, and updates the [`LinearViewCursor`] passed in. This function can be called
     /// repeatedly to get more lines of linear disassembly.
     ///
     /// # Arguments
     /// * `pos` - Position to retrieve linear disassembly lines from
     fn get_next_linear_disassembly_lines(
         &self,
         pos: &mut LinearViewCursor,
-    ) -> Array<LinearDisassemblyLine> {
-        let mut result = unsafe { Array::new(std::ptr::null_mut(), 0, ()) };
+    ) -> Array<LinearDisassemblyLine<NativeBlock>> {
+        let mut result = unsafe { Array::new(std::ptr::null_mut(), 0, NativeBlock::new()) };
 
         while result.is_empty() {
             result = pos.lines();
@@ -1242,8 +1242,8 @@ pub trait BinaryViewExt: BinaryViewBase {
     fn get_previous_linear_disassembly_lines(
         &self,
         pos: &mut LinearViewCursor,
-    ) -> Array<LinearDisassemblyLine> {
-        let mut result = unsafe { Array::new(std::ptr::null_mut(), 0, ()) };
+    ) -> Array<LinearDisassemblyLine<NativeBlock>> {
+        let mut result = unsafe { Array::new(std::ptr::null_mut(), 0, NativeBlock::new()) };
         while result.is_empty() {
             if !pos.previous() {
                 return result;

rust/src/collaboration/sync.rs

@@ -701,7 +701,7 @@ pub fn is_type_archive_snapshot_ignored<S: BnStrCompatible>(
 pub fn download_type_archive<S: BnStrCompatible>(
     file: &RemoteFile,
     location: S,
-) -> Result<Option<TypeArchive>, ()> {
+) -> Result<Option<Ref<TypeArchive>>, ()> {
     download_type_archive_with_progress(file, location, NoProgressCallback)
 }
 
@@ -711,7 +711,7 @@ pub fn download_type_archive_with_progress<S: BnStrCompatible, F: ProgressCallba
     file: &RemoteFile,
     location: S,
     mut progress: F,
-) -> Result<Option<TypeArchive>, ()> {
+) -> Result<Option<Ref<TypeArchive>>, ()> {
     let mut value = std::ptr::null_mut();
     let db_path = location.into_bytes_with_nul();
     let success = unsafe {
@@ -724,7 +724,7 @@ pub fn download_type_archive_with_progress<S: BnStrCompatible, F: ProgressCallba
         )
     };
     success
-        .then(|| NonNull::new(value).map(|handle| unsafe { TypeArchive::from_raw(handle) }))
+        .then(|| NonNull::new(value).map(|handle| unsafe { TypeArchive::ref_from_raw(handle) }))
         .ok_or(())
 }
 

rust/src/disassembly.rs

@@ -30,7 +30,7 @@ pub type DisassemblyOption = BNDisassemblyOption;
 pub type InstructionTextTokenType = BNInstructionTextTokenType;
 pub type StringType = BNStringType;
 
-#[derive(Clone, PartialEq, Debug, Default)]
+#[derive(Clone, PartialEq, Debug, Default, Eq)]
 pub struct DisassemblyTextLine {
     pub address: u64,
     pub instruction_index: usize,
@@ -234,7 +234,7 @@ impl DisassemblyTextLineTypeInfo {
     }
 }
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct InstructionTextToken {
     pub address: u64,
     pub text: String,
@@ -882,6 +882,8 @@ impl From<InstructionTextTokenKind> for BNInstructionTextTokenType {
     }
 }
 
+impl Eq for InstructionTextTokenKind {}
+
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub enum InstructionTextTokenContext {
     Normal,

rust/src/flowgraph.rs

@@ -14,12 +14,14 @@
 
 //! Interfaces for creating and displaying pretty CFGs in Binary Ninja.
 
+use crate::disassembly::DisassemblyTextLine;
 use binaryninjacore_sys::*;
 use std::slice;
-use crate::disassembly::DisassemblyTextLine;
 
 use crate::rc::*;
 
+use crate::basic_block::{BasicBlock, BlockContext};
+use crate::function::HighlightColor;
 use std::marker::PhantomData;
 
 pub type BranchType = BNBranchType;
@@ -37,15 +39,19 @@ impl FlowGraph {
         Self { handle: raw }
     }
 
+    pub(crate) unsafe fn ref_from_raw(raw: *mut BNFlowGraph) -> Ref<Self> {
+        Ref::new(Self { handle: raw })
+    }
+
     pub fn new() -> Ref<Self> {
-        unsafe { Ref::new(FlowGraph::from_raw(BNCreateFlowGraph())) }
+        unsafe { FlowGraph::ref_from_raw(BNCreateFlowGraph()) }
     }
 
     pub fn nodes<'a>(&self) -> Vec<Ref<FlowGraphNode<'a>>> {
         let mut count: usize = 0;
-        let mut nodes_ptr = unsafe { BNGetFlowGraphNodes(self.handle, &mut count as *mut usize) };
+        let nodes_ptr = unsafe { BNGetFlowGraphNodes(self.handle, &mut count as *mut usize) };
 
-        let mut nodes = unsafe { slice::from_raw_parts_mut(nodes_ptr, count) };
+        let nodes = unsafe { slice::from_raw_parts_mut(nodes_ptr, count) };
 
         let mut result = vec![];
         result.reserve(count);
@@ -131,8 +137,37 @@ impl<'a> FlowGraphNode<'a> {
         }
     }
 
-    pub fn new(graph: &FlowGraph) -> Self {
-        unsafe { FlowGraphNode::from_raw(BNCreateFlowGraphNode(graph.handle)) }
+    pub(crate) unsafe fn ref_from_raw(raw: *mut BNFlowGraphNode) -> Ref<Self> {
+        Ref::new(Self {
+            handle: raw,
+            _data: PhantomData,
+        })
+    }
+
+    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>) {
@@ -149,6 +184,30 @@ impl<'a> FlowGraphNode<'a> {
         }
     }
 
+    /// 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,

rust/src/function.rs

@@ -128,7 +128,7 @@ impl Iterator for NativeBlockIter {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct NativeBlock {
     _priv: (),
 }
@@ -2346,7 +2346,7 @@ impl Function {
     /// Flow graph of unresolved stack adjustments
     pub fn unresolved_stack_adjustment_graph(&self) -> Option<Ref<FlowGraph>> {
         let graph = unsafe { BNGetUnresolvedStackAdjustmentGraph(self.handle) };
-        (!graph.is_null()).then(|| unsafe { Ref::new(FlowGraph::from_raw(graph)) })
+        (!graph.is_null()).then(|| unsafe { FlowGraph::ref_from_raw(graph) })
     }
 
     pub fn create_graph(
@@ -2358,7 +2358,7 @@ impl Function {
         let raw_view_type = FunctionViewType::into_raw(view_type);
         let result = unsafe { BNCreateFunctionGraph(self.handle, raw_view_type, settings_raw) };
         FunctionViewType::free_raw(raw_view_type);
-        unsafe { Ref::new(FlowGraph::from_raw(result)) }
+        unsafe { FlowGraph::ref_from_raw(result) }
     }
 
     pub fn parent_components(&self) -> Array<Component> {

rust/src/lib.rs

@@ -68,6 +68,7 @@ pub mod project;
 pub mod rc;
 pub mod references;
 pub mod relocation;
+pub mod render_layer;
 pub mod section;
 pub mod segment;
 pub mod settings;