Expose shared lifted branch overrides across language bindings

Add the core API and C++, Rust, and Python context methods for applying
branch overrides to staged LLIL. Update DefaultLiftFunction to use the
shared implementation and add Rust binding coverage.

Author: zznop

architecture.cpp

@@ -666,6 +666,108 @@ const std::map<ArchAndAddr, std::map<BNBranchType, BranchOverride>>& FunctionLif
 }
 
 
+PreparedLiftedBranchOverride::PreparedLiftedBranchOverride(BNPreparedLiftedBranchOverride* object) : m_object(object)
+{}
+
+
+PreparedLiftedBranchOverride::~PreparedLiftedBranchOverride()
+{
+	BNFreePreparedLiftedBranchOverride(m_object);
+}
+
+
+bool PreparedLiftedBranchOverride::SuppressesInstruction() const
+{
+	return BNPreparedLiftedBranchOverrideSuppressesInstruction(m_object);
+}
+
+
+Ref<LowLevelILFunction> PreparedLiftedBranchOverride::GetSource()
+{
+	if (!m_source)
+		m_source = new LowLevelILFunction(BNGetPreparedLiftedBranchOverrideSource(m_object));
+	return m_source;
+}
+
+
+bool PreparedLiftedBranchOverride::Apply(uint64_t continuationAddress)
+{
+	return BNApplyPreparedLiftedBranchOverride(m_object, continuationAddress);
+}
+
+
+struct FunctionLifterContext::LiftedBranchOverrideInputs
+{
+	vector<BNOverridableBranchInfo> branches;
+	vector<BNBranchOverride> overrides;
+	vector<BNArchitectureAndAddress> indirectTargets;
+	BNLiftedBranchOverrideInfo info{};
+};
+
+
+FunctionLifterContext::LiftedBranchOverrideInputs FunctionLifterContext::GetLiftedBranchOverrideInputs(
+	const ArchAndAddr& location, const vector<OverridableBranchInfo>& branches) const
+{
+	LiftedBranchOverrideInputs result;
+	auto& rawBranches = result.branches;
+	auto& rawOverrides = result.overrides;
+	auto& indirectTargets = result.indirectTargets;
+	auto it = m_branchOverrides.find(location);
+	if (!location.arch || (it == m_branchOverrides.end()))
+		return result;
+
+	rawBranches.reserve(branches.size());
+	for (const auto& branch : branches)
+		rawBranches.push_back({branch.type, branch.target, branch.arch ? branch.arch->GetObject() : nullptr});
+
+	rawOverrides.reserve(it->second.size());
+	for (const auto& [originalType, value] : it->second)
+	{
+		rawOverrides.push_back({location.arch->GetObject(), location.address, originalType, value.type,
+			value.target.has_value(), value.targetArch ? value.targetArch->GetObject() : nullptr,
+			value.target.value_or(0)});
+	}
+
+	const set<ArchAndAddr>* targets = nullptr;
+	if (auto user = m_userIndirectBranches.find(location); user != m_userIndirectBranches.end())
+		targets = &user->second;
+	else if (auto automatic = m_autoIndirectBranches.find(location); automatic != m_autoIndirectBranches.end())
+		targets = &automatic->second;
+	if (targets)
+	{
+		indirectTargets.reserve(targets->size());
+		for (const auto& target : *targets)
+			indirectTargets.push_back({target.arch->GetObject(), target.address});
+	}
+
+	result.info = {location.arch->GetObject(), location.address, 0,
+		rawBranches.data(), rawBranches.size(), rawOverrides.data(), rawOverrides.size(),
+		m_noReturnCalls.count(location) != 0, indirectTargets.data(), indirectTargets.size()};
+	return result;
+}
+
+
+unique_ptr<PreparedLiftedBranchOverride> FunctionLifterContext::PrepareLiftedBranchOverrides(
+	LowLevelILFunction& dest, BasicBlock* block, const ArchAndAddr& location,
+	const vector<OverridableBranchInfo>& branches) const
+{
+	auto inputs = GetLiftedBranchOverrideInputs(location, branches);
+	if (!inputs.info.arch)
+		return nullptr;
+	auto result = BNPrepareLiftedBranchOverrides(dest.GetObject(), block ? block->GetObject() : nullptr, &inputs.info);
+	return result ? make_unique<PreparedLiftedBranchOverride>(result) : nullptr;
+}
+
+
+bool FunctionLifterContext::ApplyLiftedBranchOverrides(LowLevelILFunction& dest, LowLevelILFunction& source,
+	const ArchAndAddr& location, uint64_t continuationAddress, const vector<OverridableBranchInfo>& branches) const
+{
+	auto inputs = GetLiftedBranchOverrideInputs(location, branches);
+	inputs.info.continuationAddress = continuationAddress;
+	return inputs.info.arch && BNApplyLiftedBranchOverrides(dest.GetObject(), source.GetObject(), &inputs.info);
+}
+
+
 Ref<Logger>& FunctionLifterContext::GetLogger()
 {
 	return m_logger;

binaryninjaapi.h

@@ -9850,6 +9850,24 @@ namespace BinaryNinja {
 		void Finalize();
 	};
 
+	// A prepared override owns its inputs and staging IL. Use only while lifting.
+	// The architecture chooses the instruction group and its continuation address.
+	class PreparedLiftedBranchOverride
+	{
+		BNPreparedLiftedBranchOverride* m_object;
+		Ref<LowLevelILFunction> m_source;
+
+	public:
+		explicit PreparedLiftedBranchOverride(BNPreparedLiftedBranchOverride* object);
+		~PreparedLiftedBranchOverride();
+		PreparedLiftedBranchOverride(const PreparedLiftedBranchOverride&) = delete;
+		PreparedLiftedBranchOverride& operator=(const PreparedLiftedBranchOverride&) = delete;
+		bool SuppressesInstruction() const;
+		Ref<LowLevelILFunction> GetSource();
+		// A successful application consumes the prepared override; a second call returns false.
+		bool Apply(uint64_t continuationAddress);
+	};
+
 	class FunctionLifterContext
 	{
 		Ref<LowLevelILFunction> m_function;
@@ -9867,6 +9885,9 @@ namespace BinaryNinja {
 		bool* m_containsInlinedFunctions;
 		void* m_functionArchContext;
 		Ref<LifterInstructionData> m_lifterInstructionData;
+		struct LiftedBranchOverrideInputs;
+		LiftedBranchOverrideInputs GetLiftedBranchOverrideInputs(
+			const ArchAndAddr& location, const std::vector<OverridableBranchInfo>& branches) const;
 
 	public:
 		BNFunctionLifterContext* m_context;
@@ -9881,6 +9902,17 @@ namespace BinaryNinja {
 		std::map<ArchAndAddr, std::set<ArchAndAddr>>& GetUserIndirectBranches();
 		std::map<ArchAndAddr, std::set<ArchAndAddr>>& GetAutoIndirectBranches();
 		const std::map<ArchAndAddr, std::map<BNBranchType, BranchOverride>>& GetBranchOverrides() const;
+		// Returns null for no matching overrides or invalid inputs, without changing dest.
+		// The result owns a snapshot of the inputs; no staging IL is allocated until GetSource.
+		// See BNPrepareLiftedBranchOverrides and BNApplyLiftedBranchOverrides for requirements.
+		std::unique_ptr<PreparedLiftedBranchOverride> PrepareLiftedBranchOverrides(LowLevelILFunction& dest,
+			BasicBlock* block, const ArchAndAddr& location, const std::vector<OverridableBranchInfo>& branches) const;
+		// Append an isolated, unfinalized instruction group using this context's overrides.
+		// See BNApplyLiftedBranchOverrides for label, temporary-register, and ownership requirements.
+		// False leaves dest untouched, including when no original branch has an override.
+		bool ApplyLiftedBranchOverrides(LowLevelILFunction& dest, LowLevelILFunction& source,
+			const ArchAndAddr& location, uint64_t continuationAddress,
+			const std::vector<OverridableBranchInfo>& branches) const;
 		std::set<uint64_t>& GetInlinedCalls();
 		void SetContainsInlinedFunctions(bool value);
 		void* GetFunctionArchContextRaw() const { return m_functionArchContext; }

binaryninjacore.h

@@ -255,6 +255,7 @@ extern "C"
 	typedef struct BNFunction BNFunction;
 	typedef struct BNBasicBlock BNBasicBlock;
 	typedef struct BNLifterInstructionData BNLifterInstructionData;
+	typedef struct BNPreparedLiftedBranchOverride BNPreparedLiftedBranchOverride;
 	typedef struct BNDownloadProvider BNDownloadProvider;
 	typedef struct BNDownloadInstance BNDownloadInstance;
 	typedef struct BNWebsocketProvider BNWebsocketProvider;
@@ -2187,6 +2188,20 @@ extern "C"
 		uint64_t replacementTarget;
 	} BNBranchOverride;
 
+	typedef struct BNLiftedBranchOverrideInfo
+	{
+		BNArchitecture* arch;
+		uint64_t address;
+		uint64_t continuationAddress;
+		const BNOverridableBranchInfo* branches;
+		size_t branchCount;
+		const BNBranchOverride* overrides;
+		size_t overrideCount;
+		bool noReturnCall;
+		const BNArchitectureAndAddress* indirectTargets;
+		size_t indirectTargetCount;
+	} BNLiftedBranchOverrideInfo;
+
 	typedef struct BNBasicBlockAnalysisContext
 	{
 		BNFunction* function;
@@ -5894,6 +5909,17 @@ extern "C"
 		BNBasicBlockAnalysisContext* context);
 	BINARYNINJACOREAPI bool BNArchitectureSetDefaultLiftFunctionCallback(void *callback);
 	BINARYNINJACOREAPI bool BNArchitectureDefaultLiftFunction(BNLowLevelILFunction* function, BNFunctionLifterContext* context);
+	BINARYNINJACOREAPI bool BNApplyLiftedBranchOverrides(BNLowLevelILFunction* dest,
+		BNLowLevelILFunction* source, const BNLiftedBranchOverrideInfo* info);
+	BINARYNINJACOREAPI BNPreparedLiftedBranchOverride* BNPrepareLiftedBranchOverrides(
+		BNLowLevelILFunction* dest, BNBasicBlock* block, const BNLiftedBranchOverrideInfo* info);
+	BINARYNINJACOREAPI void BNFreePreparedLiftedBranchOverride(BNPreparedLiftedBranchOverride* prepared);
+	BINARYNINJACOREAPI bool BNPreparedLiftedBranchOverrideSuppressesInstruction(
+		BNPreparedLiftedBranchOverride* prepared);
+	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetPreparedLiftedBranchOverrideSource(
+		BNPreparedLiftedBranchOverride* prepared);
+	BINARYNINJACOREAPI bool BNApplyPreparedLiftedBranchOverride(
+		BNPreparedLiftedBranchOverride* prepared, uint64_t continuationAddress);
 	BINARYNINJACOREAPI bool BNArchitectureLiftFunction(BNArchitecture* arch, BNLowLevelILFunction* function,
 		BNFunctionLifterContext* context);
 	BINARYNINJACOREAPI void BNArchitectureFreeFunctionArchContext(BNArchitecture* arch, void* context);