Add API to enqueue work on the worker thread(s)

4 files changed, 82 insertions, 3 deletions

diff --git a/binaryninjaapi.cpp b/binaryninjaapi.cpp
index f98dfef7..71a5a8f7 100644
--- a/binaryninjaapi.cpp
+++ b/binaryninjaapi.cpp
@@ -24,6 +24,12 @@ using namespace BinaryNinja;
 using namespace std;
 
 
+struct WorkerThreadActionContext
+{
+	std::function<void()> action;
+};
+
+
 void BinaryNinja::InitCorePlugins()
 {
 	BNInitCorePlugins();
@@ -138,3 +144,59 @@ void BinaryNinja::AddOptionalPluginDependency(const string& name)
 {
 	BNAddOptionalPluginDependency(name.c_str());
 }
+
+
+static void WorkerActionCallback(void* ctxt)
+{
+	WorkerThreadActionContext* action = (WorkerThreadActionContext*)ctxt;
+	action->action();
+	delete action;
+}
+
+
+void BinaryNinja::WorkerEnqueue(const function<void()>& action)
+{
+	WorkerThreadActionContext* ctxt = new WorkerThreadActionContext;
+	ctxt->action = action;
+	BNWorkerEnqueue(ctxt, WorkerActionCallback);
+}
+
+
+void BinaryNinja::WorkerEnqueue(RefCountObject* owner, const function<void()>& action)
+{
+	struct
+	{
+		Ref<RefCountObject> owner;
+		function<void()> func;
+	} context;
+	context.owner = owner;
+	context.func = action;
+
+	WorkerEnqueue([=]() {
+			context.func();
+		});
+}
+
+
+void BinaryNinja::WorkerPriorityEnqueue(const function<void()>& action)
+{
+	WorkerThreadActionContext* ctxt = new WorkerThreadActionContext;
+	ctxt->action = action;
+	BNWorkerPriorityEnqueue(ctxt, WorkerActionCallback);
+}
+
+
+void BinaryNinja::WorkerPriorityEnqueue(RefCountObject* owner, const function<void()>& action)
+{
+	struct
+	{
+		Ref<RefCountObject> owner;
+		function<void()> func;
+	} context;
+	context.owner = owner;
+	context.func = action;
+
+	WorkerPriorityEnqueue([=]() {
+			context.func();
+		});
+}
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 7c417623..6902c1a3 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -371,6 +371,11 @@ namespace BinaryNinja
 	Ref<MainThreadAction> ExecuteOnMainThread(const std::function<void()>& action);
 	void ExecuteOnMainThreadAndWait(const std::function<void()>& action);
 
+	void WorkerEnqueue(const std::function<void()>& action);
+	void WorkerEnqueue(RefCountObject* owner, const std::function<void()>& action);
+	void WorkerPriorityEnqueue(const std::function<void()>& action);
+	void WorkerPriorityEnqueue(RefCountObject* owner, const std::function<void()>& action);
+
 	class DataBuffer
 	{
 		BNDataBuffer* m_buffer;
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 3aa0a1e6..a24ce9fa 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1959,6 +1959,10 @@ extern "C"
 	BINARYNINJACOREAPI BNMainThreadAction* BNExecuteOnMainThread(void* ctxt, void (*func)(void* ctxt));
 	BINARYNINJACOREAPI void BNExecuteOnMainThreadAndWait(void* ctxt, void (*func)(void* ctxt));
 
+	// Worker thread queue management
+	BINARYNINJACOREAPI void BNWorkerEnqueue(void* ctxt, void (*action)(void* ctxt));
+	BINARYNINJACOREAPI void BNWorkerPriorityEnqueue(void* ctxt, void (*action)(void* ctxt));
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/python/__init__.py b/python/__init__.py
index 82aaba32..fa8ff02c 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -9916,7 +9916,7 @@ def redirect_output_to_log():
 	global _output_to_log
 	_output_to_log = True
 
-class _MainThreadActionContext:
+class _ThreadActionContext:
 	_actions = []
 
 	def __init__(self, func):
@@ -9939,16 +9939,24 @@ class _MainThreadActionContext:
 			self.__class__._actions.remove(self)
 
 def execute_on_main_thread(func):
-	action = _MainThreadActionContext(func)
+	action = _ThreadActionContext(func)
 	obj = core.BNExecuteOnMainThread(0, action.callback)
 	if obj:
 		return MainThreadAction(obj)
 	return None
 
 def execute_on_main_thread_and_wait(func):
-	action = _MainThreadActionContext(func)
+	action = _ThreadActionContext(func)
 	core.BNExecuteOnMainThreadAndWait(0, action.callback)
 
+def worker_enqueue(func):
+	action = _ThreadActionContext(func)
+	core.BNWorkerEnqueue(0, action.callback)
+
+def worker_priority_enqueue(func):
+	action = _ThreadActionContext(func)
+	core.BNWorkerEnqueue(0, action.callback)
+
 bundled_plugin_path = core.BNGetBundledPluginDirectory()
 user_plugin_path = core.BNGetUserPluginDirectory()