[Utils] Allow update_test_checks to scrub attribute annotations

Attribute annotations on calls, e.g., #0, are not useful on their own.
This patch adds a flag to update_test_checks.py to scrub them.

Reviewed By: lebedev.ri

Differential Revision: https://reviews.llvm.org/D68851

2 files changed, 12 insertions, 1 deletions

diff --git a/llvm/utils/UpdateTestChecks/common.py b/llvm/utils/UpdateTestChecks/common.py
index f93f8bf1cc7..ccb6ba586de 100644
--- a/llvm/utils/UpdateTestChecks/common.py
+++ b/llvm/utils/UpdateTestChecks/common.py
@@ -81,6 +81,8 @@ MARCH_ARG_RE = re.compile(r'-march[= ]([^ ]+)')
 SCRUB_LEADING_WHITESPACE_RE = re.compile(r'^(\s+)')
 SCRUB_WHITESPACE_RE = re.compile(r'(?!^(|  \w))[ \t]+', flags=re.M)
 SCRUB_TRAILING_WHITESPACE_RE = re.compile(r'[ \t]+$', flags=re.M)
+SCRUB_TRAILING_WHITESPACE_TEST_RE = SCRUB_TRAILING_WHITESPACE_RE
+SCRUB_TRAILING_WHITESPACE_AND_ATTRIBUTES_RE = re.compile(r'([ \t]|(#[0-9]+))+$', flags=re.M)
 SCRUB_KILL_COMMENT_RE = re.compile(r'^ *#+ +kill:.*\n')
 SCRUB_LOOP_COMMENT_RE = re.compile(
     r'# =>This Inner Loop Header:.*|# in Loop:.*', flags=re.M)
@@ -125,7 +127,7 @@ def scrub_body(body):
   # Expand the tabs used for indentation.
   body = string.expandtabs(body, 2)
   # Strip trailing whitespace.
-  body = SCRUB_TRAILING_WHITESPACE_RE.sub(r'', body)
+  body = SCRUB_TRAILING_WHITESPACE_TEST_RE.sub(r'', body)
   return body
 
 def do_scrub(body, scrubber, scrubber_args, extra):
diff --git a/llvm/utils/update_test_checks.py b/llvm/utils/update_test_checks.py
index 8ee22654980..94e3fef1fca 100755
--- a/llvm/utils/update_test_checks.py
+++ b/llvm/utils/update_test_checks.py
@@ -64,6 +64,8 @@ def main():
                       help='Do not scrub IR names')
   parser.add_argument('--function-signature', action='store_true',
                       help='Keep function signature information around for the check line')
+  parser.add_argument('--scrub-attributes', action='store_true',
+                      help='Remove attribute annotations (#0) from the end of check line')
   parser.add_argument('tests', nargs='+')
   args = common.parse_commandline_args(parser)
 
@@ -98,6 +100,13 @@ def main():
         continue
 
     run_lines = common.find_run_lines(test, input_lines)
+
+    # If requested we scrub trailing attribute annotations, e.g., '#0', together with whitespaces
+    if args.scrub_attributes:
+      common.SCRUB_TRAILING_WHITESPACE_TEST_RE = common.SCRUB_TRAILING_WHITESPACE_AND_ATTRIBUTES_RE
+    else:
+      common.SCRUB_TRAILING_WHITESPACE_TEST_RE = common.SCRUB_TRAILING_WHITESPACE_RE
+
     prefix_list = []
     for l in run_lines:
       if '|' not in l: