From b270d7ebe1e2a406e1f7f14807987d24b0c4a846 Mon Sep 17 00:00:00 2001
From: Josh Ferrell <josh@vector35.com>
Date: Sat, 14 Jan 2023 16:01:59 -0500
Subject: Allow setting UIDF on any mlil/hlil usage

---
 function.cpp | 74 ++++++++++++++++++++++++++++++++----------------------------
 1 file changed, 40 insertions(+), 34 deletions(-)

(limited to 'function.cpp')

diff --git a/function.cpp b/function.cpp
index 298148b6..4ef91860 100644
--- a/function.cpp
+++ b/function.cpp
@@ -2497,26 +2497,29 @@ Ref<FlowGraph> Function::GetUnresolvedStackAdjustmentGraph()
 
 void Function::SetUserVariableValue(const Variable& var, uint64_t defAddr, PossibleValueSet& value)
 {
-	Ref<MediumLevelILFunction> mlil = GetMediumLevelIL();
-	const set<size_t>& varDefs = mlil->GetVariableDefinitions(var);
-	if (varDefs.size() == 0)
+	if (var.index != 0)
 	{
-		LogError("Could not get definition for Variable");
-		return;
-	}
-	bool found = false;
-	for (auto& site : varDefs)
-	{
-		const MediumLevelILInstruction& instr = mlil->GetInstruction(site);
-		if (instr.address == defAddr)
+		Ref<MediumLevelILFunction> mlil = GetMediumLevelIL();
+		const set<size_t>& varDefs = mlil->GetVariableDefinitions(var);
+		if (varDefs.size() == 0)
 		{
-			found = true;
-			break;
+			LogError("Could not get definition for Variable");
+			return;
+		}
+		bool found = false;
+		for (auto& site : varDefs)
+		{
+			const MediumLevelILInstruction& instr = mlil->GetInstruction(site);
+			if (instr.address == defAddr)
+			{
+				found = true;
+				break;
+			}
+		}
+		if (!found)
+		{
+			LogError("Could not find definition for variable at given address");
 		}
-	}
-	if (!found)
-	{
-		LogError("Could not find definition for variable at given address");
 	}
 	auto defSite = BNArchitectureAndAddress();
 	defSite.arch = GetArchitecture()->m_object;
@@ -2535,26 +2538,29 @@ void Function::SetUserVariableValue(const Variable& var, uint64_t defAddr, Possi
 
 void Function::ClearUserVariableValue(const Variable& var, uint64_t defAddr)
 {
-	Ref<MediumLevelILFunction> mlil = GetMediumLevelIL();
-	const set<size_t>& varDefs = mlil->GetVariableDefinitions(var);
-	if (varDefs.size() == 0)
+	if (var.index != 0)
 	{
-		LogError("Could not get definition for Variable");
-		return;
-	}
-	bool found = false;
-	for (auto& site : varDefs)
-	{
-		const MediumLevelILInstruction& instr = mlil->GetInstruction(site);
-		if (instr.address == defAddr)
+		Ref<MediumLevelILFunction> mlil = GetMediumLevelIL();
+		const set<size_t>& varDefs = mlil->GetVariableDefinitions(var);
+		if (varDefs.size() == 0)
 		{
-			found = true;
-			break;
+			LogError("Could not get definition for Variable");
+			return;
+		}
+		bool found = false;
+		for (auto& site : varDefs)
+		{
+			const MediumLevelILInstruction& instr = mlil->GetInstruction(site);
+			if (instr.address == defAddr)
+			{
+				found = true;
+				break;
+			}
+		}
+		if (!found)
+		{
+			LogError("Could not find definition for variable at given address");
 		}
-	}
-	if (!found)
-	{
-		LogError("Could not find definition for variable at given address");
 	}
 	auto defSite = BNArchitectureAndAddress();
 	defSite.arch = GetArchitecture()->m_object;
-- 
cgit v1.3.1