Add threading primitives.

PiperOrigin-RevId: 562649644
Change-Id: I9c35b270e4b3b50cc7eb5a152f24d71def8e2fcd

introspect/ast_nodes.py

@@ -66,6 +66,7 @@ class ValueType:
   def __init__(self, name: str, is_const: bool = False,
                is_volatile: bool = False):
     is_valid_type_name = (
+        name == 'void *(*)(void *)' or
         VALID_TYPE_NAME_PATTERN.fullmatch(name) or
         _is_valid_integral_type(name)) and name not in C_INVALID_TYPE_NAMES
     if not is_valid_type_name:

introspect/functions.py

@@ -8257,4 +8257,76 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Look up a resource provider by slot number returned by mjp_registerResourceProvider. If invalid slot number, return NULL.',  # pylint: disable=line-too-long
      )),
+    ('mju_threadPoolCreate',
+     FunctionDecl(
+         name='mju_threadPoolCreate',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjThreadPool'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='number_of_threads',
+                 type=ValueType(name='size_t'),
+             ),
+         ),
+         doc='Creates a thread pool with the specified number of threads running.',  # pylint: disable=line-too-long
+     )),
+    ('mju_threadPoolEnqueue',
+     FunctionDecl(
+         name='mju_threadPoolEnqueue',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='thread_pool',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjThreadPool'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='task',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjTask'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='start_routine',
+                 type=ValueType(name='void *(*)(void *)'),
+             ),
+             FunctionParameterDecl(
+                 name='args',
+                 type=PointerType(
+                     inner_type=ValueType(name='void'),
+                 ),
+             ),
+         ),
+         doc='Enqueues a task in a thread pool.',
+     )),
+    ('mju_taskJoin',
+     FunctionDecl(
+         name='mju_taskJoin',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='task',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjTask'),
+                 ),
+             ),
+         ),
+         doc='Waits for a task to complete.',
+     )),
+    ('mju_threadPoolDestroy',
+     FunctionDecl(
+         name='mju_threadPoolDestroy',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='thread_pool',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjThreadPool'),
+                 ),
+             ),
+         ),
+         doc='Destroys a thread pool.',
+     )),
 ])

introspect/structs.py

@@ -4475,6 +4475,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='constraint state (mjtConstraintState)             (nefc x 1)',  # pylint: disable=line-too-long
              ),
+             StructFieldDecl(
+                 name='threadpool',
+                 type=ValueType(name='uintptr_t'),
+                 doc='ThreadPool for multithreaded operations',
+             ),
          ),
      )),
     ('mjvPerturb',
@@ -6978,6 +6983,36 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              ),
          ),
      )),
+    ('mjTask',
+     StructDecl(
+         name='mjTask',
+         declname='struct mjTask_',
+         fields=(
+             StructFieldDecl(
+                 name='buffer',
+                 type=ArrayType(
+                     inner_type=ValueType(name='char'),
+                     extents=(48,),
+                 ),
+                 doc='',
+             ),
+         ),
+     )),
+    ('mjThreadPool',
+     StructDecl(
+         name='mjThreadPool',
+         declname='struct mjThreadPool_',
+         fields=(
+             StructFieldDecl(
+                 name='buffer',
+                 type=ArrayType(
+                     inner_type=ValueType(name='char'),
+                     extents=(6208,),
+                 ),
+                 doc='',
+             ),
+         ),
+     )),
     ('mjuiState',
      StructDecl(
          name='mjuiState',

introspect/type_parsing.py

@@ -68,6 +68,8 @@ def _parse_maybe_pointer(
                                                    ast_nodes.PointerType]]
 ) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
   """Internal-only helper that parses a type that may be a pointer type."""
+  if type_name == 'void *(*)(void *)':
+    return ast_nodes.ValueType(name=type_name)
   p = type_name.rfind('*')
   if p != -1:
     leftover, is_qualifier = _parse_qualifiers(
@@ -107,6 +109,9 @@ def _peel_nested_parens(input_str: str) -> MutableSequence[str]:
     A sequence of substrings enclosed with in respective parentheses. See the
     description above for the precise detail of the output.
   """
+  if input_str == 'void *(*)(void *)':
+    return ['void *(*)(void *)']
+
   start = input_str.find('(')
   end = input_str.rfind(')')
 
@@ -146,4 +151,4 @@ def parse_type(
 def parse_function_return_type(
     type_name: str
 ) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
-  return parse_type(type_name[:type_name.rfind('(')])
+  return parse_type(type_name[:type_name.find('(')])