update_test_checks: Fix preservation of meta variable names with --include-generated-funcs

llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/stable_ir_values5.ll

@@ -0,0 +1,22 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --include-generated-funcs --version 3
+; RUN: opt < %s -S | FileCheck %s
+
+; With --include-generated-funcs, check lines are separate from IR function bodies.
+; This needs to be handled separately by the UTC parser for existing check lines.
+; This test tests this case.
+
+define i32 @func({i32, i32} %x, i32 %y) {
+  %x.i34 = extractvalue {i32, i32} %x, 0
+  %1 = add i32 %y, 1
+  %2 = add i32 %x.i34, %1
+  %3 = mul i32 %2, 3
+  ret i32 %3
+}
+
+; CHECK-LABEL: define i32 @func(
+; CHECK-SAME: { i32, i32 } [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[X_I33:%.*]] = extractvalue { i32, i32 } [[X]], 0
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[X_I33]], [[Y]]
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i32 [[TMP1]], 3
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;

llvm/test/tools/UpdateTestChecks/update_test_checks/Inputs/stable_ir_values5.ll.expected

@@ -0,0 +1,23 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --include-generated-funcs --version 3
+; RUN: opt < %s -S | FileCheck %s
+
+; With --include-generated-funcs, check lines are separate from IR function bodies.
+; This needs to be handled separately by the UTC parser for existing check lines.
+; This test tests this case.
+
+define i32 @func({i32, i32} %x, i32 %y) {
+  %x.i34 = extractvalue {i32, i32} %x, 0
+  %1 = add i32 %y, 1
+  %2 = add i32 %x.i34, %1
+  %3 = mul i32 %2, 3
+  ret i32 %3
+}
+
+; CHECK-LABEL: define i32 @func(
+; CHECK-SAME: { i32, i32 } [[X:%.*]], i32 [[Y:%.*]]) {
+; 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]]
+;

llvm/test/tools/UpdateTestChecks/update_test_checks/stable_ir_values5.test

@@ -0,0 +1,2 @@
+# RUN: cp -f %S/Inputs/stable_ir_values5.ll %t.ll && %update_test_checks %t.ll
+# RUN: diff -u %t.ll %S/Inputs/stable_ir_values5.ll.expected

llvm/utils/UpdateTestChecks/common.py

@@ -430,36 +430,81 @@ 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 = {}
 
+    # We need to detect check lines, and associate them with the correct functions and prefixes.
+    # There are two formats how IR and check lines are combined. In the standard format,
+    # they are interleaved, and it suffices to parse IR function definitions in order to keep
+    # track of the current function:
+    #
+    #    define i32 @func1(i32 %x) {
+    #    ; CHECK-LABEL: define i32 @func1(
+    #    ; CHECK-NEXT:    /* check lines */
+    #    ;
+    #      %1 = /* IR body of @func1 */
+    #
+    #    define i32 @func2(i32 %x) {
+    #    ; CHECK-LABEL: define i32 @func2(
+    #    ; CHECK-NEXT:    /* check lines */
+    #    ;
+    #      %1 = /* IR body of @func2 */
+    #
+    # However, with --include-generated-funcs, check lines are separate from IR function bodies,
+    # and we also need to parse CHECK lines of function definitions:
+    #
+    #    define i32 @func1(i32 %x) {
+    #      /* IR body */
+    #
+    #    define i32 @func2(..) {
+    #      /* IR body */
+    #
+    #    ; CHECK-LABEL: define i32 @func1
+    #    ; CHECK-NEXT /* check lines for func1 */
+    #
+    #    ; CHECK-LABEL: define i32 @func2
+    #    ; CHECK-NEXT /* check lines for func2 */
+
+    result = collections.defaultdict(dict)
     current_function = None
+
+    def update_current_function(line):
+        m = IR_FUNCTION_RE.match(line)
+        if m is None:
+            return
+        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.
+            return
+
+        nonlocal current_function
+        current_function = result[func_name]
+
     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
+        if input_line == "":
+            continue
+        m = None
+        if input_line.lstrip().startswith(";"):
+            m = CHECK_RE.match(input_line)
+            if m is not None and m.group(2) == "LABEL":
+                # Update current_function if the current line is a CHECK of a function definition
+                line_remainder = input_line[len(m.group(0)) :]
+                update_current_function(line_remainder)
+        else:
             current_function = None
 
-        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 current_function is not None:
+            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
 
-            assert func_name not in result
-            current_function = result[func_name] = {}
+        # Update current_function if the current line is an IR function definition
+        update_current_function(input_line)
 
     return result