Add UI action and API to control expression folding

11 files changed, 70 insertions, 2 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 942f9d56..cba3c874 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -11423,6 +11423,9 @@ namespace BinaryNinja {
 		BNDeadStoreElimination GetVariableDeadStoreElimination(const Variable& var);
 		void SetVariableDeadStoreElimination(const Variable& var, BNDeadStoreElimination mode);
 
+		BNExprFolding GetExprFolding(uint64_t addr);
+		void SetExprFolding(uint64_t addr, BNExprFolding 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 d5cc3238..ec03c1dd 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1021,6 +1021,9 @@ extern "C"
 
 		// MLIL instruction appears to be an otherwise unused intermediate
 		MLILPossiblyUnusedIntermediate = 0x80,
+
+		// HLIL expression can be folded into other expressions or has been folded
+		HLILFoldableExpr = 0x100,
 	} BNILInstructionAttribute;
 
 	typedef enum BNIntrinsicClass
@@ -3226,6 +3229,13 @@ extern "C"
 		AllowDeadStoreElimination
 	} BNDeadStoreElimination;
 
+	typedef enum BNExprFolding
+	{
+		DefaultExprFolding,
+		PreventExprFolding,
+		AllowExprFolding
+	} BNExprFolding;
+
 	typedef struct BNDebugFunctionInfo
 	{
 		char* shortName;
@@ -4979,6 +4989,8 @@ extern "C"
 	    BNFunction* func, const BNVariable* var);
 	BINARYNINJACOREAPI void BNSetFunctionVariableDeadStoreElimination(
 	    BNFunction* func, const BNVariable* var, BNDeadStoreElimination mode);
+	BINARYNINJACOREAPI BNExprFolding BNGetExprFolding(BNFunction* func, uint64_t addr);
+	BINARYNINJACOREAPI void BNSetExprFolding(BNFunction* func, uint64_t addr, BNExprFolding 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 eb29747f..f4a4e4af 100644
--- a/docs/guide/index.md
+++ b/docs/guide/index.md
@@ -807,6 +807,28 @@ 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
+
+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.
+
+![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.
+
+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" }
+
 ## 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/docs/img/folding-after.png b/docs/img/folding-after.png
new file mode 100644
index 00000000..ac0ad68e
--- /dev/null
+++ b/docs/img/folding-after.png
diff --git a/docs/img/folding-before.png b/docs/img/folding-before.png
new file mode 100644
index 00000000..4d760dbe
--- /dev/null
+++ b/docs/img/folding-before.png
diff --git a/docs/img/folding-menu.png b/docs/img/folding-menu.png
new file mode 100644
index 00000000..bd8f95ca
--- /dev/null
+++ b/docs/img/folding-menu.png
diff --git a/function.cpp b/function.cpp
index cc9a8e13..756c9629 100644
--- a/function.cpp
+++ b/function.cpp
@@ -2813,6 +2813,18 @@ void Function::SetVariableDeadStoreElimination(const Variable& var, BNDeadStoreE
 }
 
 
+BNExprFolding Function::GetExprFolding(uint64_t addr)
+{
+	return BNGetExprFolding(m_object, addr);
+}
+
+
+void Function::SetExprFolding(uint64_t addr, BNExprFolding mode)
+{
+	BNSetExprFolding(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 429b6f69..dc7add3a 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
+	BuiltinType, ExprFolding
 )
 
 from . import associateddatastore  # Required in the main scope due to being an argument for _FunctionAssociatedDataStore
@@ -3456,6 +3456,16 @@ class Function:
 		assert isinstance(value, dict), "core.BNFunctionGetAutoMetadata did not return a dict"
 		return value
 
+	def get_expr_folding(self, addr: Union[int, highlevelil.HighLevelILInstruction]) -> ExprFolding:
+		if isinstance(addr, highlevelil.HighLevelILInstruction):
+			addr = addr.address
+		return ExprFolding(core.BNGetExprFolding(self.handle, addr))
+
+	def set_expr_folding(self, addr: Union[int, highlevelil.HighLevelILInstruction], value: ExprFolding):
+		if isinstance(addr, highlevelil.HighLevelILInstruction):
+			addr = addr.address
+		core.BNSetExprFolding(self.handle, addr, value)
+
 
 class AdvancedFunctionAnalysisDataRequestor:
 	def __init__(self, func: Optional['Function'] = None):
diff --git a/ui/commands.h b/ui/commands.h
index c8634bf7..7fb3f539 100644
--- a/ui/commands.h
+++ b/ui/commands.h
@@ -95,6 +95,9 @@ std::optional<std::string> GetVariableNameFromExpr(BinaryNinja::Function* func,
 // typedef-ed type
 TypeRef GetFunctionType(BinaryViewRef data, TypeRef type);
 
+std::optional<uint64_t> getFoldableExprAddress(
+	BinaryNinja::HighLevelILFunction* hlil, const HighlightTokenState& highlight);
+
 /*!
-	@}
+    @}
 */
diff --git a/ui/flowgraphwidget.h b/ui/flowgraphwidget.h
index 5133cd24..c7af8aa9 100644
--- a/ui/flowgraphwidget.h
+++ b/ui/flowgraphwidget.h
@@ -166,6 +166,8 @@ class BINARYNINJAUIAPI FlowGraphWidget :
 	void recenterUpdatedGraph(FlowGraphRef oldGraph, int oldXOfs, int oldYOfs);
 
 	BNDeadStoreElimination getCurrentVariableDeadStoreElimination();
+	std::optional<uint64_t> getCurrentFoldableExprAddress();
+	BNExprFolding getCurrentExprFolding();
 	std::optional<std::pair<BinaryNinja::Variable, BinaryNinja::Variable>> getMergeVariablesAtCurrentLocation();
 
   protected:
@@ -395,6 +397,7 @@ class BINARYNINJAUIAPI FlowGraphWidget :
 	void instrEditDoneEvent();
 
 	void setCurrentVariableDeadStoreElimination(BNDeadStoreElimination elimination);
+	void setCurrentExprFolding(BNExprFolding folding);
 	void splitToNewTabAndNavigateFromCursorPosition();
 	void splitToNewWindowAndNavigateFromCursorPosition();
 	void splitToNewPaneAndNavigateFromCursorPosition();
diff --git a/ui/linearview.h b/ui/linearview.h
index 588909be..e5305aa8 100644
--- a/ui/linearview.h
+++ b/ui/linearview.h
@@ -316,6 +316,8 @@ class BINARYNINJAUIAPI LinearView : public QAbstractScrollArea, public View, pub
 		TypeRef resType, uint64_t baseAddr, uint64_t& begin, uint64_t& end, bool singleLine, std::set<TypeRef>& seen);
 
 	BNDeadStoreElimination getCurrentVariableDeadStoreElimination();
+	std::optional<uint64_t> getCurrentFoldableExprAddress();
+	BNExprFolding getCurrentExprFolding();
 
 	void setDataButtonVisible(bool visible);
 	std::optional<std::pair<BinaryNinja::Variable, BinaryNinja::Variable>> getMergeVariablesAtCurrentLocation();
@@ -406,6 +408,7 @@ private Q_SLOTS:
 	std::optional<uint64_t> addressForCall();
 
 	void setCurrentVariableDeadStoreElimination(BNDeadStoreElimination elimination);
+	void setCurrentExprFolding(BNExprFolding folding);
 
 Q_SIGNALS:
 	void notifyResizeEvent(int width, int height);