Reset processingEnv for BindingGraphPlugins for each round.

Previously, plugin initialization was performed in processingSteps(), which is only called once at the time when the Processor is created. RELNOTES=Added `BindingGraphPlugin#onProcessingRoundBegin` for pre processing initialization. PiperOrigin-RevId: 618910291
google · Mar 25, 2024 · 2a6a0b4 · 2a6a0b4
1 parent fc2363d
commit 2a6a0b4
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 5 deletions.
diff --git a/java/dagger/hilt/android/processor/internal/viewmodel/ViewModelValidationPlugin.kt b/java/dagger/hilt/android/processor/internal/viewmodel/ViewModelValidationPlugin.kt
@@ -48,9 +48,14 @@ import javax.tools.Diagnostic.Kind
 class ViewModelValidationPlugin : BindingGraphPlugin {
 
   private lateinit var env: XProcessingEnv
+  private lateinit var daggerProcessingEnv: DaggerProcessingEnv
 
   override fun init(processingEnv: DaggerProcessingEnv, options: MutableMap<String, String>) {
-    env = processingEnv.toXProcessingEnv()
+    daggerProcessingEnv = processingEnv
+  }
+
+  override fun onProcessingRoundBegin() {
+    env = daggerProcessingEnv.toXProcessingEnv()
   }
 
   override fun visitGraph(bindingGraph: BindingGraph, diagnosticReporter: DiagnosticReporter) {
@@ -75,7 +80,7 @@ class ViewModelValidationPlugin : BindingGraphPlugin {
           "\nInjection of an @HiltViewModel class is prohibited since it does not create a " +
             "ViewModel instance correctly.\nAccess the ViewModel via the Android APIs " +
             "(e.g. ViewModelProvider) instead." +
-            "\nInjected ViewModel: ${target.key().type()}\n"
+            "\nInjected ViewModel: ${target.key().type()}\n",
         )
       } else if (
         isViewModelAssistedFactory(target) && !isInternalViewModelAssistedFactoryUsage(source)
@@ -86,7 +91,7 @@ class ViewModelValidationPlugin : BindingGraphPlugin {
           "\nInjection of an assisted factory for Hilt ViewModel is prohibited since it " +
             "can not be used to create a ViewModel instance correctly.\nAccess the ViewModel via " +
             "the Android APIs (e.g. ViewModelProvider) instead." +
-            "\nInjected factory: ${target.key().type()}\n"
+            "\nInjected factory: ${target.key().type()}\n",
         )
       }
     }

diff --git a/java/dagger/internal/codegen/ComponentProcessor.java b/java/dagger/internal/codegen/ComponentProcessor.java
@@ -119,6 +119,11 @@ public Iterable<XProcessingStep> processingSteps() {
     return delegate.processingSteps();
   }
 
+  @Override
+  public void preRound(XProcessingEnv env, XRoundEnv roundEnv) {
+    delegate.onProcessingRoundBegin();
+  }
+
   @Override
   public void postRound(XProcessingEnv env, XRoundEnv roundEnv) {
     delegate.postRound(env, roundEnv);

diff --git a/java/dagger/internal/codegen/DelegateComponentProcessor.java b/java/dagger/internal/codegen/DelegateComponentProcessor.java
@@ -108,15 +108,19 @@ public void initialize(
     DaggerDelegateComponentProcessor_Injector.factory()
         .create(env, plugins, legacyPlugins)
         .inject(this);
+    validationBindingGraphPlugins.initializePlugins();
+    externalBindingGraphPlugins.initializePlugins();
   }
 
   public Iterable<XProcessingStep> processingSteps() {
-    validationBindingGraphPlugins.initializePlugins();
-    externalBindingGraphPlugins.initializePlugins();
 
     return processingSteps;
   }
 
+  public void onProcessingRoundBegin() {
+    externalBindingGraphPlugins.onProcessingRoundBegin();
+  }
+
   public void postRound(XProcessingEnv env, XRoundEnv roundEnv) {
     if (!roundEnv.isProcessingOver()) {
       try {

diff --git a/java/dagger/internal/codegen/KspComponentProcessor.java b/java/dagger/internal/codegen/KspComponentProcessor.java
@@ -58,6 +58,11 @@ public Iterable<XProcessingStep> processingSteps() {
     return delegate.processingSteps();
   }
 
+  @Override
+  public void preRound(XProcessingEnv env, XRoundEnv roundEnv) {
+    delegate.onProcessingRoundBegin();
+  }
+
   @Override
   public void postRound(XProcessingEnv env, XRoundEnv roundEnv) {
     delegate.postRound(env, roundEnv);

diff --git a/java/dagger/internal/codegen/validation/ExternalBindingGraphPlugins.java b/java/dagger/internal/codegen/validation/ExternalBindingGraphPlugins.java
@@ -84,6 +84,10 @@ private void initializePlugin(BindingGraphPlugin plugin) {
     plugin.init(SpiModelBindingGraphConverter.toSpiModel(processingEnv), filteredOptions);
   }
 
+  public void onProcessingRoundBegin() {
+    plugins.forEach(BindingGraphPlugin::onProcessingRoundBegin);
+  }
+
   private void initializeLegacyPlugin(dagger.spi.BindingGraphPlugin plugin) {
     plugin.initFiler(toJavac(filer));
     plugin.initTypes(toJavac(processingEnv).getTypeUtils()); // ALLOW_TYPES_ELEMENTS

diff --git a/java/dagger/spi/model/BindingGraphPlugin.java b/java/dagger/spi/model/BindingGraphPlugin.java
@@ -59,6 +59,14 @@ default String pluginName() {
     return getClass().getCanonicalName();
   }
 
+  /**
+   * Runs before each round of Dagger annotation processing.
+   *
+   * <p>If using the plugin to process elements that need resetting at the beginning of each
+   * processing round, use this function to perform the setup.
+   */
+  default void onProcessingRoundBegin() {}
+
   /**
    * Perform any extra work after the plugin finished all its visiting. This will be called once per
    * instance of this plugin, after all graphs were {@linkplain #visitGraph(BindingGraph,