Collect Refaster rules across all provided files

Fixes google#552
PicnicSupermarket · Aug 22, 2017 · e2f4073 · e2f4073
1 parent 25b741a
commit e2f4073
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 9 deletions.
diff --git a/refaster/src/main/java/com/google/errorprone/refaster/RefasterRuleCompiler.java b/refaster/src/main/java/com/google/errorprone/refaster/RefasterRuleCompiler.java
@@ -53,8 +53,10 @@ public static void main(String[] args) {
   }
 
   private Result run(String[] argv, Context context) {
+    List<String> newArgs = new ArrayList<>(Arrays.asList(argv));
+    RefasterRuleCompilerAnalyzer analyzer;
     try {
-      argv = prepareCompilation(argv, context);
+      analyzer = prepareCompilation(newArgs, context);
     } catch (InvalidCommandLineOptionException e) {
       System.err.println(e.getMessage());
       System.err.flush();
@@ -66,8 +68,11 @@ private Result run(String[] argv, Context context) {
           new Main(
                   "RefasterRuleCompiler",
                   new PrintWriter(new OutputStreamWriter(System.err, StandardCharsets.UTF_8)))
-              .compile(argv, context);
+              .compile(newArgs.toArray(new String[0]), context);
       System.err.flush();
+      if (compileResult == Result.OK) {
+        analyzer.persist();
+      }
       return compileResult;
     } catch (InvalidCommandLineOptionException e) {
       System.err.println(e.getMessage());
@@ -76,11 +81,10 @@ private Result run(String[] argv, Context context) {
     }
   }
 
-  private String[] prepareCompilation(String[] argv, Context context)
+  private RefasterRuleCompilerAnalyzer prepareCompilation(List<String> argv, Context context)
       throws InvalidCommandLineOptionException {
     context.put(DiagnosticListener.class, new DiagnosticCollector<JavaFileObject>());
-    List<String> newArgs = new ArrayList<>(Arrays.asList(argv));
-    Iterator<String> itr = newArgs.iterator();
+    Iterator<String> itr = argv.iterator();
     String path = null;
     while (itr.hasNext()) {
       if (itr.next().equals("--out")) {
@@ -94,9 +98,10 @@ private String[] prepareCompilation(String[] argv, Context context)
 
     MaskedClassLoader.preRegisterFileManager(context);
 
-    MultiTaskListener.instance(context)
-        .add(new RefasterRuleCompilerAnalyzer(context, FileSystems.getDefault().getPath(path)));
+    RefasterRuleCompilerAnalyzer analyzer =
+        new RefasterRuleCompilerAnalyzer(context, FileSystems.getDefault().getPath(path));
+    MultiTaskListener.instance(context) .add(analyzer);
 
-    return Iterables.toArray(newArgs, String.class);
+    return analyzer;
   }
 }
diff --git a/refaster/src/main/java/com/google/errorprone/refaster/RefasterRuleCompilerAnalyzer.java b/refaster/src/main/java/com/google/errorprone/refaster/RefasterRuleCompilerAnalyzer.java
@@ -29,13 +29,15 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 /**
  * TaskListener that receives compilation of a Refaster rule class and outputs a serialized analyzer
  * to the specified path.
  */
 public class RefasterRuleCompilerAnalyzer implements TaskListener {
+  private final List<CodeTransformer> allRules = Collections.synchronizedList(new ArrayList<>());
   private final Context context;
   private final Path destinationPath;
 
@@ -67,9 +69,13 @@ public Void visitClass(ClassTree node, Context context) {
     if (rules.isEmpty()) {
       throw new IllegalArgumentException("Did not find any Refaster templates");
     }
+    allRules.addAll(rules);
+  }
+
+  public void persist() {
     try (ObjectOutputStream output =
         new ObjectOutputStream(Files.newOutputStream(destinationPath))) {
-      output.writeObject(CompositeCodeTransformer.compose(rules));
+      output.writeObject(CompositeCodeTransformer.compose(allRules));
     } catch (IOException e) {
       throw new RuntimeException(e);
     }