Merge pull request #328 from square/ralf/private-code-generators

Only call private code generators once after all other code generators are done

Author: vRallev

compiler/src/main/java/com/squareup/anvil/compiler/codegen/CodeGenerationExtension.kt

@@ -69,6 +69,9 @@ internal class CodeGenerationExtension(
     val anvilModule = RealAnvilModuleDescriptor(module)
     anvilModule.addFiles(files)
 
+    val (privateCodeGenerators, nonPrivateCodeGenerators) =
+      codeGenerators.partition { it is PrivateCodeGenerator }
+
     fun Collection<GeneratedFile>.toKtFile(): Collection<KtFile> {
       return this
         .mapNotNull { (file, content) ->
@@ -95,19 +98,19 @@ internal class CodeGenerationExtension(
           codeGenerator.flush(codeGenDir, anvilModule).toKtFile()
         }
 
-    var newFiles = codeGenerators.generateCode(files)
+    var newFiles = nonPrivateCodeGenerators.generateCode(files)
+
     while (newFiles.isNotEmpty()) {
       // Parse the KtFile for each generated file. Then feed the code generators with the new
       // parsed files until no new files are generated.
-      newFiles = codeGenerators.generateCode(newFiles)
+      newFiles = nonPrivateCodeGenerators.generateCode(newFiles)
     }
 
-    val flushedFiles = codeGenerators.flush()
+    nonPrivateCodeGenerators.flush()
 
-    // The contract is that PrivateCodeGenerators are called one last time after flush().
-    codeGenerators
-      .filterIsInstance<PrivateCodeGenerator>()
-      .generateCode(flushedFiles)
+    // PrivateCodeGenerators don't impact other code generators. Therefore, they can be called a
+    // single time at the end.
+    privateCodeGenerators.generateCode(anvilModule.allFiles)
 
     // This restarts the analysis phase and will include our files.
     return RetryWithAdditionalRoots(

compiler/src/main/java/com/squareup/anvil/compiler/codegen/RealAnvilModuleDescriptor.kt

@@ -15,11 +15,15 @@ class RealAnvilModuleDescriptor(
   delegate: ModuleDescriptor
 ) : AnvilModuleDescriptor, ModuleDescriptor by delegate {
 
+  internal val allFiles = mutableListOf<KtFile>()
+
   private val classesMap = mutableMapOf<String, List<KtClassOrObject>>()
   private val allClasses: Sequence<KtClassOrObject>
     get() = classesMap.values.asSequence().flatMap { it }
 
   fun addFiles(files: Collection<KtFile>) {
+    allFiles += files
+
     files.forEach { ktFile ->
       classesMap[ktFile.identifier] = ktFile.classesAndInnerClasses()
     }