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

6 files changed, 265 insertions, 4 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index cb33aeff..efe9a22a 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -5219,7 +5219,7 @@ namespace BinaryNinja {
 		QualifiedName GetTypeNameById(const std::string& id);
 		bool IsTypeAutoDefined(const QualifiedName& name);
 		QualifiedName DefineType(const std::string& id, const QualifiedName& defaultName, Ref<Type> type);
-		void DefineTypes(const std::vector<std::pair<std::string, QualifiedNameAndType>>& types, std::function<bool(size_t, size_t)> progress = {});
+		std::unordered_map<std::string, QualifiedName> DefineTypes(const std::vector<std::pair<std::string, QualifiedNameAndType>>& types, std::function<bool(size_t, size_t)> progress = {});
 		void DefineUserType(const QualifiedName& name, Ref<Type> type);
 		void DefineUserTypes(const std::vector<QualifiedNameAndType>& types, std::function<bool(size_t, size_t)> progress = {});
 		void DefineUserTypes(const std::vector<ParsedType>& types, std::function<bool(size_t, size_t)> progress = {});
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 246918fc..3120e4a7 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -4371,6 +4371,7 @@ extern "C"
 		BNQualifiedNameList* typesAllowRedefinition);
 	BINARYNINJACOREAPI void BNFreeQualifiedNameAndType(BNQualifiedNameAndType* obj);
 	BINARYNINJACOREAPI void BNFreeQualifiedNameAndTypeArray(BNQualifiedNameAndType* obj, size_t count);
+	BINARYNINJACOREAPI void BNFreeQualifiedNameTypeAndId(BNQualifiedNameTypeAndId* obj);
 	BINARYNINJACOREAPI char* BNEscapeTypeName(const char* name, BNTokenEscapingType escaping);
 	BINARYNINJACOREAPI char* BNUnescapeTypeName(const char* name, BNTokenEscapingType escaping);
 
@@ -4389,7 +4390,7 @@ extern "C"
 	BINARYNINJACOREAPI BNQualifiedName BNDefineAnalysisType(
 	    BNBinaryView* view, const char* id, BNQualifiedName* defaultName, BNType* type);
 	BINARYNINJACOREAPI void BNDefineUserAnalysisType(BNBinaryView* view, BNQualifiedName* name, BNType* type);
-	BINARYNINJACOREAPI void BNDefineAnalysisTypes(BNBinaryView* view, BNQualifiedNameTypeAndId* types, size_t count, bool (*progress)(void*, size_t, size_t), void* progressContext);
+	BINARYNINJACOREAPI size_t BNDefineAnalysisTypes(BNBinaryView* view, BNQualifiedNameTypeAndId* types, size_t count, bool (*progress)(void*, size_t, size_t), void* progressContext, char*** resultIds, BNQualifiedName** resultNames);
 	BINARYNINJACOREAPI void BNDefineUserAnalysisTypes(BNBinaryView* view, BNQualifiedNameAndType* types, size_t count, bool (*progress)(void*, size_t, size_t), void* progressContext);
 	BINARYNINJACOREAPI void BNUndefineAnalysisType(BNBinaryView* view, const char* id);
 	BINARYNINJACOREAPI void BNUndefineUserAnalysisType(BNBinaryView* view, BNQualifiedName* name);
diff --git a/binaryview.cpp b/binaryview.cpp
index b1bd6f4d..6120e10f 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -3731,7 +3731,7 @@ void BinaryView::DefineUserType(const QualifiedName& name, Ref<Type> type)
 }
 
 
