Explicitly model 'replace' flags (#139)

* Explicitly model 'replace' flags

Using 'delete' and 'add' for flags that must actually be present - such
as -fthinlto-index, or, later, profiles - is a bug farm. This patch
explicitly models such flags as something that cannot be 'deleted' or
'added', and must be somehow converted (contextualized to the location
of the files(presumably) used as values of those flags). The conversion
rule is a formatting string which uses a 'context' value, which will
have fields populate at the time of the conversion.

Some tests had to have the order of some flags adjusted a bit as a
NFC side-effect.

* Rebase and fix

Author: mtrofin

compiler_opt/rl/corpus.py

@@ -19,7 +19,7 @@
 
 from absl import logging
 from dataclasses import dataclass
-from typing import Iterable, List, Dict, Tuple, Any
+from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 import json
 import os
@@ -82,16 +82,36 @@ def __call__(self,
 class Corpus:
   """Represents a corpus. Comes along with some utility functions."""
 
+  @dataclass(frozen=True)
+  class ReplaceContext:
+    """Context for 'replace' rules."""
+    full_path_prefix: str
+
   def __init__(self,
                data_path: str,
                additional_flags: Tuple[str, ...] = (),
-               delete_flags: Tuple[str, ...] = ()):
+               delete_flags: Tuple[str, ...] = (),
+               replace_flags: Optional[Dict[str, str]] = None):
+    """
+    Args:
+      data_path: corpus directory.
+      additional_flags: list of flags to append to the command line
+      delete_flags: list of flags to remove (both `-flag=<value` and
+        `-flag <value>` are supported).
+      replace_flags: list of flags to be replaced. The key in the dictionary
+        is the flag. The value is a string that will be `format`-ed with a
+        `context` object - see `ReplaceContext`.
+        We verify that flags in replace_flags are present, and do not appear
+        in the additional_flags nor delete_flags.
+        Thinlto index is handled this way, too.
+    """
     self._module_specs = tuple(
         sorted(
             _build_modulespecs_from_datapath(
                 data_path=data_path,
                 additional_flags=additional_flags,
-                delete_flags=delete_flags),
+                delete_flags=delete_flags,
+                replace_flags=replace_flags),
             key=lambda m: m.size,
             reverse=True))
     self._root_dir = data_path
