Instruction starts and SSA form are a property of IL functions

6 files changed, 49 insertions, 18 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 188471b3..9d81b198 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1887,7 +1887,6 @@ namespace BinaryNinja
 		void SetCommentForAddress(uint64_t addr, const std::string& comment);
 
 		Ref<LowLevelILFunction> GetLowLevelIL() const;
-		Ref<LowLevelILFunction> GetLowLevelILSSAForm() const;
 		size_t GetLowLevelILForInstruction(Architecture* arch, uint64_t addr);
 		std::vector<size_t> GetLowLevelILExitsForInstruction(Architecture* arch, uint64_t addr);
 		RegisterValue GetRegisterValueAtInstruction(Architecture* arch, uint64_t addr, uint32_t reg);
@@ -2057,7 +2056,8 @@ namespace BinaryNinja
 		LowLevelILFunction(BNLowLevelILFunction* func);
 
 		uint64_t GetCurrentAddress() const;
-		void SetCurrentAddress(uint64_t addr);
+		void SetCurrentAddress(Architecture* arch, uint64_t addr);
+		size_t GetInstructionStart(Architecture* arch, uint64_t addr);
 
 		void ClearIndirectBranches();
 		void SetIndirectBranches(const std::vector<ArchAndAddr>& branches);
@@ -2158,6 +2158,8 @@ namespace BinaryNinja
 		uint32_t GetTemporaryFlagCount();
 
 		std::vector<Ref<BasicBlock>> GetBasicBlocks() const;
+
+		Ref<LowLevelILFunction> GetSSAForm() const;
 	};
 
 	class FunctionRecognizer
diff --git a/binaryninjacore.h b/binaryninjacore.h
index c1e7d6f9..e95f2bf3 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1694,7 +1694,6 @@ extern "C"
 	BINARYNINJACOREAPI BNBasicBlock** BNGetBasicBlocksStartingAtAddress(BNBinaryView* view, uint64_t addr, size_t* count);
 
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLowLevelIL(BNFunction* func);
-	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLowLevelILSSAForm(BNFunction* func);
 	BINARYNINJACOREAPI size_t BNGetLowLevelILForInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr);
 	BINARYNINJACOREAPI size_t* BNGetLowLevelILExitsForInstruction(BNFunction* func, BNArchitecture* arch, uint64_t addr,
 	                                                              size_t* count);
@@ -1977,7 +1976,10 @@ extern "C"
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNNewLowLevelILFunctionReference(BNLowLevelILFunction* func);
 	BINARYNINJACOREAPI void BNFreeLowLevelILFunction(BNLowLevelILFunction* func);
 	BINARYNINJACOREAPI uint64_t BNLowLevelILGetCurrentAddress(BNLowLevelILFunction* func);
-	BINARYNINJACOREAPI void BNLowLevelILSetCurrentAddress(BNLowLevelILFunction* func, uint64_t addr);
+	BINARYNINJACOREAPI void BNLowLevelILSetCurrentAddress(BNLowLevelILFunction* func,
+		BNArchitecture* arch, uint64_t addr);
+	BINARYNINJACOREAPI size_t BNLowLevelILGetInstructionStart(BNLowLevelILFunction* func,
+		BNArchitecture* arch, uint64_t addr);
 	BINARYNINJACOREAPI void BNLowLevelILClearIndirectBranches(BNLowLevelILFunction* func);
 	BINARYNINJACOREAPI void BNLowLevelILSetIndirectBranches(BNLowLevelILFunction* func, BNArchitectureAndAddress* branches,
 	                                                        size_t count);
@@ -2015,6 +2017,8 @@ extern "C"
 
 	BINARYNINJACOREAPI BNBasicBlock** BNGetLowLevelILBasicBlockList(BNLowLevelILFunction* func, size_t* count);
 
+	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetLowLevelILSSAForm(BNLowLevelILFunction* func);
+
 	// Types
 	BINARYNINJACOREAPI BNType* BNCreateVoidType(void);
 	BINARYNINJACOREAPI BNType* BNCreateBoolType(void);
diff --git a/function.cpp b/function.cpp
index d8a5a4a7..b1d008ee 100644
--- a/function.cpp
+++ b/function.cpp
@@ -147,12 +147,6 @@ Ref<LowLevelILFunction> Function::GetLowLevelIL() const
 }
 
 
