Top relative register stack access as a normal register

5 files changed, 31 insertions, 10 deletions

diff --git a/architecture.cpp b/architecture.cpp
index 56af3f80..169a8c0c 100644
--- a/architecture.cpp
+++ b/architecture.cpp
@@ -731,7 +731,9 @@ BNRegisterStackInfo Architecture::GetRegisterStackInfo(uint32_t)
 {
 	BNRegisterStackInfo result;
 	result.firstStorageReg = BN_INVALID_REGISTER;
-	result.count = 0;
+	result.topRelativeCount = BN_INVALID_REGISTER;
+	result.storageCount = 0;
+	result.topRelativeCount = 0;
 	result.stackTopReg = BN_INVALID_REGISTER;
 	return result;
 }
diff --git a/binaryninjacore.h b/binaryninjacore.h
index f50c28c8..c32d42e3 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -679,8 +679,8 @@ extern "C"
 
 	struct BNRegisterStackInfo
 	{
-		uint32_t firstStorageReg;
-		uint32_t count;
+		uint32_t firstStorageReg, firstTopRelativeReg;
+		uint32_t storageCount, topRelativeCount;
 		uint32_t stackTopReg;
 	};
 
diff --git a/lowlevelilinstruction.cpp b/lowlevelilinstruction.cpp
index 4130dd2a..91af36b1 100644
--- a/lowlevelilinstruction.cpp
+++ b/lowlevelilinstruction.cpp
@@ -2179,7 +2179,7 @@ ExprId LowLevelILFunction::SetRegisterStackTopRelativeSSA(size_t size, uint32_t
 ExprId LowLevelILFunction::SetRegisterStackAbsoluteSSA(size_t size, uint32_t regStack,
 	size_t destVersion, size_t srcVersion, uint32_t reg, ExprId val, const ILSourceLocation& loc)
 {
-	return AddExprWithLocation(LLIL_SET_REG_STACK_REL_SSA, loc, size, 0,
+	return AddExprWithLocation(LLIL_SET_REG_STACK_ABS_SSA, loc, size, 0,
 		AddExprWithLocation(LLIL_REG_STACK_DEST_SSA, loc, size, 0, regStack, destVersion, srcVersion),
 		reg, val);
 }
diff --git a/python/architecture.py b/python/architecture.py
index 289b0abd..04179e7f 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -225,10 +225,13 @@ class Architecture(object):
 				name = core.BNGetArchitectureRegisterStackName(self.handle, regs[i])
 				info = core.BNGetArchitectureRegisterStackInfo(self.handle, regs[i])
 				storage = []
-				for j in xrange(0, info.count):
+				for j in xrange(0, info.storageCount):
 					storage.append(core.BNGetArchitectureRegisterName(self.handle, info.firstStorageReg + j))
+				top_rel = []
+				for j in xrange(0, info.topRelativeCount):
+					top_rel.append(core.BNGetArchitectureRegisterName(self.handle, info.firstTopRelativeReg + j))
 				top = core.BNGetArchitectureRegisterName(self.handle, info.stackTopReg)
-				self.reg_stacks[name] = function.RegisterStackInfo(storage, top)
+				self.reg_stacks[name] = function.RegisterStackInfo(storage, top_rel, top)
 			core.BNFreeRegisterList(regs)
 		else:
 			startup._init_plugins()
@@ -310,6 +313,11 @@ class Architecture(object):
 					self._regs_by_index[reg_index] = reg
 					self.regs[reg].index = reg_index
 					reg_index += 1
+				for reg in info.top_relative_regs:
+					self._all_regs[reg] = reg_index
+					self._regs_by_index[reg_index] = reg
+					self.regs[reg].index = reg_index
+					reg_index += 1
 				if reg_stack not in self._all_reg_stacks:
 					self._all_reg_stacks[reg_stack] = reg_stack_index
 					self._reg_stacks_by_index[reg_stack_index] = reg_stack
@@ -808,17 +816,27 @@ class Architecture(object):
 		try:
 			if reg_stack not in self._reg_stacks_by_index:
 				result[0].firstStorageReg = 0
-				result[0].count = 0
+				result[0].firstTopRelativeReg = 0
+				result[0].storageCount = 0
+				result[0].topRelativeCount = 0
 				result[0].stackTopReg = 0
 				return
 			info = self.__class__.regs[self._reg_stacks_by_index[reg_stack]]
 			result[0].firstStorageReg = self._all_regs[info.storage_regs[0]]
-			result[0].count = len(info.storage_regs)
+			result[0].storageCount = len(info.storage_regs)
+			if len(info.top_relative_regs) > 0:
+				result[0].firstTopRelativeReg = self._all_regs[info.top_relative_regs[0]]
+				result[0].topRelativeCount = len(info.top_relative_regs)
+			else:
+				result[0].firstTopRelativeReg = 0
+				result[0].topRelativeCount = 0
 			result[0].stackTopReg = self._all_regs[info.stack_top_reg]
 		except KeyError:
 			log.log_error(traceback.format_exc())
 			result[0].firstStorageReg = 0
-			result[0].count = 0
+			result[0].firstTopRelativeReg = 0
+			result[0].storageCount = 0
+			result[0].topRelativeCount = 0
 			result[0].stackTopReg = 0
 
 	def _assemble(self, ctxt, code, addr, result, errors):
diff --git a/python/function.py b/python/function.py
index eb8796f1..ba00147e 100644
--- a/python/function.py
+++ b/python/function.py
@@ -1727,8 +1727,9 @@ class RegisterInfo(object):
 
 
 class RegisterStackInfo(object):
-	def __init__(self, storage_regs, stack_top_reg):
+	def __init__(self, storage_regs, top_relative_regs, stack_top_reg):
 		self.storage_regs = storage_regs
+		self.top_relative_regs = top_relative_regs
 		self.stack_top_reg = stack_top_reg
 
 	def __repr__(self):