[Utils] Add new --update-tests flag to llvm-lit (#108425)

This adds a flag to lit for detecting and updating failing tests when
possible to do so automatically. The flag uses a plugin architecture
where config files can add additional auto-updaters for the types of
tests in the test suite. When a test fails with `--update-tests` enabled
lit passes the test RUN invocation and output to each registered test
updater until one of them signals that it updated the test (or all test
updaters have been run). As such it is the responsibility of the test
updater to only update tests where it is reasonably certain that it will
actually fix the test, or come close to doing so.

Initially adds support for UpdateVerifyTests and UpdateTestChecks. The
flag is currently only implemented for lit's internal shell, so
`--update-tests` implies `LIT_USE_INTERNAL_SHELL=1`.

Builds on work in #97369
Fixes #81320

9 files changed, 103 insertions, 3 deletions

diff --git a/clang/test/lit.cfg.py b/clang/test/lit.cfg.py
index 92a3361..32ed523 100644
--- a/clang/test/lit.cfg.py
+++ b/clang/test/lit.cfg.py
@@ -362,3 +362,13 @@ if "system-aix" in config.available_features:
 # possibly be present in system and user configuration files, so disable
 # default configs for the test runs.
 config.environment["CLANG_NO_DEFAULT_CONFIG"] = "1"
+
+if lit_config.update_tests:
+    import sys
+    import os
+
+    utilspath = os.path.join(config.llvm_src_root, "utils")
+    sys.path.append(utilspath)
+    from update_any_test_checks import utc_lit_plugin
+
+    lit_config.test_updaters.append(utc_lit_plugin)
diff --git a/llvm/docs/CommandGuide/lit.rst b/llvm/docs/CommandGuide/lit.rst
index c9d5bab..dadecef 100644
--- a/llvm/docs/CommandGuide/lit.rst
+++ b/llvm/docs/CommandGuide/lit.rst
@@ -313,6 +313,11 @@ ADDITIONAL OPTIONS
 
  List all of the discovered tests and exit.
 
+.. option:: --update-tests
+
+ Pass failing tests to functions in the ``lit_config.update_tests`` list to
+ check whether any of them know how to update the test to make it pass.
+
 EXIT STATUS
 -----------
 
diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index 5a03a85..1d5b2bc 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -630,3 +630,13 @@ if "system-aix" in config.available_features:
 
 if config.has_logf128:
     config.available_features.add("has_logf128")
+
+if lit_config.update_tests:
+    import sys
+    import os
+
+    utilspath = os.path.join(config.llvm_src_root, "utils")
+    sys.path.append(utilspath)
+    from update_any_test_checks import utc_lit_plugin
+
+    lit_config.test_updaters.append(utc_lit_plugin)
diff --git a/llvm/utils/lit/lit/LitConfig.py b/llvm/utils/lit/lit/LitConfig.py
index 5dc712a..198a2bf 100644
--- a/llvm/utils/lit/lit/LitConfig.py
+++ b/llvm/utils/lit/lit/LitConfig.py
@@ -38,6 +38,7 @@ class LitConfig(object):
         parallelism_groups={},
         per_test_coverage=False,
         gtest_sharding=True,
+        update_tests=False,
     ):
         # The name of the test runner.
         self.progname = progname
@@ -89,6 +90,8 @@ class LitConfig(object):
         self.parallelism_groups = parallelism_groups
         self.per_test_coverage = per_test_coverage
         self.gtest_sharding = bool(gtest_sharding)
+        self.update_tests = update_tests
+        self.test_updaters = []
 
     @property
     def maxIndividualTestTime(self):
diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index a1785073..3a2cdc5 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1190,6 +1190,18 @@ def executeScriptInternal(
                 str(result.timeoutReached),
             )
 
+        if litConfig.update_tests:
+            for test_updater in litConfig.test_updaters:
+                try:
+                    update_output = test_updater(result, test)
+                except Exception as e:
+                    out += f"Exception occurred in test updater: {e}"
+                    continue
+                if update_output:
+                    for line in update_output.splitlines():
+                        out += f"# {line}\n"
+                    break
+
     return out, err, exitCode, timeoutInfo
 
 
diff --git a/llvm/utils/lit/lit/cl_arguments.py b/llvm/utils/lit/lit/cl_arguments.py
index ed78256..dcbe553 100644
--- a/llvm/utils/lit/lit/cl_arguments.py
+++ b/llvm/utils/lit/lit/cl_arguments.py
@@ -204,6 +204,12 @@ def parse_args():
         action="store_true",
         help="Exit with status zero even if some tests fail",
     )
