Expose Template Simplifier to the API, Adds Tests, and Adds Rust String FFI support to Python (see note)

This introduces the ability to receive and free string from Rust in Python.
For it to work, your exposed function name needs to begin with "BNRust", and the generator will create the appropriate code for freeing string return values and arguments.

6 files changed, 257 insertions, 15 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 8604e44f..fb5e77fc 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -704,8 +704,12 @@ __attribute__ ((format (printf, 1, 2)))
 	class BinaryView;
 
 	bool DemangleMS(Architecture* arch, const std::string& mangledName, Type** outType,
+		QualifiedName& outVarName, const bool simplify = false);
+	bool DemangleMS(Architecture* arch, const std::string& mangledName, Type** outType,
 		QualifiedName& outVarName, const Ref<BinaryView>& view);
 	bool DemangleGNU3(Ref<Architecture> arch, const std::string& mangledName, Type** outType,
+		QualifiedName& outVarName, const bool simplify = false);
+	bool DemangleGNU3(Ref<Architecture> arch, const std::string& mangledName, Type** outType,
 		QualifiedName& outVarName, const Ref<BinaryView>& view);
 
 	void RegisterMainThread(MainThreadActionHandler* handler);
@@ -5401,4 +5405,32 @@ __attribute__ ((format (printf, 1, 2)))
 
 		static int Compare(LinearViewCursor* a, LinearViewCursor* b);
 	};
+
+	class SimplifyName
+	{
+	public:
+		// Use these functions to interface with the simplifier
+		static std::string        to_string(const std::string& input);
+		static std::string        to_string(const QualifiedName& input);
+		static QualifiedName to_qualified_name(const std::string& input, bool simplify);
+		static QualifiedName to_qualified_name(const QualifiedName& input);
+
+		// Below is everything for the above APIs to work
+		enum SimplifierDest
+		{
+			str,
+			fqn
+		};
+
+		SimplifyName(const std::string&, const SimplifierDest, const bool);
+		~SimplifyName();
+
+		operator std::string() const;
+		operator QualifiedName();
+
+	private:
+		const char*  m_rust_string;
+		const char** m_rust_array;
+		uint64_t m_length;
+	};
 }
diff --git a/binaryninjacore.h b/binaryninjacore.h
index b13f9fb3..a34d1f7e 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -4464,6 +4464,8 @@ __attribute__ ((format (printf, 1, 2)))
 
 	//Demangler
 	BINARYNINJACOREAPI bool BNDemangleMS(BNArchitecture* arch, const char* mangledName, BNType** outType, char*** outVarName,
+		size_t* outVarNameElements, const bool simplify);
+	BINARYNINJACOREAPI bool BNDemangleMSWithOptions(BNArchitecture* arch, const char* mangledName, BNType** outType, char*** outVarName,
 		size_t* outVarNameElements, const BNBinaryView* const view);
 
 	// Download providers
@@ -4607,6 +4609,8 @@ __attribute__ ((format (printf, 1, 2)))
 
 	BINARYNINJACOREAPI bool BNIsGNU3MangledString(const char* mangledName);
 	BINARYNINJACOREAPI bool BNDemangleGNU3(BNArchitecture* arch, const char* mangledName, BNType** outType,
+		char*** outVarName, size_t* outVarNameElements, const bool simplify);
+	BINARYNINJACOREAPI bool BNDemangleGNU3WithOptions(BNArchitecture* arch, const char* mangledName, BNType** outType,
 		char*** outVarName, size_t* outVarNameElements, const BNBinaryView* const view);
 	BINARYNINJACOREAPI void BNFreeDemangledName(char*** name, size_t nameElements);
 
@@ -4872,6 +4876,11 @@ __attribute__ ((format (printf, 1, 2)))
 
 	BINARYNINJACOREAPI uint32_t BNGetAddressRenderedWidth(uint64_t addr);
 
