Corpus: let module_filter be a function (#169)

That makes it easier to filter arbitrary modules (e.g. from an
allow/deny list)

Author: mtrofin

compiler_opt/rl/corpus.py

@@ -17,11 +17,10 @@
 import concurrent.futures
 import math
 import random
-import re
 
 from absl import logging
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import json
 import os
@@ -227,7 +226,7 @@ class ReplaceContext:
   def __init__(self,
                *,
                data_path: str,
-               module_filter: Optional[re.Pattern] = None,
+               module_filter: Optional[Callable[[str], bool]] = None,
                additional_flags: Tuple[str, ...] = (),
                delete_flags: Tuple[str, ...] = (),
                replace_flags: Optional[Dict[str, str]] = None,
@@ -309,9 +308,7 @@ def __init__(self,
       raise ValueError('do not use add/delete flags to replace')
 
     if module_filter:
-      module_paths = [
-          name for name in module_paths if module_filter.match(name)
-      ]
+      module_paths = [name for name in module_paths if module_filter(name)]
 
     def get_cmdline(name: str):
       if cmd_override_was_specified:

compiler_opt/rl/corpus_test.py

@@ -230,7 +230,7 @@ def test_filter(self):
             corpus.ModuleSpec(name='largest', size=500),
             corpus.ModuleSpec(name='small', size=100)
         ],
-        module_filter=re.compile(r'.+l'))
+        module_filter=lambda name: re.compile(r'.+l').match(name))
     sample = cps.sample(999, sort=True)
     self.assertLen(sample, 3)
     self.assertEqual(sample[0].name, 'middle')

compiler_opt/tools/generate_default_trace.py

@@ -146,7 +146,8 @@ def main(_):
 
   cps = corpus.Corpus(
       data_path=_DATA_PATH.value,
-      module_filter=module_filter,
+      module_filter=lambda name: True
+      if not module_filter else module_filter.match(name),
       additional_flags=config.flags_to_add(),
       delete_flags=config.flags_to_delete(),
       replace_flags=config.flags_to_replace())