@@ -137,8 +157,8 @@ def __len__(self):
 def _build_modulespecs_from_datapath(
     data_path: str,
     additional_flags: Tuple[str, ...] = (),
-    delete_flags: Tuple[str, ...] = ()
-) -> List[ModuleSpec]:
+    delete_flags: Tuple[str, ...] = (),
+    replace_flags: Optional[Dict[str, str]] = None) -> List[ModuleSpec]:
   # TODO: (b/233935329) Per-corpus *fdo profile paths can be read into
   # {additional|delete}_flags here
   with open(
@@ -174,6 +194,7 @@ def _build_modulespecs_from_datapath(
         has_thinlto=has_thinlto,
         additional_flags=additional_flags,
         delete_flags=delete_flags,
+        replace_flags=replace_flags,
         cmd_override=cmd_override)
     size = os.path.getsize(full_module_path + '.bc')
     module_specs.append(
@@ -187,7 +208,8 @@ def _load_and_parse_command(
     has_thinlto: bool,
     additional_flags: Tuple[str, ...] = (),
     delete_flags: Tuple[str, ...] = (),
-    cmd_override: Tuple[str, ...] = ()
+    replace_flags: Optional[Dict[str, str]] = None,
+    cmd_override: Tuple[str, ...] = (),
 ) -> List[str]:
   """Cleans up base command line.
 
@@ -211,23 +233,59 @@ def _load_and_parse_command(
   else:
     with open(module_path + '.cmd', encoding='utf-8') as f:
       option_iterator = iter(f.read().split('\0'))
+
+  context = Corpus.ReplaceContext(full_path_prefix=module_path)
+  replace_flags = replace_flags.copy() if replace_flags else {}
+  if has_thinlto:
+    additional_flags = ('-mllvm', '-thinlto-assume-merged') + additional_flags
+    if cmd_override:
+      additional_flags = (
+          f'-fthinlto-index={context.full_path_prefix}.thinlto.bc',
+      ) + additional_flags
+    else:
+      fthinlto_index = '-fthinlto-index'
+      if fthinlto_index in replace_flags:
+        raise ValueError(
+            '-fthinlto-index must be handled by the infrastructure')
+      replace_flags[fthinlto_index] = '{context.full_path_prefix}.thinlto.bc'
+
+  additional_flags = ('-x', 'ir', module_path + '.bc') + additional_flags
+
+  matched_replace_flags = set()
+  # don't user add/remove for replace
+  if set(additional_flags).intersection(
+      set(replace_flags)) or set(delete_flags).intersection(set(replace_flags)):
+    raise ValueError('do not use add/delete flags to replace')
+
   option = next(option_iterator, None)
 
   while option is not None:
     if any(option.startswith(flag) for flag in delete_flags):
       if '=' not in option:
         next(option_iterator, None)
     else:
-      cmdline.append(option)
-    option = next(option_iterator, None)
-  cmdline.extend(['-x', 'ir', module_path + '.bc'])
-
-  if has_thinlto:
-    cmdline.extend([
-        f'-fthinlto-index={module_path}.thinlto.bc', '-mllvm',
-        '-thinlto-assume-merged'
-    ])
+      matching_replace = [
+          flag for flag in replace_flags if option.startswith(flag)
+      ]
+      if not matching_replace:
+        cmdline.append(option)
+      else:
+        assert len(matching_replace) == 1
+        flag = matching_replace[0]
+        if flag in matched_replace_flags:
+          raise ValueError(f'{flag} was matched twice')
+        matched_replace_flags.add(flag)
+
+        if '=' not in option:
+          next(option_iterator, None)
+          cmdline.extend([option, replace_flags[flag].format(context=context)])
+        else:
+          cmdline.append(flag + '=' +
+                         replace_flags[flag].format(context=context))
 
+    option = next(option_iterator, None)
+  if len(matched_replace_flags) != len(replace_flags):
+    raise ValueError('flags that were expected to be replaced were not found')
   cmdline.extend(additional_flags)
 
   # The options read from a .cmd file must be run with -cc1

compiler_opt/rl/corpus_test.py

@@ -33,12 +33,14 @@ def test_thinlto_file(self):
         corpus._load_and_parse_command(
             module_path=module_path, has_thinlto=False),
         ['-cc1', '-foo', '-bar=baz', '-x', 'ir', module_path + '.bc'])
+    data = ['-cc1', '-foo', '-bar=baz', '-fthinlto-index=some/path/here']
+    argfile = self.create_tempfile(content='\0'.join(data), file_path='hi.cmd')
     self.assertEqual(
         corpus._load_and_parse_command(
             module_path=module_path, has_thinlto=True), [
-                '-cc1', '-foo', '-bar=baz', '-x', 'ir', module_path + '.bc',
-                '-fthinlto-index=' + module_path + '.thinlto.bc', '-mllvm',
-                '-thinlto-assume-merged'
+                '-cc1', '-foo', '-bar=baz',
+                '-fthinlto-index=' + module_path + '.thinlto.bc', '-x', 'ir',
+                module_path + '.bc', '-mllvm', '-thinlto-assume-merged'
             ])
 
   def test_deletion(self):
@@ -70,6 +72,54 @@ def test_addition(self):
             additional_flags=additional_flags),
         ['-cc1', '-x', 'ir', module_path + '.bc', '-fix-all-bugs'])
 
+  def test_replacement(self):
+    replace_flags = {'-replace-me': '{context.full_path_prefix}.replaced'}
+    data = ['-cc1']
+    argfile = self.create_tempfile(content='\0'.join(data), file_path='hi.cmd')
+    module_path = argfile.full_path[:-4]
+    # if we expect to be able to replace a flag, and it's not in the original
+    # cmdline, raise.
+    self.assertRaises(
+        ValueError,
+        corpus._load_and_parse_command,
+        module_path=module_path,
+        has_thinlto=False,
+        replace_flags=replace_flags)
+    data = ['-cc1', '-replace-me=some_value']
+    argfile = self.create_tempfile(content='\0'.join(data), file_path='hi.cmd')
+    # the flag can be replaced in non-thinlto cases; and in the case of thinlto
+    # it does not interfere with the thinlto one
+    self.assertEqual(
+        corpus._load_and_parse_command(
+            module_path=module_path,
+            has_thinlto=False,
+            replace_flags=replace_flags), [
+                '-cc1', '-replace-me=' + module_path + '.replaced', '-x', 'ir',
+                module_path + '.bc'
+            ])
+    # variant without '='
+    data = ['-cc1', '-replace-me', 'some_value']
+    argfile = self.create_tempfile(content='\0'.join(data), file_path='hi.cmd')
+    self.assertEqual(
+        corpus._load_and_parse_command(
+            module_path=module_path,
+            has_thinlto=False,
+            replace_flags=replace_flags), [
+                '-cc1', '-replace-me', module_path + '.replaced', '-x', 'ir',
+                module_path + '.bc'
+            ])
+    data = ['-cc1', '-replace-me', 'some_value', '-fthinlto-index=blah']
+    argfile = self.create_tempfile(content='\0'.join(data), file_path='hi.cmd')
+    self.assertEqual(
+        corpus._load_and_parse_command(
+            module_path=module_path,
+            has_thinlto=True,
+            replace_flags=replace_flags), [
+                '-cc1', '-replace-me', module_path + '.replaced',
+                '-fthinlto-index=' + module_path + '.thinlto.bc', '-x', 'ir',
+                module_path + '.bc', '-mllvm', '-thinlto-assume-merged'
+            ])
+
   def test_modification(self):
     delete_compilation_flags = ('-split-dwarf-file', '-split-dwarf-output',
                                 '-fthinlto-index', '-fprofile-sample-use',
@@ -166,23 +216,27 @@ def test_get_with_thinlto(self):
     tempdir.create_file(
         '2.cmd', content='\0'.join(['-cc1', '-fthinlto-index=abc']))
 
-    ms_list = corpus._build_modulespecs_from_datapath(
+    self.assertRaises(
+        ValueError,
+        corpus._build_modulespecs_from_datapath,
         tempdir.full_path,
         additional_flags=('-add',),
         delete_flags=('-fthinlto-index',))
+    ms_list = corpus._build_modulespecs_from_datapath(
+        tempdir.full_path, additional_flags=('-add',))
     self.assertEqual(len(ms_list), 2)
     ms1 = ms_list[0]
     ms2 = ms_list[1]
     self.assertEqual(ms1.name, '1')
     self.assertEqual(ms1.exec_cmd,
-                     ('-cc1', '-x', 'ir', tempdir.full_path + '/1.bc',
-                      '-fthinlto-index=' + tempdir.full_path + '/1.thinlto.bc',
+                     ('-cc1', '-fthinlto-index=' + tempdir.full_path +
+                      '/1.thinlto.bc', '-x', 'ir', tempdir.full_path + '/1.bc',
                       '-mllvm', '-thinlto-assume-merged', '-add'))
 
     self.assertEqual(ms2.name, '2')
     self.assertEqual(ms2.exec_cmd,
-                     ('-cc1', '-x', 'ir', tempdir.full_path + '/2.bc',
-                      '-fthinlto-index=' + tempdir.full_path + '/2.thinlto.bc',
+                     ('-cc1', '-fthinlto-index=' + tempdir.full_path +
+                      '/2.thinlto.bc', '-x', 'ir', tempdir.full_path + '/2.bc',
                       '-mllvm', '-thinlto-assume-merged', '-add'))
 
   def test_get_with_override(self):
@@ -203,9 +257,7 @@ def test_get_with_override(self):
     tempdir.create_file('2.cmd', content='\0'.join(['-fthinlto-index=abc']))
 
     ms_list = corpus._build_modulespecs_from_datapath(
-        tempdir.full_path,
-        additional_flags=('-add',),
-        delete_flags=('-fthinlto-index',))
+        tempdir.full_path, additional_flags=('-add',))
     self.assertEqual(len(ms_list), 2)
     ms1 = ms_list[0]
     ms2 = ms_list[1]