Handle multiline #pragma increment_block_number (#2023)

* Handle multiline #pragma increment_block_number

* fix: align all continuation characters vertically

* simplify continuation line character alignment

* Update tools/fix_bss.py

Co-authored-by: cadmic <cadmic24@gmail.com>

---------

Co-authored-by: cadmic <cadmic24@gmail.com>

tools/fix_bss.py

@@ -11,7 +11,6 @@ from collections import Counter
 import colorama
 from dataclasses import dataclass
 import io
-import itertools
 import multiprocessing
 import multiprocessing.pool
 from pathlib import Path
@@ -19,6 +18,7 @@ import re
 import shlex
 import sys
 import time
+import traceback
 from typing import BinaryIO, Iterator
 
 from ido_block_numbers import (
@@ -501,32 +501,97 @@ def solve_bss_ordering(
     raise FixBssException("Could not find any solutions")
 
 
+# Parses #pragma increment_block_number (with line continuations already removed)
+def parse_pragma(pragma_string: str) -> dict[str, int]:
+    amounts = {}
+    for part in pragma_string.replace('"', "").split()[2:]:
+        kv = part.split(":")
+        if len(kv) != 2:
+            raise FixBssException(
+                "#pragma increment_block_number"
+                f' arguments must be version:amount pairs, not "{part}"'
+            )
+        try:
+            amount = int(kv[1])
+        except ValueError:
+            raise FixBssException(
+                "#pragma increment_block_number"
+                f' amount must be an integer, not "{kv[1]}" (in "{part}")'
+            )
+        amounts[kv[0]] = amount
+    return amounts
+
+
+# Formats #pragma increment_block_number as a list of lines
+def format_pragma(amounts: dict[str, int], max_line_length: int) -> list[str]:
+    lines = []
+    pragma_start = "#pragma increment_block_number "
+    current_line = pragma_start + '"'
+    first = True
+    for version, amount in sorted(amounts.items()):
+        part = f"{version}:{amount}"
+        if len(current_line) + len(part) + len('" \\') > max_line_length:
+            lines.append(current_line + '" ')
+            current_line = " " * len(pragma_start) + '"'
+            first = True
+        if not first:
+            current_line += " "
+        current_line += part
+        first = False
+    lines.append(current_line + '"\n')
+
+    if len(lines) >= 2:
+        # add and align vertically all continuation \ characters
+        n_align = max(map(len, lines[:-1]))
+        for i in range(len(lines) - 1):
+            lines[i] = f"{lines[i]:{n_align}}\\\n"
+
+    return lines
+
+
 def update_source_file(version_to_update: str, file: Path, new_pragmas: list[Pragma]):
     with open(file, "r", encoding="utf-8") as f:
         lines = f.readlines()
 
+    replace_lines: list[tuple[int, int, list[str]]] = []
+
     for pragma in new_pragmas:
-        line = lines[pragma.line_number - 1]
+        i = pragma.line_number - 1
-        if not line.startswith("#pragma increment_block_number "):
+        if not lines[i].startswith("#pragma increment_block_number"):
             raise FixBssException(
                 f"Expected #pragma increment_block_number on line {pragma.line_number}"
             )
 
-        # Grab pragma argument and remove quotes
+        # list the pragma line and any continuation line
-        arg = line.strip()[len("#pragma increment_block_number ") + 1 : -1]
+        pragma_lines = [lines[i]]
+        while pragma_lines[-1].endswith("\\\n"):
+            i += 1
+            pragma_lines.append(lines[i])
 
-        amounts_by_version = {}
+        # concatenate all lines into one
-        for part in arg.split():
+        pragma_string = "".join(s.replace("\\\n", "") for s in pragma_lines)
-            version, amount_str = part.split(":")
-            amounts_by_version[version] = int(amount_str)
 
-        amounts_by_version[version_to_update] = pragma.amount
+        amounts = parse_pragma(pragma_string)
-        new_arg = " ".join(
+
-            f"{version}:{amount}" for version, amount in amounts_by_version.items()
+        amounts[version_to_update] = pragma.amount
+
+        column_limit = 120  # matches .clang-format's ColumnLimit
+        new_pragma_lines = format_pragma(amounts, column_limit)
+
+        replace_lines.append(
+            (
+                pragma.line_number - 1,
+                pragma.line_number - 1 + len(pragma_lines),
+                new_pragma_lines,
+            )
         )
-        new_line = f'#pragma increment_block_number "{new_arg}"\n'
 
-        lines[pragma.line_number - 1] = new_line
+    # Replace the pragma lines starting from the end of the file, so the line numbers
+    # for pragmas earlier in the file stay accurate.
+    replace_lines.sort(key=lambda it: it[0], reverse=True)
+    for start, end, new_pragma_lines in replace_lines:
+        del lines[start:end]
+        lines[start:start] = new_pragma_lines
 
     with open(file, "w", encoding="utf-8") as f:
         f.writelines(lines)
@@ -600,9 +665,15 @@ def process_file_worker(*x):
     try:
         sys.stdout = fake_stdout
         process_file(*x)
-    except Exception as e:
+    except FixBssException as e:
+        # exception with a message for the user
         print(f"{colorama.Fore.RED}Error: {str(e)}{colorama.Fore.RESET}")
         raise
+    except Exception as e:
+        # "unexpected" exception, also print a trace for devs
+        print(f"{colorama.Fore.RED}Error: {str(e)}{colorama.Fore.RESET}")
+        traceback.print_exc(file=sys.stdout)
+        raise
     finally:
         sys.stdout = old_stdout
         print()

tools/preprocess.py

@@ -11,10 +11,11 @@
 
 import argparse
 from pathlib import Path
-import os
+import re
 import tempfile
 import subprocess
 import sys
+import typing
 
 
 def fail(message):
@@ -22,36 +23,59 @@ def fail(message):
     sys.exit(1)
 
 
-def process_file(version, filename, input, output):
+def process_file(
+    version: str,
+    filename: str,
+    input: typing.TextIO,
+    output: typing.TextIO,
+):
     output.write(f'#line 1 "{filename}"\n')
+    # whether the current line follows a #pragma increment_block_number,
+    # including continuation lines (lines after a \-ending line)
+    in_pragma_incblocknum = False
+    # the line where the #pragma increment_block_number is
+    pragma_incblocknum_first_line_num = None
+    # all the lines from the #pragma increment_block_number line to the last
+    # continuation line, as a list[str]
+    pragma_incblocknum_lines = None
     for i, line in enumerate(input, start=1):
-        if line.startswith("#pragma increment_block_number "):
+        if not in_pragma_incblocknum and line.startswith(
-            # Grab pragma argument and remove quotes
+            "#pragma increment_block_number"
-            arg = line.strip()[len("#pragma increment_block_number ") + 1 : -1]
+        ):
-            amount = 0
+            in_pragma_incblocknum = True
-            for part in arg.split():
+            pragma_incblocknum_first_line_num = i
-                kv = part.split(":")
+            pragma_incblocknum_lines = []
-                if len(kv) != 2:
-                    fail(
-                        f"{filename}:{i}: increment_block_number must be followed by a list of version:amount pairs"
-                    )
-                if kv[0] != version:
-                    continue
-                try:
-                    amount = int(kv[1])
-                except ValueError:
-                    fail(
-                        f"{filename}:{i}: increment_block_number amount must be an integer"
-                    )
 
-            # Always generate at least one struct so that fix_bss.py can know where the increment_block_number pragmas are
+        if in_pragma_incblocknum:
-            if amount == 0:
+            if line.endswith("\\\n"):
-                amount = 256
+                pragma_incblocknum_lines.append(line)
+            else:
+                in_pragma_incblocknum = False
+                pragma_incblocknum_lines.append(line)
+                amount = 0
+                for s in pragma_incblocknum_lines:
+                    # Note if we had two versions like "abc-def-version" and "def-version"
+                    # then this code would find either given "def-version", but
+                    # thankfully we don't have such nested version names.
+                    m = re.search(rf"{version}:(\d+)\b", s)
+                    if m:
+                        amount = int(m.group(1))
+                        break
 
-            # Write fake structs for BSS ordering
+                # Always generate at least one struct,
-            for j in range(amount):
+                # so that fix_bss.py can know where the increment_block_number pragmas are
-                output.write(f"struct increment_block_number_{i:05}_{j:03};\n")
+                if amount == 0:
-            output.write(f'#line {i + 1} "{filename}"\n')
+                    amount = 256
+
+                # Write fake structs for BSS ordering
+                # pragma_incblocknum_first_line_num is used for symbol uniqueness, and
+                # also by fix_bss.py to locate the pragma these symbols originate from.
+                for j in range(amount):
+                    output.write(
+                        "struct increment_block_number_"
+                        f"{pragma_incblocknum_first_line_num:05}_{j:03};\n"
+                    )
+                output.write(f'#line {i + 1} "{filename}"\n')
         else:
             output.write(line)