+	BINARYNINJACOREAPI void BNRustFreeString(const char* const);
+	BINARYNINJACOREAPI void BNRustFreeStringArray(const char** const, uint64_t);
+	BINARYNINJACOREAPI const char** const BNRustSimplifyStrToFQN(const char* const, bool);
+	BINARYNINJACOREAPI const char* const BNRustSimplifyStrToStr(const char* const);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/demangle.cpp b/demangle.cpp
index 9ea2bfa0..2cdc0288 100644
--- a/demangle.cpp
+++ b/demangle.cpp
@@ -1,5 +1,5 @@
 #include "binaryninjaapi.h"
-
+#include <string>
 using namespace std;
 
 namespace BinaryNinja
@@ -7,10 +7,17 @@ namespace BinaryNinja
 	bool DemangleMS(Architecture* arch, const std::string& mangledName, Type** outType,
 		QualifiedName& outVarName, const Ref<BinaryView>& view)
 	{
+		const bool simplify = Settings::Instance()->Get<bool>("analysis.types.TemplateSimplifier", view);
+		return DemangleMS(arch, mangledName, outType, outVarName, simplify);
+	}
+
+	bool DemangleMS(Architecture* arch, const std::string& mangledName, Type** outType,
+		QualifiedName& outVarName, const bool simplify)
+	{
 		BNType* localType = nullptr;
 		char** localVarName = nullptr;
 		size_t localSize = 0;
-		if (!BNDemangleMS(arch->GetObject(), mangledName.c_str(), &localType, &localVarName, &localSize, view->GetObject()))
+		if (!BNDemangleMS(arch->GetObject(), mangledName.c_str(), &localType, &localVarName, &localSize, simplify))
 			return false;
 		if (!localType)
 			return false;
@@ -24,14 +31,20 @@ namespace BinaryNinja
 		return true;
 	}
 
-
 	bool DemangleGNU3(Ref<Architecture> arch, const std::string& mangledName, Type** outType,
 		QualifiedName& outVarName, const Ref<BinaryView>& view)
 	{
+		const bool simplify = Settings::Instance()->Get<bool>("analysis.types.TemplateSimplifier", view);
+		return DemangleGNU3(arch, mangledName, outType, outVarName, simplify);
+	}
+
+	bool DemangleGNU3(Ref<Architecture> arch, const std::string& mangledName, Type** outType,
+		QualifiedName& outVarName, const bool simplify)
+	{
 		BNType* localType;
 		char** localVarName = nullptr;
 		size_t localSize = 0;
-		if (!BNDemangleGNU3(arch->GetObject(), mangledName.c_str(), &localType, &localVarName, &localSize, view->GetObject()))
+		if (!BNDemangleGNU3(arch->GetObject(), mangledName.c_str(), &localType, &localVarName, &localSize, simplify))
 			return false;
 		if (!localType)
 			return false;
@@ -44,4 +57,77 @@ namespace BinaryNinja
 		delete [] localVarName;
 		return true;
 	}
+
+
+	string SimplifyName::to_string(const string& input)
+	{
+		return (string)SimplifyName(input, SimplifierDest::str, true);
+	}
+
+
+	string SimplifyName::to_string(const QualifiedName& input)
+	{
+		return (string)SimplifyName(input.GetString(), SimplifierDest::str, true);
+	}
+
+
+	QualifiedName SimplifyName::to_qualified_name(const string& input, bool simplify)
+	{
+		return (QualifiedName)SimplifyName(input, SimplifierDest::fqn, simplify);
+	}
+
+
+	QualifiedName SimplifyName::to_qualified_name(const QualifiedName& input)
+	{
+		return (QualifiedName)SimplifyName(input.GetString(), SimplifierDest::fqn, true);
+	}
+
+
+	SimplifyName::SimplifyName(const string& input, const SimplifierDest dest, const bool simplify) :
+			m_rust_string(nullptr), m_rust_array(nullptr), m_length(0)
+	{
+		if (dest == SimplifierDest::str)
+			m_rust_string = BNRustSimplifyStrToStr(input.c_str());
+		else
+			m_rust_array = BNRustSimplifyStrToFQN(input.c_str(), simplify);
+	}
+
+
+	SimplifyName::~SimplifyName()
+	{
+		if (m_rust_string)
+			BNRustFreeString(m_rust_string);
+		if (m_rust_array)
+		{
+			if (m_length == 0)
+			{
+				// Should never reach here
+				LogWarn("Deallocating SimplifyName without having been used; Likely misuse of API.\n");
+				uint64_t index = 0;
+				while (m_rust_array[index][0] != 0x0)
+					++index;
+				m_length = index + 1;
+			}
+			BNRustFreeStringArray(m_rust_array, m_length);
+		}
+	}
+
+
+	SimplifyName::operator string() const
+	{
+		return string(m_rust_string);
+	}
+
+
+	SimplifyName::operator QualifiedName()
+	{
+		QualifiedName result;
+		uint64_t      index = 0;
+		while (m_rust_array[index][0] != 0x0)
+		{
+			result.push_back(string(m_rust_array[index++]));
+		}
+		m_length = index;
+		return result;
+	}
 }
