[llvm-branch-commits] [llvm] [lit] Add lit.llvm.fn_param shared helper for --param fn= substitutions (PR #200350)

[via llvm-branch-commits]

https://github.com/jofrn updated https://github.com/llvm/llvm-project/pull/200350

>From e02f075dbcc6c1be589fbe576e2e01bb400687ee Mon Sep 17 00:00:00 2001
From: jofrn <165626406+jofrn at users.noreply.github.com>
Date: Fri, 29 May 2026 01:09:16 -0700
Subject: [PATCH] [lit] Add lit.llvm.fn_param shared helper for --param fn=
 substitutions

---
 llvm/test/lit.cfg.py                |  6 ++++
 llvm/utils/lit/lit/llvm/fn_param.py | 46 +++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 llvm/utils/lit/lit/llvm/fn_param.py

diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index 09df1e3fd6281..f4b9e8aaed6ef 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -39,6 +39,12 @@
 )
 config.test_format = lit.formats.ShTest(not use_lit_shell, extra_substitutions)
 
+# --param fn=NAMES narrows compilation to the named functions and their
+# transitive dependencies. See lit.llvm.fn_param.install for dispatch details.
+from lit.llvm import fn_param
+
+fn_param.install(config, lit_config)
+
 # suffixes: A list of file extensions to treat as test files. This is overriden
 # by individual lit.local.cfg files in the test subdirectories.
 config.suffixes = [".ll", ".c", ".test", ".txt", ".s", ".mir", ".yaml", ".spv"]
diff --git a/llvm/utils/lit/lit/llvm/fn_param.py b/llvm/utils/lit/lit/llvm/fn_param.py
new file mode 100644
index 0000000000000..48d8400b80155
--- /dev/null
+++ b/llvm/utils/lit/lit/llvm/fn_param.py
@@ -0,0 +1,46 @@
+"""Shared building blocks for `--param fn=NAMES`-driven lit substitutions.
+
+Used by `lit.llvm.fn_selection` and `lit.llvm.fn_extract` to narrow compilation
+to a subset of functions. Kept here so the two helpers stay short and share
+parsing + capture-substitution wiring."""
+
+from lit.TestingConfig import SubstituteCaptures
+
+
+def parse_fn_names(lit_config, param="fn"):
+    """Return the comma-separated list passed via `--param <param>=NAMES`,
+    or an empty list when the param is absent or empty."""
+    val = lit_config.params.get(param)
+    if not val:
+        return []
+    return [n.strip() for n in val.split(",") if n.strip()]
+
+
+def add_capture_sub(config, pattern, replacement):
+    """Append a substitution that preserves regex backreferences in `replacement`."""
+    config.substitutions.append((pattern, SubstituteCaptures(replacement)))
+
+
+def install(config, lit_config):
+    """Dispatch `--param fn=NAMES` to the right helper, and ask FileCheck to
+    drop CHECKs outside the selected CHECK-LABEL sections.
+
+    `--param fn-pass=1` opts into `lit.llvm.fn_selection` (the select-function
+    pass, opt-only); otherwise `lit.llvm.fn_extract` is used (prepends
+    llvm-extract, tool-agnostic)."""
+    names = parse_fn_names(lit_config)
+    if not names:
+        return
+    # Splice `--filter-label=NAMES` after any FileCheck invocation so the
+    # downstream FileCheck only checks the CHECK-LABEL sections we kept.
+    add_capture_sub(
+        config, r"(\S*FileCheck)\b", r"\1 --filter-label=" + ",".join(names)
+    )
+    if lit_config.params.get("fn-pass"):
+        # from lit.llvm import fn_selection
+        # fn_selection.install(config, lit_config)
+        pass
+    else:
+        # from lit.llvm import fn_extract
+        # fn_extract.install(config, lit_config)
+        pass