Add nullable attribute to AST nodes for function parameters in API introspection. Fixes #309

This change introduces an `nullable` flag to `ArrayType`, `PointerType`, and `FunctionParameterDecl`. The code generation process now parses function comments from `mujoco.h` to identify parameters marked as "Nullable" and sets this flag accordingly.

PiperOrigin-RevId: 797870883
Change-Id: Ic7cac5a7a9c14177fbf97b8655c5c4809d4ef683

Author: Google DeepMind

python/mujoco/introspect/ast_nodes.py

@@ -62,9 +62,11 @@ class ValueType:
   name: str
   is_const: bool = False
   is_volatile: bool = False
+  nullable: bool = False
 
   def __init__(self, name: str, is_const: bool = False,
-               is_volatile: bool = False):
+               is_volatile: bool = False,
+               nullable: bool = False):
     is_valid_type_name = (
         name == 'void *(*)(void *)' or
         VALID_TYPE_NAME_PATTERN.fullmatch(name) or
@@ -74,6 +76,7 @@ def __init__(self, name: str, is_const: bool = False,
     self.name = name
     self.is_const = is_const
     self.is_volatile = is_volatile
+    self.nullable = nullable
 
   def decl(self, name_or_decl: Optional[str] = None) -> str:
     parts = []
@@ -96,9 +99,11 @@ class ArrayType:
 
   inner_type: Union[ValueType, 'PointerType']
   extents: Tuple[int, ...]
+  nullable: bool = False
 
-  def __init__(self, inner_type: Union[ValueType, 'PointerType'],
-               extents: Sequence[int]):
+  def __init__(
+      self, inner_type: Union[ValueType, 'PointerType'], extents: Sequence[int]
+  ):
     self.inner_type = inner_type
     self.extents = tuple(extents)
 
@@ -119,13 +124,16 @@ class PointerType:
   """Represents a C pointer type."""
 
   inner_type: Union[ValueType, ArrayType, 'PointerType']
+  nullable: bool = False
   is_const: bool = False
   is_volatile: bool = False
   is_restrict: bool = False
 
   def decl(self, name_or_decl: Optional[str] = None) -> str:
     """Creates a string that declares an object of this type."""
     parts = ['*']
+    if self.nullable:
+      parts.append('nullable')
     if self.is_const:
       parts.append('const')
     if self.is_volatile:
@@ -155,6 +163,7 @@ class FunctionParameterDecl:
 
   name: str
   type: Union[ValueType, ArrayType, PointerType]
+  nullable: bool = False
 
   def __str__(self):
     return self.type.decl(self.name)

python/mujoco/introspect/ast_nodes_test.py

@@ -174,6 +174,29 @@ def test_anonymous_union_decl(self):
     self.assertEqual(str(union), 'union {int foo; float bar[3];}')
     self.assertEqual(union.decl('var'), 'union {int foo; float bar[3];} var')
 
+  def test_function_parameter_decl_nullable(self):
+    param_ptr_nullable = ast_nodes.FunctionParameterDecl(
+        name='ptr_param',
+        type=ast_nodes.PointerType(ast_nodes.ValueType('int')),
+        nullable=True
+    )
+    self.assertTrue(param_ptr_nullable.nullable)
+    self.assertEqual(str(param_ptr_nullable), 'int * ptr_param')
+
+    param_ptr_not_nullable = ast_nodes.FunctionParameterDecl(
+        name='ptr_param',
+        type=ast_nodes.PointerType(ast_nodes.ValueType('int')),
+    )
+    self.assertFalse(param_ptr_not_nullable.nullable)
+    self.assertEqual(str(param_ptr_not_nullable), 'int * ptr_param')
+
+    param_array_nullable = ast_nodes.FunctionParameterDecl(
+        name='array_param',
+        type=ast_nodes.ArrayType(ast_nodes.ValueType('float'), (10,)),
+        nullable=True
+    )
+    self.assertTrue(param_array_nullable.nullable)
+    self.assertEqual(str(param_array_nullable), 'float array_param[10]')
 
 if __name__ == '__main__':
   absltest.main()

python/mujoco/introspect/codegen/generate_functions.py

@@ -58,21 +58,43 @@ def _make_function(self, node: ClangJsonNode) -> ast_nodes.FunctionDecl:
         node['type']['qualType'])
     parameters = []
     comments = []
+    nullable_params = set()
+
     for child in node['inner']:
       child_kind = child.get('kind')
-      if child_kind == 'ParmVarDecl':
-        parameters.append(self._make_parameter(child))
       if child_kind == 'FullComment':
         comments.append(self._make_comment(child))
+        nullable_params.update(self._find_nullable_params(child))
     comment = ' '.join(comments).strip()
+
+    for child in node['inner']:
+      child_kind = child.get('kind')
+      if child_kind == 'ParmVarDecl':
+        parameters.append(self._make_parameter(child, nullable_params))
+
     return ast_nodes.FunctionDecl(
         name=name, return_type=return_type, parameters=parameters, doc=comment)
 
+  def _find_nullable_params(self, node: ClangJsonNode) -> set[str]:
+    """Finds the names of parameters that are marked as nullable."""
+    nullable_params = set()
+    for child in node['inner']:
+      child_kind = child.get('kind')
+      if child_kind == 'ParagraphComment':
+        nullable_params.update(self._find_nullable_params(child))
+      if child_kind == 'TextComment':
+        if 'Nullable' in child['text']:
+          for param in child['text'].split(':')[1].split(','):
+            nullable_params.add(param.strip())
+    return nullable_params
+
   def _make_parameter(
-      self, node: ClangJsonNode) -> ast_nodes.FunctionParameterDecl:
+      self, node: ClangJsonNode, nullable_params: set[str]
+  ) -> ast_nodes.FunctionParameterDecl:
     """Makes a ParameterDecl from a Clang AST ParmVarDecl node."""
     name = node['name']
     type_name = node['type']['qualType']
+    nullable = name in nullable_params
 
     # For a pointer parameters, look up in the original header to see if
     # n array extent was declared there.
@@ -85,7 +107,10 @@ def _make_parameter(
       type_name = decl[:name_begin] + decl[name_end:]
 
     return ast_nodes.FunctionParameterDecl(
-        name=name, type=type_parsing.parse_type(type_name))
+        nullable=nullable,
+        name=name,
+        type=type_parsing.parse_type(type_name),
+    )
 
   def _make_comment(self, node: ClangJsonNode) -> str:
     """Makes a comment string from a Clang AST FullComment node."""