Allow setting UIDF on any mlil/hlil usage

2 files changed, 48 insertions, 36 deletions

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;
diff --git a/ui/commands.h b/ui/commands.h
index fb4a1ba2..62016bdb 100644
--- a/ui/commands.h
+++ b/ui/commands.h
@@ -30,8 +30,8 @@ bool BINARYNINJAUIAPI createInferredMember(QWidget* parent, BinaryViewRef data,
 bool BINARYNINJAUIAPI createStructMembers(
     QWidget* parent, BinaryViewRef data, HighlightTokenState& highlight, FunctionRef func);
 
-bool BINARYNINJAUIAPI inputPossibleValueSet(QWidget* parent, BinaryViewRef data, FunctionRef currentFunction,
-    HighlightTokenState& highlight, uint64_t defSiteAddress, size_t ilInstructionIndex = BN_INVALID_EXPR);
+bool BINARYNINJAUIAPI inputPossibleValueSet(QWidget* parent, BinaryViewRef data, FunctionRef currentFunction, BNFunctionGraphType funcType,
+    const BinaryNinja::Variable& var, size_t ilInstructionIndex = BN_INVALID_EXPR);
 
 bool BINARYNINJAUIAPI getEnumSelection(QWidget* parent, BinaryViewRef data, FunctionRef func, uint64_t constValue,
 	TypeRef& selectedEnum, bool checkValue, bool canTruncate);
@@ -56,6 +56,12 @@ std::string BINARYNINJAUIAPI createStructureName(BinaryViewRef data);
 std::optional<BinaryNinja::Variable> BINARYNINJAUIAPI getSplitVariableForAssignment(
 	FunctionRef func, BNFunctionGraphType ilType, uint64_t location, const BinaryNinja::Variable& var);
 
+std::optional<size_t> getVariableDefinitionInstructionIndex(
+    FunctionRef func, BNFunctionGraphType funcType, const BinaryNinja::Variable& var, size_t ilInstructionIndex);
+std::optional<size_t> getVariableDefinitionAddress(
+    FunctionRef func, BNFunctionGraphType funcType, const BinaryNinja::Variable& var, size_t ilInstructionIndex);
+
+
 /*!
 	@}
 */