Add UI action and API to control switch recovery in HLIL

8 files changed, 65 insertions, 8 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index c36e21a2..275ec979 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -11436,6 +11436,9 @@ namespace BinaryNinja {
 		BNEarlyReturn GetEarlyReturn(uint64_t addr);
 		void SetEarlyReturn(uint64_t addr, BNEarlyReturn mode);
 
+		BNSwitchRecovery GetSwitchRecovery(uint64_t addr);
+		void SetSwitchRecovery(uint64_t addr, BNSwitchRecovery mode);
+
 		std::map<Variable, std::set<Variable>> GetMergedVariables();
 		void MergeVariables(const Variable& target, const std::set<Variable>& sources);
 		void UnmergeVariables(const Variable& target, const std::set<Variable>& sources);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index c4ac0911..1210fcca 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -37,7 +37,7 @@
 // Current ABI version for linking to the core. This is incremented any time
 // there are changes to the API that affect linking, including new functions,
 // new types, or modifications to existing functions or types.
-#define BN_CURRENT_CORE_ABI_VERSION 110
+#define BN_CURRENT_CORE_ABI_VERSION 111
 
 // Minimum ABI version that is supported for loading of plugins. Plugins that
 // are linked to an ABI version less than this will not be able to load and
@@ -1030,6 +1030,9 @@ extern "C"
 
 		// HLIL condition can be rewritten as an early return
 		HLILEarlyReturnPossible = 0x400,
+
+		// HLIL condition chain can be rewritten as a switch statement
+		HLILSwitchRecoveryPossible = 0x800,
 	} BNILInstructionAttribute;
 
 	typedef enum BNIntrinsicClass
@@ -3254,6 +3257,13 @@ extern "C"
 		FalseSideEarlyReturn
 	} BNEarlyReturn;
 
