[llvm] update_test_checks: keep names stable with generated functions (PR #87988)

[Nicolai Hähnle via llvm-commits]

https://github.com/nhaehnle updated https://github.com/llvm/llvm-project/pull/87988

>From cba180eda6d315588877ee774aac6848f73255be Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicolai=20H=C3=A4hnle?= <nicolai.haehnle at amd.com>
Date: Mon, 8 Apr 2024 15:10:39 +0200
Subject: [PATCH] update_test_checks: keep names stable with generated
 functions

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 at amd.com>
---
 .../Inputs/stable_ir_values_funcs.ll.expected | 10 ++--
 llvm/utils/UpdateTestChecks/common.py         | 59 +++++++++++--------
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/stable_ir_values_funcs.ll.expected b/llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/stable_ir_values_funcs.ll.expected
index 1559319ac013a2..86f929ffe36af6 100644
--- a/llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/stable_ir_values_funcs.ll.expected
+++ b/llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/stable_ir_values_funcs.ll.expected
@@ -16,9 +16,9 @@ define i32 @func({i32, i32} %x, i32 %y) {
 
 ; CHECK-LABEL: define i32 @func(
 ; CHECK-SAME: { i32, i32 } [[X:%.*]], i32 [[Y:%.*]]) {
-; CHECK-NEXT:    [[X_I34:%.*]] = extractvalue { i32, i32 } [[X]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[Y]], 1
-; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[X_I34]], [[TMP1]]
-; CHECK-NEXT:    [[TMP3:%.*]] = mul i32 [[TMP2]], 3
-; CHECK-NEXT:    ret i32 [[TMP3]]
+; CHECK-NEXT:    [[X_I33:%.*]] = extractvalue { i32, i32 } [[X]], 0
+; CHECK-NEXT:    [[TMP3:%.*]] = add i32 [[Y]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[X_I33]], [[TMP3]]
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i32 [[TMP1]], 3
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
diff --git a/llvm/utils/UpdateTestChecks/common.py b/llvm/utils/UpdateTestChecks/common.py
index ecb19d233a8d1a..66ea3761745cd2 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