Sort switch blocks in graph view by case number

7 files changed, 50 insertions, 0 deletions

diff --git a/basicblock.cpp b/basicblock.cpp
index 8025a9fc..9a9f3b5b 100644
--- a/basicblock.cpp
+++ b/basicblock.cpp
@@ -593,6 +593,12 @@ vector<vector<InstructionTextToken>> BasicBlock::GetAnnotations()
 }
 
 
+std::optional<int64_t> BasicBlock::GetSortHint()
+{
+	return GetFunction()->GetBlockSortHint(GetArchitecture(), GetStart());
+}
+
+
 vector<DisassemblyTextLine> BasicBlock::GetDisassemblyText(DisassemblySettings* settings)
 {
 	size_t count;
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 5a8398d2..bc36f633 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -12641,6 +12641,11 @@ namespace BinaryNinja {
 			\return List of automatic annotations for the start of this block
 		*/
 		std::vector<std::vector<InstructionTextToken>> GetAnnotations();
+		/*! Hint for sorting this block in graph layouts
+
+			\return Integer for sorting this block, if defined
+		 */
+		std::optional<int64_t> GetSortHint();
 
 		/*! property which returns a list of DisassemblyTextLine objects for the current basic block.
 
@@ -13432,6 +13437,7 @@ namespace BinaryNinja {
 		bool IsCallInstruction(Architecture* arch, uint64_t addr);
 
 		std::vector<std::vector<InstructionTextToken>> GetBlockAnnotations(Architecture* arch, uint64_t addr);
+		std::optional<int64_t> GetBlockSortHint(Architecture* arch, uint64_t addr);
 
 		BNIntegerDisplayType GetIntegerConstantDisplayType(
 		    Architecture* arch, uint64_t instrAddr, uint64_t value, size_t operand);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 0f820e06..fe9b3dbe 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -5653,6 +5653,8 @@ extern "C"
 
 	BINARYNINJACOREAPI BNInstructionTextLine* BNGetFunctionBlockAnnotations(
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
+	BINARYNINJACOREAPI bool BNGetFunctionBlockSortHint(
+	    BNFunction* func, BNArchitecture* arch, uint64_t addr, int64_t* result);
 
 	BINARYNINJACOREAPI BNIntegerDisplayType BNGetIntegerConstantDisplayType(
 	    BNFunction* func, BNArchitecture* arch, uint64_t instrAddr, uint64_t value, size_t operand);
diff --git a/function.cpp b/function.cpp
index a8ab3b99..094a50db 100644
--- a/function.cpp
+++ b/function.cpp
@@ -1798,6 +1798,18 @@ vector<vector<InstructionTextToken>> Function::GetBlockAnnotations(Architecture*
 }
 
 
+std::optional<int64_t> Function::GetBlockSortHint(Architecture* arch, uint64_t addr)
+{
+	int64_t result;
+	if (!BNGetFunctionBlockSortHint(m_object, arch->GetObject(), addr, &result))
+	{
+		return std::nullopt;
+	}
+
+	return result;
+}
+
+
 BNIntegerDisplayType Function::GetIntegerConstantDisplayType(
     Architecture* arch, uint64_t instrAddr, uint64_t value, size_t operand)
 {
diff --git a/python/basicblock.py b/python/basicblock.py
index fc8b89e1..69f9f996 100644
--- a/python/basicblock.py
+++ b/python/basicblock.py
@@ -705,6 +705,14 @@ class BasicBlock:
 		return self.function.get_block_annotations(self.start, self.arch)
 
 	@property
+	def sort_hint(self) -> Optional[int]:
+		"""Graph edge sorting hint for this block (read-only)"""
+		if self.function is None:
+			raise ValueError("Attempting to call BasicBlock.sort_hint when Function is None")
+
+		return self.function.get_block_sort_hint(self.start, self.arch)
+
+	@property
 	def disassembly_text(self) -> List['_function.DisassemblyTextLine']:
 		"""
 		``disassembly_text`` property which returns a list of function.DisassemblyTextLine objects for the current basic block.
diff --git a/python/function.py b/python/function.py
index 9f5e89bb..8c9eaf1d 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2437,6 +2437,14 @@ class Function:
 		finally:
 			core.BNFreeInstructionTextLines(lines, count.value)
 
+	def get_block_sort_hint(self, addr: int, arch: Optional['architecture.Architecture'] = None) -> Optional[int]:
+		if arch is None:
+			arch = self.arch
+		result = ctypes.c_int64()
+		if not core.BNGetFunctionBlockSortHint(self.handle, arch.handle, addr, ctypes.byref(result)):
+			return None
+		return result.value
+
 	def set_auto_type(self, value: StringOrType) -> None:
 		if isinstance(value, str):
 			(value, _) = self.view.parse_type_string(value)
diff --git a/rust/src/function.rs b/rust/src/function.rs
index 90432af6..320382de 100644
--- a/rust/src/function.rs
+++ b/rust/src/function.rs
@@ -510,6 +510,14 @@ impl Function {
         unsafe { Array::new(lines, count, ()) }
     }
 
+    pub fn block_sort_hint(&self, addr: u64, arch: Option<CoreArchitecture>) -> Option<i64> {
+        let arch = arch.unwrap_or_else(|| self.arch());
+        let mut result = 0;
+        unsafe {
+            BNGetFunctionBlockSortHint(self.handle, arch.handle, addr, &mut result).then_some(result)
+        }
+    }
+
     pub fn variable_name(&self, var: &Variable) -> String {
         unsafe {
             let raw_var = BNVariable::from(var);