Add convenience methods to CodeBuilder and cleanup some WASM bindings code generation

PiperOrigin-RevId: 827916844
Change-Id: Ie382f2a5af3d7e3ec2c81cde5d6b768835da2365

wasm/codegen/helpers/code_builder.py

@@ -56,12 +56,16 @@ class CodeBuilder:
       line += " {" if line else "{"
       self._builder.line(line)
       self._builder._indent_level += 1
+      self._line_count_enter = len(self._builder._lines)
       return self._builder
 
     def __exit__(self, exc_type, exc_val, exc_tb):
       if self._builder._indent_level > 0:
         self._builder._indent_level -= 1
-      self._builder.line("}")
+      if self._line_count_enter == len(self._builder._lines):
+        self._builder._lines[-1] += "}"
+      else:
+        self._builder.line("}")
 
   def block(self, header_line="") -> IndentBlock:
     """Creates a block including braces and an optional header before the opening brace.
@@ -75,3 +79,21 @@ class CodeBuilder:
         An IndentBlock instance that manages the indentation.
     """
     return self.IndentBlock(self, header_line)
+
+  def function(self, signature="") -> IndentBlock:
+    """Creates a function."""
+    return self.block(signature)
+
+  def struct(self, name="") -> IndentBlock:
+    """Creates a struct."""
+    return self.block(f"struct {name}")
+
+  def private(self) -> None:
+    """Creates a private section."""
+    self.newline()
+    self.line("private:")
+
+  def public(self) -> None:
+    """Creates a public section."""
+    self.newline()
+    self.line("public:")

wasm/codegen/helpers/functions.py

@@ -114,17 +114,14 @@ def should_be_wrapped(func: ast_nodes.FunctionDecl) -> bool:
 def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str:
   """Generates C++ code for a wrapper function."""
 
-  builder = code_builder.CodeBuilder()
-  # Build function header
   params_unpack_statements = get_params_unpack_statements(func.parameters)
   wrapper_params_list = get_params_string(func.parameters)
   not_nullable_params = get_params_notnullable(func.parameters)
-  wrapper_params_str = ", ".join(wrapper_params_list)
+  wrapper_params = ", ".join(wrapper_params_list)
   ret_type = get_compatible_return_type(func)
-  builder.line(f"{ret_type} {func.name}_wrapper({wrapper_params_str})")
 
-  # Build function body
+  builder = code_builder.CodeBuilder()
-  with builder.block():
+  with builder.function(f"{ret_type} {func.name}_wrapper({wrapper_params})"):
     invoker_params_list = get_params_string_maybe_with_conversion(
         func.parameters
     )
@@ -136,6 +133,7 @@ def generate_function_wrapper(func: ast_nodes.FunctionDecl) -> str:
     for unpack_statement in params_unpack_statements:
       builder.line(unpack_statement)
     builder.line(f"{invoker_statement};")
+
   return builder.to_string()
 
 

wasm/codegen/helpers/helpers_test.py

