Allow @AndroidEntryPoint base class to be @AndroidEntryPoint class when using plugin.

This fixes a bug when using the plugin where the child generated class uses `@Overrides void inject()` which fails because that method doesn't exist in an ancestor until after bytecode injection.

Fixes #1910

RELNOTES=Fix #1910: Allow usage of an @AndroidEntryPoint base class when using short-form notation with the hilt gradle plugin.

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

3 files changed, 89 insertions, 25 deletions

diff --git a/java/dagger/hilt/android/processor/internal/androidentrypoint/AndroidEntryPointMetadata.java b/java/dagger/hilt/android/processor/internal/androidentrypoint/AndroidEntryPointMetadata.java
index ae6aaa817..d1f9ae5ec 100644
--- a/java/dagger/hilt/android/processor/internal/androidentrypoint/AndroidEntryPointMetadata.java
+++ b/java/dagger/hilt/android/processor/internal/androidentrypoint/AndroidEntryPointMetadata.java
@@ -22,6 +22,7 @@ import static dagger.internal.codegen.extension.DaggerStreams.toImmutableSet;
 import com.google.auto.common.MoreElements;
 import com.google.auto.common.MoreTypes;
 import com.google.auto.value.AutoValue;
+import com.google.auto.value.extension.memoized.Memoized;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
@@ -61,6 +62,9 @@ public abstract class AndroidEntryPointMetadata {
   /** The name of the generated base class, beginning with 'Hilt_'. */
   public abstract ClassName generatedClassName();
 
+  /** Returns {@code true} if the class requires bytecode injection to replace the base class. */
+  public abstract boolean requiresBytecodeInjection();
+
   /** Returns the {@link AndroidType} for the annotated element. */
   public abstract AndroidType androidType();
 
@@ -76,6 +80,16 @@ public abstract class AndroidEntryPointMetadata {
   /** Returns the initialization arguments for the component manager. */
   public abstract Optional<CodeBlock> componentManagerInitArgs();
 
+  /**
+   * Returns the metadata for the root most class in the hierarchy.
+   *
+   * <p>If this is the only metadata in the class hierarchy, it returns this.
+   */
+  @Memoized
+  public AndroidEntryPointMetadata rootMetadata() {
+    return baseMetadata().map(AndroidEntryPointMetadata::rootMetadata).orElse(this);
+  }
+
   /** Returns true if this class allows optional injection. */
   public boolean allowsOptionalInjection() {
     return Processors.hasAnnotation(element(), AndroidClassNames.OPTIONAL_INJECT);
@@ -149,8 +163,7 @@ public abstract class AndroidEntryPointMetadata {
 
   private static ImmutableSet<? extends AnnotationMirror> hiltAnnotations(Element element) {
     return element.getAnnotationMirrors().stream()
-        .filter(
-            mirror -> HILT_ANNOTATION_NAMES.contains(ClassName.get(mirror.getAnnotationType())))
+        .filter(mirror -> HILT_ANNOTATION_NAMES.contains(ClassName.get(mirror.getAnnotationType())))
         .collect(toImmutableSet());
   }
 
@@ -170,6 +183,7 @@ public abstract class AndroidEntryPointMetadata {
       TypeElement element,
       TypeElement baseElement,
       ClassName generatedClassName,
+      boolean requiresBytecodeInjection,
       AndroidType androidType,
       Optional<AndroidEntryPointMetadata> baseMetadata,
       ImmutableSet<ClassName> installInComponents,
@@ -179,6 +193,7 @@ public abstract class AndroidEntryPointMetadata {
         element,
         baseElement,
         generatedClassName,
+        requiresBytecodeInjection,
         androidType,
         baseMetadata,
         installInComponents,
@@ -201,8 +216,7 @@ public abstract class AndroidEntryPointMetadata {
         hiltAnnotations);
     ClassName annotationClassName =
         ClassName.get(
-            MoreTypes.asTypeElement(
-                Iterables.getOnlyElement(hiltAnnotations).getAnnotationType()));
+            MoreTypes.asTypeElement(Iterables.getOnlyElement(hiltAnnotations).getAnnotationType()));
 
     ProcessorErrors.checkState(
         element.getKind() == ElementKind.CLASS,
@@ -218,8 +232,10 @@ public abstract class AndroidEntryPointMetadata {
             "value");
     final TypeElement baseElement;
     final ClassName generatedClassName;
-    if (DISABLE_ANDROID_SUPERCLASS_VALIDATION.get(env)
-        && MoreTypes.isTypeOf(Void.class, androidEntryPointClassValue.asType())) {
+    boolean requiresBytecodeInjection =
+        DISABLE_ANDROID_SUPERCLASS_VALIDATION.get(env)
+            && MoreTypes.isTypeOf(Void.class, androidEntryPointClassValue.asType());
+    if (requiresBytecodeInjection) {
       baseElement = MoreElements.asType(env.getTypeUtils().asElement(androidEntryPointElement.getSuperclass()));
       // If this AndroidEntryPoint is a Kotlin class and its base type is also Kotlin and has
       // default values declared in its constructor then error out because for the short-form
@@ -249,7 +265,10 @@ public abstract class AndroidEntryPointMetadata {
 
       // Check that the root $CLASS extends Hilt_$CLASS
       String extendsName =
-          env.getTypeUtils().asElement(androidEntryPointElement.getSuperclass()).getSimpleName().toString();
+          env.getTypeUtils()
+              .asElement(androidEntryPointElement.getSuperclass())
+              .getSimpleName()
+              .toString();
       generatedClassName = generatedClassName(androidEntryPointElement);
       ProcessorErrors.checkState(
           extendsName.contentEquals(generatedClassName.simpleName()),
@@ -268,6 +287,7 @@ public abstract class AndroidEntryPointMetadata {
           androidEntryPointElement,
           baseElement,
           generatedClassName,
+          requiresBytecodeInjection,
           baseMetadata.get().androidType(),
           baseMetadata,
           baseMetadata.get().installInComponents(),
@@ -279,6 +299,7 @@ public abstract class AndroidEntryPointMetadata {
           androidEntryPointElement,
           baseElement,
           generatedClassName,
+          requiresBytecodeInjection,
           type.androidType,
           Optional.empty(),
           ImmutableSet.of(type.component),
@@ -307,8 +328,7 @@ public abstract class AndroidEntryPointMetadata {
     // None type is returned if this is an interface or Object
     if (superClass.getKind() != TypeKind.NONE && superClass.getKind() != TypeKind.ERROR) {
       Preconditions.checkState(superClass.getKind() == TypeKind.DECLARED);
-      return baseMetadata(
-          env, element, MoreTypes.asTypeElement(superClass), inheritanceTrace);
+      return baseMetadata(env, element, MoreTypes.asTypeElement(superClass), inheritanceTrace);
     }
 
     return Optional.empty();
@@ -453,7 +473,7 @@ public abstract class AndroidEntryPointMetadata {
       }
       throw new BadInputException(
           "@AndroidEntryPoint base class must extend ComponentActivity, (support) Fragment, "
-          + "View, Service, or BroadcastReceiver.",
+              + "View, Service, or BroadcastReceiver.",
           element);
     }
   }
@@ -482,7 +502,7 @@ public abstract class AndroidEntryPointMetadata {
         isBaseAnnotated
             ? "Classes that extend an @%1$s base class must also be annotated @%1$s"
             : "Classes that extend a @AndroidEntryPoint base class must not use @%1$s when the "
-                + "base class "
+                    + "base class "
                 + "does not use @%1$s",
         annotationName.simpleName());
   }
diff --git a/java/dagger/hilt/android/processor/internal/androidentrypoint/ApplicationGenerator.java b/java/dagger/hilt/android/processor/internal/androidentrypoint/ApplicationGenerator.java
index 0af1dccbe..08a72ed86 100644
--- a/java/dagger/hilt/android/processor/internal/androidentrypoint/ApplicationGenerator.java
+++ b/java/dagger/hilt/android/processor/internal/androidentrypoint/ApplicationGenerator.java
@@ -28,7 +28,6 @@ import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
 import com.squareup.javapoet.TypeVariableName;
 import dagger.hilt.android.processor.internal.AndroidClassNames;
-import dagger.hilt.processor.internal.ClassNames;
 import dagger.hilt.processor.internal.ComponentNames;
 import dagger.hilt.processor.internal.Processors;
 import java.io.IOException;
@@ -154,11 +153,10 @@ public final class ApplicationGenerator {
         .add("// This is a known unsafe cast, but is safe in the only correct use case:\n")
         .add("// $T extends $T\n", metadata.elementClassName(), metadata.generatedClassName())
         .addStatement(
-            "(($T) generatedComponent()).$L($T.<$T>unsafeCast(this))",
+            "(($T) generatedComponent()).$L($L)",
             metadata.injectorClassName(),
             metadata.injectMethodName(),
-            ClassNames.UNSAFE_CASTS,
-            metadata.elementClassName())
+            Generators.unsafeCastThisTo(metadata.elementClassName()))
         .build();
   }
 }
diff --git a/java/dagger/hilt/android/processor/internal/androidentrypoint/Generators.java b/java/dagger/hilt/android/processor/internal/androidentrypoint/Generators.java
index 553eb81c1..11e674e6a 100644
--- a/java/dagger/hilt/android/processor/internal/androidentrypoint/Generators.java
+++ b/java/dagger/hilt/android/processor/internal/androidentrypoint/Generators.java
@@ -160,7 +160,7 @@ final class Generators {
                 .addAnnotation(Override.class)
                 .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                 .returns(TypeName.OBJECT)
-                .addStatement("return componentManager().generatedComponent()")
+                .addStatement("return $L.generatedComponent()", componentManagerCallBlock(metadata))
                 .build());
   }
 
@@ -314,8 +314,13 @@ final class Generators {
         if (metadata.overridesAndroidEntryPointClass()) {
           typeSpecBuilder.addField(injectedField(metadata));
 
+          // Only add @Override if an ancestor extends a generated Hilt class.
+          // When using bytecode injection, this isn't always guaranteed.
+          if (ancestorExtendsGeneratedHiltClass(metadata)) {
+            methodSpecBuilder.addAnnotation(Override.class);
+          }
+
           methodSpecBuilder
-              .addAnnotation(Override.class)
               .beginControlFlow("if (!injected)")
               .addStatement("injected = true");
         } else if (metadata.allowsOptionalInjection()) {
@@ -323,11 +328,11 @@ final class Generators {
           methodSpecBuilder.addStatement("injected = true");
         }
         methodSpecBuilder.addStatement(
-            "(($T) generatedComponent()).$L($T.<$T>unsafeCast(this))",
+            "(($T) $L).$L($L)",
             metadata.injectorClassName(),
+            generatedComponentCallBlock(metadata),
             metadata.injectMethodName(),
-            ClassNames.UNSAFE_CASTS,
-            metadata.elementClassName());
+            unsafeCastThisTo(metadata.elementClassName()));
         if (metadata.overridesAndroidEntryPointClass()) {
           methodSpecBuilder.endControlFlow();
         }
@@ -342,12 +347,11 @@ final class Generators {
             .beginControlFlow("synchronized (injectedLock)")
             .beginControlFlow("if (!injected)")
             .addStatement(
-                "(($T) $T.generatedComponent(context)).$L($T.<$T>unsafeCast(this))",
+                "(($T) $T.generatedComponent(context)).$L($L)",
                 metadata.injectorClassName(),
                 metadata.componentManager(),
                 metadata.injectMethodName(),
-                ClassNames.UNSAFE_CASTS,
-                metadata.elementClassName())
+                unsafeCastThisTo(metadata.elementClassName()))
             .addStatement("injected = true")
             .endControlFlow()
             .endControlFlow()
@@ -389,8 +393,7 @@ final class Generators {
         return CodeBlock.of("getHost()");
       case VIEW:
         return CodeBlock.of(
-            "componentManager().maybeGetParentComponentManager()",
-            metadata.componentManagerParam());
+            "$L.maybeGetParentComponentManager()", componentManagerCallBlock(metadata));
       case BROADCAST_RECEIVER:
         // Broadcast receivers receive a "context" parameter
         return CodeBlock.of("context.getApplicationContext()");
@@ -399,6 +402,49 @@ final class Generators {
     }
   }
 
+  /**
+   * Returns the call to {@code generatedComponent()} with casts if needed.
+   *
+   * <p>A cast is required when the root generated Hilt class uses bytecode injection because
+   * subclasses won't have access to the {@code generatedComponent()} method in that case.
+   */
+  private static CodeBlock generatedComponentCallBlock(AndroidEntryPointMetadata metadata) {
+    return CodeBlock.of(
+        "$L.generatedComponent()",
+        metadata.rootMetadata().requiresBytecodeInjection()
+            ? unsafeCastThisTo(metadata.rootMetadata().generatedClassName())
+            : "this");
+  }
+
+  /**
+   * Returns the call to {@code componentManager()} with casts if needed.
+   *
+   * <p>A cast is required when the root generated Hilt class uses bytecode injection because
+   * subclasses won't have access to the {@code componentManager()} method in that case.
+   */
+  private static CodeBlock componentManagerCallBlock(AndroidEntryPointMetadata metadata) {
+    return CodeBlock.of(
+        "$L.componentManager()",
+        metadata.rootMetadata().requiresBytecodeInjection()
+            ? unsafeCastThisTo(metadata.rootMetadata().generatedClassName())
+            : "this");
+  }
+
+  static CodeBlock unsafeCastThisTo(ClassName castType) {
+    return CodeBlock.of("$T.<$T>unsafeCast(this)", ClassNames.UNSAFE_CASTS, castType);
+  }
+
+  /** Returns {@code true} if the an ancestor annotated class extends the generated class */
+  private static boolean ancestorExtendsGeneratedHiltClass(AndroidEntryPointMetadata metadata) {
+    while (metadata.baseMetadata().isPresent()) {
+      metadata = metadata.baseMetadata().get();
+      if (!metadata.requiresBytecodeInjection()) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   // private boolean injected = false;
   private static FieldSpec injectedField(AndroidEntryPointMetadata metadata) {
     FieldSpec.Builder builder = FieldSpec.builder(TypeName.BOOLEAN, "injected")