auto-sync: fix patch StreamOperation.py

imbillow · Jul 6, 2024 · c75ea7a · c75ea7a
1 parent f081e30
commit c75ea7a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/suite/auto-sync/src/autosync/cpptranslator/CppTranslator.py b/suite/auto-sync/src/autosync/cpptranslator/CppTranslator.py
@@ -444,6 +444,8 @@ def translate(self) -> None:
                 cb: [(Node, str)]
                 for cb in captures_bundle:
                     patch_kwargs = self.get_patch_kwargs(patch)
+                    patch_kwargs["tree"] = self.tree
+                    patch_kwargs["ts_cpp_lang"] = self.ts_cpp_lang
                     bytes_patch: bytes = patch.get_patch(cb, self.src, **patch_kwargs)
                     p_list.append((bytes_patch, cb[0][0]))
                 self.patch_src(p_list)

diff --git a/suite/auto-sync/src/autosync/cpptranslator/patches/StreamOperation.py b/suite/auto-sync/src/autosync/cpptranslator/patches/StreamOperation.py
@@ -3,7 +3,7 @@
 
 from tree_sitter import Node
 
-from autosync.cpptranslator.patches.Helper import get_text
+from autosync.cpptranslator.patches.Helper import get_text, get_text_from_node
 from autosync.cpptranslator.patches.Patch import Patch
 
 
@@ -105,6 +105,26 @@ def get_patch(self, captures: [(Node, str)], src: bytes, **kwargs) -> bytes:
                 + last_op_text.replace(b"'", b'"')
                 + b");\n"
             )
+        elif last_op.type == 'identifier':
+            queue_str = f'''
+                (declaration (
+                    (primitive_type) @typ
+                    (init_declarator 
+                        (identifier) @ident (#eq? @ident "{last_op_text.decode('utf8')}")
+                    )
+                )) @decl
+'''
+            query = kwargs["ts_cpp_lang"].query(queue_str)
+            root_node = kwargs["tree"].root_node
+            cap = list(filter(lambda x: 'typ' in x[1], query.matches(root_node, end_byte=last_op.start_byte)))[-1]
+            typ = get_text_from_node(src, cap[1]["typ"])
+            match typ:
+                case b'int':
+                    res += b"printInt32(" + s_name + b", " + last_op_text + b");"
+                case b'int64_t':
+                    res += b"printInt64(" + s_name + b", " + last_op_text + b");"
+                case _:
+                    res += b"SStream_concat0(" + s_name + b", " + last_op_text + b");"
         else:
             res += b"SStream_concat0(" + s_name + b", " + last_op_text + b");"
         stream = captures[0][0]