From 7dce9142f6de231a0d19403841e82bfb67e49b47 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Tupynamb=C3=A1=20Dutra?=
 <rafael.dutra@cispa.saarland>
Date: Mon, 7 Mar 2022 10:49:42 +0100
Subject: [PATCH] Fix TODO field for recursive structs

---
 pfp/interp.py | 53 +++++++++++++++++++++++++++------------------------
 1 file changed, 28 insertions(+), 25 deletions(-)

diff --git a/pfp/interp.py b/pfp/interp.py
index b83fde4..a2d7280 100644
--- a/pfp/interp.py
+++ b/pfp/interp.py
@@ -691,12 +691,27 @@ def add_decl(self, classname, classnode, node, is_union):
         else:
             if classname not in self._to_define:
                 self._to_define[classname] = []
-            self._to_define[classname].append((node.name, node, len(self._call_stack) > 1 and not self._call_stack[-1]))
-            node.cpp = "/*TODO field " + node.name + "("
+            is_var = len(self._call_stack) > 1 and not self._call_stack[-1]
+            self._to_define[classname].append((node.name, node, is_var))
+
+            name = node.name
+            if hasattr(node, "originalname"):
+                name = node.originalname
+            node.cpp = "GENERATE"
+            if is_var:
+                node.cpp += "_VAR"
+            self._variable_types[node.name] = classname
+            node.cpp += "(" + name + ", ::g->" + node.name + ".generate("
+            arg_num = 0
             if hasattr(node.type, "args") and node.type.args:
                 for arg in node.type.args.exprs:
+                    arg_num += 1
                     node.cpp += arg.cpp + ", "
-            node.cpp += ")*/"
+            if arg_num > 0:
+                node.cpp = node.cpp[:-2]
+            node.cpp += "/*TODO class " + classname + "*/"
+            node.cpp += "))"
+
             node.type.cpp = classname
             if not self._incomplete:
                 self.add_class(classname, classnode, is_union)
@@ -982,33 +997,21 @@ def add_class(self, classname, classnode, is_union=False):
                     self._defined[field_name] = classname
                     self._globals.append((field_name, classname + " " + node.name + "(" + classname + "_" + node.name + "_instances);\n"))
                     self._instances += "std::vector<" + classname + "*> " + classname + "_" + node.name + "_instances;\n"
-                name = field_name
-                if hasattr(node, "originalname"):
-                    name = node.originalname
-                node.cpp = "GENERATE"
-                if is_var:
-                    node.cpp += "_VAR"
-                self._variable_types[field_name] = classname
-                node.cpp += "(" + name + ", ::g->" + field_name + ".generate("
+
                 arg_num = 0
-                todofield = "/*TODO field " + field_name + "("
                 if hasattr(node.type, "args") and node.type.args:
-                    for arg in node.type.args.exprs:
-                        arg_num += 1
-                        node.cpp += arg.cpp + ", "
-                        todofield += arg.cpp + ", "
-                todofield += ")*/"
+                    arg_num = len(node.type.args.exprs)
+                todoclass = ", " if arg_num else ""
                 if hasattr(classnode, "args") and classnode.args is not None:
                     for i in range(arg_num, len(classnode.args.params)):
-                        arg_num += 1
-                        node.cpp += classnode.args.params[i].name + ", "
-                if arg_num > 0:
-                    node.cpp = node.cpp[:-2]
-                node.cpp += "))"
+                        todoclass += classnode.args.params[i].name + ", "
+                if todoclass:
+                    todoclass = todoclass[:-2]
+                todo = "/*TODO class " + classname + "*/"
+                node.cpp = node.cpp.replace(todo, todoclass)
                 for decl in classnode.decls:
-                    decl.cpp = decl.cpp.replace(todofield, node.cpp)
-                self._to_replace.append((todofield, node.cpp))
-                node.type.cpp = classname
+                    decl.cpp = decl.cpp.replace(todo, todoclass)
+                self._to_replace.append((todo, todoclass))
 
     def add_class_generate(self, classname, classnode, is_union=False):
         if classname + "::generate" in self._defined: