From 5fd8ec1a39f568d67bd6644abe821729ba6a0725 Mon Sep 17 00:00:00 2001
From: Brad Corso <bcorso@google.com>
Date: Wed, 4 Dec 2024 14:47:02 -0800
Subject: [PATCH] Add the originating element in LazyMapKeyProxyGenerator.

Forgetting the originating element was an oversight. Gradle's incremental processing requires originating elements to determine when to reprocess files. This CL fixes `LazyMapKeyProxyGenerator` by extending `SourceFileGenerator`, which takes care of assigning the originating element. In addition, it removes some duplicate code that's already handled in `SourceFileGenerator` (e.g. adding the generated and suppress annotations).

Fixes #4526

RELNOTES=Fixes #4526: Add the originating element in LazyMapKeyProxyGenerator.
PiperOrigin-RevId: 702869366
---
 .../writing/LazyMapKeyProxyGenerator.java     | 64 ++++++-------------
 1 file changed, 20 insertions(+), 44 deletions(-)

diff --git a/java/dagger/internal/codegen/writing/LazyMapKeyProxyGenerator.java b/java/dagger/internal/codegen/writing/LazyMapKeyProxyGenerator.java
index 623e13e1439..a5b15f2bd73 100644
--- a/java/dagger/internal/codegen/writing/LazyMapKeyProxyGenerator.java
+++ b/java/dagger/internal/codegen/writing/LazyMapKeyProxyGenerator.java
@@ -20,73 +20,49 @@
 import static dagger.internal.codegen.binding.MapKeys.KEEP_FIELD_TYPE_FIELD;
 import static dagger.internal.codegen.binding.MapKeys.LAZY_CLASS_KEY_NAME_FIELD;
 import static dagger.internal.codegen.binding.MapKeys.lazyClassKeyProxyClassName;
-import static dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression.CAST;
-import static dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression.DEPRECATION;
-import static dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression.KOTLIN_INTERNAL;
-import static dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression.RAWTYPES;
-import static dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression.UNCHECKED;
 import static javax.lang.model.element.Modifier.FINAL;
 import static javax.lang.model.element.Modifier.PUBLIC;
 import static javax.lang.model.element.Modifier.STATIC;
 
+import androidx.room.compiler.processing.XElement;
 import androidx.room.compiler.processing.XFiler;
 import androidx.room.compiler.processing.XMethodElement;
 import androidx.room.compiler.processing.XProcessingEnv;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
-import com.squareup.javapoet.AnnotationSpec;
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.FieldSpec;
-import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.TypeSpec;
-import dagger.internal.DaggerGenerated;
-import dagger.internal.codegen.javapoet.AnnotationSpecs;
-import dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression;
+import dagger.internal.codegen.base.SourceFileGenerator;
 import dagger.internal.codegen.javapoet.TypeNames;
-import java.util.Optional;
 import javax.inject.Inject;
 
 /**
  * Generate a class containing fields that works with proguard rules to support @LazyClassKey
  * usages.
  */
-public final class LazyMapKeyProxyGenerator {
-  private static final String GENERATED_COMMENTS = "https://dagger.dev";
-  private final XProcessingEnv processingEnv;
-  private final XFiler filer;
+public final class LazyMapKeyProxyGenerator extends SourceFileGenerator<XMethodElement> {
 
   @Inject
-  LazyMapKeyProxyGenerator(XProcessingEnv processingEnv, XFiler filer) {
-    this.processingEnv = processingEnv;
-    this.filer = filer;
+  LazyMapKeyProxyGenerator(XFiler filer, XProcessingEnv processingEnv) {
+    super(filer, processingEnv);
   }
 
-  public void generate(XMethodElement element) {
-    ClassName lazyClassKeyProxyClassName = lazyClassKeyProxyClassName(element);
-    TypeSpec.Builder typeSpecBuilder =
-        classBuilder(lazyClassKeyProxyClassName)
-            .addModifiers(PUBLIC, FINAL)
-            .addAnnotation(TypeNames.IDENTIFIER_NAME_STRING)
-            .addAnnotation(DaggerGenerated.class)
-            .addFields(lazyClassKeyFields(element));
-    Optional<AnnotationSpec> generatedAnnotation =
-        Optional.ofNullable(processingEnv.findGeneratedAnnotation())
-            .map(
-                annotation ->
-                    AnnotationSpec.builder(annotation.getClassName())
-                        .addMember("value", "$S", "dagger.internal.codegen.LazyClassKeyProcessor")
-                        .addMember("comments", "$S", GENERATED_COMMENTS)
-                        .build());
-    generatedAnnotation.ifPresent(typeSpecBuilder::addAnnotation);
-    typeSpecBuilder.addAnnotation(
-        AnnotationSpecs.suppressWarnings(
-            ImmutableSet.<Suppression>builder()
-                .add(UNCHECKED, RAWTYPES, KOTLIN_INTERNAL, CAST, DEPRECATION)
-                .build()));
+  @Override
+  public XElement originatingElement(XMethodElement input) {
+    return input;
+  }
+
+  @Override
+  public ImmutableList<TypeSpec.Builder> topLevelTypes(XMethodElement input) {
+    return ImmutableList.of(lazyClassKeyProxyTypeSpec(input).toBuilder());
+  }
 
-    filer.write(
-        JavaFile.builder(lazyClassKeyProxyClassName.packageName(), typeSpecBuilder.build()).build(),
-        XFiler.Mode.Isolating);
+  private TypeSpec lazyClassKeyProxyTypeSpec(XMethodElement element) {
+    return classBuilder(lazyClassKeyProxyClassName(element))
+        .addModifiers(PUBLIC, FINAL)
+        .addAnnotation(TypeNames.IDENTIFIER_NAME_STRING)
+        .addFields(lazyClassKeyFields(element))
+        .build();
   }
 
   private static ImmutableList<FieldSpec> lazyClassKeyFields(XMethodElement element) {