From 73b649e79d5cd219cb56866abf2ee5e6a64d976a Mon Sep 17 00:00:00 2001
From: Chinmay <chinmay1dd@gmail.com>
Date: Mon, 17 Aug 2020 22:53:09 -0700
Subject: UI changes for user-informed dataflow

---
 function.cpp | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

(limited to 'function.cpp')

diff --git a/function.cpp b/function.cpp
index ecb81a67..168decc2 100644
--- a/function.cpp
+++ b/function.cpp
@@ -1773,17 +1773,18 @@ Ref<FlowGraph> Function::GetUnresolvedStackAdjustmentGraph()
 
 void Function::SetUserVariableValue(const Variable& var, uint64_t defAddr, PossibleValueSet& value)
 {
-	auto mlil = GetMediumLevelIL();
-	auto varDefs = mlil->GetVariableDefinitions(var);
+	Ref<MediumLevelILFunction> mlil = GetMediumLevelIL();
+	const set<size_t>& varDefs = mlil->GetVariableDefinitions(var);
 	if (varDefs.size() == 0)
 	{
 		LogError("Could not get definition for Variable");
 		return;
 	} 
 	bool found = false;
-	for (const auto& site : varDefs)
+	for (auto& site : varDefs)
 	{
-		if (site == defAddr)
+		const MediumLevelILInstruction& instr = mlil->GetInstruction(site);
+		if (instr.address == defAddr)
 		{
 			found = true;
 			break;
@@ -1810,16 +1811,18 @@ void Function::SetUserVariableValue(const Variable& var, uint64_t defAddr, Possi
 
 void Function::ClearUserVariableValue(const Variable& var, uint64_t defAddr)
 {
-	auto mlil = GetMediumLevelIL();
-	auto varDefs = mlil->GetVariableDefinitions(var);
+	Ref<MediumLevelILFunction> mlil = GetMediumLevelIL();
+	const set<size_t>& varDefs = mlil->GetVariableDefinitions(var);
 	if (varDefs.size() == 0)
 	{
 		LogError("Could not get definition for Variable");
+		return;
 	} 
 	bool found = false;
-	for (auto site : varDefs)
+	for (auto& site : varDefs)
 	{
-		if (site == defAddr)
+		const MediumLevelILInstruction& instr = mlil->GetInstruction(site);
+		if (instr.address == defAddr)
 		{
 			found = true;
 			break;
-- 
cgit v1.3.1