+	typedef enum BNSwitchRecovery
+	{
+		DefaultSwitchRecovery,
+		PreventSwitchRecovery,
+		AllowSwitchRecovery
+	} BNSwitchRecovery;
+
 	typedef struct BNDebugFunctionInfo
 	{
 		char* shortName;
@@ -5013,6 +5023,8 @@ extern "C"
 	BINARYNINJACOREAPI void BNSetConditionInverted(BNFunction* func, uint64_t addr, bool invert);
 	BINARYNINJACOREAPI BNEarlyReturn BNGetEarlyReturn(BNFunction* func, uint64_t addr);
 	BINARYNINJACOREAPI void BNSetEarlyReturn(BNFunction* func, uint64_t addr, BNEarlyReturn mode);
+	BINARYNINJACOREAPI BNSwitchRecovery BNGetSwitchRecovery(BNFunction* func, uint64_t addr);
+	BINARYNINJACOREAPI void BNSetSwitchRecovery(BNFunction* func, uint64_t addr, BNSwitchRecovery mode);
 	BINARYNINJACOREAPI BNMergedVariable* BNGetMergedVariables(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI void BNFreeMergedVariableList(BNMergedVariable* vars, size_t count);
 	BINARYNINJACOREAPI void BNMergeVariables(BNFunction* func, const BNVariable* target, const BNVariable* sources,
diff --git a/docs/guide/index.md b/docs/guide/index.md
index f4a4e4af..1ba08f8c 100644
--- a/docs/guide/index.md
+++ b/docs/guide/index.md
@@ -807,7 +807,7 @@ Performing this action on both variables in the example results in the following
 
 ![Dead Store Elimination Results](../img/dead-store-after.png "Dead Store Elimination Results"){ width="500" }
 
-## Expression Folding
+## High Level Optimization Overrides
 
 Binary Ninja automatically performs optimization passes on high level code. One of these optimizations is to fold
 assignments with a single use into the statement that uses it. An example of a function call being folded into another
@@ -815,20 +815,33 @@ is shown below:
 
 ![Expression Folding](../img/folding-before.png "Expression Folding"){ width="500" }
 
-Binary Ninja uses heuristics to determine if this will improve readability, but sometimes it doesn't make the preferred
-choice. You can override the heuristic by right-clicking an expression and choosing "Allow" or "Prevent" from the
-"Expression Folding" submenu.
+Binary Ninja uses heuristics to determine if optimizatoins will improve readability, but sometimes it doesn't make the
+preferred choice. In the case above, you can override the heuristic by right-clicking the inner call expression and
+choosing "Allow" or "Prevent" from the "Expression Folding" submenu.
 
 ![Expression Folding Menu](../img/folding-menu.png "Expression Folding Menu"){ width="500" }
 
-This option will only appear if the expression can be folded. If Binary Ninja's analysis determines that it is not sound
-to fold an expression, the submenu will not be present.
+These options will only appear if the given optimizations are applied or can be applied. If Binary Ninja's analysis
+determines that it is not sound to perform an optimization, the submenus will not be present.
 
 Choosing "Prevent" from the "Expression Folding" menu on the call to `_strlen` in the example above results in the
 following output:
 
 ![Expression Folding Results](../img/folding-after.png "Expression Folding Results"){ width="500" }
 
+The heuristics can sometimes choose to not apply an optimization. You can override the heuristic to apply the
+optimization using the same right-click submenus.
+
+Several optimizations offer the option to override heuristics. The optimizations that support
+this feature are shown in the table below:
+
+| Optimization               | Description                                                                                                                                                    | Valid Locations                                                                                    |
+|----------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------|
+| Early Return               | Rewrites trailing `if`/`else` constructs to return early from one side. Avoids large indented `if` blocks.                                                     | `if` statements                                                                                    |
+| Expression Folding         | Folds a store of an expression into another expression. Heuristics try to reduce lines of code and number of active variables.                                 | Inner expressions (when optimization is applied) or assignments (when optimization is not applied) |
+| Show Condition as Inverted | Heuristics try to pick the best condition for `if`/`else` constructs, but an override can invert the chosen condition and rearrange the `if`/`else` construct. | `if` statements                                                                                    |
+| Switch Recovery            | Heuristics try to avoid tiny switch constructs. Overrides can choose to display the code as a `switch` or a chain of `if`/`else` constructs.                   | `if` or `switch` statements                                                                        |
+
 ## Merging and Splitting Variables
 
 Binary Ninja automatically splits all variables that the analysis determines to be safely splittable. This allows the user to assign different types to different uses of the same register or memory location. Sometimes, however, the code is more clear if two or more of these variables are merged together into a single variable.
diff --git a/function.cpp b/function.cpp
index 5ce3a843..6d026359 100644
--- a/function.cpp
+++ b/function.cpp
@@ -2849,6 +2849,18 @@ void Function::SetEarlyReturn(uint64_t addr, BNEarlyReturn mode)
 }
 
 
+BNSwitchRecovery Function::GetSwitchRecovery(uint64_t addr)
+{
+	return BNGetSwitchRecovery(m_object, addr);
+}
+
+
+void Function::SetSwitchRecovery(uint64_t addr, BNSwitchRecovery mode)
+{
+	BNSetSwitchRecovery(m_object, addr, mode);
+}
+
+
 std::map<Variable, std::set<Variable>> Function::GetMergedVariables()
 {
 	size_t count;
diff --git a/python/function.py b/python/function.py
index 7d14c2bb..ccecb992 100644
--- a/python/function.py
+++ b/python/function.py
@@ -29,7 +29,7 @@ from . import _binaryninjacore as core
 from .enums import (
 	AnalysisSkipReason, FunctionGraphType, SymbolType, InstructionTextTokenType, HighlightStandardColor,
 	HighlightColorStyle, DisassemblyOption, IntegerDisplayType, FunctionAnalysisSkipOverride, FunctionUpdateType,
-	BuiltinType, ExprFolding, EarlyReturn
+	BuiltinType, ExprFolding, EarlyReturn, SwitchRecovery
 )
 
 from . import associateddatastore  # Required in the main scope due to being an argument for _FunctionAssociatedDataStore
@@ -3538,6 +3538,16 @@ class Function:
 			addr = addr.address
 		core.BNSetEarlyReturn(self.handle, addr, value)
 
+	def get_switch_recovery(self, addr: Union[int, highlevelil.HighLevelILInstruction]) -> SwitchRecovery:
+		if isinstance(addr, highlevelil.HighLevelILInstruction):
+			addr = addr.address
+		return SwitchRecovery(core.BNGetSwitchRecovery(self.handle, addr))
+
+	def set_switch_recovery(self, addr: Union[int, highlevelil.HighLevelILInstruction], value: SwitchRecovery):
+		if isinstance(addr, highlevelil.HighLevelILInstruction):
+			addr = addr.address
+		core.BNSetSwitchRecovery(self.handle, addr, value)
+
 
 class AdvancedFunctionAnalysisDataRequestor:
 	def __init__(self, func: Optional['Function'] = None):
diff --git a/ui/commands.h b/ui/commands.h
index 754d99a7..e8f12786 100644
--- a/ui/commands.h
+++ b/ui/commands.h
@@ -99,6 +99,7 @@ std::optional<uint64_t> getFoldableExprAddress(
 	BinaryNinja::HighLevelILFunction* hlil, const HighlightTokenState& highlight);
 std::optional<uint64_t> getInvertableConditionAddress(BinaryNinja::HighLevelILFunction* hlil, size_t instrIndex);
 std::optional<uint64_t> getEarlyReturnAddress(BinaryNinja::HighLevelILFunction* hlil, size_t instrIndex);
+std::optional<uint64_t> getSwitchRecoveryAddress(BinaryNinja::HighLevelILFunction* hlil, size_t instrIndex);
 
 /*!
     @}
diff --git a/ui/flowgraphwidget.h b/ui/flowgraphwidget.h
index d988da12..90b86000 100644
--- a/ui/flowgraphwidget.h
+++ b/ui/flowgraphwidget.h
@@ -172,6 +172,8 @@ class BINARYNINJAUIAPI FlowGraphWidget :
 	bool getCurrentConditionInverted();
 	std::optional<uint64_t> getCurrentEarlyReturnAddress();
 	BNEarlyReturn getCurrentEarlyReturn();
+	std::optional<uint64_t> getCurrentSwitchRecoveryAddress();
+	BNSwitchRecovery getCurrentSwitchRecovery();
 	std::optional<std::pair<BinaryNinja::Variable, BinaryNinja::Variable>> getMergeVariablesAtCurrentLocation();
 
   protected:
@@ -404,6 +406,7 @@ class BINARYNINJAUIAPI FlowGraphWidget :
 	void setCurrentExprFolding(BNExprFolding folding);
 	void toggleConditionInverted();
 	void setCurrentEarlyReturn(BNEarlyReturn earlyReturn);
+	void setCurrentSwitchRecovery(BNSwitchRecovery recovery);
 	void splitToNewTabAndNavigateFromCursorPosition();
 	void splitToNewWindowAndNavigateFromCursorPosition();
 	void splitToNewPaneAndNavigateFromCursorPosition();
diff --git a/ui/linearview.h b/ui/linearview.h
index eea40d7f..9ba28407 100644
--- a/ui/linearview.h
+++ b/ui/linearview.h
@@ -322,6 +322,8 @@ class BINARYNINJAUIAPI LinearView : public QAbstractScrollArea, public View, pub
 	bool getCurrentConditionInverted();
 	std::optional<uint64_t> getCurrentEarlyReturnAddress();
 	BNEarlyReturn getCurrentEarlyReturn();
+	std::optional<uint64_t> getCurrentSwitchRecoveryAddress();
+	BNSwitchRecovery getCurrentSwitchRecovery();
 
 	void setDataButtonVisible(bool visible);
 	std::optional<std::pair<BinaryNinja::Variable, BinaryNinja::Variable>> getMergeVariablesAtCurrentLocation();
@@ -415,6 +417,7 @@ private Q_SLOTS:
 	void setCurrentExprFolding(BNExprFolding folding);
 	void toggleConditionInverted();
 	void setCurrentEarlyReturn(BNEarlyReturn earlyReturn);
+	void setCurrentSwitchRecovery(BNSwitchRecovery recovery);
 
 Q_SIGNALS:
 	void notifyResizeEvent(int width, int height);