[Bug] Fix memory OOM in Genesis push_differentiable example

python/gstaichi/lang/ast/ast_transformers/call_transformer.py

@@ -166,12 +166,13 @@ def _canonicalize_formatted_string(raw_string: str, *raw_args: list, **raw_keywo
         return args
 
     @staticmethod
-    def _expand_Call_dataclass_args(args: tuple[ast.stmt]) -> tuple[ast.stmt]:
+    def _expand_Call_dataclass_args(args: tuple[ast.stmt, ...]) -> tuple[tuple[ast.stmt, ...], tuple[ast.stmt, ...]]:
         """
         We require that each node has a .ptr attribute added to it, that contains
         the associated Python object
         """
         args_new = []
+        added_args = []
         for arg in args:
             val = arg.ptr
             if dataclasses.is_dataclass(val):
@@ -192,20 +193,24 @@ def _expand_Call_dataclass_args(args: tuple[ast.stmt]) -> tuple[ast.stmt]:
                     )
                     if dataclasses.is_dataclass(field.type):
                         arg_node.ptr = field.type
-                        args_new.extend(CallTransformer._expand_Call_dataclass_args((arg_node,)))
+                        _added_args, _args_new = CallTransformer._expand_Call_dataclass_args((arg_node,))
+                        args_new.extend(_args_new)
+                        added_args.extend(_added_args)
                     else:
                         args_new.append(arg_node)
+                        added_args.append(arg_node)
             else:
                 args_new.append(arg)
-        return tuple(args_new)
+        return tuple(added_args), tuple(args_new)
 
     @staticmethod
-    def _expand_Call_dataclass_kwargs(kwargs: list[ast.keyword]) -> list[ast.keyword]:
+    def _expand_Call_dataclass_kwargs(kwargs: list[ast.keyword]) -> tuple[list[ast.keyword], list[ast.keyword]]:
         """
         We require that each node has a .ptr attribute added to it, that contains
         the associated Python object
         """
         kwargs_new = []
+        added_kwargs = []
         for i, kwarg in enumerate(kwargs):
             val = kwarg.ptr[kwarg.arg]
             if dataclasses.is_dataclass(val):
@@ -233,12 +238,15 @@ def _expand_Call_dataclass_kwargs(kwargs: list[ast.keyword]) -> list[ast.keyword
                     )
                     if dataclasses.is_dataclass(field.type):
                         kwarg_node.ptr = {child_name: field.type}
-                        kwargs_new.extend(CallTransformer._expand_Call_dataclass_kwargs([kwarg_node]))
+                        _added_kwargs, _kwargs_new = extend(CallTransformer._expand_Call_dataclass_kwargs([kwarg_node]))
+                        kwargs_new.extend(_kwargs_new)
+                        added_kwargs.extend(_added_kwargs)
                     else:
                         kwargs_new.append(kwarg_node)
+                        added_kwargs.append(kwarg_node)
             else:
                 kwargs_new.append(kwarg)
-        return kwargs_new
+        return added_kwargs, kwargs_new
 
     @staticmethod
     def build_Call(ctx: ASTTransformerContext, node: ast.Call, build_stmt, build_stmts) -> Any | None:
@@ -257,11 +265,17 @@ def build_Call(ctx: ASTTransformerContext, node: ast.Call, build_stmt, build_stm
             # not related to dataclasses). Necessary for calling further child functions
             build_stmts(ctx, node.args)
             build_stmts(ctx, node.keywords)
-            node.args = CallTransformer._expand_Call_dataclass_args(node.args)
-            node.keywords = CallTransformer._expand_Call_dataclass_kwargs(node.keywords)
+
+            added_args, node.args = CallTransformer._expand_Call_dataclass_args(node.args)
+            added_keywords, node.keywords = CallTransformer._expand_Call_dataclass_kwargs(node.keywords)
+
             # create variables for the now-expanded dataclass members
-            build_stmts(ctx, node.args)
-            build_stmts(ctx, node.keywords)
+            for arg in added_args:
+                assert not hasattr(arg, "ptr")
+                build_stmt(ctx, arg)
+            for arg in added_keywords:
+                assert not hasattr(arg, "ptr")
+                build_stmt(ctx, arg)
 
         args = []
         for arg in node.args:

tests/python/gstaichi/lang/ast/test_call_transformer.py

@@ -29,7 +29,7 @@ class MyStructEF:
     my_struct_cd: MyStructCD
 
 
-def dump_ast_list(nodes: list[ast.AST]) -> str:
+def dump_ast_list(nodes: tuple[ast.stmt, ...]) -> str:
     res_l = []
     for node in nodes:
         res_l.append(ast.dump(node, indent=2))
@@ -57,26 +57,24 @@ def dump_ast_list(nodes: list[ast.AST]) -> str:
         ),
         (
             [ast.Name(id="my_struct_ef", ctx=ast.Load(), ptr=MyStructEF)],
-            [
+            (
                 ast.Name(id="__ti_my_struct_ef__ti_e", ctx=ast.Load()),
                 ast.Name(id="__ti_my_struct_ef__ti_f", ctx=ast.Load()),
                 ast.Name(id="__ti_my_struct_ef__ti_my_struct_cd__ti_c", ctx=ast.Load()),
                 ast.Name(id="__ti_my_struct_ef__ti_my_struct_cd__ti_d", ctx=ast.Load()),
                 ast.Name(id="__ti_my_struct_ef__ti_my_struct_cd__ti_my_struct_ab__ti_a", ctx=ast.Load()),
                 ast.Name(id="__ti_my_struct_ef__ti_my_struct_cd__ti_my_struct_ab__ti_b", ctx=ast.Load()),
-            ],
+            ),
         ),
     ],
 )
 @test_utils.test()
-def test_expand_Call_dataclass_args(args_in: list[ast.AST], expected_args_out: list[ast.AST]) -> None:
+def test_expand_Call_dataclass_args(args_in: tuple[ast.stmt, ...], expected_args_out: tuple[ast.stmt, ...]) -> None:
     for arg in args_in:
         arg.lineno = 1
-        arg.endlineno = 2
+        arg.end_lineno = 2
         arg.col_offset = 1
         arg.end_col_offset = 2
-    args_out = CallTransformer._expand_Call_dataclass_args(args_in)
-    print("args_in", dump_ast_list(args_in))
-    print("args_out", dump_ast_list(args_out))
-    print("expected_args_out", dump_ast_list(expected_args_out))
+    args_added, args_out = CallTransformer._expand_Call_dataclass_args(args_in)
     assert dump_ast_list(expected_args_out) == dump_ast_list(args_out)
+    assert dump_ast_list(expected_args_out) == dump_ast_list(args_added)