diff --git a/python/demangle.py b/python/demangle.py
index 67b211df..912d1c62 100644
--- a/python/demangle.py
+++ b/python/demangle.py
@@ -48,12 +48,14 @@ def get_qualified_name(names):
 	return "::".join(names)
 
 
-def demangle_ms(arch, mangled_name, view = None):
+def demangle_ms(arch, mangled_name, options = False):
 	"""
 	``demangle_ms`` demangles a mangled Microsoft Visual Studio C++ name to a Type object.
 
 	:param Architecture arch: Architecture for the symbol. Required for pointer and integer sizes.
 	:param str mangled_name: a mangled Microsoft Visual Studio C++ name
+	:param options: (optional) Whether to simplify demangled names : None falls back to user settings, a BinaryView uses that BinaryView's settings, or a boolean to set it directally
+	:type options: Tuple[bool, BinaryView, None]
 	:return: returns tuple of (Type, demangled_name) or (None, mangled_name) on error
 	:rtype: Tuple
 	:Example:
@@ -66,9 +68,9 @@ def demangle_ms(arch, mangled_name, view = None):
 	outName = ctypes.POINTER(ctypes.c_char_p)()
 	outSize = ctypes.c_ulonglong()
 	names = []
-	if isinstance(view, BinaryView):
-		view = view.handle
-	if core.BNDemangleMS(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize), view):
+	if (isinstance(options, BinaryView) and core.BNDemangleMSWithOptions(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize), options)) or \
+		(isinstance(options, bool) and core.BNDemangleMS(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize), options)) or \
+		(options is None and core.BNDemangleMSWithOptions(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize), None)):
 		for i in range(outSize.value):
 			names.append(pyNativeStr(outName[i]))
 		core.BNFreeDemangledName(ctypes.byref(outName), outSize.value)
@@ -76,14 +78,24 @@ def demangle_ms(arch, mangled_name, view = None):
 	return (None, mangled_name)
 
 
-def demangle_gnu3(arch, mangled_name, view = None):
+def demangle_gnu3(arch, mangled_name, options = None):
+	"""
+	``demangle_gnu3`` demangles a mangled name to a Type object.
+
+	:param Architecture arch: Architecture for the symbol. Required for pointer and integer sizes.
+	:param str mangled_name: a mangled GNU3 name
+	:param options: (optional) Whether to simplify demangled names : None falls back to user settings, a BinaryView uses that BinaryView's settings, or a boolean to set it directally
+	:type options: Tuple[bool, BinaryView, None]
+	:return: returns tuple of (Type, demangled_name) or (None, mangled_name) on error
+	:rtype: Tuple
+	"""
 	handle = ctypes.POINTER(core.BNType)()
 	outName = ctypes.POINTER(ctypes.c_char_p)()
 	outSize = ctypes.c_ulonglong()
 	names = []
-	if isinstance(view, BinaryView):
-		view = view.handle
-	if core.BNDemangleGNU3(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize), view):
+	if (isinstance(options, BinaryView) and core.BNDemangleGNU3WithOptions(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize), options)) or \
+		(isinstance(options, bool) and core.BNDemangleGNU3(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize), options)) or \
+		(options is None and core.BNDemangleGNU3WithOptions(arch.handle, mangled_name, ctypes.byref(handle), ctypes.byref(outName), ctypes.byref(outSize), None)):
 		for i in range(outSize.value):
 			names.append(pyNativeStr(outName[i]))
 		core.BNFreeDemangledName(ctypes.byref(outName), outSize.value)
@@ -91,3 +103,62 @@ def demangle_gnu3(arch, mangled_name, view = None):
 			return (None, names)
 		return (types.Type(handle), names)
 	return (None, mangled_name)
+
+
+def simplify_name_to_string(input_name):
+	"""
+	``simplify_name_to_string`` simplifies a templated C++ name with default arguments and returns a string
+
+	:param input_name: String or qualified name to be simplified
+	:type input_name: Union[str, QualifiedName]
+	:return: simplified name (or original name if simplifier fails/cannot simplify)
+	:rtype: str
+	:Example:
+
+		>>> bdemangle.simplify_name_to_string("std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >")
+		'std::string'
+		>>>
+	"""
+	result = None
+	if isinstance(input_name, str):
+		result = core.BNRustSimplifyStrToStr(input_name)
+	elif isinstance(input_name, types.QualifiedName):
+		result = core.BNRustSimplifyStrToStr(str(input_name))
+	else:
+		raise TypeError("Parameter must be of type `str` or `types.QualifiedName`")
+	return result
+
+
+def simplify_name_to_qualified_name(input_name, simplify = True):
+	"""
+	``simplify_name_to_qualified_name`` simplifies a templated C++ name with default arguments and returns a qualified name.  This can also tokenize a string to a qualified name with/without simplifying it
+
+	:param input_name: String or qualified name to be simplified
+	:type input_name: Union[str, QualifiedName]
+	:param bool simplify_name: (optional) Whether to simplify input string (no effect if given a qualified name; will always simplify)
+	:return: simplified name (or one-element array containing the input if simplifier fails/cannot simplify)
+	:rtype: QualifiedName
+	:Example:
+
+		>>> demangle.simplify_name_to_qualified_name(QualifiedName(["std", "__cxx11", "basic_string<wchar, std::char_traits<wchar>, std::allocator<wchar> >"]), True)
+		'std::wstring'
+		>>>
+	"""
+	result = None
+	if isinstance(input_name, str):
+		result = core.BNRustSimplifyStrToFQN(input_name, simplify)
+	elif isinstance(input_name, types.QualifiedName):
+		result = core.BNRustSimplifyStrToFQN(str(input_name), True)
+	else:
+		raise TypeError("Parameter must be of type `str` or `types.QualifiedName`")
+
+	native_result = []
+	for name in result:
+		if name == b'':
+			break
+		native_result.append(name)
+	name_count = len(native_result)
+
+	native_result = types.QualifiedName(native_result)
+	core.BNRustFreeStringArray(result, name_count+1)
+	return native_result
diff --git a/python/generator.cpp b/python/generator.cpp
index ff1b153d..4f329299 100644
--- a/python/generator.cpp
+++ b/python/generator.cpp
@@ -350,9 +350,12 @@ int main(int argc, char* argv[])
 					break;
 				}
 		}
-		if (name == "BNFreeString")
+		if (name == "BNFreeString" || name == "BNRustFreeString")
 			stringArgument = false;
 
+		// Rust-allocated strings are deallocated differently
+		bool rustFFI = name.rfind("BNRust", 0) == 0;
+
 		bool callbackConvention = false;
 		if (name == "BNAllocString")
 		{
@@ -377,7 +380,7 @@ int main(int argc, char* argv[])
 			for (auto& j : i.second->GetParameters())
 			{
 				fprintf(out, "\t\t");
-				if (name == "BNFreeString")
+				if (name == "BNFreeString" || name == "BNRustFreeString")
 				{
 					// BNFreeString expects a pointer to a string allocated by the core, so do not use
 					// a c_char_p here, as that would be allocated by the Python runtime.  This can
@@ -440,7 +443,10 @@ int main(int argc, char* argv[])
 			else
 				fprintf(out, "\tresult = %s(*args)\n", funcName.c_str());
 			fprintf(out, "\tstring = str(pyNativeStr(ctypes.cast(result, ctypes.c_char_p).value))\n");
-			fprintf(out, "\tBNFreeString(result)\n");
+			if (rustFFI)
+				fprintf(out, "\tBNRustFreeString(result)\n");
+			else
+				fprintf(out, "\tBNFreeString(result)\n");
 			fprintf(out, "\treturn string\n");
 		}
 		else if (pointerResult)
diff --git a/suite/testcommon.py b/suite/testcommon.py
index 09f9d589..b7251d42 100644
--- a/suite/testcommon.py
+++ b/suite/testcommon.py
@@ -590,6 +590,44 @@ class TestBuilder(Builder):
         testfunction = binja.Type.function(inttype, [inttype, inttype, inttype])
         return ["Test_function params: " + str(testfunction.parameters), "Test_function pointer: " + str(testfunction.pointer(binja.Architecture["x86"], testfunction))]
 
+    def test_Simplifier(self):
+        """Template Simplification"""
+        result = [binja.demangle.simplify_name_to_string(s) for s in [
+            # Simple
+            "std::__cxx11::basic_string<T, std::char_traits<T>, std::allocator<T> >",
+            "std::vector<T, std::allocator<T> >",
+            "std::vector<T, std::allocator<T>, std::lessthan<T> >",
+            "std::deque<T, std::allocator<T> >",
+            "std::forward_list<T, std::allocator<T> >",
+            "std::list<T, std::allocator<T> >",
+            "std::stack<T, std::deque<T> >",
+            "std::queue<T, std::deque<T> >",
+            "std::set<T, std::less<T>, std::allocator<T> >",
+            "std::multiset<T, std::less<T>, std::allocator<T> >",
+            "std::map<T1, T2, std::less<T1>, std::allocator<std::pair<const T1, T2> > >",
+            "std::multimap<T1, T2, std::less<T1>, std::allocator<std::pair<const T1, T2> > >",
+            "std::unordered_set<T, std::hash<T>, std::equal_to<T>, std::allocator<T> >",
+            "std::unordered_multiset<T, std::hash<T>, std::equal_to<T>, std::allocator<T> >",
+            "std::unordered_map<T1, T2, std::hash<T1>, std::equal_to<T1>, std::allocator<std::pair<const T1, T2> > >",
+            "std::unordered_multimap<T1, T2, std::hash<T1>, std::equal_to<T1>, std::allocator<std::pair<const T1, T2> > >",
+
+            # More complex
+            "std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >::basic_string",
+            "std::vector<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::array<uint32_t, 5ul> >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::array<uint32_t, 5ul> > > >::_M_default_append(uint64_t)",
+            "std::vector<std::vector<T, std::allocator<T> >, std::allocator<std::vector<T, std::allocator<T> > > >::_M_check_len(uint64_t, char const*) const",
+        ]]
+
+        # Test all the APIs
+        qName = binja.types.QualifiedName(["std", "__cxx11", "basic_string<T, std::char_traits<T>, std::allocator<T> >"])
+        result.append(binja.demangle.simplify_name_to_string(qName))
+        result.append(str(binja.demangle.simplify_name_to_qualified_name(qName)))
+        result.append(str(binja.demangle.simplify_name_to_qualified_name(str(qName))))
+        result.append(str(binja.demangle.simplify_name_to_qualified_name(str(qName), False).name))
+        result.append("::".join(binja.demangle_gnu3(binja.Architecture['x86_64'], "_ZNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEE9_M_createERmm", False)[1]))
+        result.append("::".join(binja.demangle_gnu3(binja.Architecture['x86_64'], "_ZNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEE9_M_createERmm", True)[1]))
+
+        return result
+
     def test_Struct(self):
         """Struct produced different result"""
         retinfo = []