Add support for ternary simplification in Psuedo C and disable in graph view

2 files changed, 118 insertions, 0 deletions

diff --git a/lang/c/pseudoc.cpp b/lang/c/pseudoc.cpp
index 86a7b9db..4de9f447 100644
--- a/lang/c/pseudoc.cpp
+++ b/lang/c/pseudoc.cpp
@@ -442,6 +442,103 @@ PseudoCFunction::FieldDisplayType PseudoCFunction::GetFieldDisplayType(
 		return FieldDisplayNone;
 }
 
+std::optional<PseudoCFunction::TernaryInfo> PseudoCFunction::CanSimplifyToTernary(const BinaryNinja::HighLevelILInstruction &instr) const
+{
+    // Only handle if-statements
+    if (instr.operation != HLIL_IF)
+        return std::nullopt;
+
+    auto conditionExpr = instr.GetConditionExpr<HLIL_IF>();
+    auto trueExpr = instr.GetTrueExpr<HLIL_IF>();
+    auto falseExpr = instr.GetFalseExpr<HLIL_IF>();
+
+	if (GetHighLevelILFunction()->HasSideEffects(conditionExpr))
+		return std::nullopt;
+    // Both branches must be assignment operations
+    if (trueExpr.operation != HLIL_ASSIGN || falseExpr.operation != HLIL_ASSIGN)
+        return std::nullopt;
+
+    // Get the destination expressions of the assignments
+    auto trueDestExpr = trueExpr.GetDestExpr<HLIL_ASSIGN>();
+    auto falseDestExpr = falseExpr.GetDestExpr<HLIL_ASSIGN>();
+
+    // Verify that the destination expressions are variable references
+    if (trueDestExpr.operation != HLIL_VAR || falseDestExpr.operation != HLIL_VAR)
+        return std::nullopt;
+
+	auto trueExprDestExpr = trueExpr.GetDestExpr<HLIL_ASSIGN>();
+	auto falseExprDestExpr = falseExpr.GetDestExpr<HLIL_ASSIGN>();
+	if (trueExprDestExpr.operation != HLIL_VAR || falseExprDestExpr.operation != HLIL_VAR)
+		return std::nullopt;
+
+	auto trueExprDestVar = trueExprDestExpr.GetVariable<HLIL_VAR>();
+	auto falseExprDestVar = falseExprDestExpr.GetVariable<HLIL_VAR>();
+	if (trueExprDestVar != falseExprDestVar)
+		return std::nullopt;
+
+	auto trueExprSourceExpr = trueExpr.GetSourceExpr<HLIL_ASSIGN>();
+	auto falseExprSourceExpr = falseExpr.GetSourceExpr<HLIL_ASSIGN>();
+	if (GetHighLevelILFunction()->HasSideEffects(trueExprSourceExpr) || GetHighLevelILFunction()->HasSideEffects(falseExprSourceExpr))
+		return std::nullopt;
+
+    // Avoid folding for "else if" cases
+    for (auto parent = instr; parent.HasParent(); parent = parent.GetParent())
+    {
+        if (parent.operation != HLIL_IF)
+            break;
+        auto parentFalse = parent.GetFalseExpr<HLIL_IF>();
+        if (parentFalse.operation == HLIL_IF)
+            return std::nullopt;
+    }
+
+    TernaryInfo info;
+    info.conditional = conditionExpr;
+    info.assignDest  = trueDestExpr;
+    info.trueAssign  = trueExprSourceExpr;
+    info.falseAssign = falseExprSourceExpr;
+    return info;
+}
+
+bool PseudoCFunction::TryEmitSimplifiedTernary(
+	const BinaryNinja::HighLevelILInstruction &instr,
+	DisassemblySettings* settings,
+	BinaryNinja::HighLevelILTokenEmitter &emitter
+)
+{
+	auto ternaryOpt = CanSimplifyToTernary(instr);
+	if (!ternaryOpt.has_value())
+		return false;
+
+	auto &[conditional, assignDest, trueAssign, falseAssign] = ternaryOpt.value();
+
+	std::vector<InstructionTextToken> tokens = emitter.GetCurrentTokens();
+	size_t originalTokenCount = tokens.size();
+
+	// Emit the destination expression
+	GetExprText(assignDest, emitter, settings, AssignmentOperatorPrecedence);
+	emitter.Append(OperationToken, " = ");
+
+	// Emit the condition expression
+	GetExprText(conditional, emitter, settings, TernaryOperatorPrecedence);
+	emitter.Append(OperationToken, " ? ");
+
+	// Emit the true-source expression
+	GetExprText(trueAssign, emitter, settings, TernaryOperatorPrecedence);
+	emitter.Append(OperationToken, " : ");
+
+	// Emit the false-source expression
+	GetExprText(falseAssign, emitter, settings, TernaryOperatorPrecedence);
+	emitter.Append(KeywordToken, ";");
+
+	// If the ternary expression is too complex (i.e. too many tokens), revert back.
+	if (tokens.size() - originalTokenCount > emitter.GetMaxTernarySimplificationTokens())
+	{
+		tokens.resize(originalTokenCount);
+		return false;
+	}
+	return true;
+}
+
 
 void PseudoCFunction::AppendDefaultSplitExpr(const BinaryNinja::HighLevelILInstruction& instr,
 	BinaryNinja::HighLevelILTokenEmitter& tokens, DisassemblySettings* settings, BNOperatorPrecedence precedence)
@@ -672,6 +769,9 @@ void PseudoCFunction::GetExprTextInternal(const HighLevelILInstruction& instr, H
 	case HLIL_IF:
 		[&]()
 		{
+			if (instr.ast && TryEmitSimplifiedTernary(instr, settings, tokens))
+				return;
+
 			const auto condExpr = instr.GetConditionExpr<HLIL_IF>();
 			const auto trueExpr = instr.GetTrueExpr<HLIL_IF>();
 			const auto falseExpr = instr.GetFalseExpr<HLIL_IF>();
diff --git a/lang/c/pseudoc.h b/lang/c/pseudoc.h
index d983b2d9..ca723272 100644
--- a/lang/c/pseudoc.h
+++ b/lang/c/pseudoc.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include "binaryninjaapi.h"
+#include "highlevelilinstruction.h"
 
 class PseudoCFunction: public BinaryNinja::LanguageRepresentationFunction
 {
@@ -15,6 +16,23 @@ class PseudoCFunction: public BinaryNinja::LanguageRepresentationFunction
 		FieldDisplayNone
 	};
 
+	struct TernaryInfo
+	{
+		BinaryNinja::HighLevelILInstruction conditional;
+		BinaryNinja::HighLevelILInstruction assignDest;
+		BinaryNinja::HighLevelILInstruction trueAssign;
+		BinaryNinja::HighLevelILInstruction falseAssign;
+	};
+
+	std::optional<PseudoCFunction::TernaryInfo> CanSimplifyToTernary(
+		const BinaryNinja::HighLevelILInstruction& instr
+	) const;
+	bool TryEmitSimplifiedTernary(
+		const BinaryNinja::HighLevelILInstruction& instr,
+		BinaryNinja::DisassemblySettings* settings,
+		BinaryNinja::HighLevelILTokenEmitter& emitter
+	);
+
 	BinaryNinja::Ref<BinaryNinja::Type> GetFieldType(const BinaryNinja::HighLevelILInstruction& var, bool deref);
 	FieldDisplayType GetFieldDisplayType(BinaryNinja::Ref<BinaryNinja::Type> type, uint64_t offset, size_t memberIndex, bool deref);