update_test_checks: keep names stable with generated functions (#87988)

Collect the original check lines in a manner that is independent of
where the check lines appear in the file. This is so that we keep
FileCheck variable names stable even when --include-generated-funcs is
used.

Reported-by: Ruiling Song <ruiling.song@amd.com>

1 files changed, 35 insertions, 24 deletions

diff --git a/llvm/utils/UpdateTestChecks/common.py b/llvm/utils/UpdateTestChecks/common.py
index eed36a0..15d3d5e 100644
--- a/llvm/utils/UpdateTestChecks/common.py
+++ b/llvm/utils/UpdateTestChecks/common.py
@@ -430,36 +430,47 @@ def collect_original_check_lines(ti: TestInfo, prefix_set: set):
     result[func_name][prefix] is filled with a list of right-hand-sides of check
     lines.
     """
-    result = {}
+    result = collections.defaultdict(lambda: {})
 
+    current_prefix = None
     current_function = None
     for input_line_info in ti.ro_iterlines():
         input_line = input_line_info.line
-        if current_function is not None:
-            if input_line == "":
-                continue
-            if input_line.lstrip().startswith(";"):
-                m = CHECK_RE.match(input_line)
-                if (
-                    m is not None
-                    and m.group(1) in prefix_set
-                    and m.group(2) not in ["LABEL", "SAME"]
-                ):
-                    if m.group(1) not in current_function:
-                        current_function[m.group(1)] = []
-                    current_function[m.group(1)].append(input_line[m.end() :].strip())
-                continue
-            current_function = None
+        if input_line.lstrip().startswith(";"):
+            m = CHECK_RE.match(input_line)
+            if m is not None:
+                prefix = m.group(1)
+                check_kind = m.group(2)
+                line = input_line[m.end() :].strip()
+
+                if prefix != current_prefix:
+                    current_function = None
+                    current_prefix = None
+
+                if check_kind not in ["LABEL", "SAME"]:
+                    if current_function is not None:
+                        current_function.append(line)
+                    continue
 
-        m = IR_FUNCTION_RE.match(input_line)
-        if m is not None:
-            func_name = m.group(1)
-            if ti.args.function is not None and func_name != ti.args.function:
-                # When filtering on a specific function, skip all others.
-                continue
+                if check_kind == "SAME":
+                    continue
+
+                if check_kind == "LABEL":
+                    m = IR_FUNCTION_RE.match(line)
+                    if m is not None:
+                        func_name = m.group(1)
+                        if (
+                            ti.args.function is not None
+                            and func_name != ti.args.function
+                        ):
+                            # When filtering on a specific function, skip all others.
+                            continue
+
+                        current_prefix = prefix
+                        current_function = result[func_name][prefix] = []
+                        continue
 
-            assert func_name not in result
-            current_function = result[func_name] = {}
+        current_function = None
 
     return result