1 files changed, 62 insertions, 36 deletions

diff --git a/rust/src/references.rs b/rust/src/references.rs
index d12c33c5..52fb6713 100644
--- a/rust/src/references.rs
+++ b/rust/src/references.rs
@@ -1,53 +1,70 @@
+#![allow(dead_code)]
 use crate::architecture::CoreArchitecture;
 use crate::function::Function;
-use crate::rc::{CoreArrayProvider, CoreArrayProviderInner, Guard, Ref};
+use crate::rc::{CoreArrayProvider, CoreArrayProviderInner, Ref};
 use binaryninjacore_sys::{BNFreeCodeReferences, BNFreeDataReferences, BNReferenceSource};
-use std::mem::ManuallyDrop;
 
 /// A struct representing a single code cross-reference.
-/// Taking a cue from [`crate::linearview::LinearDisassemblyLine`], this struct uses [ManuallyDrop] to
-/// prevent destructors from being run on the [`Function`] object allocated by
-/// the core in `BNGetCodeReferences` (et al). The reference is cleaned up on [Drop] of
-/// the enclosing array object.
 #[derive(Debug)]
 pub struct CodeReference {
-    arch: CoreArchitecture,
-    func: ManuallyDrop<Ref<Function>>,
-    pub address: u64,
-}
-
-/// A struct representing a single data cross-reference.
-/// Data references have no associated metadata, so this object has only
-/// a single [`DataReference::address`] attribute.
-pub struct DataReference {
+    pub arch: Option<CoreArchitecture>,
+    pub func: Option<Ref<Function>>,
     pub address: u64,
 }
 
 impl CodeReference {
-    pub(crate) unsafe fn new(handle: &BNReferenceSource) -> Self {
-        let func = ManuallyDrop::new(Function::from_raw(handle.func));
-        let arch = CoreArchitecture::from_raw(handle.arch);
-        let address = handle.addr;
+    pub(crate) fn from_raw(value: &BNReferenceSource) -> Self {
         Self {
-            func,
-            arch,
-            address,
+            func: match value.func.is_null() {
+                false => Some(unsafe { Function::from_raw(value.func) }.to_owned()),
+                true => None,
+            },
+            arch: match value.func.is_null() {
+                false => Some(unsafe { CoreArchitecture::from_raw(value.arch) }),
+                true => None,
+            },
+            address: value.addr,
+        }
+    }
+
+    pub(crate) fn from_owned_raw(value: BNReferenceSource) -> Self {
+        let owned = Self::from_raw(&value);
+        Self::free_raw(value);
+        owned
+    }
+
+    pub(crate) fn into_raw(value: Self) -> BNReferenceSource {
+        BNReferenceSource {
+            func: match value.func {
+                Some(func) => unsafe { Ref::into_raw(func) }.handle,
+                None => std::ptr::null_mut(),
+            },
+            arch: value.arch.map(|a| a.handle).unwrap_or(std::ptr::null_mut()),
+            addr: value.address,
         }
     }
-}
 
-impl<'a> CodeReference {
-    /// A handle to the referenced function bound by the [CodeReference] object's lifetime.
-    /// A user can call `.to_owned()` to promote this into its own ref-counted struct
-    /// and use it after the lifetime of the [CodeReference].
-    pub fn function(&'a self) -> &'a Function {
-        self.func.as_ref()
+    pub(crate) fn into_owned_raw(value: &Self) -> BNReferenceSource {
+        BNReferenceSource {
+            func: match &value.func {
+                Some(func) => func.handle,
+                None => std::ptr::null_mut(),
+            },
+            arch: value.arch.map(|a| a.handle).unwrap_or(std::ptr::null_mut()),
+            addr: value.address,
+        }
+    }
+
+    pub(crate) fn free_raw(value: BNReferenceSource) {
+        let _ = unsafe { Function::ref_from_raw(value.func) };
     }
 
-    /// A handle to the [CodeReference]'s [CoreArchitecture]. This type is [Copy] so reference
-    /// shenanigans are not needed here.
-    pub fn architecture(&self) -> CoreArchitecture {
-        self.arch
+    pub fn new(address: u64, func: Option<Ref<Function>>, arch: Option<CoreArchitecture>) -> Self {
+        Self {
+            func,
+            arch,
+            address,
+        }
     }
 }
 
@@ -56,19 +73,27 @@ impl<'a> CodeReference {
 impl CoreArrayProvider for CodeReference {
     type Raw = BNReferenceSource;
     type Context = ();
-    type Wrapped<'a> = Guard<'a, CodeReference>;
+    type Wrapped<'a> = Self;
 }
 
 unsafe impl CoreArrayProviderInner for CodeReference {
     unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
         BNFreeCodeReferences(raw, count)
     }
+
     unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
-        Guard::new(CodeReference::new(raw), &())
+        CodeReference::from_raw(raw)
     }
 }
 
-// Data Reference Array<T> boilerplate
+// TODO: This only exists so that Array can free.
+// TODO: Is there any way we can have this instead be Array<Location> of some sort?
+/// A struct representing a single data cross-reference.
+/// Data references have no associated metadata, so this object has only
+/// a single [`DataReference::address`] attribute.
+pub struct DataReference {
+    pub address: u64,
+}
 
 impl CoreArrayProvider for DataReference {
     type Raw = u64;
@@ -80,6 +105,7 @@ unsafe impl CoreArrayProviderInner for DataReference {
     unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
         BNFreeDataReferences(raw)
     }
+
     unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
         DataReference { address: *raw }
     }