Adding return value registers to function type

4 files changed, 83 insertions, 3 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index af067ae9..d4624ca5 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2204,6 +2204,7 @@ namespace BinaryNinja
 
 		Ref<Type> GetType() const;
 		Confidence<Ref<Type>> GetReturnType() const;
+		Confidence<std::vector<uint32_t>> GetReturnRegisters() const;
 		Confidence<Ref<CallingConvention>> GetCallingConvention() const;
 		Confidence<std::vector<Variable>> GetParameterVariables() const;
 		Confidence<bool> HasVariableArguments() const;
@@ -2213,6 +2214,7 @@ namespace BinaryNinja
 
 		void SetAutoType(Type* type);
 		void SetAutoReturnType(const Confidence<Ref<Type>>& type);
+		void SetAutoReturnRegisters(const Confidence<std::vector<uint32_t>>& returnRegs);
 		void SetAutoCallingConvention(const Confidence<Ref<CallingConvention>>& convention);
 		void SetAutoParameterVariables(const Confidence<std::vector<Variable>>& vars);
 		void SetAutoHasVariableArguments(const Confidence<bool>& varArgs);
@@ -2223,6 +2225,7 @@ namespace BinaryNinja
 
 		void SetUserType(Type* type);
 		void SetReturnType(const Confidence<Ref<Type>>& type);
+		void SetReturnRegisters(const Confidence<std::vector<uint32_t>>& returnRegs);
 		void SetCallingConvention(const Confidence<Ref<CallingConvention>>& convention);
 		void SetParameterVariables(const Confidence<std::vector<Variable>>& vars);
 		void SetHasVariableArguments(const Confidence<bool>& varArgs);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 278299bc..7f107c60 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2180,6 +2180,7 @@ extern "C"
 
 	BINARYNINJACOREAPI BNType* BNGetFunctionType(BNFunction* func);
 	BINARYNINJACOREAPI BNTypeWithConfidence BNGetFunctionReturnType(BNFunction* func);
+	BINARYNINJACOREAPI BNRegisterSetWithConfidence BNGetFunctionReturnRegisters(BNFunction* func);
 	BINARYNINJACOREAPI BNCallingConventionWithConfidence BNGetFunctionCallingConvention(BNFunction* func);
 	BINARYNINJACOREAPI BNParameterVariablesWithConfidence BNGetFunctionParameterVariables(BNFunction* func);
 	BINARYNINJACOREAPI void BNFreeParameterVariables(BNParameterVariablesWithConfidence* vars);
@@ -2188,9 +2189,10 @@ extern "C"
 	BINARYNINJACOREAPI BNRegisterStackAdjustment* BNGetFunctionRegisterStackAdjustments(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI void BNFreeRegisterStackAdjustments(BNRegisterStackAdjustment* adjustments);
 	BINARYNINJACOREAPI BNRegisterSetWithConfidence BNGetFunctionClobberedRegisters(BNFunction* func);
-	BINARYNINJACOREAPI void BNFreeClobberedRegisters(BNRegisterSetWithConfidence* regs);
+	BINARYNINJACOREAPI void BNFreeRegisterSet(BNRegisterSetWithConfidence* regs);
 
 	BINARYNINJACOREAPI void BNSetAutoFunctionReturnType(BNFunction* func, BNTypeWithConfidence* type);
+	BINARYNINJACOREAPI void BNSetAutoFunctionReturnRegisters(BNFunction* func, BNRegisterSetWithConfidence* regs);
 	BINARYNINJACOREAPI void BNSetAutoFunctionCallingConvention(BNFunction* func, BNCallingConventionWithConfidence* convention);
 	BINARYNINJACOREAPI void BNSetAutoFunctionParameterVariables(BNFunction* func, BNParameterVariablesWithConfidence* vars);
 	BINARYNINJACOREAPI void BNSetAutoFunctionHasVariableArguments(BNFunction* func, BNBoolWithConfidence* varArgs);
@@ -2201,6 +2203,7 @@ extern "C"
 	BINARYNINJACOREAPI void BNSetAutoFunctionClobberedRegisters(BNFunction* func, BNRegisterSetWithConfidence* regs);
 
 	BINARYNINJACOREAPI void BNSetUserFunctionReturnType(BNFunction* func, BNTypeWithConfidence* type);
+	BINARYNINJACOREAPI void BNSetUserFunctionReturnRegisters(BNFunction* func, BNRegisterSetWithConfidence* regs);
 	BINARYNINJACOREAPI void BNSetUserFunctionCallingConvention(BNFunction* func, BNCallingConventionWithConfidence* convention);
 	BINARYNINJACOREAPI void BNSetUserFunctionParameterVariables(BNFunction* func, BNParameterVariablesWithConfidence* vars);
 	BINARYNINJACOREAPI void BNSetUserFunctionHasVariableArguments(BNFunction* func, BNBoolWithConfidence* varArgs);
diff --git a/function.cpp b/function.cpp
index 879c7156..ba2b419e 100644
--- a/function.cpp
+++ b/function.cpp
@@ -490,6 +490,18 @@ Confidence<Ref<Type>> Function::GetReturnType() const
 }
 
 
