Add UI action and API to control switch recovery in HLIL

Author: D0ntPanic

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);

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,

docs/guide/index.md

@@ -807,28 +807,41 @@ 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
 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.

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;

python/function.py

@@ -29,7 +29,7 @@
 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 @@ def set_early_return(self, addr: Union[int, highlevelil.HighLevelILInstruction],
 			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):

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);
 
 /*!
     @}

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();

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);