@@ -190,8 +190,7 @@ class FunctionUtilsTest(absltest.TestCase):
     result = functions.generate_function_wrapper(func)
     self.assertEqual(
         result,
-        """int get_id_wrapper()
+        """int get_id_wrapper() {
-{
   return get_id();
 }""",
     )
@@ -218,8 +217,7 @@ class FunctionUtilsTest(absltest.TestCase):
     result = functions.generate_function_wrapper(func)
     self.assertEqual(
         result,
-        """int get_id_wrapper(const NumberArray& mat, int nr)
+        """int get_id_wrapper(const NumberArray& mat, int nr) {
-{
   UNPACK_ARRAY(mjtNum, mat);
   return get_id(mat_.data(), nr);
 }""",

wasm/codegen/helpers/structs.py

@@ -73,15 +73,14 @@ class WrappedStructData:
 
 
 def build_primitive_type_definition(field: ast_nodes.StructFieldDecl) -> str:
-  """Builds the C++ code for a primitive type field wrapper."""
+  """Builds the C++ getter/setter code for a primitive type field wrapper."""
   if not isinstance(field.type, ast_nodes.ValueType):
     raise ValueError(f"{field.type} must be ValueType.")
+
   builder = code_builder.CodeBuilder()
-  # build getter for primitive type field
+  with builder.function(f"{field.type.name} {field.name}() const"):
-  with builder.block(f"{field.type.name} {field.name}() const"):
     builder.line(f"return ptr_->{field.name};")
-  # build setter for primitive type field
+  with builder.function(f"void set_{field.name}({field.type.name} value)"):
-  with builder.block(f"void set_{field.name}({field.type.name} value)"):
     builder.line(f"ptr_->{field.name} = value;")
   return builder.to_string()
 
@@ -91,7 +90,7 @@ def build_memory_view_definition(
 ) -> str:
   """Builds the C++ code for a pointer type field wrapper."""
   builder = code_builder.CodeBuilder()
-  with builder.block(f"emscripten::val {field.name}() const"):
+  with builder.function(f"emscripten::val {field.name}() const"):
     builder.line(
         "return"
         f" emscripten::val(emscripten::typed_memory_view({array_size_str},"
@@ -101,13 +100,13 @@ def build_memory_view_definition(
 
 
 def build_string_field_definition(field: ast_nodes.StructFieldDecl) -> str:
-  """Builds the C++ code for a string type field wrapper."""
+  """Builds the C++ code getter/setter for a string type field wrapper."""
   builder = code_builder.CodeBuilder()
-  with builder.block(f"mjString {field.name}() const"):
+  with builder.function(f"mjString {field.name}() const"):
     builder.line(
         f'return (ptr_ && ptr_->{field.name}) ? *(ptr_->{field.name}) : "";'
     )
-  with builder.block(f"void set_{field.name}(const mjString& value)"):
+  with builder.function(f"void set_{field.name}(const mjString& value)"):
     with builder.block(f"if (ptr_ && ptr_->{field.name})"):
       builder.line(f"*(ptr_->{field.name}) = value;")
   return builder.to_string()
@@ -124,7 +123,7 @@ def build_mjvec_pointer_definition(
         f"*(reinterpret_cast<std::vector<uint8_t>*>(ptr_->{field.name}))"
     )
   builder = code_builder.CodeBuilder()
-  with builder.block(f"{vector_type} &{field.name}() const"):
+  with builder.function(f"{vector_type} &{field.name}() const"):
     builder.line(f"return {ptr_field_expr};")
   return builder.to_string()
 
@@ -453,40 +452,42 @@ def _build_struct_header_internal(
     is_mjs: bool = False,
 ):
   """Builds the C++ header file code for a struct."""
+  s = struct_name
+  w = common.uppercase_first_letter(s)
 
   shallow_copy = use_shallow_copy(wrapped_fields)
 
-  wrapper_name = common.uppercase_first_letter(struct_name)
   builder = code_builder.CodeBuilder()
-  with builder.block(f"struct {wrapper_name}"):
+  with builder.struct(f"{w}"):
     if not is_mjs:
-      builder.line(f"{wrapper_name}();")
+      builder.line(f"{w}();")
-      builder.line(f"{wrapper_name}(const {wrapper_name} &);")
+      builder.line(f"{w}(const {w} &);")
-      builder.line(f"{wrapper_name} &operator=(const {wrapper_name} &);")
+      builder.line(f"{w} &operator=(const {w} &);")
 
-    builder.line(f"explicit {wrapper_name}({struct_name} *ptr);")
+    builder.line(f"explicit {w}({s} *ptr);")
-    builder.line(f"~{wrapper_name}();")
+    builder.line(f"~{w}();")
 
     if shallow_copy:
-      builder.line(f"std::unique_ptr<{wrapper_name}> copy();")
+      builder.line(f"std::unique_ptr<{w}> copy();")
 
     for field in wrapped_fields:
       if field.definition and field not in fields_with_init:
         for line in field.definition.splitlines():
           builder.line(line)
 
-    builder.line(f"{struct_name}* get() const {{ return ptr_; }}")
+    with builder.function(f"{s}* get() const"):
-    builder.line(f"void set({struct_name}* ptr) {{ ptr_ = ptr; }}")
+      builder.line("return ptr_;")
 
-    builder.newline()
+    with builder.function(f"void set({s}* ptr)"):
-    builder.line("private:")
+      builder.line("ptr_ = ptr;")
-    builder.line(f"{struct_name}* ptr_;")
+
+    builder.private()
+    builder.line(f"{s}* ptr_;")
     if not is_mjs:
       builder.line("bool owned_ = false;")
 
     if is_mjs and fields_with_init:
-      builder.newline()
+      builder.public()
-      builder.line("public:")
       for field in fields_with_init:
         if field.definition:
           builder.line(f"{field.definition}")
@@ -580,9 +581,9 @@ def build_struct_source(
     wrapped_fields: List[WrappedFieldData],
 ):
   """Builds the C++ .cc file code for a struct."""
-  wrapper_name = common.uppercase_first_letter(struct_name)
+  s = struct_name
-  is_mjs_struct = "Mjs" in wrapper_name
+  w = common.uppercase_first_letter(s)
-  builder = code_builder.CodeBuilder()
+  is_mjs = "Mjs" in w
 
   fields_with_init = _find_fields_with_init(wrapped_fields)
   shallow_copy = use_shallow_copy(wrapped_fields)
@@ -593,65 +594,49 @@ def build_struct_source(
         field_with_init.initialization for field_with_init in fields_with_init
     )
 
+  builder = code_builder.CodeBuilder()
+
   # constructor passing native ptr
-  builder.line(
+  with builder.function(f"{w}::{w}({s} *ptr) : ptr_(ptr){fields_init}"):
-      f"{wrapper_name}::{wrapper_name}({struct_name} *ptr) :"
+    pass
-      f" ptr_(ptr){fields_init} {{}}"
-  )
 
   # constructor with default values
-  if not is_mjs_struct:
+  if not is_mjs:
-    with builder.block(
+    with builder.function(f"{w}::{w}() : ptr_(new {s}){fields_init}"):
-        f"{wrapper_name}::{wrapper_name}() : ptr_(new"
-        f" {struct_name}){fields_init}"
-    ):
       builder.line("owned_ = true;")
-      default_func = _default_function_statement(struct_name)
+      default_func = _default_function_statement(s)
       if default_func:
         builder.line(default_func)
 
-  if shallow_copy and not is_mjs_struct:
+  if shallow_copy and not is_mjs:
     # copy constructor
-    with builder.block(
+    with builder.function(f"{w}::{w}(const {w} &other) : {w}()"):
-        f"{wrapper_name}::{wrapper_name}(const {wrapper_name} &other)"
-        + (f" : {wrapper_name}()" if not is_mjs_struct else "")
-    ):
       builder.line("*ptr_ = *other.get();")
-      if fields_with_init:
+      for field_with_init in fields_with_init:
-        for field_with_init in fields_with_init:
+        if field_with_init.ptr_copy_reset is not None:
-          if field_with_init.ptr_copy_reset is not None:
+          builder.line(field_with_init.ptr_copy_reset)
-            builder.line(field_with_init.ptr_copy_reset)
 
     # assignment operator
-    with builder.block(
+    with builder.function(f"{w}& {w}::operator=(const {w} &other)"):
-        f"{wrapper_name}&"
-        f" {wrapper_name}::operator=(const"
-        f" {wrapper_name} &other)"
-    ):
       with builder.block("if (this == &other)"):
         builder.line("return *this;")
       builder.line("*ptr_ = *other.get();")
-      if fields_with_init:
+      for field_with_init in fields_with_init:
-        for field_with_init in fields_with_init:
+        if field_with_init.ptr_copy_reset is not None:
-          if field_with_init.ptr_copy_reset is not None:
+          builder.line(field_with_init.ptr_copy_reset)
-            builder.line(field_with_init.ptr_copy_reset)
       builder.line("return *this;")
 
   # destructor
-  if is_mjs_struct:
+  with builder.function(f"{w}::~{w}()"):
-    builder.line(f"{wrapper_name}::~{wrapper_name}() {{}}")
+    if not is_mjs:
-  else:
-    with builder.block(f"{wrapper_name}::~{wrapper_name}()"):
       with builder.block("if (owned_ && ptr_)"):
-        delete_ptr = _delete_ptr_statement(struct_name)
+        delete_ptr = _delete_ptr_statement(s)
         builder.line(delete_ptr)
 
   # copy function
   if shallow_copy:
-    with builder.block(
+    with builder.function(f"std::unique_ptr<{w}> {w}::copy()"):
-        f"std::unique_ptr<{wrapper_name}> {wrapper_name}::copy()"
+      builder.line(f"return std::make_unique<{w}>(*this);")
-    ):
-      builder.line(f"return std::make_unique<{wrapper_name}>(*this);")
 
   return builder.to_string()