From f976d6d228612c6dd863ee4e094ae9d7452fcfa9 Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Tue, 5 Sep 2023 06:28:58 -0700
Subject: [PATCH] Use a raw function pointer instead of `std::function` objects
 in thread pool.

std::function can have an allocation cost when it's created, and also adds overhead to function calls. Since this API must be used from C, using a raw function pointer is good enough.
PiperOrigin-RevId: 562760324
Change-Id: I2b714a13e07a2b71566d013f21f76c506e817a52
---
 doc/includes/references.h            | 2 +-
 include/mujoco/mjthread.h            | 2 +-
 introspect/structs.py                | 2 +-
 src/thread/task.h                    | 6 +++---
 src/thread/thread_pool.h             | 3 +--
 unity/Runtime/Bindings/MjBindings.cs | 2 +-
 6 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/doc/includes/references.h b/doc/includes/references.h
index 34d9f187..d91fa3f2 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -1440,7 +1440,7 @@ struct mjrContext_ {              // custom OpenGL context
 };
 typedef struct mjrContext_ mjrContext;
 struct mjTask_ {
-  char buffer[48];
+  char buffer[24];
 };
 typedef struct mjTask_ mjTask;
 struct mjThreadPool_ {
diff --git a/include/mujoco/mjthread.h b/include/mujoco/mjthread.h
index 1538e954..66290370 100644
--- a/include/mujoco/mjthread.h
+++ b/include/mujoco/mjthread.h
@@ -27,7 +27,7 @@ extern "C" {
 // These types are implemented in C++, they're just used as opaque pointers in C
 // to provide type safety for functions.
 struct mjTask_ {
-  char buffer[48];
+  char buffer[24];
 };
 typedef struct mjTask_ mjTask;
 
diff --git a/introspect/structs.py b/introspect/structs.py
index 4bdc211c..d399c6bc 100644
--- a/introspect/structs.py
+++ b/introspect/structs.py
@@ -6992,7 +6992,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  name='buffer',
                  type=ArrayType(
                      inner_type=ValueType(name='char'),
-                     extents=(48,),
+                     extents=(24,),
                  ),
                  doc='',
              ),
diff --git a/src/thread/task.h b/src/thread/task.h
index 390a4571..296592a5 100644
--- a/src/thread/task.h
+++ b/src/thread/task.h
@@ -20,13 +20,13 @@
 #ifdef __cplusplus
 
 #include <atomic>
-#include <functional>
 #include <thread>
 
 namespace mujoco {
 
 class Task {
  public:
+  using FunctionPtr = void* (*)(void*);
   enum Status {
     QUEUED,
     COMPLETE,
@@ -34,7 +34,7 @@ class Task {
 
   static void Initialize(
       Task* task,
-      std::function<void*(void*)> start_routine,
+      FunctionPtr start_routine,
       void* args) {
     // instantiate a task at the pointer passed in
     new(task) Task();
@@ -55,7 +55,7 @@ class Task {
   }
 
  private:
-  std::function<void*(void*)> start_routine_;
+  FunctionPtr start_routine_;
 
   void* args_;
 
diff --git a/src/thread/thread_pool.h b/src/thread/thread_pool.h
index e987318e..11b51145 100644
--- a/src/thread/thread_pool.h
+++ b/src/thread/thread_pool.h
@@ -21,7 +21,6 @@
 
 #include <atomic>
 #include <cstddef>
-#include <functional>
 #include <thread>
 
 #include "thread/lockless_queue.h"
@@ -43,7 +42,7 @@ class ThreadPool {
 
   // start a task in the threadpool
   void Enqueue(
-      Task* task, std::function<void*(void*)> start_routine, void* args) {
+      Task* task, Task::FunctionPtr start_routine, void* args) {
     Task::Initialize(task, start_routine, args);
     lockless_queue_.push(static_cast<void*>(task));
   }
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index 37a1278a..0a646168 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -2358,7 +2358,7 @@ public unsafe struct mjrContext_ {
 
 [StructLayout(LayoutKind.Sequential)]
 public unsafe struct mjTask_ {
-  public fixed sbyte buffer[48];
+  public fixed sbyte buffer[24];
 }
 
 [StructLayout(LayoutKind.Sequential)]