Allow flag roles to be dependent on semantic class

5 files changed, 60 insertions, 33 deletions

diff --git a/architecture.cpp b/architecture.cpp
index 571794c6..78a67105 100644
--- a/architecture.cpp
+++ b/architecture.cpp
@@ -318,10 +318,10 @@ uint32_t* Architecture::GetAllSemanticFlagGroupsCallback(void* ctxt, size_t* cou
 }
 
 
-BNFlagRole Architecture::GetFlagRoleCallback(void* ctxt, uint32_t flag)
+BNFlagRole Architecture::GetFlagRoleCallback(void* ctxt, uint32_t flag, uint32_t semClass)
 {
 	Architecture* arch = (Architecture*)ctxt;
-	return arch->GetFlagRole(flag);
+	return arch->GetFlagRole(flag, semClass);
 }
 
 
@@ -780,7 +780,7 @@ vector<uint32_t> Architecture::GetAllSemanticFlagGroups()
 }
 
 
-BNFlagRole Architecture::GetFlagRole(uint32_t)
+BNFlagRole Architecture::GetFlagRole(uint32_t, uint32_t)
 {
 	return SpecialFlagRole;
 }
@@ -819,8 +819,7 @@ uint32_t Architecture::GetSemanticClassForFlagWriteType(uint32_t)
 size_t Architecture::GetFlagWriteLowLevelIL(BNLowLevelILOperation op, size_t size, uint32_t flagWriteType,
 	uint32_t flag, BNRegisterOrConstant* operands, size_t operandCount,LowLevelILFunction& il)
 {
-	(void)flagWriteType;
-	BNFlagRole role = GetFlagRole(flag);
+	BNFlagRole role = GetFlagRole(flag, GetSemanticClassForFlagWriteType(flagWriteType));
 	return BNGetDefaultArchitectureFlagWriteLowLevelIL(m_object, op, size, role, operands,
 		operandCount, il.GetObject());
 }
@@ -834,15 +833,17 @@ size_t Architecture::GetDefaultFlagWriteLowLevelIL(BNLowLevelILOperation op, siz
 }
 
 
-ExprId Architecture::GetFlagConditionLowLevelIL(BNLowLevelILFlagCondition cond, uint32_t, LowLevelILFunction& il)
+ExprId Architecture::GetFlagConditionLowLevelIL(BNLowLevelILFlagCondition cond,
+	uint32_t semClass, LowLevelILFunction& il)
 {
-	return BNGetDefaultArchitectureFlagConditionLowLevelIL(m_object, cond, il.GetObject());
+	return BNGetDefaultArchitectureFlagConditionLowLevelIL(m_object, cond, semClass, il.GetObject());
 }
 
 
-ExprId Architecture::GetDefaultFlagConditionLowLevelIL(BNLowLevelILFlagCondition cond, LowLevelILFunction& il)
+ExprId Architecture::GetDefaultFlagConditionLowLevelIL(BNLowLevelILFlagCondition cond,
+	uint32_t semClass, LowLevelILFunction& il)
 {
-	return BNGetDefaultArchitectureFlagConditionLowLevelIL(m_object, cond, il.GetObject());
+	return BNGetDefaultArchitectureFlagConditionLowLevelIL(m_object, cond, semClass, il.GetObject());
 }
 
 
@@ -1325,9 +1326,9 @@ vector<uint32_t> CoreArchitecture::GetAllSemanticFlagGroups()
 }
 
 
-BNFlagRole CoreArchitecture::GetFlagRole(uint32_t flag)
+BNFlagRole CoreArchitecture::GetFlagRole(uint32_t flag, uint32_t semClass)
 {
-	return BNGetArchitectureFlagRole(m_object, flag);
+	return BNGetArchitectureFlagRole(m_object, flag, semClass);
 }
 
 
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index cbd51d05..b74fa971 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1609,7 +1609,7 @@ namespace BinaryNinja
 		static uint32_t* GetAllFlagWriteTypesCallback(void* ctxt, size_t* count);
 		static uint32_t* GetAllSemanticFlagClassesCallback(void* ctxt, size_t* count);
 		static uint32_t* GetAllSemanticFlagGroupsCallback(void* ctxt, size_t* count);
-		static BNFlagRole GetFlagRoleCallback(void* ctxt, uint32_t flag);
+		static BNFlagRole GetFlagRoleCallback(void* ctxt, uint32_t flag, uint32_t semClass);
 		static uint32_t* GetFlagsRequiredForFlagConditionCallback(void* ctxt, BNLowLevelILFlagCondition cond,
 			uint32_t semClass, size_t* count);
 		static uint32_t* GetFlagsRequiredForSemanticFlagGroupCallback(void* ctxt, uint32_t semGroup, size_t* count);
