update for main handle

Co-Authored-By: Vo Luu Tuong Anh <[email protected]>

Author: shiroinekotfs

jcppkernel/__main__.py

@@ -123,10 +123,8 @@ def do_execute(
         self, code, silent, store_history=True, user_expressions=None, allow_stdin=True
     ):
         cpp_res_path = f'"{self.resDir}/gcpph.hpp"'
-        code = CPPCodeProcessingUnit._add_code_compat(
-            CPPCodeProcessingUnit, code, cpp_res_path
-        )
-
+        code = CPPCodeProcessingUnit()._add_code_compat(code, cpp_res_path)
+        
         with CPPTempFileProcessing._new_temp_file(
             CPPTempFileProcessing, self.files, suffix=".cpp"
         ) as source_file, CPPTempFileProcessing._new_temp_file(

jcppkernel/code_processing.py

@@ -9,26 +9,73 @@ class CPPCodeProcessingUnit:
     comments = r'\/\/.*?(?=\r?\n)|\/\*[\s\S]*?\*\/'
     main_method = r'\b(?:[a-zA-Z_]\w*(?:\s*\*?\s*)?\s+)+main\s*\('
     function_define = r'^[\w\s:<>&*,]+\s+\**\s*\w+\s*\(.*\)\s*{'
-
+    
     def _prep_main_function(self, code):
         lines = code.splitlines()
-        defs_lines = []
-        main_body_lines = []
+        defs_lines = []       # Collect lines that belong to function definitions.
+        main_body_lines = []  # Collect lines that are not part of any function.
         in_function = False
         brace_count = 0
-        func_def_regex = code_compile(self.function_define)
-        for line in lines:
-            if not in_function and func_def_regex.search(line):
-                in_function = True
-                brace_count = line.count('{') - line.count('}')
-                defs_lines.append(line)
-            elif in_function:
+
+        # Regex for a complete function definition with the opening brace on the same line.
+        func_def_regex = code_compile(
+            r'^\s*(?:(?:'
+            r'inline|virtual|static|constexpr|_Noreturn'
+            r')\s+)*(?!\b(?:'
+            r'break|case|continue|default|do|else|for|goto|if|return|switch|while|else if'
+            r')\b)(.*?)\s+(\w+)\s*\([^)]*\)\s*{'
+        )
+        
+        # Regex for a potential function header without the opening brace.
+        func_header_regex = code_compile(
+            r'^\s*(?:(?:'
+            r'inline|virtual|static|constexpr|_Noreturn'
+            r')\s+)*(?!\b(?:'
+            r'break|case|continue|default|do|else|for|goto|if|return|switch|while|else if'
+            r')\b)(.*?)\s+(\w+)\s*\([^)]*\)\s*$'
+        )
+        
+        header_buffer = []  # Buffer to accumulate lines for multi-line function headers.
+        i = 0
+        while i < len(lines):
+            line = lines[i]
+            
+            if not in_function:
+                if func_def_regex.search(line):
+                    in_function = True
+                    brace_count = line.count('{') - line.count('}')
+                    defs_lines.append(line)
+                elif func_header_regex.search(line):
+                    header_buffer = [line]
+                    j = i + 1
+                    found_brace = False
+                    while j < len(lines):
+                        next_line = lines[j]
+                        header_buffer.append(next_line)
+                        if '{' in next_line:
+                            found_brace = True
+                            in_function = True
+                            brace_count = next_line.count('{') - next_line.count('}')
+                            break
+                        if next_line.strip() == "":
+                            break
+                        j += 1
+                    if found_brace:
+                        defs_lines.extend(header_buffer)
+                        i = j
+                        header_buffer = []
+                    else:
+                        main_body_lines.extend(header_buffer)
+                        header_buffer = []
+                else:
+                    main_body_lines.append(line)
+            else:
                 defs_lines.append(line)
                 brace_count += line.count('{') - line.count('}')
                 if brace_count <= 0:
                     in_function = False
-            else:
-                main_body_lines.append(line)
+            i += 1
+
         defs = "\n".join(defs_lines).strip()
         main_body = "\n".join(main_body_lines).strip()
         if main_body:
@@ -37,11 +84,11 @@ def _prep_main_function(self, code):
             wrapped_main = ""
         code = (defs + "\n\n" if defs else "") + wrapped_main
         return code
-        
+
     def _add_code_compat(self, code, cpp_res_path):
         code = code_sub(self.comments, r'', code)
         if not code_search(self.main_method, code):
-            code = self._prep_main_function(self, code)
+            code = self._prep_main_function(code)
 
         code = f"#include {cpp_res_path}\n{code}"
         code = CPPExternalHeader._support_external_header(CPPExternalHeader, code)