+    execution_group.add_argument(
+        "--update-tests",
+        dest="update_tests",
+        action="store_true",
+        help="Try to update regression tests to reflect current behavior, if possible",
+    )
     execution_test_time_group = execution_group.add_mutually_exclusive_group()
     execution_test_time_group.add_argument(
         "--skip-test-time-recording",
diff --git a/llvm/utils/lit/lit/llvm/config.py b/llvm/utils/lit/lit/llvm/config.py
index 5f762ec..c05ec16 100644
--- a/llvm/utils/lit/lit/llvm/config.py
+++ b/llvm/utils/lit/lit/llvm/config.py
@@ -64,12 +64,17 @@ class LLVMConfig(object):
             self.with_environment("_TAG_REDIR_ERR", "TXT")
             self.with_environment("_CEE_RUNOPTS", "FILETAG(AUTOCVT,AUTOTAG) POSIX(ON)")
 
+        if lit_config.update_tests:
+            self.use_lit_shell = True
+
         # Choose between lit's internal shell pipeline runner and a real shell.
         # If LIT_USE_INTERNAL_SHELL is in the environment, we use that as an
         # override.
         lit_shell_env = os.environ.get("LIT_USE_INTERNAL_SHELL")
         if lit_shell_env:
             self.use_lit_shell = lit.util.pythonize_bool(lit_shell_env)
+            if not self.use_lit_shell and lit_config.update_tests:
+                print("note: --update-tests is not supported when using external shell")
 
         if not self.use_lit_shell:
             features.add("shell")
diff --git a/llvm/utils/lit/lit/main.py b/llvm/utils/lit/lit/main.py
index 24ba804..745e376 100755
--- a/llvm/utils/lit/lit/main.py
+++ b/llvm/utils/lit/lit/main.py
@@ -42,6 +42,7 @@ def main(builtin_params={}):
         config_prefix=opts.configPrefix,
         per_test_coverage=opts.per_test_coverage,
         gtest_sharding=opts.gtest_sharding,
+        update_tests=opts.update_tests,
     )
 
     discovered_tests = lit.discovery.find_tests_for_inputs(
diff --git a/llvm/utils/update_any_test_checks.py b/llvm/utils/update_any_test_checks.py
index e8eef1a..76fe336 100755
--- a/llvm/utils/update_any_test_checks.py
+++ b/llvm/utils/update_any_test_checks.py
@@ -34,9 +34,12 @@ def find_utc_tool(search_path, utc_name):
     return None
 
 
-def run_utc_tool(utc_name, utc_tool, testname):
+def run_utc_tool(utc_name, utc_tool, testname, environment):
     result = subprocess.run(
-        [utc_tool, testname], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        [utc_tool, testname],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=environment,
     )
     return (result.returncode, result.stdout, result.stderr)
 
@@ -60,6 +63,42 @@ def expand_listfile_args(arg_list):
     return exp_arg_list
 
 
+def utc_lit_plugin(result, test):
+    testname = test.getFilePath()
+    if not testname:
+        return None
+
+    script_name = os.path.abspath(__file__)
+    utc_search_path = os.path.join(os.path.dirname(script_name), os.path.pardir)
+
+    with open(testname, "r") as f:
+        header = f.readline().strip()
+
+    m = RE_ASSERTIONS.search(header)
+    if m is None:
+        return None
+
+    utc_name = m.group(1)
+    utc_tool = find_utc_tool([utc_search_path], utc_name)
+    if not utc_tool:
+        return f"update-utc-tests: {utc_name} not found"
+
+    return_code, stdout, stderr = run_utc_tool(
+        utc_name, utc_tool, testname, test.config.environment
+    )
+
+    stderr = stderr.decode(errors="replace")
+    if return_code != 0:
+        if stderr:
+            return f"update-utc-tests: {utc_name} exited with return code {return_code}\n{stderr.rstrip()}"
+        return f"update-utc-tests: {utc_name} exited with return code {return_code}"
+
+    stdout = stdout.decode(errors="replace")
+    if stdout:
+        return f"update-utc-tests: updated {testname}\n{stdout.rstrip()}"
+    return f"update-utc-tests: updated {testname}"
+
+
 def main():
     from argparse import RawTextHelpFormatter
 
@@ -78,6 +117,11 @@ def main():
         nargs="*",
         help="Additional directories to scan for update_*_test_checks scripts",
     )
+    parser.add_argument(
+        "--path",
+        help="""Additional directories to scan for executables invoked by the update_*_test_checks scripts,
+separated by the platform path separator""",
+    )
     parser.add_argument("tests", nargs="+")
     config = parser.parse_args()
 
@@ -88,6 +132,10 @@ def main():
     script_name = os.path.abspath(__file__)
     utc_search_path.append(os.path.join(os.path.dirname(script_name), os.path.pardir))
 
+    local_env = os.environ.copy()
+    if config.path:
+        local_env["PATH"] = config.path + os.pathsep + local_env["PATH"]
+
     not_autogenerated = []
     utc_tools = {}
     have_error = False
@@ -117,7 +165,7 @@ def main():
                         continue
 
                 future = executor.submit(
-                    run_utc_tool, utc_name, utc_tools[utc_name], testname
+                    run_utc_tool, utc_name, utc_tools[utc_name], testname, local_env
                 )
                 jobs.append((testname, future))