1 files changed, 78 insertions, 6 deletions

diff --git a/llvm/utils/UpdateTestChecks/common.py b/llvm/utils/UpdateTestChecks/common.py
index d49fe50..488425d 100644
--- a/llvm/utils/UpdateTestChecks/common.py
+++ b/llvm/utils/UpdateTestChecks/common.py
@@ -18,6 +18,8 @@ else:
 _verbose = False
 
 def parse_commandline_args(parser):
+  parser.add_argument('--include-generated-funcs', action='store_true',
+                      help='Output checks for functions not in source')
   parser.add_argument('-v', '--verbose', action='store_true',
                       help='Show verbose output')
   parser.add_argument('-u', '--update-only', action='store_true',
@@ -62,19 +64,25 @@ class TestInfo(object):
     self.test_autogenerated_note = self.autogenerated_note_prefix + script_name
     self.test_autogenerated_note += get_autogennote_suffix(parser, self.args)
 
+  def ro_iterlines(self):
+    for line_num, input_line in enumerate(self.input_lines):
+      args, argv = check_for_command(input_line, self.parser,
+                                     self.args, self.argv, self.argparse_callback)
+      yield InputLineInfo(input_line, line_num, args, argv)
+
   def iterlines(self, output_lines):
     output_lines.append(self.test_autogenerated_note)
-    for line_num, input_line in enumerate(self.input_lines):
+    for line_info in self.ro_iterlines():
+      input_line = line_info.line
       # Discard any previous script advertising.
       if input_line.startswith(self.autogenerated_note_prefix):
         continue
-      self.args, self.argv = check_for_command(input_line, self.parser,
-                                               self.args, self.argv, self.argparse_callback)
+      self.args = line_info.args
+      self.argv = line_info.argv
       if not self.args.enabled:
         output_lines.append(input_line)
         continue
-      yield InputLineInfo(input_line, line_num, self.args, self.argv)
-
+      yield line_info
 
 def itertests(test_patterns, parser, script_name, comment_prefix=None, argparse_callback=None):
   for pattern in test_patterns:
@@ -250,7 +258,7 @@ class function_body(object):
   def __str__(self):
     return self.scrub
 
-def build_function_body_dictionary(function_re, scrubber, scrubber_args, raw_tool_output, prefixes, func_dict, verbose, record_args, check_attributes):
+def build_function_body_dictionary(function_re, scrubber, scrubber_args, raw_tool_output, prefixes, func_dict, func_order, verbose, record_args, check_attributes):
   for m in function_re.finditer(raw_tool_output):
     if not m:
       continue
@@ -292,6 +300,7 @@ def build_function_body_dictionary(function_re, scrubber, scrubber_args, raw_too
               continue
 
       func_dict[prefix][func] = function_body(scrubbed_body, scrubbed_extra, args_and_sig, attrs)
+      func_order[prefix].append(func)
 
 ##### Generator of LLVM IR CHECK lines
 
@@ -594,3 +603,66 @@ def check_for_command(line, parser, args, argv, argparse_callback):
         if argparse_callback is not None:
           argparse_callback(args)
     return args, argv
+
+def find_arg_in_test(test_info, get_arg_to_check, arg_string, is_global):
+  result = get_arg_to_check(test_info.args)
+  if not result and is_global:
+    # See if this has been specified via UTC_ARGS.  This is a "global" option
+    # that affects the entire generation of test checks.  If it exists anywhere
+    # in the test, apply it to everything.
+    saw_line = False
+    for line_info in test_info.ro_iterlines():
+      line = line_info.line
+      if not line.startswith(';') and line.strip() != '':
+        saw_line = True
+      result = get_arg_to_check(line_info.args)
+      if result:
+        if warn and saw_line:
+          # We saw the option after already reading some test input lines.
+          # Warn about it.
+          print('WARNING: Found {} in line following test start: '.format(arg_string)
+                + line, file=sys.stderr)
+          print('WARNING: Consider moving {} to top of file'.format(arg_string),
+                file=sys.stderr)
+        break
+  return result
+
+def dump_input_lines(output_lines, test_info, prefix_set, comment_string):
+  for input_line_info in test_info.iterlines(output_lines):
+    line = input_line_info.line
+    args = input_line_info.args
+    if line.strip() == comment_string:
+      continue
+    if line.lstrip().startswith(comment_string):
+      m = CHECK_RE.match(line)
+      if m and m.group(1) in prefix_set:
+        continue
+    output_lines.append(line.rstrip('\n'))
+
+def add_checks_at_end(output_lines, prefix_list, func_order,
+                      comment_string, check_generator):
+  added = set()
+  for prefixes, tool_args, *rest in prefix_list:
+    for prefix in prefixes:
+      for func in func_order[prefix]:
+        if added:
+          output_lines.append(comment_string)
+        added.add(func)
+
+        # The add_*_checks routines expect a run list whose items are
+        # tuples that have a list of prefixes as their first element and
+        # tool command args string as their second element.  They output
+        # checks for each prefix in the list of prefixes.  By doing so, it
+        # implicitly assumes that for each function every run line will
+        # generate something for that function.  That is not the case for
+        # generated functions as some run lines might not generate them
+        # (e.g. -fopenmp vs. no -fopenmp).
+        #
+        # Therefore, pass just the prefix we're interested in.  This has
+        # the effect of generating all of the checks for functions of a
+        # single prefix before moving on to the next prefix.  So checks
+        # are ordered by prefix instead of by function as in "normal"
+        # mode.
+        check_generator(output_lines,
+                        [([prefix], tool_args)],
+                        func)