@@ -1687,7 +1687,7 @@ namespace BinaryNinja
 		virtual std::vector<uint32_t> GetAllFlagWriteTypes();
 		virtual std::vector<uint32_t> GetAllSemanticFlagClasses();
 		virtual std::vector<uint32_t> GetAllSemanticFlagGroups();
-		virtual BNFlagRole GetFlagRole(uint32_t flag);
+		virtual BNFlagRole GetFlagRole(uint32_t flag, uint32_t semClass = 0);
 		virtual std::vector<uint32_t> GetFlagsRequiredForFlagCondition(BNLowLevelILFlagCondition cond,
 			uint32_t semClass = 0);
 		virtual std::vector<uint32_t> GetFlagsRequiredForSemanticFlagGroup(uint32_t semGroup);
@@ -1699,7 +1699,7 @@ namespace BinaryNinja
 		ExprId GetDefaultFlagWriteLowLevelIL(BNLowLevelILOperation op, size_t size, BNFlagRole role,
 			BNRegisterOrConstant* operands, size_t operandCount, LowLevelILFunction& il);
 		virtual ExprId GetFlagConditionLowLevelIL(BNLowLevelILFlagCondition cond, uint32_t semClass, LowLevelILFunction& il);
-		ExprId GetDefaultFlagConditionLowLevelIL(BNLowLevelILFlagCondition cond, LowLevelILFunction& il);
+		ExprId GetDefaultFlagConditionLowLevelIL(BNLowLevelILFlagCondition cond, uint32_t semClass, LowLevelILFunction& il);
 		virtual ExprId GetSemanticFlagGroupLowLevelIL(uint32_t semGroup, LowLevelILFunction& il);
 		virtual BNRegisterInfo GetRegisterInfo(uint32_t reg);
 		virtual uint32_t GetStackPointerRegister();
@@ -1832,7 +1832,7 @@ namespace BinaryNinja
 		virtual std::vector<uint32_t> GetAllFlagWriteTypes() override;
 		virtual std::vector<uint32_t> GetAllSemanticFlagClasses() override;
 		virtual std::vector<uint32_t> GetAllSemanticFlagGroups() override;
-		virtual BNFlagRole GetFlagRole(uint32_t flag) override;
+		virtual BNFlagRole GetFlagRole(uint32_t flag, uint32_t semClass = 0) override;
 		virtual std::vector<uint32_t> GetFlagsRequiredForFlagCondition(BNLowLevelILFlagCondition cond,
 			uint32_t semClass = 0) override;
 		virtual std::vector<uint32_t> GetFlagsRequiredForSemanticFlagGroup(uint32_t semGroup) override;
@@ -2526,7 +2526,8 @@ namespace BinaryNinja
 			const ILSourceLocation& loc = ILSourceLocation());
 		ExprId RegisterStackTopRelative(size_t size, uint32_t regStack, ExprId entry,
 			const ILSourceLocation& loc = ILSourceLocation());
-		ExprId RegisterStackPop(size_t size, uint32_t regStack, const ILSourceLocation& loc = ILSourceLocation());
+		ExprId RegisterStackPop(size_t size, uint32_t regStack, uint32_t flags = 0,
+			const ILSourceLocation& loc = ILSourceLocation());
 		ExprId RegisterStackTopRelativeSSA(size_t size, const SSARegisterStack& regStack, ExprId entry,
 			const SSARegister& top, const ILSourceLocation& loc = ILSourceLocation());
 		ExprId RegisterStackAbsoluteSSA(size_t size, const SSARegisterStack& regStack, uint32_t reg,
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 388d3643..bc58092d 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -441,7 +441,9 @@ extern "C"
 		OverflowFlagRole = 5,
 		HalfCarryFlagRole = 6,
 		EvenParityFlagRole = 7,
-		OddParityFlagRole = 8
+		OddParityFlagRole = 8,
+		OrderedFlagRole = 9,
+		UnorderedFlagRole = 10
 	};
 
 	enum BNFunctionGraphType
@@ -1102,7 +1104,7 @@ extern "C"
 		uint32_t* (*getAllFlagWriteTypes)(void* ctxt, size_t* count);
 		uint32_t* (*getAllSemanticFlagClasses)(void* ctxt, size_t* count);
 		uint32_t* (*getAllSemanticFlagGroups)(void* ctxt, size_t* count);