-void BinaryView::DefineTypes(const vector<pair<string, QualifiedNameAndType>>& types, std::function<bool(size_t, size_t)> progress)
+std::unordered_map<std::string, QualifiedName> BinaryView::DefineTypes(const vector<pair<string, QualifiedNameAndType>>& types, std::function<bool(size_t, size_t)> progress)
 {
 	BNQualifiedNameTypeAndId* apiTypes = new BNQualifiedNameTypeAndId[types.size()];
 	for (size_t i = 0; i < types.size(); i++)
@@ -3743,7 +3743,20 @@ void BinaryView::DefineTypes(const vector<pair<string, QualifiedNameAndType>>& t
 
 	ProgressContext cb;
 	cb.callback = progress;
-	BNDefineAnalysisTypes(m_object, apiTypes, types.size(), ProgressCallback, &cb);
+	char** resultIds;
+	BNQualifiedName* resultNames;
+	size_t resultCount = BNDefineAnalysisTypes(m_object, apiTypes, types.size(), ProgressCallback, &cb, &resultIds, &resultNames);
+
+	unordered_map<string, QualifiedName> result;
+	for (size_t i = 0; i < resultCount; i ++)
+	{
+		string id = resultIds[i];
+		QualifiedName name = QualifiedName::FromAPIObject(&resultNames[i]);
+		result.insert({id, name});
+	}
+
+	BNFreeStringList(resultIds, resultCount);
+	BNFreeTypeNameList(resultNames, resultCount);
 
 	for (size_t i = 0; i < types.size(); i++)
 	{
@@ -3751,6 +3764,8 @@ void BinaryView::DefineTypes(const vector<pair<string, QualifiedNameAndType>>& t
 		BNFreeString(apiTypes[i].id);
 	}
 	delete [] apiTypes;
+
+	return result;
 }
 
 
diff --git a/python/binaryview.py b/python/binaryview.py
index e84b40d1..5573b81f 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -7061,6 +7061,89 @@ class BinaryView:
 		_name = _types.QualifiedName(name)._to_core_struct()
 		core.BNDefineUserAnalysisType(self.handle, _name, type_obj.handle)
 
+	def define_types(self, types: List[Tuple[str, Optional['_types.QualifiedNameType'], StringOrType]], progress_func: Optional[ProgressFuncType]) -> Mapping[str, '_types.QualifiedName']:
+		"""
+		``define_types`` registers multiple types as though calling :py:func:`define_type` multiple times.
+		The difference with this plural version is that it is optimized for adding many types
+		at the same time, using knowledge of all types at add-time to improve runtime.
+		There is an optional ``progress_func`` callback function in case you want updates for a long-running call.
+
+		.. warning:: This method should only be used for automatically generated types, see :py:func:`define_user_types` for interactive plugin uses.
+
+		The return values of this function provide a map of each type id and which name was chosen for that type
+		(which may be different from the requested name).
+
+		:param types: List of type ids/names/definitions for the new types. Check :py:func:`define_type` for more details.
+		:param progress: Function to call for progress updates
+		:return: A map of all the chosen names for the defined types with their ids.
+		"""
+		api_types = (core.BNQualifiedNameTypeAndId * len(types))()
+		for i, (type_id, default_name, type_obj) in enumerate(types):
+			if isinstance(type_obj, str):
+				(type_obj, new_name) = self.parse_type_string(type_obj)
+				if default_name is None:
+					default_name = new_name
+			assert default_name is not None, "default_name can only be None if named type is derived from string passed to type_obj"
+			api_types[i].name = _types.QualifiedName(default_name)._to_core_struct()
+			api_types[i].id = core.cstr(type_id)
+			api_types[i].type = type_obj.handle
+
+		if progress_func:
+			progress_func_obj = ctypes.CFUNCTYPE(
+				ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong
+			)(lambda ctxt, cur, total: progress_func(cur, total))
+		else:
+			progress_func_obj = ctypes.CFUNCTYPE(
+				ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong
+			)(lambda ctxt, cur, total: True)
+
+		result_ids = ctypes.POINTER(ctypes.c_char_p)()
+		result_names = ctypes.POINTER(core.BNQualifiedName)()
+
+		result_count = core.BNDefineAnalysisTypes(self.handle, api_types, len(types), progress_func_obj, None, result_ids, result_names)
+
+		try:
+			result = {}
+			for i in range(result_count):
+				id = core.pyNativeStr(result_ids[i])
+				name = _types.QualifiedName._from_core_struct(result_names[i])
+				result[id] = name
+			return result
+		finally:
+			core.BNFreeStringList(result_ids, result_count)
+			core.BNFreeTypeNameList(result_names, result_count)
+
+	def define_user_types(self, types: List[Tuple[Optional['_types.QualifiedNameType'], StringOrType]], progress_func: Optional[ProgressFuncType]):
+		"""
+		``define_user_types`` registers multiple types as though calling :py:func:`define_user_type` multiple times.
+		The difference with this plural version is that it is optimized for adding many types
+		at the same time, using knowledge of all types at add-time to improve runtime.
+		There is an optional ``progress_func`` callback function in case you want updates for a long-running call.
+
+		:param types: List of type names/definitions for the new types. Check :py:func:`define_user_type` for more details.
+		:param progress: Function to call for progress updates
+		"""
+		api_types = (core.BNQualifiedNameAndType * len(types))()
+		for i, (default_name, type_obj) in enumerate(types):
+			if isinstance(type_obj, str):
+				(type_obj, new_name) = self.parse_type_string(type_obj)
+				if default_name is None:
+					default_name = new_name
+			assert default_name is not None, "default_name can only be None if named type is derived from string passed to type_obj"
+			api_types[i].name = _types.QualifiedName(default_name)._to_core_struct()
+			api_types[i].type = type_obj.handle
+
+		if progress_func:
+			progress_func_obj = ctypes.CFUNCTYPE(
+				ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong
+			)(lambda ctxt, cur, total: progress_func(cur, total))
+		else:
+			progress_func_obj = ctypes.CFUNCTYPE(
+				ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong
+			)(lambda ctxt, cur, total: True)
+
+		core.BNDefineUserAnalysisTypes(self.handle, api_types, len(types), progress_func_obj, None)
+
 	def undefine_type(self, type_id: str) -> None:
 		"""
 		``undefine_type`` removes a :py:class:`Type` from the global list of types for the current :py:class:`BinaryView`
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> {