[Ahead-of-time Subcomponents] Invoke code generation when compiling c…

…lasses

annotated as Subcomponents. This is the initial step in generating
separate generated implementation classes for each subcomponent for the purpose
of improving (incremental) build performance.

RELNOTES=n/a

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

java/dagger/internal/codegen/BindingGraphConverter.java

@@ -47,10 +47,13 @@
 final class BindingGraphConverter {
 
   private final BindingDeclarationFormatter bindingDeclarationFormatter;
+  private final CompilerOptions compilerOptions;
 
   @Inject
-  BindingGraphConverter(BindingDeclarationFormatter bindingDeclarationFormatter) {
+  BindingGraphConverter(
+      BindingDeclarationFormatter bindingDeclarationFormatter, CompilerOptions compilerOptions) {
     this.bindingDeclarationFormatter = bindingDeclarationFormatter;
+    this.compilerOptions = compilerOptions;
   }
 
   /**
@@ -72,7 +75,7 @@ private final class Traverser extends ComponentTreeTraverser {
     private ComponentNode currentComponent;
 
     Traverser(BindingGraph graph) {
-      super(graph);
+      super(graph, compilerOptions);
     }
 
     @Override

java/dagger/internal/codegen/BindingGraphFactory.java

@@ -79,22 +79,26 @@ final class BindingGraphFactory {
   private final InjectBindingRegistry injectBindingRegistry;
   private final KeyFactory keyFactory;
   private final BindingFactory bindingFactory;
+  private final CompilerOptions compilerOptions;
 
   @Inject
   BindingGraphFactory(
       DaggerElements elements,
       InjectBindingRegistry injectBindingRegistry,
       KeyFactory keyFactory,
-      BindingFactory bindingFactory) {
+      BindingFactory bindingFactory,
+      CompilerOptions compilerOptions) {
     this.elements = elements;
     this.injectBindingRegistry = injectBindingRegistry;
     this.keyFactory = keyFactory;
     this.bindingFactory = bindingFactory;
+    this.compilerOptions = compilerOptions;
   }
 
   /** Creates a binding graph for a root component. */
   BindingGraph create(ComponentDescriptor componentDescriptor) {
-    checkArgument(componentDescriptor.kind().isTopLevel());
+    checkArgument(
+        componentDescriptor.kind().isTopLevel() || compilerOptions.aheadOfTimeSubcomponents());
     return create(Optional.empty(), componentDescriptor);
   }
 

java/dagger/internal/codegen/ComponentBindingExpressions.java

@@ -84,7 +84,7 @@ final class ComponentBindingExpressions {
         generatedComponentModel,
         subcomponentNames,
         componentRequirementFields,
-        new ReferenceReleasingManagerFields(graph, generatedComponentModel),
+        new ReferenceReleasingManagerFields(graph, generatedComponentModel, compilerOptions),
         new StaticSwitchingProviders(generatedComponentModel, types),
         optionalFactories,
         types,

java/dagger/internal/codegen/ComponentDescriptor.java

@@ -442,30 +442,38 @@ static final class Factory {
     private final Types types;
     private final DependencyRequestFactory dependencyRequestFactory;
     private final ModuleDescriptor.Factory moduleDescriptorFactory;
+    private final CompilerOptions compilerOptions;
 
     @Inject
     Factory(
         DaggerElements elements,
         Types types,
         DependencyRequestFactory dependencyRequestFactory,
-        ModuleDescriptor.Factory moduleDescriptorFactory) {
+        ModuleDescriptor.Factory moduleDescriptorFactory,
+        CompilerOptions compilerOptions) {
       this.elements = elements;
       this.types = types;
       this.dependencyRequestFactory = dependencyRequestFactory;
       this.moduleDescriptorFactory = moduleDescriptorFactory;
+      this.compilerOptions = compilerOptions;
     }
 
     /**
      * Returns a component descriptor for a type annotated with either {@link Component @Component}
-     * or {@link ProductionComponent @ProductionComponent}.
+     * or {@link ProductionComponent @ProductionComponent}. This is also compatible with {@link
+     * Subcomponent @Subcomponent} or {@link ProductionSubcomponent @ProductionSubcomponent} when
+     * generating ahead-of-time subcomponents.
      */
-    ComponentDescriptor forComponent(TypeElement componentDefinitionType) {
-      Optional<Kind> kind = Kind.forAnnotatedElement(componentDefinitionType);
+    ComponentDescriptor forComponent(TypeElement componentType) {
+      Optional<Kind> kind = Kind.forAnnotatedElement(componentType);
       checkArgument(
-          kind.isPresent() && kind.get().isTopLevel(),
-          "%s must be annotated with @Component or @ProductionComponent",
-          componentDefinitionType);
-      return create(componentDefinitionType, kind.get(), Optional.empty());
+          kind.isPresent(), "%s must have a component or subcomponent annotation", componentType);
+      if (!compilerOptions.aheadOfTimeSubcomponents()) {
+        checkArgument(kind.get().isTopLevel(),
+            "%s must be annotated with @Component or @ProductionComponent.",
+            componentType);
+      }
+      return create(componentType, kind.get(), Optional.empty());
     }
 
     private ComponentDescriptor create(

java/dagger/internal/codegen/ComponentModelBuilder.java

@@ -79,7 +79,24 @@ static GeneratedComponentModel buildComponentModel(
             types,
             elements,
             compilerOptions);
-    return new RootComponentModelBuilder(
+    if (graph.componentDescriptor().kind().isTopLevel()) {
+      return new RootComponentModelBuilder(
+              types,
+              elements,
+              graph,
+              generatedComponentModel,
+              subcomponentNames,
+              optionalFactories,
+              bindingExpressions,
+              componentRequirementFields,
+              builder)
+          .build();
+    }
+    checkState(
+        compilerOptions.aheadOfTimeSubcomponents(),
+        "Calling 'buildComponentModel()' on %s when not generating ahead-of-time subcomponents.",
+        graph.componentDescriptor().componentDefinitionType());
+    return new BaseSubcomponentModelBuilder(
             types,
             elements,
             graph,
@@ -153,7 +170,7 @@ private ComponentModelBuilder(
    * called once (and will throw on successive invocations). If the component must be regenerated,
    * use a new instance.
    */
-  protected final GeneratedComponentModel build() {
+  protected GeneratedComponentModel build() {
     checkState(
         !done,
         "ComponentModelBuilder has already built the GeneratedComponentModel for [%s].",
@@ -387,6 +404,46 @@ private DeclaredType parentType() {
     }
   }
 
+  /** Builds the model for a top-level abstract base implementation of a subcomponent. */
+  // TODO(b/72748365): Implement.
+  private static final class BaseSubcomponentModelBuilder extends ComponentModelBuilder {
+    private final GeneratedComponentModel generatedComponentModel;
+
+    BaseSubcomponentModelBuilder(
+        DaggerTypes types,
+        Elements elements,
+        BindingGraph graph,
+        GeneratedComponentModel generatedComponentModel,
+        SubcomponentNames subcomponentNames,
+        OptionalFactories optionalFactories,
+        ComponentBindingExpressions bindingExpressions,
+        ComponentRequirementFields componentRequirementFields,
+        Optional<ComponentBuilder> builder) {
+      super(
+          types,
+          elements,
+          graph,
+          generatedComponentModel,
+          subcomponentNames,
+          optionalFactories,
+          bindingExpressions,
+          componentRequirementFields,
+          builder);
+      this.generatedComponentModel = generatedComponentModel;
+    }
+
+    @Override
+    protected GeneratedComponentModel build() {
+      return generatedComponentModel;
+    }
+
+    @Override
+    protected void addBuilderClass(TypeSpec builder) {}
+
+    @Override
+    protected void addFactoryMethods() {}
+  }
+
   /** Returns the list of {@link ParameterSpec}s for the corresponding graph's factory method. */
   private static ImmutableList<ParameterSpec> getFactoryMethodParameterSpecs(BindingGraph graph) {
     return graph

java/dagger/internal/codegen/ComponentProcessingStep.java

@@ -53,6 +53,7 @@ final class ComponentProcessingStep implements ProcessingStep {
   private final BindingGraphConverter bindingGraphConverter;
   private final BindingGraphPlugins validationPlugins;
   private final BindingGraphPlugins spiPlugins;
+  private final CompilerOptions compilerOptions;
 
   @Inject
   ComponentProcessingStep(
@@ -65,7 +66,8 @@ final class ComponentProcessingStep implements ProcessingStep {
       ComponentGenerator componentGenerator,
       BindingGraphConverter bindingGraphConverter,
       @Validation BindingGraphPlugins validationPlugins,
-      BindingGraphPlugins spiPlugins) {
+      BindingGraphPlugins spiPlugins,
+      CompilerOptions compilerOptions) {
     this.messager = messager;
     this.componentValidator = componentValidator;
     this.builderValidator = builderValidator;
@@ -76,6 +78,7 @@ final class ComponentProcessingStep implements ProcessingStep {
     this.bindingGraphConverter = bindingGraphConverter;
     this.validationPlugins = validationPlugins;
     this.spiPlugins = spiPlugins;
+    this.compilerOptions = compilerOptions;
   }
 
   @Override
@@ -146,6 +149,25 @@ public ImmutableSet<Element> process(
         rejectedElements.add(componentTypeElement);
       }
     }
+
+    if (compilerOptions.aheadOfTimeSubcomponents()) {
+      for (TypeElement subcomponentTypeElement : typesIn(subcomponentElements)) {
+        if (!subcomponentIsClean(
+            subcomponentTypeElement, reportsBySubcomponent, builderReportsBySubcomponent)) {
+          continue;
+        }
+        try {
+          ComponentDescriptor componentDescriptor =
+              componentDescriptorFactory.forComponent(subcomponentTypeElement);
+          BindingGraph bindingGraph = bindingGraphFactory.create(componentDescriptor);
+          // TODO(b/72748365): Do subgraph validation.
+          generateComponent(bindingGraph);
+        } catch (TypeNotPresentException e) {
+          rejectedElements.add(subcomponentTypeElement);
+        }
+      }
+    }
+
     return rejectedElements.build();
   }
 
@@ -212,15 +234,27 @@ private boolean isClean(
       return false;
     }
     for (Element element : report.referencedSubcomponents()) {
-      ValidationReport<?> subcomponentBuilderReport = builderReportsBySubcomponent.get(element);
-      if (subcomponentBuilderReport != null && !subcomponentBuilderReport.isClean()) {
-        return false;
-      }
-      ValidationReport<?> subcomponentReport = reportsBySubcomponent.get(element);
-      if (subcomponentReport != null && !subcomponentReport.isClean()) {
+      if (!subcomponentIsClean(element, reportsBySubcomponent, builderReportsBySubcomponent)) {
         return false;
       }
     }
     return true;
   }
+
+  /** Returns true if the reports associated with the subcomponent are clean. */
+  private boolean subcomponentIsClean(
+      Element subcomponentElement,
+      Map<Element, ValidationReport<TypeElement>> reportsBySubcomponent,
+      Map<Element, ValidationReport<TypeElement>> builderReportsBySubcomponent) {
+    ValidationReport<?> subcomponentBuilderReport =
+        builderReportsBySubcomponent.get(subcomponentElement);
+    if (subcomponentBuilderReport != null && !subcomponentBuilderReport.isClean()) {
+      return false;
+    }
+    ValidationReport<?> subcomponentReport = reportsBySubcomponent.get(subcomponentElement);
+    if (subcomponentReport != null && !subcomponentReport.isClean()) {
+      return false;
+    }
+    return true;
+  }
 }

java/dagger/internal/codegen/ComponentTreeTraverser.java

@@ -74,9 +74,10 @@ public class ComponentTreeTraverser {
   private final Deque<BindingGraph> bindingGraphPath = new ArrayDeque<>();
 
   /** Constructs a traverser for a root (component, not subcomponent) binding graph. */
-  public ComponentTreeTraverser(BindingGraph rootGraph) {
+  public ComponentTreeTraverser(BindingGraph rootGraph, CompilerOptions compilerOptions) {
     checkArgument(
-        rootGraph.componentDescriptor().kind().isTopLevel(),
+        rootGraph.componentDescriptor().kind().isTopLevel()
+            || compilerOptions.aheadOfTimeSubcomponents(),
         "only top-level graphs can be traversed, not %s",
         rootGraph.componentDescriptor().componentDefinitionType().getQualifiedName());
     bindingGraphPath.add(rootGraph);

java/dagger/internal/codegen/DaggerKythePlugin.java

@@ -155,6 +155,7 @@ public void run(
       DaggerDaggerKythePlugin_PluginComponent.builder()
           .types(JavacTypes.instance(javaContext))
           .elements(JavacElements.instance(javaContext))
+          .compilerOptions(KytheBindingGraphFactory.createCompilerOptions())
           .build()
           .inject(this);
     }
@@ -168,8 +169,15 @@ interface PluginComponent {
 
     @Component.Builder
     interface Builder {
-      @BindsInstance Builder types(Types types);
-      @BindsInstance Builder elements(Elements elements);
+      @BindsInstance
+      Builder types(Types types);
+
+      @BindsInstance
+      Builder elements(Elements elements);
+
+      @BindsInstance
+      Builder compilerOptions(CompilerOptions compilerOptions);
+
       PluginComponent build();
     }
   }