Allow JavacPluginProcessingEnvironment to be created with Elements and Types.

This allows the JavacPluginProcessingEnvironment to be shared with KeyFactoryTest.

This CL also reorganize the javac plugin related classes to only expose JavacPluginModule and makes other classes private.

RELNOTES=N/A
PiperOrigin-RevId: 404025485
diff --git a/java/dagger/internal/codegen/compileroption/JavacPluginCompilerOptions.java b/java/dagger/internal/codegen/javac/JavacPluginCompilerOptions.java
similarity index 91%
rename from java/dagger/internal/codegen/compileroption/JavacPluginCompilerOptions.java
rename to java/dagger/internal/codegen/javac/JavacPluginCompilerOptions.java
index a86cc1b..d3f5262 100644
--- a/java/dagger/internal/codegen/compileroption/JavacPluginCompilerOptions.java
+++ b/java/dagger/internal/codegen/javac/JavacPluginCompilerOptions.java
@@ -14,17 +14,19 @@
  * limitations under the License.
  */
 
-package dagger.internal.codegen.compileroption;
+package dagger.internal.codegen.javac;
 
 import static dagger.internal.codegen.compileroption.ValidationType.NONE;
 import static javax.tools.Diagnostic.Kind.NOTE;
 
+import dagger.internal.codegen.compileroption.CompilerOptions;
+import dagger.internal.codegen.compileroption.ValidationType;
 import javax.inject.Inject;
 import javax.lang.model.element.TypeElement;
 import javax.tools.Diagnostic;
 
 /** {@link CompilerOptions} for Javac plugins (e.g. for Dagger statistics or Kythe). */
