Adding APIs to ensure IL is available for upcoming caching methods

4 files changed, 134 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 342dcc48..72e13567 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1611,8 +1611,11 @@ namespace BinaryNinja
 
 	class Function: public CoreRefCountObject<BNFunction, BNNewFunctionReference, BNFreeFunction>
 	{
+		int m_advancedAnalysisRequests;
+
 	public:
 		Function(BNFunction* func);
+		virtual ~Function();
 
 		Ref<Architecture> GetArchitecture() const;
 		Ref<Platform> GetPlatform() const;
@@ -1684,6 +1687,24 @@ namespace BinaryNinja
 			BNIntegerDisplayType type);
 
 		void Reanalyze();
+
+		void RequestAdvancedAnalysisData();
+		void ReleaseAdvancedAnalysisData();
+		void ReleaseAdvancedAnalysisData(size_t count);
+	};
+
+	class AdvancedFunctionAnalysisDataRequestor
+	{
+		Ref<Function> m_func;
+
+	public:
+		AdvancedFunctionAnalysisDataRequestor(Function* func = nullptr);
+		AdvancedFunctionAnalysisDataRequestor(const AdvancedFunctionAnalysisDataRequestor& req);
+		~AdvancedFunctionAnalysisDataRequestor();
+		AdvancedFunctionAnalysisDataRequestor& operator=(const AdvancedFunctionAnalysisDataRequestor& req);
+
+		Ref<Function> GetFunction() { return m_func; }
+		void SetFunction(Function* func);
 	};
 
 	struct FunctionGraphEdge
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 6eff4d46..0c2738cb 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1384,6 +1384,9 @@ extern "C"
 	BINARYNINJACOREAPI void BNUpdateAnalysis(BNBinaryView* view);
 	BINARYNINJACOREAPI void BNAbortAnalysis(BNBinaryView* view);
 	BINARYNINJACOREAPI bool BNIsFunctionUpdateNeeded(BNFunction* func);
+	BINARYNINJACOREAPI void BNRequestAdvancedFunctionAnalysisData(BNFunction* func);
+	BINARYNINJACOREAPI void BNReleaseAdvancedFunctionAnalysisData(BNFunction* func);
+	BINARYNINJACOREAPI void BNReleaseAdvancedFunctionAnalysisDataMultiple(BNFunction* func, size_t count);
 
 	BINARYNINJACOREAPI BNFunction* BNNewFunctionReference(BNFunction* func);
 	BINARYNINJACOREAPI void BNFreeFunction(BNFunction* func);
diff --git a/function.cpp b/function.cpp
index cbd9ec40..de12918f 100644
--- a/function.cpp
+++ b/function.cpp
@@ -27,6 +27,14 @@ using namespace std;
 Function::Function(BNFunction* func)
 {
 	m_object = func;
+	m_advancedAnalysisRequests = 0;
+}
+
+
+Function::~Function()
+{
+	if (m_advancedAnalysisRequests > 0)
+		BNReleaseAdvancedFunctionAnalysisDataMultiple(m_object, (size_t)m_advancedAnalysisRequests);
 }
 
 
@@ -578,3 +586,67 @@ void Function::Reanalyze()
 {
 	BNReanalyzeFunction(m_object);
 }
+
+
+void Function::RequestAdvancedAnalysisData()
+{
+	BNRequestAdvancedFunctionAnalysisData(m_object);
+#ifdef WIN32
+	InterlockedIncrement((LONG*)&m_advancedAnalysisRequests);
+#else
+	__sync_fetch_and_add(&m_advancedAnalysisRequests, 1);
+#endif
+}
+
+
+void Function::ReleaseAdvancedAnalysisData()
+{
+	BNReleaseAdvancedFunctionAnalysisData(m_object);
+#ifdef WIN32
+	InterlockedDecrement((LONG*)&m_advancedAnalysisRequests);
+#else
+	__sync_fetch_and_add(&m_advancedAnalysisRequests, -1);
+#endif
+}
+
+
+AdvancedFunctionAnalysisDataRequestor::AdvancedFunctionAnalysisDataRequestor(Function* func): m_func(func)
+{
+	if (m_func)
+		m_func->RequestAdvancedAnalysisData();
+}
+
+
+AdvancedFunctionAnalysisDataRequestor::AdvancedFunctionAnalysisDataRequestor(const AdvancedFunctionAnalysisDataRequestor& req)
+{
+	m_func = req.m_func;
+	if (m_func)
+		m_func->RequestAdvancedAnalysisData();
+}
+
+
+AdvancedFunctionAnalysisDataRequestor::~AdvancedFunctionAnalysisDataRequestor()
+{
+	if (m_func)
+		m_func->ReleaseAdvancedAnalysisData();
+}
+
+
+AdvancedFunctionAnalysisDataRequestor& AdvancedFunctionAnalysisDataRequestor::operator=(
+	const AdvancedFunctionAnalysisDataRequestor& req)
+{
+	SetFunction(req.m_func);
+	return *this;
+}
+
+
+void AdvancedFunctionAnalysisDataRequestor::SetFunction(Function* func)
+{
+	if (m_func)
+		m_func->ReleaseAdvancedAnalysisData();
+
+	m_func = func;
+
+	if (m_func)
+		m_func->RequestAdvancedAnalysisData();
+}
diff --git a/python/__init__.py b/python/__init__.py
index d199385a..cd559fe5 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -4330,8 +4330,11 @@ class Function(object):
 	def __init__(self, view, handle):
 		self._view = view
 		self.handle = core.handle_of_type(handle, core.BNFunction)
+		self._advanced_analysis_requests = 0
 
 	def __del__(self):
+		if self._advanced_analysis_requests > 0:
+			core.BNReleaseAdvancedFunctionAnalysisDataMultiple(self.handle, self._advanced_analysis_requests)
 		core.BNFreeFunction(self.handle)
 
 	@property
@@ -4726,6 +4729,41 @@ class Function(object):
 		"""
 		core.BNReanalyzeFunction(self.handle)
 
+	def request_advanced_analysis_data(self):
+		core.BNRequestAdvancedFunctionAnalysisData(self.handle)
+		self._advanced_analysis_requests += 1
+
+	def release_advanced_analysis_data(self):
+		core.BNReleaseAdvancedFunctionAnalysisData(self.handle)
+		self._advanced_analysis_requests -= 1
+
+class AdvancedFunctionAnalysisDataRequestor(object):
+	def __init__(self, func = None):
+		self._function = func
+		if self._function is not None:
+			self._function.request_advanced_analysis_data()
+
+	def __del__(self):
+		if self._function is not None:
+			self._function.release_advanced_analysis_data()
+
+	@property
+	def function(self):
+		return self._function
+
+	@function.setter
+	def function(self, func):
+		if self._function is not None:
+			self._function.release_advanced_analysis_data()
+		self._function = func
+		if self._function is not None:
+			self._function.request_advanced_analysis_data()
+
+	def close(self):
+		if self._function is not None:
+			self._function.release_advanced_analysis_data()
+		self._function = None
+
 class BasicBlockEdge:
 	def __init__(self, branch_type, target, arch):
 		self.type = core.BNBranchType_names[branch_type]