[UpdateTestChecks] Split core functionality of add_ir_checks into add_checks

Cherry picked from D45272, also added some setup for add_asm_checks to use add_checks as well.

llvm-svn: 329266

1 files changed, 13 insertions, 7 deletions

diff --git a/llvm/utils/UpdateTestChecks/common.py b/llvm/utils/UpdateTestChecks/common.py
index 10dade6c50a..12b88a98f1e 100644
--- a/llvm/utils/UpdateTestChecks/common.py
+++ b/llvm/utils/UpdateTestChecks/common.py
@@ -156,22 +156,23 @@ def genericize_check_lines(lines):
   return lines
 
 
-def add_ir_checks(output_lines, comment_marker, prefix_list, func_dict, func_name):
-  # Label format is based on IR string.
-  check_label_format = '{} %s-LABEL: @%s('.format(comment_marker)
-
+def add_checks(output_lines, comment_marker, prefix_list, func_dict, func_name, check_label_format):
   printed_prefixes = []
   for p in prefix_list:
     checkprefixes = p[0]
     for checkprefix in checkprefixes:
       if checkprefix in printed_prefixes:
         break
-      if not func_dict[checkprefix][func_name]:
+      # TODO func_dict[checkprefix] may be None, '' or not exist.
+      # Fix the call sites.
+      if func_name not in func_dict[checkprefix] or not func_dict[checkprefix][func_name]:
         continue
+
       # Add some space between different check prefixes, but not after the last
       # check line (before the test code).
       #if len(printed_prefixes) != 0:
-      #  output_lines.append(';')
+      #  output_lines.append(comment_marker)
+
       printed_prefixes.append(checkprefix)
       output_lines.append(check_label_format % (checkprefix, func_name))
       func_body = func_dict[checkprefix][func_name].splitlines()
@@ -188,7 +189,7 @@ def add_ir_checks(output_lines, comment_marker, prefix_list, func_dict, func_nam
       #if func_body[0].startswith("#") or func_body[0].startswith("entry:"):
       #  is_blank_line = True
       #else:
-      #  output_lines.append('; %s:       %s' % (checkprefix, func_body[0]))
+      #  output_lines.append('%s %s:       %s' % (comment_marker, checkprefix, func_body[0]))
       #  is_blank_line = False
 
       is_blank_line = False
@@ -213,3 +214,8 @@ def add_ir_checks(output_lines, comment_marker, prefix_list, func_dict, func_nam
       # line of code in the test function.
       output_lines.append(comment_marker)
       break
+
+def add_ir_checks(output_lines, comment_marker, prefix_list, func_dict, func_name):
+  # Label format is based on IR string.
+  check_label_format = '{} %s-LABEL: @%s('.format(comment_marker)
+  add_checks(output_lines, comment_marker, prefix_list, func_dict, func_name, check_label_format)