-		BNFlagRole (*getFlagRole)(void* ctxt, uint32_t flag);
+		BNFlagRole (*getFlagRole)(void* ctxt, uint32_t flag, uint32_t semClass);
 		uint32_t* (*getFlagsRequiredForFlagCondition)(void* ctxt, BNLowLevelILFlagCondition cond,
 			uint32_t semClass, size_t* count);
 		uint32_t* (*getFlagsRequiredForSemanticFlagGroup)(void* ctxt, uint32_t semGroup, size_t* count);
@@ -2078,7 +2080,7 @@ extern "C"
 	BINARYNINJACOREAPI uint32_t* BNGetAllArchitectureFlagWriteTypes(BNArchitecture* arch, size_t* count);
 	BINARYNINJACOREAPI uint32_t* BNGetAllArchitectureSemanticFlagClasses(BNArchitecture* arch, size_t* count);
 	BINARYNINJACOREAPI uint32_t* BNGetAllArchitectureSemanticFlagGroups(BNArchitecture* arch, size_t* count);
-	BINARYNINJACOREAPI BNFlagRole BNGetArchitectureFlagRole(BNArchitecture* arch, uint32_t flag);
+	BINARYNINJACOREAPI BNFlagRole BNGetArchitectureFlagRole(BNArchitecture* arch, uint32_t flag, uint32_t semClass);
 	BINARYNINJACOREAPI uint32_t* BNGetArchitectureFlagsRequiredForFlagCondition(BNArchitecture* arch, BNLowLevelILFlagCondition cond,
 		uint32_t semClass, size_t* count);
 	BINARYNINJACOREAPI uint32_t* BNGetArchitectureFlagsRequiredForSemanticFlagGroup(BNArchitecture* arch,
@@ -2097,7 +2099,7 @@ extern "C"
 	BINARYNINJACOREAPI size_t BNGetArchitectureFlagConditionLowLevelIL(BNArchitecture* arch, BNLowLevelILFlagCondition cond,
 		uint32_t semClass, BNLowLevelILFunction* il);
 	BINARYNINJACOREAPI size_t BNGetDefaultArchitectureFlagConditionLowLevelIL(BNArchitecture* arch, BNLowLevelILFlagCondition cond,
-		BNLowLevelILFunction* il);
+		uint32_t semClass, BNLowLevelILFunction* il);
 	BINARYNINJACOREAPI size_t BNGetArchitectureSemanticFlagGroupLowLevelIL(BNArchitecture* arch,
 		uint32_t semGroup, BNLowLevelILFunction* il);
 	BINARYNINJACOREAPI uint32_t* BNGetModifiedArchitectureRegistersOnWrite(BNArchitecture* arch, uint32_t reg, size_t* count);
diff --git a/lowlevelilinstruction.cpp b/lowlevelilinstruction.cpp
index 4b1acff3..9909e8ee 100644
--- a/lowlevelilinstruction.cpp
+++ b/lowlevelilinstruction.cpp
@@ -1680,7 +1680,7 @@ ExprId LowLevelILInstruction::CopyTo(LowLevelILFunction* dest,
 		return dest->RegisterStackTopRelative(size, GetSourceRegisterStack<LLIL_REG_STACK_REL>(),
 			subExprHandler(GetSourceExpr<LLIL_REG_STACK_REL>()), *this);
 	case LLIL_REG_STACK_POP:
-		return dest->RegisterStackPop(size, GetSourceRegisterStack<LLIL_REG_STACK_POP>(), *this);
+		return dest->RegisterStackPop(size, GetSourceRegisterStack<LLIL_REG_STACK_POP>(), flags, *this);
 	case LLIL_REG_STACK_REL_SSA:
 		return dest->RegisterStackTopRelativeSSA(size, GetSourceSSARegisterStack<LLIL_REG_STACK_REL_SSA>(),
 			subExprHandler(GetSourceExpr<LLIL_REG_STACK_REL_SSA>()),
@@ -1827,6 +1827,7 @@ ExprId LowLevelILInstruction::CopyTo(LowLevelILFunction* dest,
 	case LLIL_FCMP_LE:
 	case LLIL_FCMP_GE:
 	case LLIL_FCMP_GT:
+	case LLIL_FCMP_O:
 	case LLIL_FCMP_UO:
 		return dest->AddExprWithLocation(operation, *this, size, flags,
 			subExprHandler(AsTwoOperand().GetLeftExpr()), subExprHandler(AsTwoOperand().GetRightExpr()));
@@ -2435,9 +2436,9 @@ ExprId LowLevelILFunction::RegisterStackTopRelative(size_t size, uint32_t regSta
 }
 
 
-ExprId LowLevelILFunction::RegisterStackPop(size_t size, uint32_t regStack, const ILSourceLocation& loc)
+ExprId LowLevelILFunction::RegisterStackPop(size_t size, uint32_t regStack, uint32_t flags, const ILSourceLocation& loc)
 {
-	return AddExprWithLocation(LLIL_REG_STACK_POP, loc, size, 0, regStack);
+	return AddExprWithLocation(LLIL_REG_STACK_POP, loc, size, flags, regStack);
 }
 
 
diff --git a/python/architecture.py b/python/architecture.py
index d2ab586e..2b422962 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -208,7 +208,7 @@ class Architecture(object):
 			self._flag_roles = {}
 			self.__dict__["flag_roles"] = {}
 			for flag in self.__dict__["flags"]:
-				role = FlagRole(core.BNGetArchitectureFlagRole(self.handle, self._flags[flag]))
+				role = FlagRole(core.BNGetArchitectureFlagRole(self.handle, self._flags[flag], 0))
 				self.__dict__["flag_roles"][flag] = role
 				self._flag_roles[self._flags[flag]] = role
 
@@ -840,21 +840,28 @@ class Architecture(object):
 			count[0] = 0
 			return None
 
-	def _get_flag_role(self, ctxt, flag):
+	def _get_flag_role(self, ctxt, flag, sem_class):
 		try:
-			if flag in self._flag_roles:
-				return self._flag_roles[flag]
-			return FlagRole.SpecialFlagRole
+			if sem_class in self._semantic_flag_classes_by_index:
+				sem_class = self._semantic_flag_classes_by_index[sem_class]
+			else:
+				sem_class = None
+			return self.perform_get_flag_role(flag, sem_class)
 		except KeyError:
 			log.log_error(traceback.format_exc())
-			return None
+			return FlagRole.SpecialFlagRole
+
+	def perform_get_flag_role(self, flag, sem_class):
+		if flag in self._flag_roles:
+			return self._flag_roles[flag]
+		return FlagRole.SpecialFlagRole
 
 	def _get_flags_required_for_flag_condition(self, ctxt, cond, sem_class, count):
 		try:
 			if sem_class in self._semantic_flag_classes_by_index:
 				sem_class = self._semantic_flag_classes_by_index[sem_class]
 			else:
-				sem_class = 0
+				sem_class = None
 			flag_names = self.perform_get_flags_required_for_flag_condition(cond, sem_class)
 			flags = []
 			for name in flag_names:
@@ -1322,7 +1329,7 @@ class Architecture(object):
 		:param LowLevelILFunction il: LowLevelILFunction object to append LowLevelILExpr objects to
 		:rtype: LowLevelILExpr
 		"""
-		return self.get_default_flag_condition_low_level_il(cond, il)
+		return self.get_default_flag_condition_low_level_il(cond, sem_class, il)
 
 	@abc.abstractmethod
 	def perform_get_semantic_flag_group_low_level_il(self, sem_group, il):
@@ -1744,6 +1751,19 @@ class Architecture(object):
 		"""
 		return self._semantic_flag_groups[sem_group]
 
+	def get_flag_role(self, flag, sem_class = None):
+		"""
+		``get_flag_role`` gets the role of a given flag.
+
+		:param int flag: flag
+		:param int sem_class: optional semantic flag class
+		:return: flag role
+		:rtype: FlagRole
+		"""
+		flag = self.get_flag_index(flag)
+		sem_class = self.get_semantic_flag_class_index(sem_class)
+		return FlagRole(core.BNGetArchitectureFlagRole(self.handle, flag, sem_class))
+
 	def get_flag_write_low_level_il(self, op, size, write_type, flag, operands, il):
 		"""
 		:param LowLevelILOperation op:
@@ -1801,13 +1821,15 @@ class Architecture(object):
 		"""
 		return lowlevelil.LowLevelILExpr(core.BNGetArchitectureFlagConditionLowLevelIL(self.handle, cond, il.handle))
 
-	def get_default_flag_condition_low_level_il(self, cond, il):
+	def get_default_flag_condition_low_level_il(self, cond, sem_class, il):
 		"""
 		:param LowLevelILFlagCondition cond:
 		:param LowLevelILFunction il:
+		:param str sem_class:
 		:rtype: LowLevelILExpr
 		"""
-		return lowlevelil.LowLevelILExpr(core.BNGetDefaultArchitectureFlagConditionLowLevelIL(self.handle, cond, il.handle))
+		class_index = self.get_semantic_flag_class_index(sem_class)
+		return lowlevelil.LowLevelILExpr(core.BNGetDefaultArchitectureFlagConditionLowLevelIL(self.handle, cond, class_index, il.handle))
 
 	def get_semantic_flag_group_low_level_il(self, sem_group, il):
 		"""