Move stack var refs and constant refs behind advanced analysis requests for significant memory savings

4 files changed, 86 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 411c9c9c..a58d42fe 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -3639,7 +3639,11 @@ namespace BinaryNinja {
 		std::vector<uint32_t> GetRegistersReadByInstruction(Architecture* arch, uint64_t addr);
 		std::vector<uint32_t> GetRegistersWrittenByInstruction(Architecture* arch, uint64_t addr);
 		std::vector<StackVariableReference> GetStackVariablesReferencedByInstruction(Architecture* arch, uint64_t addr);
+		std::vector<StackVariableReference> GetStackVariablesReferencedByInstructionIfAvailable(
+			Architecture* arch, uint64_t addr);
 		std::vector<BNConstantReference> GetConstantsReferencedByInstruction(Architecture* arch, uint64_t addr);
+		std::vector<BNConstantReference> GetConstantsReferencedByInstructionIfAvailable(
+			Architecture* arch, uint64_t addr);
 
 		std::vector<ILReferenceSource> GetMediumLevelILVariableReferences(const Variable& var);
 		std::vector<VariableReferenceSource> GetMediumLevelILVariableReferencesFrom(Architecture* arch, uint64_t addr);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 4783fed0..5687fcf2 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -3572,9 +3572,13 @@ extern "C"
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI BNStackVariableReference* BNGetStackVariablesReferencedByInstruction(
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
+	BINARYNINJACOREAPI BNStackVariableReference* BNGetStackVariablesReferencedByInstructionIfAvailable(
+		BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI void BNFreeStackVariableReferenceList(BNStackVariableReference* refs, size_t count);
 	BINARYNINJACOREAPI BNConstantReference* BNGetConstantsReferencedByInstruction(
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
+	BINARYNINJACOREAPI BNConstantReference* BNGetConstantsReferencedByInstructionIfAvailable(
+		BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
 	BINARYNINJACOREAPI void BNFreeConstantReferenceList(BNConstantReference* refs);
 
 	BINARYNINJACOREAPI BNLowLevelILFunction* BNGetFunctionLiftedIL(BNFunction* func);
diff --git a/function.cpp b/function.cpp
index fd6ea568..712c8155 100644
--- a/function.cpp
+++ b/function.cpp
@@ -570,6 +570,33 @@ vector<StackVariableReference> Function::GetStackVariablesReferencedByInstructio
 }
 
 
+vector<StackVariableReference> Function::GetStackVariablesReferencedByInstructionIfAvailable(
+	Architecture* arch, uint64_t addr)
+{
+	size_t count;
+	BNStackVariableReference* refs =
+		BNGetStackVariablesReferencedByInstructionIfAvailable(m_object, arch->GetObject(), addr, &count);
+
+	vector<StackVariableReference> result;
+	result.reserve(count);
+	for (size_t i = 0; i < count; i++)
+	{
+		StackVariableReference ref;
+		ref.sourceOperand = refs[i].sourceOperand;
+		ref.type = Confidence<Ref<Type>>(
+			refs[i].type ? new Type(BNNewTypeReference(refs[i].type)) : nullptr, refs[i].typeConfidence);
+		ref.name = refs[i].name;
+		ref.var = Variable::FromIdentifier(refs[i].varIdentifier);
+		ref.referencedOffset = refs[i].referencedOffset;
+		ref.size = refs[i].size;
+		result.push_back(ref);
+	}
+
+	BNFreeStackVariableReferenceList(refs, count);
+	return result;
+}
+
+
 vector<BNConstantReference> Function::GetConstantsReferencedByInstruction(Architecture* arch, uint64_t addr)
 {
 	size_t count;
@@ -583,6 +610,20 @@ vector<BNConstantReference> Function::GetConstantsReferencedByInstruction(Archit
 }
 
 
+vector<BNConstantReference> Function::GetConstantsReferencedByInstructionIfAvailable(Architecture* arch, uint64_t addr)
+{
+	size_t count;
+	BNConstantReference* refs =
+		BNGetConstantsReferencedByInstructionIfAvailable(m_object, arch->GetObject(), addr, &count);
+
+	vector<BNConstantReference> result;
+	result.insert(result.end(), &refs[0], &refs[count]);
+
+	BNFreeConstantReferenceList(refs);
+	return result;
+}
+
+
 Ref<LowLevelILFunction> Function::GetLiftedIL() const
 {
 	return new LowLevelILFunction(BNGetFunctionLiftedIL(m_object));
diff --git a/python/function.py b/python/function.py
index 006677b8..2bd64f4a 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2062,6 +2062,28 @@ class Function:
 		core.BNFreeStackVariableReferenceList(refs, count.value)
 		return result
 
+	def get_stack_vars_referenced_by_address_if_available(
+	    self, addr: int, arch: Optional['architecture.Architecture'] = None
+	) -> List['variable.StackVariableReference']:
+		if arch is None:
+			arch = self.arch
+		count = ctypes.c_ulonglong()
+		refs = core.BNGetStackVariablesReferencedByInstructionIfAvailable(self.handle, arch.handle, addr, count)
+		assert refs is not None, "core.BNGetStackVariablesReferencedByInstructionIfAvailable returned None"
+		result = []
+		for i in range(0, count.value):
+			var_type = types.Type.create(
+			    core.BNNewTypeReference(refs[i].type), platform=self.platform, confidence=refs[i].typeConfidence
+			)
+			var = variable.Variable.from_identifier(self, refs[i].varIdentifier)
+			result.append(
+			    variable.StackVariableReference(
+			        refs[i].sourceOperand, var_type, refs[i].name, var, refs[i].referencedOffset, refs[i].size
+			    )
+			)
+		core.BNFreeStackVariableReferenceList(refs, count.value)
+		return result
+
 	@property
 	def auto_function_tags(self):
 		"""
@@ -2218,6 +2240,21 @@ class Function:
 		core.BNFreeConstantReferenceList(refs)
 		return result
 
+	def get_constants_referenced_by_address_if_available(self, addr: int,
+	                                arch: 'architecture.Architecture' = None) -> List[variable.ConstantReference]:
+		if arch is None:
+			arch = self.arch
+		count = ctypes.c_ulonglong()
+		refs = core.BNGetConstantsReferencedByInstructionIfAvailable(self.handle, arch.handle, addr, count)
+		assert refs is not None, "core.BNGetConstantsReferencedByInstructionIfAvailable returned None"
+		result = []
+		for i in range(0, count.value):
+			result.append(
+			    variable.ConstantReference(refs[i].value, refs[i].size, refs[i].pointer, refs[i].intermediate)
+			)
+		core.BNFreeConstantReferenceList(refs)
+		return result
+
 	def remove_auto_function_tag(self, tag: 'binaryview.Tag') -> None:
 		"""
 		``remove_user_function_tag`` removes a Tag object as a function tag.