-public final class JavacPluginCompilerOptions extends CompilerOptions {
+final class JavacPluginCompilerOptions extends CompilerOptions {
 
   @Inject
   JavacPluginCompilerOptions() {}
diff --git a/java/dagger/internal/codegen/javac/JavacPluginModule.java b/java/dagger/internal/codegen/javac/JavacPluginModule.java
index f292a9b..aca9f23 100644
--- a/java/dagger/internal/codegen/javac/JavacPluginModule.java
+++ b/java/dagger/internal/codegen/javac/JavacPluginModule.java
@@ -19,6 +19,8 @@
 import androidx.room.compiler.processing.XMessager;
 import androidx.room.compiler.processing.XProcessingEnv;
 import androidx.room.compiler.processing.compat.XConverters;
+import com.sun.tools.javac.model.JavacElements;
+import com.sun.tools.javac.model.JavacTypes;
 import com.sun.tools.javac.util.Context;
 import dagger.Binds;
 import dagger.Module;
@@ -26,44 +28,54 @@
 import dagger.internal.codegen.binding.BindingGraphFactory;
 import dagger.internal.codegen.binding.ComponentDescriptorFactory;
 import dagger.internal.codegen.compileroption.CompilerOptions;
-import dagger.internal.codegen.compileroption.JavacPluginCompilerOptions;
 import dagger.internal.codegen.langmodel.DaggerElements;
 import dagger.internal.codegen.langmodel.DaggerTypes;
 import javax.annotation.processing.ProcessingEnvironment;
-import javax.inject.Singleton;
+import javax.lang.model.util.Elements; // ALLOW_TYPES_ELEMENTS
+import javax.lang.model.util.Types; // ALLOW_TYPES_ELEMENTS
 
 /**
  * A module that provides a {@link BindingGraphFactory} and {@link ComponentDescriptorFactory} for
  * use in {@code javac} plugins. Requires a binding for the {@code javac} {@link Context}.
  */
-@Module
-public abstract class JavacPluginModule {
-  @Binds
-  abstract CompilerOptions compilerOptions(JavacPluginCompilerOptions compilerOptions);
+@Module(includes = JavacPluginModule.BindsModule.class)
+public final class JavacPluginModule {
+  @Module
+  interface BindsModule {
+    @Binds
+    CompilerOptions compilerOptions(JavacPluginCompilerOptions compilerOptions);
+  }
 
-  @Provides
-  static XMessager messager() {
-    return XConverters.toXProcessing(JavacPlugins.getNullMessager());
+  private final XProcessingEnv processingEnv;
+
+  public JavacPluginModule(Context context) {
+    this(JavacElements.instance(context), JavacTypes.instance(context));
+  }
+
+  public JavacPluginModule(Elements elements, Types types) {
+    this.processingEnv =
+        XProcessingEnv.create(new JavacPluginProcessingEnvironment(elements, types));
   }
 
   @Provides
-  static DaggerElements daggerElements(XProcessingEnv xProcessingEnv) {
-    ProcessingEnvironment env = XConverters.toJavac(xProcessingEnv);
-    return new DaggerElements(
-        env.getElementUtils(), env.getTypeUtils());  // ALLOW_TYPES_ELEMENTS
+  XMessager messager() {
+    return processingEnv.getMessager();
   }
 
   @Provides
-  static DaggerTypes daggerTypes(XProcessingEnv xProcessingEnv, DaggerElements elements) {
-    return new DaggerTypes(
-        XConverters.toJavac(xProcessingEnv).getTypeUtils(), elements);  // ALLOW_TYPES_ELEMENTS
+  DaggerElements daggerElements() {
+    ProcessingEnvironment env = XConverters.toJavac(processingEnv);
+    return new DaggerElements(env.getElementUtils(), env.getTypeUtils()); // ALLOW_TYPES_ELEMENTS
   }
 
   @Provides
-  @Singleton
-  static XProcessingEnv xProcessingEnv(Context javaContext) {
-    return JavacPlugins.getXProcessingEnv(javaContext);
+  DaggerTypes daggerTypes(DaggerElements elements) {
+    ProcessingEnvironment env = XConverters.toJavac(processingEnv);
+    return new DaggerTypes(env.getTypeUtils(), elements); // ALLOW_TYPES_ELEMENTS
   }
 
-  private JavacPluginModule() {}
+  @Provides
+  XProcessingEnv xProcessingEnv() {
+    return processingEnv;
+  }
 }
diff --git a/java/dagger/internal/codegen/javac/JavacPluginProcessingEnvironment.java b/java/dagger/internal/codegen/javac/JavacPluginProcessingEnvironment.java
new file mode 100644
index 0000000..3f6a575
--- /dev/null
+++ b/java/dagger/internal/codegen/javac/JavacPluginProcessingEnvironment.java
@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2021 The Dagger Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package dagger.internal.codegen.javac;
+
+import com.google.common.collect.ImmutableMap;
+import java.util.Locale;
+import javax.annotation.processing.Filer;
+import javax.annotation.processing.Messager;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.SourceVersion;
+import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.AnnotationValue;
+import javax.lang.model.element.Element;
+import javax.lang.model.util.Elements;  // ALLOW_TYPES_ELEMENTS since in interface API
+import javax.lang.model.util.Types;  // ALLOW_TYPES_ELEMENTS since in interface API
+import javax.tools.Diagnostic;
+import javax.tools.FileObject;
+import javax.tools.JavaFileManager.Location;
+import javax.tools.JavaFileObject;
+
+/**
+ * An implementation of {@link ProcessingEnvironment} that runs in a javac plugin environment.
+ *
+ * <p>This environment runs after the classes are already compiled, so parts of the {@link
+ * ProcessingEnvironment} API like {@link Filer}, {@link Messager} don't make sense in this
+ * environment, so they've been replaced with throwing and no-op implementations respectively.
+ */
+final class JavacPluginProcessingEnvironment implements ProcessingEnvironment {
+  private final Elements elements;
+  private final Types types;
+  private final Filer filer = new ThrowingFiler();
+  private final Messager messager = new NoopMessager();
+
+  JavacPluginProcessingEnvironment(Elements elements, Types types) {
+    this.elements = elements;
+    this.types = types;
+  }
+
+  @Override
+  public Elements getElementUtils() {
+    return elements;
+  }
+
+  @Override
+  public Types getTypeUtils() {
+    return types;
+  }
+
+  @Override
+  public Filer getFiler() {
+    return filer;
+  }
+
+  @Override
+  public Locale getLocale() {
+    // Null means there's no locale in effect
+    return null;
+  }
+
+  @Override
+  public Messager getMessager() {
+    return messager;
+  }
+
+  @Override
+  public ImmutableMap<String, String> getOptions() {
+    // TODO(erichang): You can technically parse options out of the context, but it is internal
+    // implementation and unclear that any of the tools will ever be passing an option.
+    return ImmutableMap.of();
+  }
+
+  @Override
+  public SourceVersion getSourceVersion() {
+    // This source version doesn't really matter because it is saying what version generated code
+    // should adhere to, which there shouldn't be any because the Filer doesn't work.
+    return SourceVersion.latestSupported();
+  }
+
+  private static final class ThrowingFiler implements Filer {
+    @Override
+    public JavaFileObject createClassFile(CharSequence name, Element... originatingElements) {
+      throw new UnsupportedOperationException("Cannot use a Filer in this context");
+    }
+
+    @Override
+    public FileObject createResource(
+        Location location,
+        CharSequence pkg,
+        CharSequence relativeName,
+        Element... originatingElements) {
+      throw new UnsupportedOperationException("Cannot use a Filer in this context");
+    }
+
+    @Override
+    public JavaFileObject createSourceFile(CharSequence name, Element... originatingElements) {
+      throw new UnsupportedOperationException("Cannot use a Filer in this context");
+    }
+
+    @Override
+    public FileObject getResource(Location location, CharSequence pkg, CharSequence relativeName) {
+      throw new UnsupportedOperationException("Cannot use a Filer in this context");
+    }
+  }
+
+  private static final class NoopMessager implements Messager {
+    @Override
+    public void printMessage(Diagnostic.Kind kind, CharSequence charSequence) {}
+
+    @Override
+    public void printMessage(Diagnostic.Kind kind, CharSequence charSequence, Element element) {}
+
+    @Override
+    public void printMessage(
+        Diagnostic.Kind kind,
+        CharSequence charSequence,
+        Element element,
+        AnnotationMirror annotationMirror) {}
+
+    @Override
+    public void printMessage(
+        Diagnostic.Kind kind,
+        CharSequence charSequence,
+        Element element,
+        AnnotationMirror annotationMirror,
+        AnnotationValue annotationValue) {}
+  }
+}
diff --git a/java/dagger/internal/codegen/javac/JavacPlugins.java b/java/dagger/internal/codegen/javac/JavacPlugins.java
deleted file mode 100644
index 7cd7a27..0000000
--- a/java/dagger/internal/codegen/javac/JavacPlugins.java
+++ /dev/null
@@ -1,140 +0,0 @@
-/*
- * Copyright (C) 2021 The Dagger Authors.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package dagger.internal.codegen.javac;
-
-import androidx.room.compiler.processing.XProcessingEnv;
-import com.google.common.collect.ImmutableMap;
-import com.sun.tools.javac.model.JavacElements;
-import com.sun.tools.javac.model.JavacTypes;
-import com.sun.tools.javac.util.Context;
-import java.util.Locale;
-import javax.annotation.processing.Filer;
-import javax.annotation.processing.Messager;
-import javax.annotation.processing.ProcessingEnvironment;
-import javax.lang.model.SourceVersion;
-import javax.lang.model.element.AnnotationMirror;
-import javax.lang.model.element.AnnotationValue;
-import javax.lang.model.element.Element;
-import javax.lang.model.util.Elements;  // ALLOW_TYPES_ELEMENTS since in interface API
-import javax.lang.model.util.Types;  // ALLOW_TYPES_ELEMENTS since in interface API
-import javax.tools.Diagnostic;
-import javax.tools.FileObject;
-import javax.tools.JavaFileManager.Location;
-import javax.tools.JavaFileObject;
-
-/** Helper methods for working use in javac plugins. */
-public final class JavacPlugins {
-
-  /**
-   * Gets an {@link XProcessingEnv} from a Javac {@link Context}. This processing environment is
-   * limited in that it uses a no-op Filer and Messager.
-   */
-  public static XProcessingEnv getXProcessingEnv(Context context) {
-    return XProcessingEnv.create(new JavacPluginProcessingEnvironment(context));
-  }
-
-  private static final class JavacPluginProcessingEnvironment implements ProcessingEnvironment {
-    private final Context context;
-
-    private JavacPluginProcessingEnvironment(Context context) {
-      this.context = context;
-    }
-
-    @Override public Elements getElementUtils() {
-      return JavacElements.instance(context);
-    }
-
-    @Override public Filer getFiler() {
-      return new Filer() {
-        @Override public JavaFileObject createClassFile(
-            CharSequence name, Element... originatingElements) {
-          throw new UnsupportedOperationException("Cannot use a Filer in this context");
-        }
-
-        @Override public FileObject createResource(
-            Location location,
-            CharSequence pkg,
-            CharSequence relativeName,
-            Element... originatingElements) {
-          throw new UnsupportedOperationException("Cannot use a Filer in this context");
-        }
-
-        @Override public JavaFileObject createSourceFile(
-            CharSequence name, Element... originatingElements) {
-          throw new UnsupportedOperationException("Cannot use a Filer in this context");
-        }
-
-        @Override public FileObject getResource(
-            Location location, CharSequence pkg, CharSequence relativeName) {
-          throw new UnsupportedOperationException("Cannot use a Filer in this context");
-        }
-      };
-    }
-
-    @Override public Locale getLocale() {
-      // Null means there's no locale in effect
-      return null;
-    }
-
-    @Override public Messager getMessager() {
-      return getNullMessager();
-    }
-
-    @Override public ImmutableMap<String, String> getOptions() {
-      // TODO(erichang): You can technically parse options out of the context, but it is internal
-      // implementation and unclear that any of the tools will ever be passing an option.
-      return ImmutableMap.of();
-    }
-
-    @Override public SourceVersion getSourceVersion() {
-      // This source version doesn't really matter because it is saying what version generated code
-      // should adhere to, which there shouldn't be any because the Filer doesn't work.
-      return SourceVersion.latestSupported();
-    }
-
-    @Override public Types getTypeUtils() {
-      return JavacTypes.instance(context);
-    }
-  }
-
-  public static Messager getNullMessager() {
-    return new Messager() {
-      @Override
-      public void printMessage(Diagnostic.Kind kind, CharSequence charSequence) {}
-
-      @Override
-      public void printMessage(Diagnostic.Kind kind, CharSequence charSequence, Element element) {}
-
-      @Override
-      public void printMessage(
-          Diagnostic.Kind kind,
-          CharSequence charSequence,
-          Element element,
-          AnnotationMirror annotationMirror) {}
-
-      @Override
-      public void printMessage(
-          Diagnostic.Kind kind,
-          CharSequence charSequence,
-          Element element,
-          AnnotationMirror annotationMirror,
-          AnnotationValue annotationValue) {}
-    };
-  }
-
-  private JavacPlugins() {}
-}
diff --git a/java/dagger/internal/codegen/kythe/DaggerKythePlugin.java b/java/dagger/internal/codegen/kythe/DaggerKythePlugin.java
index 424363d..d15f7d4 100644
--- a/java/dagger/internal/codegen/kythe/DaggerKythePlugin.java
+++ b/java/dagger/internal/codegen/kythe/DaggerKythePlugin.java
@@ -21,7 +21,6 @@
 
 import static dagger.internal.codegen.langmodel.DaggerElements.isAnyAnnotationPresent;
 
-
 import androidx.room.compiler.processing.XProcessingEnv;
 import androidx.room.compiler.processing.compat.XConverters;
 import com.google.auto.service.AutoService;
@@ -34,8 +33,6 @@
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.tree.JCTree.JCClassDecl;
 import com.sun.tools.javac.tree.JCTree.JCCompilationUnit;
-import com.sun.tools.javac.util.Context;
-import dagger.BindsInstance;
 import dagger.Component;
 import dagger.internal.codegen.binding.Binding;
 import dagger.internal.codegen.binding.BindingDeclaration;
@@ -180,7 +177,7 @@
     if (bindingGraphFactory == null) {
       emitter = entrySets.getEmitter();
       DaggerDaggerKythePlugin_PluginComponent.builder()
-          .context(kytheGraph.getJavaContext())
+          .javacPluginModule(new JavacPluginModule(kytheGraph.getJavaContext()))
           .build()
           .inject(this);
     }
@@ -191,13 +188,5 @@
   @Component(modules = {InjectBindingRegistryModule.class, JavacPluginModule.class})
   interface PluginComponent {
     void inject(DaggerKythePlugin plugin);
-
-    @Component.Builder
-    interface Builder {
-      @BindsInstance
-      Builder context(Context context);
-
-      PluginComponent build();
-    }
   }
 }
diff --git a/javatests/dagger/internal/codegen/BUILD b/javatests/dagger/internal/codegen/BUILD
index 74f1f74..bcfa824 100644
--- a/javatests/dagger/internal/codegen/BUILD
+++ b/javatests/dagger/internal/codegen/BUILD
@@ -75,6 +75,7 @@
         "//java/dagger/internal/codegen/binding",
         "//java/dagger/internal/codegen/bindinggraphvalidation",
         "//java/dagger/internal/codegen/compileroption",
+        "//java/dagger/internal/codegen/javac",
         "//java/dagger/internal/codegen/javapoet",
         "//java/dagger/internal/codegen/kotlin",
         "//java/dagger/internal/codegen/langmodel",
diff --git a/javatests/dagger/internal/codegen/KeyFactoryTest.java b/javatests/dagger/internal/codegen/KeyFactoryTest.java
index 03dd1d4..6bb4c63 100644
--- a/javatests/dagger/internal/codegen/KeyFactoryTest.java
+++ b/javatests/dagger/internal/codegen/KeyFactoryTest.java
@@ -26,11 +26,11 @@
 import com.google.common.collect.Iterables;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.testing.compile.CompilationRule;
-import dagger.BindsInstance;
 import dagger.Component;
 import dagger.Module;
 import dagger.Provides;
 import dagger.internal.codegen.binding.KeyFactory;
+import dagger.internal.codegen.javac.JavacPluginModule;
 import dagger.internal.codegen.langmodel.DaggerElements;
 import dagger.internal.codegen.langmodel.DaggerTypes;
 import dagger.multibindings.ElementsIntoSet;
@@ -68,7 +68,11 @@
   @Inject KeyFactory keyFactory;
 
   @Before public void setUp() {
-    DaggerKeyFactoryTest_TestComponent.factory().create(compilationRule).inject(this);
+    DaggerKeyFactoryTest_TestComponent.builder()
+        .javacPluginModule(
+            new JavacPluginModule(compilationRule.getElements(), compilationRule.getTypes()))
+        .build()
+        .inject(this);
   }
 
   @Test public void forInjectConstructorWithResolvedType() {
@@ -333,26 +337,8 @@
   }
 
   @Singleton
-  @Component(modules = {TestModule.class})
+  @Component(modules = JavacPluginModule.class)
   interface TestComponent {
     void inject(KeyFactoryTest test);
-
-    @Component.Factory
-    interface Factory {
-      TestComponent create(@BindsInstance CompilationRule compilationRule);
-    }
-  }
-
-  @Module
-  static class TestModule {
-    @Provides
-    static DaggerElements elements(CompilationRule compilationRule) {
-      return new DaggerElements(compilationRule.getElements(), compilationRule.getTypes());
-    }
-
-    @Provides
-    static DaggerTypes types(CompilationRule compilationRule, DaggerElements elements) {
-      return new DaggerTypes(compilationRule.getTypes(), elements);
-    }
   }
 }