Add DefineUserTypes to Rust/Python (+ fix return type)

2 files changed, 162 insertions, 0 deletions

diff --git a/rust/src/binaryview.rs b/rust/src/binaryview.rs
index b6a8f16c..31a12105 100644
--- a/rust/src/binaryview.rs
+++ b/rust/src/binaryview.rs
@@ -21,6 +21,8 @@ use binaryninjacore_sys::*;
 
 pub use binaryninjacore_sys::BNModificationStatus as ModificationStatus;
 
+use std::collections::HashMap;
+use std::ffi::c_void;
 use std::ops;
 use std::ops::Range;
 use std::os::raw::c_char;
@@ -142,6 +144,19 @@ pub trait BinaryViewBase: AsRef<BinaryView> {
     }
 }
 
+// TODO: Copied from debuginfo.rs, this should be consolidated
+struct ProgressContext(Option<Box<dyn Fn(usize, usize) -> Result<()>>>);
+
+extern "C" fn cb_progress(ctxt: *mut c_void, cur: usize, max: usize) -> bool {
+    ffi_wrap!("BinaryViewExt::cb_progress", unsafe {
+        let progress = ctxt as *mut ProgressContext;
+        match &(*progress).0 {
+            Some(func) => (func)(cur, max).is_ok(),
+            None => true,
+        }
+    })
+}
+
 pub trait BinaryViewExt: BinaryViewBase {
     fn file(&self) -> Ref<FileMetadata> {
         unsafe {
@@ -531,6 +546,89 @@ pub trait BinaryViewExt: BinaryViewBase {
         }
     }
 
+    fn define_auto_types<S: BnStrCompatible>(
+        &self,
+        names_sources_and_types: Vec<(S, S, &Type)>,
+        progress: Option<Box<dyn Fn(usize, usize) -> Result<()>>>,
+    ) -> HashMap<String, QualifiedName> {
+        let mut names = vec![];
+        let mut ids = vec![];
+        let mut types = vec![];
+        let mut api_types =
+            Vec::<BNQualifiedNameTypeAndId>::with_capacity(names_sources_and_types.len());
+        for (name, source, type_obj) in names_sources_and_types.into_iter() {
+            names.push(QualifiedName::from(name));
+            ids.push(source.into_bytes_with_nul());
+            types.push(type_obj);
+        }
+
+        for ((name, source), type_obj) in names.iter().zip(ids.iter()).zip(types.iter()) {
+            api_types.push(BNQualifiedNameTypeAndId {
+                name: name.0,
+                id: source.as_ref().as_ptr() as *mut _,
+                type_: type_obj.handle,
+            });
+        }
+
+        let mut progress_raw = ProgressContext(progress);
+        let mut result_ids: *mut *mut c_char = ptr::null_mut();
+        let mut result_names: *mut BNQualifiedName = ptr::null_mut();
+        let result_count = unsafe {
+            BNDefineAnalysisTypes(
+                self.as_ref().handle,
+                api_types.as_mut_ptr(),
+                api_types.len(),
+                Some(cb_progress),
+                &mut progress_raw as *mut _ as *mut c_void,
+                &mut result_ids as *mut _,
+                &mut result_names as *mut _,
+            )
+        };
+
+        let mut result = HashMap::with_capacity(result_count);
+
+        let id_array = unsafe { Array::<BnString>::new(result_ids, result_count, ()) };
+        let name_array = unsafe { Array::<QualifiedName>::new(result_names, result_count, ()) };
+
+        for (id, name) in id_array.iter().zip(name_array.iter()) {
+            result.insert(id.as_str().to_owned(), name.clone());
+        }
+
+        result
+    }
+
+    fn define_user_types<S: BnStrCompatible>(
+        &self,
+        names_and_types: Vec<(S, &Type)>,
+        progress: Option<Box<dyn Fn(usize, usize) -> Result<()>>>,
+    ) {
+        let mut names = vec![];
+        let mut types = vec![];
+        let mut api_types = Vec::<BNQualifiedNameAndType>::with_capacity(names_and_types.len());
+        for (name, type_obj) in names_and_types.into_iter() {
+            names.push(QualifiedName::from(name));
+            types.push(type_obj);
+        }
+
+        for (name, type_obj) in names.iter().zip(types.iter()) {
+            api_types.push(BNQualifiedNameAndType {
+                name: name.0,
+                type_: type_obj.handle,
+            });
+        }
+
+        let mut progress_raw = ProgressContext(progress);
+        unsafe {
+            BNDefineUserAnalysisTypes(
+                self.as_ref().handle,
+                api_types.as_mut_ptr(),
+                api_types.len(),
+                Some(cb_progress),
+                &mut progress_raw as *mut _ as *mut c_void,
+            )
+        };
+    }
+
     fn undefine_auto_type<S: BnStrCompatible>(&self, id: S) {
         let id_str = id.into_bytes_with_nul();
         unsafe {
diff --git a/rust/src/types.rs b/rust/src/types.rs
index 9d9e25cb..b2c15ab8 100644
--- a/rust/src/types.rs
+++ b/rust/src/types.rs
@@ -2283,6 +2283,24 @@ impl Drop for QualifiedName {
     }
 }
 
+impl CoreArrayProvider for QualifiedName {
+    type Raw = BNQualifiedName;
+    type Context = ();
+}
+unsafe impl CoreOwnedArrayProvider for QualifiedName {
+    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
+        BNFreeTypeNameList(raw, count);
+    }
+}
+
+unsafe impl<'a> CoreArrayWrapper<'a> for QualifiedName {
+    type Wrapped = &'a QualifiedName;
+
+    unsafe fn wrap_raw(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped {
+        mem::transmute(raw)
+    }
+}
+
 //////////////////////////
 // QualifiedNameAndType
 
@@ -2326,6 +2344,52 @@ unsafe impl<'a> CoreArrayWrapper<'a> for QualifiedNameAndType {
 }
 
 //////////////////////////
+// QualifiedNameTypeAndId
+
+#[repr(transparent)]
+pub struct QualifiedNameTypeAndId(pub(crate) BNQualifiedNameTypeAndId);
+
+impl QualifiedNameTypeAndId {
+    pub fn name(&self) -> &QualifiedName {
+        unsafe { mem::transmute(&self.0.name) }
+    }
+
+    pub fn id(&self) -> &BnStr {
+        unsafe { BnStr::from_raw(self.0.id) }
+    }
+
+    pub fn type_object(&self) -> Guard<Type> {
+        unsafe { Guard::new(Type::from_raw(self.0.type_), self) }
+    }
+}
+
+impl Drop for QualifiedNameTypeAndId {
+    fn drop(&mut self) {
+        unsafe {
+            BNFreeQualifiedNameTypeAndId(&mut self.0);
+        }
+    }
+}
+
+impl CoreArrayProvider for QualifiedNameTypeAndId {
+    type Raw = BNQualifiedNameTypeAndId;
+    type Context = ();
+}
+unsafe impl CoreOwnedArrayProvider for QualifiedNameTypeAndId {
+    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
+        BNFreeTypeIdList(raw, count);
+    }
+}
+
+unsafe impl<'a> CoreArrayWrapper<'a> for QualifiedNameTypeAndId {
+    type Wrapped = &'a QualifiedNameTypeAndId;
+
+    unsafe fn wrap_raw(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped {
+        mem::transmute(raw)
+    }
+}
+
+//////////////////////////
 // NameAndType
 
 pub struct NameAndType<S: BnStrCompatible> {