+Confidence<vector<uint32_t>> Function::GetReturnRegisters() const
+{
+	BNRegisterSetWithConfidence regs = BNGetFunctionReturnRegisters(m_object);
+	vector<uint32_t> regList;
+	for (size_t i = 0; i < regs.count; i++)
+		regList.push_back(regs.regs[i]);
+	Confidence<vector<uint32_t>> result(regList, regs.confidence);
+	BNFreeRegisterSet(&regs);
+	return result;
+}
+
+
 Confidence<Ref<CallingConvention>> Function::GetCallingConvention() const
 {
 	BNCallingConventionWithConfidence cc = BNGetFunctionCallingConvention(m_object);
@@ -549,7 +561,7 @@ Confidence<set<uint32_t>> Function::GetClobberedRegisters() const
 	for (size_t i = 0; i < regs.count; i++)
 		regSet.insert(regs.regs[i]);
 	Confidence<set<uint32_t>> result(regSet, regs.confidence);
-	BNFreeClobberedRegisters(&regs);
+	BNFreeRegisterSet(&regs);
 	return result;
 }
 
@@ -569,6 +581,19 @@ void Function::SetAutoReturnType(const Confidence<Ref<Type>>& type)
 }
 
 
+void Function::SetAutoReturnRegisters(const Confidence<std::vector<uint32_t>>& returnRegs)
+{
+	BNRegisterSetWithConfidence regs;
+	regs.regs = new uint32_t[returnRegs.GetValue().size()];
+	regs.count = returnRegs.GetValue().size();
+	for (size_t i = 0; i < regs.count; i++)
+		regs.regs[i] = returnRegs.GetValue()[i];
+	regs.confidence = returnRegs.GetConfidence();
+	BNSetAutoFunctionReturnRegisters(m_object, &regs);
+	delete[] regs.regs;
+}
+
+
 void Function::SetAutoCallingConvention(const Confidence<Ref<CallingConvention>>& convention)
 {
 	BNCallingConventionWithConfidence cc;
@@ -668,6 +693,19 @@ void Function::SetReturnType(const Confidence<Ref<Type>>& type)
 }
 
 
+void Function::SetReturnRegisters(const Confidence<std::vector<uint32_t>>& returnRegs)
+{
+	BNRegisterSetWithConfidence regs;
+	regs.regs = new uint32_t[returnRegs.GetValue().size()];
+	regs.count = returnRegs.GetValue().size();
+	for (size_t i = 0; i < regs.count; i++)
+		regs.regs[i] = returnRegs.GetValue()[i];
+	regs.confidence = returnRegs.GetConfidence();
+	BNSetUserFunctionReturnRegisters(m_object, &regs);
+	delete[] regs.regs;
+}
+
+
 void Function::SetCallingConvention(const Confidence<Ref<CallingConvention>>& convention)
 {
 	BNCallingConventionWithConfidence cc;
diff --git a/python/function.py b/python/function.py
index eabf86b3..db4ca0f5 100644
--- a/python/function.py
+++ b/python/function.py
@@ -593,6 +593,30 @@ class Function(object):
 		core.BNSetUserFunctionReturnType(self.handle, type_conf)
 
 	@property
+	def return_regs(self):
+		"""Registers that are used for the return value"""
+		result = core.BNGetFunctionReturnRegisters(self.handle)
+		reg_set = []
+		for i in xrange(0, result.count):
+			reg_set.append(self.arch.get_reg_name(result.regs[i]))
+		regs = types.RegisterSet(reg_set, confidence = result.confidence)
+		core.BNFreeRegisterSet(result)
+		return regs
+
+	@return_regs.setter
+	def return_regs(self, value):
+		regs = core.BNRegisterSetWithConfidence()
+		regs.regs = (ctypes.c_uint * len(value))()
+		regs.count = len(value)
+		for i in xrange(0, len(value)):
+			regs.regs[i] = self.arch.get_reg_index(value[i])
+		if hasattr(value, 'confidence'):
+			regs.confidence = value.confidence
+		else:
+			regs.confidence = types.max_confidence
+		core.BNSetUserFunctionReturnRegisters(self.handle, regs)
+
+	@property
 	def calling_convention(self):
 		"""Calling convention used by the function"""
 		result = core.BNGetFunctionCallingConvention(self.handle)
@@ -712,7 +736,7 @@ class Function(object):
 		for i in xrange(0, result.count):
 			reg_set.append(self.arch.get_reg_name(result.regs[i]))
 		regs = types.RegisterSet(reg_set, confidence = result.confidence)
-		core.BNFreeClobberedRegisters(result)
+		core.BNFreeRegisterSet(result)
 		return regs
 
 	@clobbered_regs.setter
@@ -1072,6 +1096,18 @@ class Function(object):
 			type_conf.confidence = value.confidence
 		core.BNSetAutoFunctionReturnType(self.handle, type_conf)
 
+	def set_auto_return_regs(self, value):
+		regs = core.BNRegisterSetWithConfidence()
+		regs.regs = (ctypes.c_uint * len(value))()
+		regs.count = len(value)
+		for i in xrange(0, len(value)):
+			regs.regs[i] = self.arch.get_reg_index(value[i])
+		if hasattr(value, 'confidence'):
+			regs.confidence = value.confidence
+		else:
+			regs.confidence = types.max_confidence
+		core.BNSetAutoFunctionReturnRegisters(self.handle, regs)
+
 	def set_auto_calling_convention(self, value):
 		conv_conf = core.BNCallingConventionWithConfidence()
 		if value is None: