Allow annotation processors to record build statistics

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=271012418

5 files changed, 87 insertions, 29 deletions

diff --git a/java/com/google/turbine/binder/Binder.java b/java/com/google/turbine/binder/Binder.java
index 07fd7b6..4f3126a 100644
--- a/java/com/google/turbine/binder/Binder.java
+++ b/java/com/google/turbine/binder/Binder.java
@@ -171,7 +171,13 @@ public class Binder {
     }
 
     return new BindingResult(
-        result.build(), boundModules, classPathEnv, tli, generatedSources, generatedClasses);
+        result.build(),
+        boundModules,
+        classPathEnv,
+        tli,
+        generatedSources,
+        generatedClasses,
+        /* statistics= */ ImmutableMap.of());
   }
 
   /** Records enclosing declarations of member classes, and group classes by compilation unit. */
@@ -442,6 +448,7 @@ public class Binder {
     private final TopLevelIndex tli;
     private final ImmutableList<SourceFile> generatedSources;
     private final ImmutableMap<String, byte[]> generatedClasses;
+    private final ImmutableMap<String, byte[]> statistics;
 
     public BindingResult(
         ImmutableMap<ClassSymbol, SourceTypeBoundClass> units,
@@ -449,13 +456,15 @@ public class Binder {
         CompoundEnv<ClassSymbol, BytecodeBoundClass> classPathEnv,
         TopLevelIndex tli,
         ImmutableList<SourceFile> generatedSources,
-        ImmutableMap<String, byte[]> generatedClasses) {
+        ImmutableMap<String, byte[]> generatedClasses,
+        ImmutableMap<String, byte[]> statistics) {
       this.units = units;
       this.modules = modules;
       this.classPathEnv = classPathEnv;
       this.tli = tli;
       this.generatedSources = generatedSources;
       this.generatedClasses = generatedClasses;
+      this.statistics = statistics;
     }
 
     /** Bound nodes for sources in the compilation. */
@@ -483,5 +492,9 @@ public class Binder {
     public ImmutableMap<String, byte[]> generatedClasses() {
       return generatedClasses;
     }
+
+    public ImmutableMap<String, byte[]> statistics() {
+      return statistics;
+    }
   }
 }
diff --git a/java/com/google/turbine/binder/Processing.java b/java/com/google/turbine/binder/Processing.java
index c2ccb31..89925de 100644
--- a/java/com/google/turbine/binder/Processing.java
+++ b/java/com/google/turbine/binder/Processing.java
@@ -39,9 +39,12 @@ import com.google.turbine.diag.SourceFile;
 import com.google.turbine.diag.TurbineLog;
 import com.google.turbine.parse.Parser;
 import com.google.turbine.processing.ModelFactory;
+import com.google.turbine.processing.TurbineElements;
 import com.google.turbine.processing.TurbineFiler;
+import com.google.turbine.processing.TurbineMessager;
 import com.google.turbine.processing.TurbineProcessingEnvironment;
 import com.google.turbine.processing.TurbineRoundEnvironment;
+import com.google.turbine.processing.TurbineTypes;
 import com.google.turbine.tree.Tree.CompUnit;
 import com.google.turbine.type.AnnoInfo;
 import java.net.MalformedURLException;
@@ -106,15 +109,22 @@ public class Processing {
         CompoundEnv.<ClassSymbol, TypeBoundClass>of(result.classPathEnv()).append(tenv);
     ModelFactory factory = new ModelFactory(env, processorInfo.loader(), result.tli());
 
+    Map<String, byte[]> statistics = new LinkedHashMap<>();
+
+    TurbineTypes turbineTypes = new TurbineTypes(factory);
+    TurbineProcessingEnvironment processingEnv =
+        new TurbineProcessingEnvironment(
+            filer,
+            turbineTypes,
+            new TurbineElements(factory, turbineTypes),
+            new TurbineMessager(factory, log),
+            processorInfo.options(),
+            processorInfo.sourceVersion(),
+            processorInfo.loader(),
+            statistics);
     for (Processor processor : processorInfo.processors()) {
-      processor.init(
-          new TurbineProcessingEnvironment(
-              factory,
-              filer,
-              log,
-              processorInfo.options(),
-              processorInfo.sourceVersion(),
-              processorInfo.loader()));
+
+      processor.init(processingEnv);
     }
 
     Map<Processor, Pattern> wanted = new HashMap<>();
@@ -245,8 +255,22 @@ public class Processing {
               result.classPathEnv(),
               result.tli(),
               result.generatedSources(),
-              filer.generatedClasses());
+              filer.generatedClasses(),
+              /* statistics= */ ImmutableMap.of());
+    }
+
+    if (!statistics.isEmpty()) {
+      result =
+          new BindingResult(
+              result.units(),
+              result.modules(),
+              result.classPathEnv(),
+              result.tli(),
+              result.generatedSources(),
+              result.generatedClasses(),
+              ImmutableMap.copyOf(statistics));
     }
+
     return result;
   }
 
diff --git a/java/com/google/turbine/main/Main.java b/java/com/google/turbine/main/Main.java
index 874b915..97b6247 100644
--- a/java/com/google/turbine/main/Main.java
+++ b/java/com/google/turbine/main/Main.java
@@ -21,6 +21,7 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.hash.Hashing;
 import com.google.common.io.MoreFiles;
 import com.google.turbine.binder.Binder;
@@ -106,12 +107,18 @@ public class Main {
      */
     public abstract int reducedClasspathLength();
 
+    public abstract ImmutableMap<String, byte[]> processorStatistics();
+
     static Result create(
         boolean transitiveClasspathFallback,
         int transitiveClasspathLength,
-        int reducedClasspathLength) {
+        int reducedClasspathLength,
+        ImmutableMap<String, byte[]> processorStatistics) {
       return new AutoValue_Main_Result(
-          transitiveClasspathFallback, transitiveClasspathLength, reducedClasspathLength);
+          transitiveClasspathFallback,
+          transitiveClasspathLength,
+          reducedClasspathLength,
+          processorStatistics);
     }
   }
 
@@ -167,7 +174,8 @@ public class Main {
           return Result.create(
               /* transitiveClasspathFallback= */ true,
               /* transitiveClasspathLength= */ transitiveClasspathLength,
-              /* reducedClasspathLength= */ reducedClasspathLength);
+              /* reducedClasspathLength= */ reducedClasspathLength,
+              /* processorStatistics= */ ImmutableMap.of());
         }
         break;
       default:
@@ -193,7 +201,8 @@ public class Main {
     return Result.create(
         /* transitiveClasspathFallback= */ transitiveClasspathFallback,
         /* transitiveClasspathLength= */ transitiveClasspathLength,
-        /* reducedClasspathLength= */ reducedClasspathLength);
+        /* reducedClasspathLength= */ reducedClasspathLength,
+        /* processorStatistics= */ bound.statistics());
   }
 
   /**
diff --git a/java/com/google/turbine/processing/TurbineMessager.java b/java/com/google/turbine/processing/TurbineMessager.java
index ab18a28..9c333b2 100644
--- a/java/com/google/turbine/processing/TurbineMessager.java
+++ b/java/com/google/turbine/processing/TurbineMessager.java
@@ -45,7 +45,7 @@ import javax.tools.Diagnostic;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Turbine's implementation of {@link Messager}. */
-class TurbineMessager implements Messager {
+public class TurbineMessager implements Messager {
   private final ModelFactory factory;
   private final TurbineLog log;
 
diff --git a/java/com/google/turbine/processing/TurbineProcessingEnvironment.java b/java/com/google/turbine/processing/TurbineProcessingEnvironment.java
index f92909d..726d075 100644
--- a/java/com/google/turbine/processing/TurbineProcessingEnvironment.java
+++ b/java/com/google/turbine/processing/TurbineProcessingEnvironment.java
@@ -16,39 +16,44 @@
 
 package com.google.turbine.processing;
 
-import com.google.turbine.diag.TurbineLog;
 import java.util.Locale;
 import java.util.Map;
 import javax.annotation.processing.Filer;
 import javax.annotation.processing.Messager;
 import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.SourceVersion;
+import javax.lang.model.util.Elements;
+import javax.lang.model.util.Types;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Turbine's {@link ProcessingEnvironment). */
 public class TurbineProcessingEnvironment implements ProcessingEnvironment {
 
   private final Filer filer;
-  private final TurbineTypes turbineTypes;
+  private final Types types;
   private final Map<String, String> processorOptions;
-  private final TurbineElements turbineElements;
+  private final Elements elements;
+  private final Map<String, byte[]> statistics;
   private final SourceVersion sourceVersion;
   private final Messager messager;
   private final ClassLoader processorLoader;
 
   public TurbineProcessingEnvironment(
-      ModelFactory factory,
       Filer filer,
-      TurbineLog log,
+      Types types,
+      Elements elements,
+      Messager messager,
       Map<String, String> processorOptions,
       SourceVersion sourceVersion,
-      @Nullable ClassLoader processorLoader) {
+      @Nullable ClassLoader processorLoader,
+      Map<String, byte[]> statistics) {
     this.filer = filer;
-    this.turbineTypes = new TurbineTypes(factory);
+    this.types = types;
     this.processorOptions = processorOptions;
     this.sourceVersion = sourceVersion;
-    this.turbineElements = new TurbineElements(factory, turbineTypes);
-    this.messager = new TurbineMessager(factory, log);
+    this.elements = elements;
+    this.statistics = statistics;
+    this.messager = messager;
     this.processorLoader = processorLoader;
   }
 
@@ -68,13 +73,13 @@ public class TurbineProcessingEnvironment implements ProcessingEnvironment {
   }
 
   @Override
-  public TurbineElements getElementUtils() {
-    return turbineElements;
+  public Elements getElementUtils() {
+    return elements;
   }
 
   @Override
-  public TurbineTypes getTypeUtils() {
-    return turbineTypes;
+  public Types getTypeUtils() {
+    return types;
   }
 
   @Override
@@ -90,4 +95,11 @@ public class TurbineProcessingEnvironment implements ProcessingEnvironment {
   public ClassLoader processorLoader() {
     return processorLoader;
   }
+
+  public void addStatistics(String key, byte[] extension) {
+    byte[] existing = statistics.put(key, extension);
+    if (existing != null) {
+      throw new IllegalStateException("duplicate statistics reported for " + key);
+    }
+  }
 }