[Rust] Make TypeLibrary ref counted

And some other misc cleanup

2 files changed, 57 insertions, 92 deletions

diff --git a/rust/src/binary_view.rs b/rust/src/binary_view.rs
index 4379e7f5..cae53e66 100644
--- a/rust/src/binary_view.rs
+++ b/rust/src/binary_view.rs
@@ -1634,10 +1634,10 @@ pub trait BinaryViewExt: BinaryViewBase {
         unsafe { BNAddBinaryViewTypeLibrary(self.as_ref().handle, library.as_raw()) }
     }
 
-    fn type_library_by_name(&self, name: &str) -> Option<TypeLibrary> {
+    fn type_library_by_name(&self, name: &str) -> Option<Ref<TypeLibrary>> {
         let name = name.to_cstr();
         let result = unsafe { BNGetBinaryViewTypeLibrary(self.as_ref().handle, name.as_ptr()) };
-        NonNull::new(result).map(|h| unsafe { TypeLibrary::from_raw(h) })
+        NonNull::new(result).map(|h| unsafe { TypeLibrary::ref_from_raw(h) })
     }
 
     /// Should be called by custom py:py:class:`BinaryView` implementations
@@ -1782,7 +1782,7 @@ pub trait BinaryViewExt: BinaryViewBase {
         &self,
         addr: u64,
         platform: &Platform,
-    ) -> Option<(TypeLibrary, QualifiedName)> {
+    ) -> Option<(Ref<TypeLibrary>, QualifiedName)> {
         let mut result_lib = std::ptr::null_mut();
         let mut result_name = BNQualifiedName::default();
         let success = unsafe {
@@ -1797,7 +1797,7 @@ pub trait BinaryViewExt: BinaryViewBase {
         if !success {
             return None;
         }
-        let lib = unsafe { TypeLibrary::from_raw(NonNull::new(result_lib)?) };
+        let lib = unsafe { TypeLibrary::ref_from_raw(NonNull::new(result_lib)?) };
         let name = QualifiedName::from_owned_raw(result_name);
         Some((lib, name))
     }
@@ -1808,7 +1808,7 @@ pub trait BinaryViewExt: BinaryViewBase {
     fn lookup_imported_type_library<T: Into<QualifiedName>>(
         &self,
         name: T,
-    ) -> Option<(TypeLibrary, QualifiedName)> {
+    ) -> Option<(Ref<TypeLibrary>, QualifiedName)> {
         let raw_name = QualifiedName::into_raw(name.into());
         let mut result_lib = std::ptr::null_mut();
         let mut result_name = BNQualifiedName::default();
@@ -1824,7 +1824,7 @@ pub trait BinaryViewExt: BinaryViewBase {
         if !success {
             return None;
         }
-        let lib = unsafe { TypeLibrary::from_raw(NonNull::new(result_lib)?) };
+        let lib = unsafe { TypeLibrary::ref_from_raw(NonNull::new(result_lib)?) };
         let name = QualifiedName::from_owned_raw(result_name);
         Some((lib, name))
     }
diff --git a/rust/src/type_library.rs b/rust/src/type_library.rs
index 6d09abfa..7f94652d 100644
--- a/rust/src/type_library.rs
+++ b/rust/src/type_library.rs
@@ -1,5 +1,6 @@
 use binaryninjacore_sys::*;
 
+use crate::rc::{Guard, RefCountable};
 use crate::{
     architecture::CoreArchitecture,
     metadata::Metadata,
@@ -8,22 +9,21 @@ use crate::{
     string::{BnString, IntoCStr},
     types::{QualifiedName, QualifiedNameAndType, Type},
 };
-use core::{ffi, mem, ptr};
 use std::path::Path;
+use std::ptr::NonNull;
 
 #[repr(transparent)]
 pub struct TypeLibrary {
-    handle: ptr::NonNull<BNTypeLibrary>,
+    handle: NonNull<BNTypeLibrary>,
 }
 
 impl TypeLibrary {
-    pub(crate) unsafe fn from_raw(handle: ptr::NonNull<BNTypeLibrary>) -> Self {
+    pub(crate) unsafe fn from_raw(handle: NonNull<BNTypeLibrary>) -> Self {
         Self { handle }
     }
 
-    pub(crate) unsafe fn ref_from_raw(handle: &*mut BNTypeLibrary) -> &Self {
-        assert!(!handle.is_null());
-        mem::transmute(handle)
+    pub(crate) unsafe fn ref_from_raw(handle: NonNull<BNTypeLibrary>) -> Ref<Self> {
+        Ref::new(Self { handle })
     }
 
     #[allow(clippy::mut_from_ref)]
@@ -31,22 +31,15 @@ impl TypeLibrary {
         &mut *self.handle.as_ptr()
     }
 
-    pub fn new_reference(&self) -> Self {
-        unsafe {
-            Self::from_raw(ptr::NonNull::new(BNNewTypeLibraryReference(self.as_raw())).unwrap())
-        }
-    }
-
-    pub fn new_duplicated(&self) -> Self {
-        unsafe { Self::from_raw(ptr::NonNull::new(BNDuplicateTypeLibrary(self.as_raw())).unwrap()) }
+    pub fn new_duplicated(&self) -> Ref<Self> {
+        unsafe { Self::ref_from_raw(NonNull::new(BNDuplicateTypeLibrary(self.as_raw())).unwrap()) }
     }
 
     /// Creates an empty type library object with a random GUID and the provided name.
-    pub fn new(arch: CoreArchitecture, name: &str) -> TypeLibrary {
+    pub fn new(arch: CoreArchitecture, name: &str) -> Ref<TypeLibrary> {
         let name = name.to_cstr();
-        let new_lib =
-            unsafe { BNNewTypeLibrary(arch.handle, name.as_ref().as_ptr() as *const ffi::c_char) };
-        unsafe { TypeLibrary::from_raw(ptr::NonNull::new(new_lib).unwrap()) }
+        let new_lib = unsafe { BNNewTypeLibrary(arch.handle, name.as_ptr()) };
+        unsafe { TypeLibrary::ref_from_raw(NonNull::new(new_lib).unwrap()) }
     }
 
     pub fn all(arch: CoreArchitecture) -> Array<TypeLibrary> {
@@ -60,47 +53,35 @@ impl TypeLibrary {
     pub fn decompress_to_file(path: &Path, output_path: &Path) -> bool {
         let path = path.to_cstr();
         let output = output_path.to_cstr();
-        unsafe {
-            BNTypeLibraryDecompressToFile(
-                path.as_ref().as_ptr() as *const ffi::c_char,
-                output.as_ref().as_ptr() as *const ffi::c_char,
-            )
-        }
+        unsafe { BNTypeLibraryDecompressToFile(path.as_ptr(), output.as_ptr()) }
     }
 
     /// Loads a finalized type library instance from file
-    pub fn load_from_file(path: &Path) -> Option<TypeLibrary> {
+    pub fn load_from_file(path: &Path) -> Option<Ref<TypeLibrary>> {
         let path = path.to_cstr();
-        let handle =
-            unsafe { BNLoadTypeLibraryFromFile(path.as_ref().as_ptr() as *const ffi::c_char) };
-        ptr::NonNull::new(handle).map(|h| unsafe { TypeLibrary::from_raw(h) })
+        let handle = unsafe { BNLoadTypeLibraryFromFile(path.as_ptr()) };
+        NonNull::new(handle).map(|h| unsafe { TypeLibrary::ref_from_raw(h) })
     }
 
     /// Saves a finalized type library instance to file
     pub fn write_to_file(&self, path: &Path) -> bool {
         let path = path.to_cstr();
-        unsafe {
-            BNWriteTypeLibraryToFile(self.as_raw(), path.as_ref().as_ptr() as *const ffi::c_char)
-        }
+        unsafe { BNWriteTypeLibraryToFile(self.as_raw(), path.as_ptr()) }
     }
 
     /// Looks up the first type library found with a matching name. Keep in mind that names are not
     /// necessarily unique.
-    pub fn from_name(arch: CoreArchitecture, name: &str) -> Option<TypeLibrary> {
+    pub fn from_name(arch: CoreArchitecture, name: &str) -> Option<Ref<TypeLibrary>> {
         let name = name.to_cstr();
-        let handle = unsafe {
-            BNLookupTypeLibraryByName(arch.handle, name.as_ref().as_ptr() as *const ffi::c_char)
-        };
-        ptr::NonNull::new(handle).map(|h| unsafe { TypeLibrary::from_raw(h) })
+        let handle = unsafe { BNLookupTypeLibraryByName(arch.handle, name.as_ptr()) };
+        NonNull::new(handle).map(|h| unsafe { TypeLibrary::ref_from_raw(h) })
     }
 
     /// Attempts to grab a type library associated with the provided Architecture and GUID pair
-    pub fn from_guid(arch: CoreArchitecture, guid: &str) -> Option<TypeLibrary> {
+    pub fn from_guid(arch: CoreArchitecture, guid: &str) -> Option<Ref<TypeLibrary>> {
         let guid = guid.to_cstr();
-        let handle = unsafe {
-            BNLookupTypeLibraryByGuid(arch.handle, guid.as_ref().as_ptr() as *const ffi::c_char)
-        };
-        ptr::NonNull::new(handle).map(|h| unsafe { TypeLibrary::from_raw(h) })
+        let handle = unsafe { BNLookupTypeLibraryByGuid(arch.handle, guid.as_ptr()) };
+        NonNull::new(handle).map(|h| unsafe { TypeLibrary::ref_from_raw(h) })
     }
 
     /// The Architecture this type library is associated with
@@ -119,9 +100,7 @@ impl TypeLibrary {
     /// Sets the name of a type library instance that has not been finalized
     pub fn set_name(&self, value: &str) {
         let value = value.to_cstr();
-        unsafe {
-            BNSetTypeLibraryName(self.as_raw(), value.as_ref().as_ptr() as *const ffi::c_char)
-        }
+        unsafe { BNSetTypeLibraryName(self.as_raw(), value.as_ptr()) }
     }
 
     /// The `dependency_name` of a library is the name used to record dependencies across
@@ -137,12 +116,7 @@ impl TypeLibrary {
     /// Sets the dependency name of a type library instance that has not been finalized
     pub fn set_dependency_name(&self, value: &str) {
         let value = value.to_cstr();
-        unsafe {
-            BNSetTypeLibraryDependencyName(
-                self.as_raw(),
-                value.as_ref().as_ptr() as *const ffi::c_char,
-            )
-        }
+        unsafe { BNSetTypeLibraryDependencyName(self.as_raw(), value.as_ptr()) }
     }
 
     /// Returns the GUID associated with the type library
@@ -154,9 +128,7 @@ impl TypeLibrary {
     /// Sets the GUID of a type library instance that has not been finalized
     pub fn set_guid(&self, value: &str) {
         let value = value.to_cstr();
-        unsafe {
-            BNSetTypeLibraryGuid(self.as_raw(), value.as_ref().as_ptr() as *const ffi::c_char)
-        }
+        unsafe { BNSetTypeLibraryGuid(self.as_raw(), value.as_ptr()) }
     }
 
     /// A list of extra names that will be considered a match by [Platform::get_type_libraries_by_name]
@@ -170,12 +142,7 @@ impl TypeLibrary {
     /// Adds an extra name to this type library used during library lookups and dependency resolution
     pub fn add_alternate_name(&self, value: &str) {
         let value = value.to_cstr();
-        unsafe {
-            BNAddTypeLibraryAlternateName(
-                self.as_raw(),
-                value.as_ref().as_ptr() as *const ffi::c_char,
-            )
-        }
+        unsafe { BNAddTypeLibraryAlternateName(self.as_raw(), value.as_ptr()) }
     }
 
     /// Returns a list of all platform names that this type library will register with during platform
@@ -214,9 +181,7 @@ impl TypeLibrary {
     /// Retrieves a metadata associated with the given key stored in the type library
     pub fn query_metadata(&self, key: &str) -> Option<Metadata> {
         let key = key.to_cstr();
-        let result = unsafe {
-            BNTypeLibraryQueryMetadata(self.as_raw(), key.as_ref().as_ptr() as *const ffi::c_char)
-        };
+        let result = unsafe { BNTypeLibraryQueryMetadata(self.as_raw(), key.as_ptr()) };
         (!result.is_null()).then(|| unsafe { Metadata::from_raw(result) })
     }
 
@@ -233,21 +198,13 @@ impl TypeLibrary {
     /// * `md` - object to store.
     pub fn store_metadata(&self, key: &str, md: &Metadata) {
         let key = key.to_cstr();
-        unsafe {
-            BNTypeLibraryStoreMetadata(
-                self.as_raw(),
-                key.as_ref().as_ptr() as *const ffi::c_char,
-                md.handle,
-            )
-        }
+        unsafe { BNTypeLibraryStoreMetadata(self.as_raw(), key.as_ptr(), md.handle) }
     }
 
     /// Removes the metadata associated with key from the current type library.
     pub fn remove_metadata(&self, key: &str) {
         let key = key.to_cstr();
-        unsafe {
-            BNTypeLibraryRemoveMetadata(self.as_raw(), key.as_ref().as_ptr() as *const ffi::c_char)
-        }
+        unsafe { BNTypeLibraryRemoveMetadata(self.as_raw(), key.as_ptr()) }
     }
 
     /// Retrieves the metadata associated with the current type library.
@@ -262,7 +219,7 @@ impl TypeLibrary {
     // /// The Type Container's Platform will be the first platform associated with the Type Library.
     // pub fn type_container(&self) -> TypeContainer {
     //     let result = unsafe{ BNGetTypeLibraryTypeContainer(self.as_raw())};
-    //     unsafe{TypeContainer::from_raw(ptr::NonNull::new(result).unwrap())}
+    //     unsafe{TypeContainer::from_raw(NonNull::new(result).unwrap())}
     // }
 
     /// Directly inserts a named object into the type library's object store.
@@ -302,13 +259,7 @@ impl TypeLibrary {
     pub fn add_type_source(&self, name: QualifiedName, source: &str) {
         let source = source.to_cstr();
         let mut raw_name = QualifiedName::into_raw(name);
-        unsafe {
-            BNAddTypeLibraryNamedTypeSource(
-                self.as_raw(),
-                &mut raw_name,
-                source.as_ref().as_ptr() as *const ffi::c_char,
-            )
-        }
+        unsafe { BNAddTypeLibraryNamedTypeSource(self.as_raw(), &mut raw_name, source.as_ptr()) }
         QualifiedName::free_raw(raw_name);
     }
 
@@ -349,16 +300,30 @@ impl TypeLibrary {
     }
 }
 
-impl Drop for TypeLibrary {
-    fn drop(&mut self) {
-        unsafe { BNFreeTypeLibrary(self.as_raw()) }
+unsafe impl RefCountable for TypeLibrary {
+    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
+        Ref::new(Self {
+            handle: NonNull::new(BNNewTypeLibraryReference(handle.handle.as_ptr())).unwrap(),
+        })
+    }
+
+    unsafe fn dec_ref(handle: &Self) {
+        BNFreeTypeLibrary(handle.handle.as_ptr());
+    }
+}
+
+impl ToOwned for TypeLibrary {
+    type Owned = Ref<Self>;
+
+    fn to_owned(&self) -> Self::Owned {
+        unsafe { RefCountable::inc_ref(self) }
     }
 }
 
 impl CoreArrayProvider for TypeLibrary {
     type Raw = *mut BNTypeLibrary;
     type Context = ();
-    type Wrapped<'a> = &'a Self;
+    type Wrapped<'a> = Guard<'a, Self>;
 }
 
 unsafe impl CoreArrayProviderInner for TypeLibrary {
@@ -366,7 +331,7 @@ unsafe impl CoreArrayProviderInner for TypeLibrary {
         BNFreeTypeLibraryList(raw, count)
     }
 
-    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
-        Self::ref_from_raw(raw)
+    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> {
+        Guard::new(Self::from_raw(NonNull::new(*raw).unwrap()), context)
     }
 }