Cache getLocalAndInheritedMethods invocations in Dagger.

This calculation (in particular, override checking) is expensive for large component interfaces. I've migrated other callers as well for consistency, even though they're not performed in hotspots.

This saves almost [1s per build[(https://pprof.corp.google.com/user-profile?id=984561c1b75e48a2b80e610430171bc3&id0=123e032498e52ee0bf5df2a0ec7f5729&tab=flame&path=02pxzzs&showfrom=getLocalAndInherited) for GMM

RELNOTES=Build performance improvements
PiperOrigin-RevId: 360710395

java/dagger/internal/codegen/AssistedFactoryProcessingStep.java

@@ -137,7 +137,7 @@ ValidationReport<TypeElement> validate(TypeElement factory) {
       }
 
       ImmutableSet<ExecutableElement> abstractFactoryMethods =
-          AssistedInjectionAnnotations.assistedFactoryMethods(factory, elements, types);
+          AssistedInjectionAnnotations.assistedFactoryMethods(factory, elements);
 
       if (abstractFactoryMethods.isEmpty()) {
         report.addError(

java/dagger/internal/codegen/AssistedProcessingStep.java

@@ -27,7 +27,6 @@
 import dagger.internal.codegen.binding.InjectionAnnotations;
 import dagger.internal.codegen.kotlin.KotlinMetadataUtil;
 import dagger.internal.codegen.langmodel.DaggerElements;
-import dagger.internal.codegen.langmodel.DaggerTypes;
 import dagger.internal.codegen.validation.TypeCheckingProcessingStep;
 import dagger.internal.codegen.validation.ValidationReport;
 import java.lang.annotation.Annotation;
@@ -47,21 +46,18 @@ final class AssistedProcessingStep extends TypeCheckingProcessingStep<VariableEl
   private final KotlinMetadataUtil kotlinMetadataUtil;
   private final InjectionAnnotations injectionAnnotations;
   private final DaggerElements elements;
-  private final DaggerTypes types;
   private final Messager messager;
 
   @Inject
   AssistedProcessingStep(
       KotlinMetadataUtil kotlinMetadataUtil,
       InjectionAnnotations injectionAnnotations,
       DaggerElements elements,
-      DaggerTypes types,
       Messager messager) {
     super(MoreElements::asVariable);
     this.kotlinMetadataUtil = kotlinMetadataUtil;
     this.injectionAnnotations = injectionAnnotations;
     this.elements = elements;
-    this.types = types;
     this.messager = messager;
   }
 
@@ -113,7 +109,7 @@ private boolean isAssistedFactoryCreateMethod(Element element) {
       TypeElement enclosingElement = closestEnclosingTypeElement(element);
       return AssistedInjectionAnnotations.isAssistedFactoryType(enclosingElement)
           // This assumes we've already validated AssistedFactory and that a valid method exists.
-          && AssistedInjectionAnnotations.assistedFactoryMethod(enclosingElement, elements, types)
+          && AssistedInjectionAnnotations.assistedFactoryMethod(enclosingElement, elements)
               .equals(element);
     }
     return false;

java/dagger/internal/codegen/ProcessingEnvironmentModule.java

@@ -22,15 +22,18 @@
 import dagger.Provides;
 import dagger.Reusable;
 import dagger.internal.codegen.SpiModule.ProcessorClassLoader;
+import dagger.internal.codegen.base.ClearableCache;
 import dagger.internal.codegen.compileroption.CompilerOptions;
 import dagger.internal.codegen.compileroption.ProcessingEnvironmentCompilerOptions;
 import dagger.internal.codegen.compileroption.ProcessingOptions;
 import dagger.internal.codegen.langmodel.DaggerElements;
+import dagger.multibindings.IntoSet;
 import dagger.spi.BindingGraphPlugin;
 import java.util.Map;
 import javax.annotation.processing.Filer;
 import javax.annotation.processing.Messager;
 import javax.annotation.processing.ProcessingEnvironment;
+import javax.inject.Singleton;
 import javax.lang.model.SourceVersion;
 import javax.lang.model.util.Types;
 
@@ -73,10 +76,15 @@ static SourceVersion sourceVersion(ProcessingEnvironment processingEnvironment)
   }
 
   @Provides
+  @Singleton
   static DaggerElements daggerElements(ProcessingEnvironment processingEnvironment) {
     return new DaggerElements(processingEnvironment);
   }
 
+  @Binds
+  @IntoSet
+  ClearableCache daggerElementAsClearableCache(DaggerElements elements);
+
   @Provides
   @ProcessorClassLoader
   static ClassLoader processorClassloader(ProcessingEnvironment processingEnvironment) {

java/dagger/internal/codegen/binding/AssistedInjectionAnnotations.java

@@ -59,14 +59,14 @@
 public final class AssistedInjectionAnnotations {
   /** Returns the factory method for the given factory {@link TypeElement}. */
   public static ExecutableElement assistedFactoryMethod(
-      TypeElement factory, DaggerElements elements, DaggerTypes types) {
-    return getOnlyElement(assistedFactoryMethods(factory, elements, types));
+      TypeElement factory, DaggerElements elements) {
+    return getOnlyElement(assistedFactoryMethods(factory, elements));
   }
 
   /** Returns the list of abstract factory methods for the given factory {@link TypeElement}. */
   public static ImmutableSet<ExecutableElement> assistedFactoryMethods(
-      TypeElement factory, DaggerElements elements, DaggerTypes types) {
-    return MoreElements.getLocalAndInheritedMethods(factory, types, elements).stream()
+      TypeElement factory, DaggerElements elements) {
+    return elements.getLocalAndInheritedMethods(factory).stream()
         .filter(method -> method.getModifiers().contains(ABSTRACT))
         .filter(method -> !method.isDefault())
         .collect(toImmutableSet());
@@ -170,7 +170,7 @@ public static AssistedFactoryMetadata create(
         TypeMirror factory, DaggerElements elements, DaggerTypes types) {
       DeclaredType factoryType = asDeclared(factory);
       TypeElement factoryElement = asTypeElement(factoryType);
-      ExecutableElement factoryMethod = assistedFactoryMethod(factoryElement, elements, types);
+      ExecutableElement factoryMethod = assistedFactoryMethod(factoryElement, elements);
       ExecutableType factoryMethodType = asExecutable(types.asMemberOf(factoryType, factoryMethod));
       DeclaredType assistedInjectType = asDeclared(factoryMethodType.getReturnType());
       return new AutoValue_AssistedInjectionAnnotations_AssistedFactoryMetadata(

java/dagger/internal/codegen/binding/BindingFactory.java

@@ -191,7 +191,7 @@ public ProvisionBinding assistedFactoryBinding(
     }
 
     ExecutableElement factoryMethod =
-        AssistedInjectionAnnotations.assistedFactoryMethod(factory, elements, types);
+        AssistedInjectionAnnotations.assistedFactoryMethod(factory, elements);
     ExecutableType factoryMethodType =
         MoreTypes.asExecutable(types.asMemberOf(factoryType, factoryMethod));
     return ProvisionBinding.builder()

java/dagger/internal/codegen/binding/BindsTypeChecker.java

@@ -18,7 +18,6 @@
 
 import static com.google.common.collect.Iterables.getOnlyElement;
 
-import com.google.auto.common.MoreElements;
 import com.google.auto.common.MoreTypes;
 import com.google.common.collect.ImmutableList;
 import dagger.internal.codegen.base.ContributionType;
@@ -83,8 +82,8 @@ private ImmutableList<TypeMirror> methodParameterTypes(DeclaredType type, String
         // type.asElement().getEnclosedElements() is not used because some non-standard JDKs (e.g.
         // J2CL) don't redefine Set.add() (whose only purpose of being redefined in the standard JDK
         // is documentation, and J2CL's implementation doesn't declare docs for JDK types).
-        // MoreElements.getLocalAndInheritedMethods ensures that the method will always be present.
-        MoreElements.getLocalAndInheritedMethods(MoreTypes.asTypeElement(type), types, elements)) {
+        // getLocalAndInheritedMethods ensures that the method will always be present.
+        elements.getLocalAndInheritedMethods(MoreTypes.asTypeElement(type))) {
       if (method.getSimpleName().contentEquals(methodName)) {
         methodsForName.add(method);
       }

java/dagger/internal/codegen/binding/ComponentRequirement.java

@@ -16,7 +16,6 @@
 
 package dagger.internal.codegen.binding;
 
-import static com.google.auto.common.MoreElements.getLocalAndInheritedMethods;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static dagger.internal.codegen.binding.SourceFiles.simpleVariableName;
@@ -38,7 +37,6 @@
 import dagger.Provides;
 import dagger.internal.codegen.kotlin.KotlinMetadataUtil;
 import dagger.internal.codegen.langmodel.DaggerElements;
-import dagger.internal.codegen.langmodel.DaggerTypes;
 import dagger.model.BindingKind;
 import dagger.model.Key;
 import dagger.multibindings.Multibinds;
@@ -115,24 +113,20 @@ public enum NullPolicy {
    * of the default behavior in {@link #nullPolicy}.
    *
    * <p>Some implementations' null policy can be determined upon construction (e.g., for binding
-   * instances), but others' require Elements and Types, which must wait until {@link #nullPolicy}
-   * is called.
+   * instances), but others' require Elements which must wait until {@link #nullPolicy} is called.
    */
   abstract Optional<NullPolicy> overrideNullPolicy();
 
   /** The requirement's null policy. */
-  public NullPolicy nullPolicy(
-      DaggerElements elements, DaggerTypes types, KotlinMetadataUtil metadataUtil) {
+  public NullPolicy nullPolicy(DaggerElements elements, KotlinMetadataUtil metadataUtil) {
     if (overrideNullPolicy().isPresent()) {
       return overrideNullPolicy().get();
     }
     switch (kind()) {
       case MODULE:
         return componentCanMakeNewInstances(typeElement(), metadataUtil)
             ? NullPolicy.NEW
-            : requiresAPassedInstance(elements, types, metadataUtil)
-                ? NullPolicy.THROW
-                : NullPolicy.ALLOW;
+            : requiresAPassedInstance(elements, metadataUtil) ? NullPolicy.THROW : NullPolicy.ALLOW;
       case DEPENDENCY:
       case BOUND_INSTANCE:
         return NullPolicy.THROW;
@@ -144,13 +138,12 @@ public NullPolicy nullPolicy(
    * Returns true if the passed {@link ComponentRequirement} requires a passed instance in order to
    * be used within a component.
    */
-  public boolean requiresAPassedInstance(
-      DaggerElements elements, DaggerTypes types, KotlinMetadataUtil metadataUtil) {
+  public boolean requiresAPassedInstance(DaggerElements elements, KotlinMetadataUtil metadataUtil) {
     if (!kind().isModule()) {
       // Bound instances and dependencies always require the user to provide an instance.
       return true;
     }
-    return requiresModuleInstance(elements, types, metadataUtil)
+    return requiresModuleInstance(elements, metadataUtil)
         && !componentCanMakeNewInstances(typeElement(), metadataUtil);
   }
 
@@ -164,17 +157,15 @@ public boolean requiresAPassedInstance(
    * <p>Alternatively, if the module is a Kotlin Object then the binding methods are considered
    * {@code static}, requiring no module instance.
    */
-  private boolean requiresModuleInstance(
-      DaggerElements elements, DaggerTypes types, KotlinMetadataUtil metadataUtil) {
+  private boolean requiresModuleInstance(DaggerElements elements, KotlinMetadataUtil metadataUtil) {
     boolean isKotlinObject =
         metadataUtil.isObjectClass(typeElement())
             || metadataUtil.isCompanionObjectClass(typeElement());
     if (isKotlinObject) {
       return false;
     }
 
-    ImmutableSet<ExecutableElement> methods =
-        getLocalAndInheritedMethods(typeElement(), types, elements);
+    ImmutableSet<ExecutableElement> methods = elements.getLocalAndInheritedMethods(typeElement());
     return methods.stream()
         .filter(this::isBindingMethod)
         .map(ExecutableElement::getModifiers)

java/dagger/internal/codegen/componentgenerator/ComponentCreatorImplementationFactory.java

@@ -237,7 +237,7 @@ private MethodSpec normalSetterMethod(ComponentRequirement requirement) {
       method.addStatement(
           "this.$N = $L",
           fields.get(requirement),
-          requirement.nullPolicy(elements, types, metadataUtil).equals(NullPolicy.ALLOW)
+          requirement.nullPolicy(elements, metadataUtil).equals(NullPolicy.ALLOW)
               ? CodeBlock.of("$N", parameter)
               : CodeBlock.of("$T.checkNotNull($N)", Preconditions.class, parameter));
       return maybeReturnThis(method);
@@ -310,7 +310,7 @@ MethodSpec factoryMethod() {
 
     private void addNullHandlingForField(
         ComponentRequirement requirement, FieldSpec field, MethodSpec.Builder factoryMethod) {
-      switch (requirement.nullPolicy(elements, types, metadataUtil)) {
+      switch (requirement.nullPolicy(elements, metadataUtil)) {
         case NEW:
           checkState(requirement.kind().isModule());
           factoryMethod
@@ -334,7 +334,7 @@ private void addNullHandlingForField(
 
     private void addNullHandlingForParameter(
         ComponentRequirement requirement, String parameter, MethodSpec.Builder factoryMethod) {
-      if (!requirement.nullPolicy(elements, types, metadataUtil).equals(NullPolicy.ALLOW)) {
+      if (!requirement.nullPolicy(elements, metadataUtil).equals(NullPolicy.ALLOW)) {
         // Factory method parameters are always required unless they are a nullable
         // binds-instance (i.e. ALLOW)
         factoryMethod.addStatement("$T.checkNotNull($L)", Preconditions.class, parameter);

java/dagger/internal/codegen/componentgenerator/ComponentHjarGenerator.java

@@ -148,8 +148,7 @@ public Optional<TypeSpec.Builder> write(ComponentDescriptor componentDescriptor)
         && !hasBindsInstanceMethods(componentDescriptor)
         && componentRequirements(componentDescriptor)
             .noneMatch(
-                requirement ->
-                    requirement.requiresAPassedInstance(elements, types, metadataUtil))) {
+                requirement -> requirement.requiresAPassedInstance(elements, metadataUtil))) {
       generatedComponent.addMethod(createMethod(componentDescriptor));
     }
 

java/dagger/internal/codegen/componentgenerator/ComponentImplementationBuilder.java

@@ -16,7 +16,6 @@
 
 package dagger.internal.codegen.componentgenerator;
 
-import static com.google.auto.common.MoreElements.getLocalAndInheritedMethods;
 import static com.google.auto.common.MoreTypes.asDeclared;
 import static com.google.common.base.Preconditions.checkState;
 import static com.squareup.javapoet.MethodSpec.constructorBuilder;
@@ -146,7 +145,8 @@ ComponentImplementation build() {
         .map(ComponentCreatorImplementation::spec)
         .ifPresent(this::addCreatorClass);
 
-    getLocalAndInheritedMethods(graph.componentTypeElement(), types, elements)
+    elements
+        .getLocalAndInheritedMethods(graph.componentTypeElement())
         .forEach(method -> componentImplementation.claimMethodName(method.getSimpleName()));
 
     addFactoryMethods();
@@ -495,7 +495,7 @@ private void createRootComponentFactoryMethod() {
   private boolean canInstantiateAllRequirements() {
     return !Iterables.any(
         graph.componentRequirements(),
-        dependency -> dependency.requiresAPassedInstance(elements, types, metadataUtil));
+        dependency -> dependency.requiresAPassedInstance(elements, metadataUtil));
   }
 
   private void createSubcomponentFactoryMethod(ExecutableElement factoryMethod) {

java/dagger/internal/codegen/langmodel/BUILD

@@ -26,6 +26,7 @@ java_library(
     tags = ["maven:merged"],
     deps = [
         "//java/dagger:core",
+        "//java/dagger/internal/codegen/base:shared",
         "//java/dagger/internal/guava:base",
         "//java/dagger/internal/guava:collect",
         "//java/dagger/internal/guava:concurrent",

java/dagger/internal/codegen/langmodel/DaggerElements.java

@@ -17,7 +17,6 @@
 package dagger.internal.codegen.langmodel;
 
 import static com.google.auto.common.MoreElements.asExecutable;
-import static com.google.auto.common.MoreElements.getLocalAndInheritedMethods;
 import static com.google.auto.common.MoreElements.hasModifiers;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.Lists.asList;
@@ -34,10 +33,12 @@
 import com.google.common.graph.Traverser;
 import com.squareup.javapoet.ClassName;
 import dagger.Reusable;
+import dagger.internal.codegen.base.ClearableCache;
 import java.io.Writer;
 import java.lang.annotation.Annotation;
 import java.util.Collection;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -75,8 +76,9 @@
 
 /** Extension of {@link Elements} that adds Dagger-specific methods. */
 @Reusable
-public final class DaggerElements implements Elements {
-
+public final class DaggerElements implements Elements, ClearableCache {
+  private final Map<TypeElement, ImmutableSet<ExecutableElement>> getLocalAndInheritedMethodsCache =
+      new HashMap<>();
   private final Elements elements;
   private final Types types;
 
@@ -100,8 +102,13 @@ public static boolean elementEncloses(TypeElement encloser, Element enclosed) {
   private static final Traverser<Element> GET_ENCLOSED_ELEMENTS =
       Traverser.forTree(Element::getEnclosedElements);
 
+  public ImmutableSet<ExecutableElement> getLocalAndInheritedMethods(TypeElement type) {
+    return getLocalAndInheritedMethodsCache.computeIfAbsent(
+        type, k -> MoreElements.getLocalAndInheritedMethods(type, types, elements));
+  }
+
   public ImmutableSet<ExecutableElement> getUnimplementedMethods(TypeElement type) {
-    return FluentIterable.from(getLocalAndInheritedMethods(type, types, elements))
+    return FluentIterable.from(getLocalAndInheritedMethods(type))
         .filter(hasModifiers(ABSTRACT))
         .toSet();
   }
@@ -492,4 +499,9 @@ public Name getName(CharSequence cs) {
   public boolean isFunctionalInterface(TypeElement type) {
     return elements.isFunctionalInterface(type);
   }
+
+  @Override
+  public void clearCache() {
+    getLocalAndInheritedMethodsCache.clear();
+  }
 }

java/dagger/internal/codegen/validation/ComponentDescriptorValidator.java

@@ -318,7 +318,7 @@ private void validateCreators(ComponentDescriptor component) {
       Set<ComponentRequirement> mustBePassed =
           Sets.filter(
               componentModuleAndDependencyRequirements,
-              input -> input.nullPolicy(elements, types, metadataUtil).equals(NullPolicy.THROW));
+              input -> input.nullPolicy(elements, metadataUtil).equals(NullPolicy.THROW));
       // Component requirements that the creator must be able to set, but can't
       Set<ComponentRequirement> missingRequirements =
           Sets.difference(mustBePassed, creatorModuleAndDependencyRequirements);

java/dagger/internal/codegen/validation/ComponentValidator.java

@@ -17,7 +17,6 @@
 package dagger.internal.codegen.validation;
 
 import static com.google.auto.common.MoreElements.asType;
-import static com.google.auto.common.MoreElements.getLocalAndInheritedMethods;
 import static com.google.auto.common.MoreElements.isAnnotationPresent;
 import static com.google.auto.common.MoreTypes.asDeclared;
 import static com.google.auto.common.MoreTypes.asExecutable;
@@ -243,8 +242,7 @@ private void validateNoReusableAnnotation() {
     }
 
     private void validateComponentMethods() {
-      getLocalAndInheritedMethods(component, types, elements).stream()
-          .filter(method -> method.getModifiers().contains(ABSTRACT))
+      elements.getUnimplementedMethods(component).stream()
           .map(ComponentMethodValidator::new)
           .forEachOrdered(ComponentMethodValidator::validateMethod);
     }

java/dagger/internal/codegen/writing/AssistedFactoryBindingExpression.java

@@ -83,7 +83,7 @@ Expression getDependencyExpression(ClassName requestingClass) {
   private TypeSpec anonymousfactoryImpl(Expression assistedInjectionExpression) {
     TypeElement factory = asType(binding.bindingElement().get());
     DeclaredType factoryType = asDeclared(binding.key().type());
-    ExecutableElement factoryMethod = assistedFactoryMethod(factory, elements, types);
+    ExecutableElement factoryMethod = assistedFactoryMethod(factory, elements);
 
     // We can't use MethodSpec.overriding directly because we need to control the parameter names.
     MethodSpec factoryOverride = MethodSpec.overriding(factoryMethod, factoryType, types).build();