-Ref<LowLevelILFunction> Function::GetLowLevelILSSAForm() const
-{
-	return new LowLevelILFunction(BNGetFunctionLowLevelILSSAForm(m_object));
-}
-
-
 size_t Function::GetLowLevelILForInstruction(Architecture* arch, uint64_t addr)
 {
 	return BNGetLowLevelILForInstruction(m_object, arch->GetObject(), addr);
diff --git a/lowlevelil.cpp b/lowlevelil.cpp
index bf3b980e..1a353b51 100644
--- a/lowlevelil.cpp
+++ b/lowlevelil.cpp
@@ -48,9 +48,15 @@ uint64_t LowLevelILFunction::GetCurrentAddress() const
 }
 
 
-void LowLevelILFunction::SetCurrentAddress(uint64_t addr)
+void LowLevelILFunction::SetCurrentAddress(Architecture* arch, uint64_t addr)
 {
-	BNLowLevelILSetCurrentAddress(m_object, addr);
+	BNLowLevelILSetCurrentAddress(m_object, arch ? arch->GetObject() : nullptr, addr);
+}
+
+
+size_t LowLevelILFunction::GetInstructionStart(Architecture* arch, uint64_t addr)
+{
+	return BNLowLevelILGetInstructionStart(m_object, arch ? arch->GetObject() : nullptr, addr);
 }
 
 
@@ -643,3 +649,12 @@ vector<Ref<BasicBlock>> LowLevelILFunction::GetBasicBlocks() const
 	BNFreeBasicBlockList(blocks, count);
 	return result;
 }
+
+
+Ref<LowLevelILFunction> LowLevelILFunction::GetSSAForm() const
+{
+	BNLowLevelILFunction* func = BNGetLowLevelILSSAForm(m_object);
+	if (!func)
+		return nullptr;
+	return new LowLevelILFunction(func);
+}
diff --git a/python/function.py b/python/function.py
index 47099395..86c93bf7 100644
--- a/python/function.py
+++ b/python/function.py
@@ -293,11 +293,6 @@ class Function(object):
 		return lowlevelil.LowLevelILFunction(self.arch, core.BNGetFunctionLowLevelIL(self.handle), self)
 
 	@property
-	def low_level_il_ssa_form(self):
-		"""Function low level IL in SSA form (read-only)"""
-		return lowlevelil.LowLevelILFunction(self.arch, core.BNGetFunctionLowLevelILSSAForm(self.handle), self)
-
-	@property
 	def lifted_il(self):
 		"""Function lifted IL (read-only)"""
 		return lowlevelil.LowLevelILFunction(self.arch, core.BNGetFunctionLiftedIL(self.handle), self)
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 3ce55793..390e4210 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -314,7 +314,12 @@ class LowLevelILFunction(object):
 
 	@current_address.setter
 	def current_address(self, value):
-		core.BNLowLevelILSetCurrentAddress(self.handle, value)
+		core.BNLowLevelILSetCurrentAddress(self.handle, self.arch.handle, value)
+
+	def set_current_address(self, value, arch = None):
+		if arch is None:
+			arch = self.arch
+		core.BNLowLevelILSetCurrentAddress(self.handle, arch.handle, value)
 
 	@property
 	def temp_reg_count(self):
@@ -340,6 +345,14 @@ class LowLevelILFunction(object):
 		core.BNFreeBasicBlockList(blocks, count.value)
 		return result
 
+	@property
+	def ssa_form(self):
+		"""Low level IL in SSA form (read-only)"""
+		result = core.BNGetLowLevelILSSAForm(self.handle)
+		if not result:
+			return None
+		return LowLevelILFunction(self.arch, result, self.source_function)
+
 	def __setattr__(self, name, value):
 		try:
 			object.__setattr__(self, name, value)
@@ -373,6 +386,14 @@ class LowLevelILFunction(object):
 		finally:
 			core.BNFreeBasicBlockList(blocks, count.value)
 
+	def get_instruction_start(self, addr, arch = None):
+		if arch is None:
+			arch = self.arch
+		result = core.BNLowLevelILGetInstructionStart(self.handle, arch.handle, addr)
+		if result >= core.BNGetLowLevelILInstructionCount(self.handle):
+			return None
+		return result
+
 	def clear_indirect_branches(self):
 		core.BNLowLevelILClearIndirectBranches(self.handle)