[introspect] Extract struct array extents from comments.

PiperOrigin-RevId: 687254268
Change-Id: Ia70dca073eaa4a99c60f76b7c69c730856193512

introspect/ast_nodes.py

@@ -234,6 +234,7 @@ class StructFieldDecl:
       'AnonymousUnionDecl',
   ]
   doc: str
+  array_extent: Optional[Tuple[Union[str, int], ...]] = None
 
   def __str__(self):
     return self.type.decl(self.name)

introspect/codegen/generate_structs.py

@@ -47,6 +47,8 @@ _EXCLUDED = (
     'mjResource_',
 )
 
+_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s\s+\((.+) x (.+)\)\Z')
+
 
 def traverse(node, visitor):
   visitor.visit(node)
@@ -112,8 +114,24 @@ class MjStructVisitor:
         doc = self._make_comment(child)
     if 'name' in node:
       field_type = self._normalize_type(node['type']['qualType'])
+      m = _ARRAY_COMMENT_PATTERN.match(doc)
+      if m is None:
+        array_extent = None
+      else:
+        doc = m.group(1)
+        array_extent_0 = m.group(2)
+        array_extent_1 = m.group(3)
+        try:
+          array_extent_1 = int(array_extent_1)
+        except ValueError:
+          pass
+        if array_extent_1 == 1:
+          array_extent = (array_extent_0,)
+        else:
+          array_extent = (array_extent_0, array_extent_1)
       return ast_nodes.StructFieldDecl(
-          name=node['name'], type=field_type, doc=doc)
+          name=node['name'], type=field_type, doc=doc,
+          array_extent=array_extent)
     else:
       return _AnonymousTypePlaceholder(self._make_anonymous_key(node))
 

introspect/structs_test.py

@@ -40,7 +40,8 @@ class StructsTest(absltest.TestCase):
         self.assertEqual(field.doc, 'warning statistics')
       elif field.name == 'qpos':
         self.assertEqual(field.type, type_parsing.parse_type('mjtNum*'))
-        self.assertEqual(re.sub(r'\s+', ' ', field.doc), 'position (nq x 1)')
+        self.assertEqual(field.doc, 'position')
+        self.assertEqual(field.array_extent, ('nq',))
 
     self.assertIn('warning', field_names)
     self.assertIn('qpos', field_names)