[Rust] Misc cleanup for ABB

2 files changed, 206 insertions, 119 deletions

diff --git a/rust/src/architecture/basic_block.rs b/rust/src/architecture/basic_block.rs
index 1867f50a..f092e927 100644
--- a/rust/src/architecture/basic_block.rs
+++ b/rust/src/architecture/basic_block.rs
@@ -4,6 +4,7 @@ use crate::function::{Function, Location, NativeBlock};
 use crate::rc::Ref;
 use binaryninjacore_sys::*;
 use std::collections::{HashMap, HashSet};
+use std::fmt::Debug;
 
 pub struct BasicBlockAnalysisContext {
     pub(crate) handle: *mut BNBasicBlockAnalysisContext,
@@ -36,65 +37,79 @@ impl BasicBlockAnalysisContext {
 
         let ctx_ref = &*handle;
 
-        let indirect_branches = (0..ctx_ref.indirectBranchesCount)
-            .map(|i| {
-                let raw: BNIndirectBranchInfo =
-                    unsafe { std::ptr::read(ctx_ref.indirectBranches.add(i)) };
-                IndirectBranchInfo::from(raw)
-            })
-            .collect::<Vec<_>>();
+        let raw_indirect_branches: &[BNIndirectBranchInfo] =
+            std::slice::from_raw_parts(ctx_ref.indirectBranches, ctx_ref.indirectBranchesCount);
+        let indirect_branches: Vec<IndirectBranchInfo> = raw_indirect_branches
+            .iter()
+            .map(IndirectBranchInfo::from)
+            .collect();
 
-        let indirect_no_return_calls = (0..ctx_ref.indirectNoReturnCallsCount)
-            .map(|i| {
-                let raw = unsafe { std::ptr::read(ctx_ref.indirectNoReturnCalls.add(i)) };
-                Location::from(raw)
-            })
-            .collect::<HashSet<_>>();
+        let raw_indirect_no_return_calls: &[BNArchitectureAndAddress] = std::slice::from_raw_parts(
+            ctx_ref.indirectNoReturnCalls,
+            ctx_ref.indirectNoReturnCallsCount,
+        );
+        let indirect_no_return_calls: HashSet<Location> = raw_indirect_no_return_calls
+            .iter()
+            .map(Location::from)
+            .collect();
 
-        let contextual_returns = (0..ctx_ref.contextualFunctionReturnCount)
-            .map(|i| {
-                let loc = unsafe {
-                    let raw = std::ptr::read(ctx_ref.contextualFunctionReturnLocations.add(i));
-                    Location::from(raw)
-                };
-                let val = unsafe { *ctx_ref.contextualFunctionReturnValues.add(i) };
-                (loc, val)
-            })
-            .collect::<HashMap<_, _>>();
+        let raw_contextual_return_locs: &[BNArchitectureAndAddress] = unsafe {
+            std::slice::from_raw_parts(
+                ctx_ref.contextualFunctionReturnLocations,
+                ctx_ref.contextualFunctionReturnCount,
+            )
+        };
+        let raw_contextual_return_vals: &[bool] = unsafe {
+            std::slice::from_raw_parts(
+                ctx_ref.contextualFunctionReturnValues,
+                ctx_ref.contextualFunctionReturnCount,
+            )
+        };
+        let contextual_returns: HashMap<Location, bool> = raw_contextual_return_locs
+            .iter()
+            .map(Location::from)
+            .zip(raw_contextual_return_vals.iter().copied())
+            .collect();
 
-        let direct_code_references = (0..ctx_ref.directRefCount)
-            .map(|i| {
-                let src = unsafe {
-                    let raw = std::ptr::read(ctx_ref.directRefSources.add(i));
-                    Location::from(raw)
-                };
-                let tgt = unsafe { *ctx_ref.directRefTargets.add(i) };
-                (tgt, src)
-            })
-            .collect::<HashMap<_, _>>();
+        let raw_direct_ref_sources: &[BNArchitectureAndAddress] =
+            unsafe { std::slice::from_raw_parts(ctx_ref.directRefSources, ctx_ref.directRefCount) };
+        let raw_direct_ref_targets: &[u64] =
+            unsafe { std::slice::from_raw_parts(ctx_ref.directRefTargets, ctx_ref.directRefCount) };
+        let direct_code_references: HashMap<u64, Location> = raw_direct_ref_targets
+            .iter()
+            .copied()
+            .zip(raw_direct_ref_sources.iter().map(Location::from))
+            .collect();
 
-        let direct_no_return_calls = (0..ctx_ref.directNoReturnCallsCount)
-            .map(|i| {
-                let raw = unsafe { std::ptr::read(ctx_ref.directNoReturnCalls.add(i)) };
-                Location::from(raw)
-            })
-            .collect::<HashSet<_>>();
+        let raw_direct_no_return_calls: &[BNArchitectureAndAddress] = std::slice::from_raw_parts(
+            ctx_ref.directNoReturnCalls,
+            ctx_ref.directNoReturnCallsCount,
+        );
+        let direct_no_return_calls: HashSet<Location> = raw_direct_no_return_calls
+            .iter()
+            .map(Location::from)
+            .collect();
 
-        let halted_disassembly_addresses = (0..ctx_ref.haltedDisassemblyAddressesCount)
-            .map(|i| {
-                let raw = unsafe { std::ptr::read(ctx_ref.haltedDisassemblyAddresses.add(i)) };
-                Location::from(raw)
-            })
-            .collect::<HashSet<_>>();
+        let raw_halted_disassembly_address: &[BNArchitectureAndAddress] =
+            std::slice::from_raw_parts(
+                ctx_ref.haltedDisassemblyAddresses,
+                ctx_ref.haltedDisassemblyAddressesCount,
+            );
+        let halted_disassembly_addresses: HashSet<Location> = raw_halted_disassembly_address
+            .iter()
+            .map(Location::from)
+            .collect();
 
-        let inlined_unresolved_indirect_branches = (0..ctx_ref
-            .inlinedUnresolvedIndirectBranchCount)
-            .map(|i| {
-                let raw =
-                    unsafe { std::ptr::read(ctx_ref.inlinedUnresolvedIndirectBranches.add(i)) };
-                Location::from(raw)
-            })
-            .collect::<HashSet<_>>();
+        let raw_inlined_unresolved_indirect_branches: &[BNArchitectureAndAddress] =
+            std::slice::from_raw_parts(
+                ctx_ref.inlinedUnresolvedIndirectBranches,
+                ctx_ref.inlinedUnresolvedIndirectBranchCount,
+            );
+        let inlined_unresolved_indirect_branches: HashSet<Location> =
+            raw_inlined_unresolved_indirect_branches
+                .iter()
+                .map(Location::from)
+                .collect();
 
         BasicBlockAnalysisContext {
             handle,
@@ -116,6 +131,7 @@ impl BasicBlockAnalysisContext {
         }
     }
 
+    /// Adds a contextual function return location and its value to the current function.
     pub fn add_contextual_return(&mut self, loc: impl Into<Location>, value: bool) {
         let loc = loc.into();
         if !self.contextual_returns.contains_key(&loc) {
@@ -125,16 +141,19 @@ impl BasicBlockAnalysisContext {
         self.contextual_returns.insert(loc, value);
     }
 
+    /// Adds a direct code reference to the current function.
     pub fn add_direct_code_reference(&mut self, target: u64, src: impl Into<Location>) {
         self.direct_code_references
             .entry(target)
             .or_insert(src.into());
     }
 
+    /// Adds a direct no-return call location to the current function.
     pub fn add_direct_no_return_call(&mut self, loc: impl Into<Location>) {
         self.direct_no_return_calls.insert(loc.into());
     }
 
+    /// Adds an address to the set of halted disassembly addresses.
     pub fn add_halted_disassembly_address(&mut self, loc: impl Into<Location>) {
         self.halted_disassembly_addresses.insert(loc.into());
     }
@@ -143,6 +162,9 @@ impl BasicBlockAnalysisContext {
         self.inlined_unresolved_indirect_branches.insert(loc.into());
     }
 
+    /// Creates a new [`BasicBlock`] at the specified address for the given [`CoreArchitecture`].
+    ///
+    /// After creating, you can add using [`BasicBlockAnalysisContext::add_basic_block`].
     pub fn create_basic_block(
         &self,
         arch: CoreArchitecture,
@@ -158,80 +180,127 @@ impl BasicBlockAnalysisContext {
         unsafe { Some(BasicBlock::ref_from_raw(raw_block, NativeBlock::new())) }
     }
 
+    /// Adds a [`BasicBlock`] to the current function.
+    ///
+    /// You can create a [`BasicBlock`] via [`BasicBlockAnalysisContext::create_basic_block`].
     pub fn add_basic_block(&self, block: Ref<BasicBlock<NativeBlock>>) {
         unsafe {
             BNAnalyzeBasicBlocksContextAddBasicBlockToFunction(self.handle, block.handle);
         }
     }
 
+    /// Adds a temporary outgoing reference to the specified function.
     pub fn add_temp_outgoing_reference(&self, target: &Function) {
         unsafe {
             BNAnalyzeBasicBlocksContextAddTempReference(self.handle, target.handle);
         }
     }
 
+    /// To be called before finalizing the basic block analysis.
+    fn update_direct_code_references(&mut self) {
+        let total = self.direct_code_references.len();
+        let mut sources: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
+        let mut targets: Vec<u64> = Vec::with_capacity(total);
+        for (target, src) in &self.direct_code_references {
+            sources.push(src.into());
+            targets.push(*target);
+        }
+        unsafe {
+            BNAnalyzeBasicBlocksContextSetDirectCodeReferences(
+                self.handle,
+                sources.as_mut_ptr(),
+                targets.as_mut_ptr(),
+                total,
+            );
+        }
+    }
+
+    /// To be called before finalizing the basic block analysis.
+    fn update_direct_no_return_calls(&mut self) {
+        let total = self.direct_no_return_calls.len();
+        let mut raw_locations: Vec<_> = self
+            .direct_no_return_calls
+            .iter()
+            .map(BNArchitectureAndAddress::from)
+            .collect();
+        unsafe {
+            BNAnalyzeBasicBlocksContextSetDirectNoReturnCalls(
+                self.handle,
+                raw_locations.as_mut_ptr(),
+                total,
+            );
+        }
+    }
+
+    /// To be called before finalizing the basic block analysis.
+    fn update_inlined_unresolved_indirect_branches(&mut self) {
+        let total = self.inlined_unresolved_indirect_branches.len();
+        let mut raw_locations: Vec<_> = self
+            .inlined_unresolved_indirect_branches
+            .iter()
+            .map(BNArchitectureAndAddress::from)
+            .collect();
+        unsafe {
+            BNAnalyzeBasicBlocksContextSetInlinedUnresolvedIndirectBranches(
+                self.handle,
+                raw_locations.as_mut_ptr(),
+                total,
+            );
+        }
+    }
+
+    /// To be called before finalizing the basic block analysis.
+    fn update_halted_disassembly_addresses(&mut self) {
+        let total = self.halted_disassembly_addresses.len();
+        let mut raw_locations: Vec<_> = self
+            .halted_disassembly_addresses
+            .iter()
+            .map(BNArchitectureAndAddress::from)
+            .collect();
+        unsafe {
+            BNAnalyzeBasicBlocksContextSetHaltedDisassemblyAddresses(
+                self.handle,
+                raw_locations.as_mut_ptr(),
+                total,
+            );
+        }
+    }
+
+    /// To be called before finalizing the basic block analysis.
+    fn update_contextual_returns(&mut self) {
+        let total = self.contextual_returns.len();
+        let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
+        let mut values: Vec<bool> = Vec::with_capacity(total);
+        for (loc, value) in &self.contextual_returns {
+            locations.push(loc.into());
+            values.push(*value);
+        }
+        unsafe {
+            BNAnalyzeBasicBlocksContextSetContextualFunctionReturns(
+                self.handle,
+                locations.as_mut_ptr(),
+                values.as_mut_ptr(),
+                total,
+            );
+        }
+    }
+
+    /// Finalizes the function's basic block analysis.
     pub fn finalize(&mut self) {
         if !self.direct_code_references.is_empty() {
-            let total = self.direct_code_references.len();
-            let mut sources: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            let mut targets: Vec<u64> = Vec::with_capacity(total);
-            for (target, src) in &self.direct_code_references {
-                sources.push(BNArchitectureAndAddress::from(src));
-                targets.push(*target);
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetDirectCodeReferences(
-                    self.handle,
-                    sources.as_mut_ptr(),
-                    targets.as_mut_ptr(),
-                    total,
-                );
-            }
+            self.update_direct_code_references();
         }
 
         if !self.direct_no_return_calls.is_empty() {
-            let total = self.direct_no_return_calls.len();
-            let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            for loc in &self.direct_no_return_calls {
-                locations.push(BNArchitectureAndAddress::from(loc));
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetDirectNoReturnCalls(
-                    self.handle,
-                    locations.as_mut_ptr(),
-                    total,
-                );
-            }
+            self.update_direct_no_return_calls();
         }
 
         if !self.halted_disassembly_addresses.is_empty() {
-            let total = self.halted_disassembly_addresses.len();
-            let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            for loc in &self.halted_disassembly_addresses {
-                locations.push(BNArchitectureAndAddress::from(loc));
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetHaltedDisassemblyAddresses(
-                    self.handle,
-                    locations.as_mut_ptr(),
-                    total,
-                );
-            }
+            self.update_halted_disassembly_addresses();
         }
 
         if !self.inlined_unresolved_indirect_branches.is_empty() {
-            let total = self.inlined_unresolved_indirect_branches.len();
-            let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            for loc in &self.inlined_unresolved_indirect_branches {
-                locations.push(BNArchitectureAndAddress::from(loc));
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetInlinedUnresolvedIndirectBranches(
-                    self.handle,
-                    locations.as_mut_ptr(),
-                    total,
-                );
-            }
+            self.update_inlined_unresolved_indirect_branches();
         }
 
         unsafe {
@@ -239,23 +308,35 @@ impl BasicBlockAnalysisContext {
         }
 
         if self.contextual_returns_dirty {
-            let total = self.contextual_returns.len();
-            let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            let mut values: Vec<bool> = Vec::with_capacity(total);
-            for (loc, value) in &self.contextual_returns {
-                locations.push(BNArchitectureAndAddress::from(loc));
-                values.push(*value);
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetContextualFunctionReturns(
-                    self.handle,
-                    locations.as_mut_ptr(),
-                    values.as_mut_ptr(),
-                    total,
-                );
-            }
+            self.update_contextual_returns();
         }
 
         unsafe { BNAnalyzeBasicBlocksContextFinalize(self.handle) };
     }
 }
+
+impl Debug for BasicBlockAnalysisContext {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BasicBlockAnalysisContext")
+            .field("indirect_branches", &self.indirect_branches)
+            .field("indirect_no_return_calls", &self.indirect_no_return_calls)
+            .field("analysis_skip_override", &self.analysis_skip_override)
+            .field("translate_tail_calls", &self.translate_tail_calls)
+            .field("disallow_branch_to_string", &self.disallow_branch_to_string)
+            .field("max_function_size", &self.max_function_size)
+            .field("guided_analysis_mode", &self.guided_analysis_mode)
+            .field(
+                "trigger_guided_on_invalid_instruction",
+                &self.trigger_guided_on_invalid_instruction,
+            )
+            .field("max_size_reached", &self.max_size_reached)
+            .field("contextual_returns", &self.contextual_returns)
+            .field("direct_code_references", &self.direct_code_references)
+            .field("direct_no_return_calls", &self.direct_no_return_calls)
+            .field(
+                "halted_disassembly_addresses",
+                &self.halted_disassembly_addresses,
+            )
+            .finish()
+    }
+}
diff --git a/rust/src/architecture/branches.rs b/rust/src/architecture/branches.rs
index ff0ab516..80baa1b6 100644
--- a/rust/src/architecture/branches.rs
+++ b/rust/src/architecture/branches.rs
@@ -136,6 +136,12 @@ impl From<IndirectBranchInfo> for BNIndirectBranchInfo {
     }
 }
 
+impl From<&BNIndirectBranchInfo> for IndirectBranchInfo {
+    fn from(value: &BNIndirectBranchInfo) -> Self {
+        Self::from(*value)
+    }
+}
+
 impl CoreArrayProvider for IndirectBranchInfo {
     type Raw = BNIndirectBranchInfo;
     type Context = ();