Generate keep rules for each LazyClassKey referenced class, potential fix for #4323

fixes #4323

RELNOTES=n/a
PiperOrigin-RevId: 686240258

java/dagger/internal/codegen/processingstep/LazyClassKeyProcessingStep.java

diff --git a/java/dagger/internal/codegen/processingstep/LazyClassKeyProcessingStep.java b/java/dagger/internal/codegen/processingstep/LazyClassKeyProcessingStep.java
new file mode 100644
index 0000000..54c406a
--- /dev/null
+++ b/java/dagger/internal/codegen/processingstep/LazyClassKeyProcessingStep.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright (C) 2024 The Dagger Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package dagger.internal.codegen.processingstep;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import androidx.room.compiler.processing.XElement;
+import androidx.room.compiler.processing.XFiler;
+import androidx.room.compiler.processing.XProcessingEnv;
+import androidx.room.compiler.processing.XTypeElement;
+import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.LinkedHashMultimap;
+import com.google.common.collect.SetMultimap;
+import com.squareup.javapoet.ClassName;
+import dagger.internal.codegen.javapoet.TypeNames;
+import dagger.internal.codegen.xprocessing.XElements;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.io.OutputStreamWriter;
+import java.nio.file.Path;
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+import javax.inject.Inject;
+
+/** Generate keep rules for LazyClassKey referenced classes to prevent class merging. */
+final class LazyClassKeyProcessingStep extends TypeCheckingProcessingStep<XElement> {
+  private static final String PROGUARD_KEEP_RULE = "-keep,allowobfuscation,allowshrinking class ";
+  private final SetMultimap<ClassName, ClassName> processedElements = LinkedHashMultimap.create();
+
+  @Inject
+  LazyClassKeyProcessingStep() {}
+
+  @Override
+  public ImmutableSet<ClassName> annotationClassNames() {
+    return ImmutableSet.of(TypeNames.LAZY_CLASS_KEY);
+  }
+
+  @Override
+  protected void process(XElement element, ImmutableSet<ClassName> annotations) {
+    ClassName lazyClassKey =
+        element
+            .getAnnotation(TypeNames.LAZY_CLASS_KEY)
+            .getAsType("value")
+            .getTypeElement()
+            .getClassName();
+    XTypeElement moduleElement = XElements.asTypeElement(element.getEnclosingElement());
+    processedElements.put(moduleElement.getClassName(), lazyClassKey);
+  }
+
+  @Override
+  public void processOver(
+      XProcessingEnv env, Map<String, ? extends Set<? extends XElement>> elementsByAnnotation) {
+    super.processOver(env, elementsByAnnotation);
+    StringBuilder proguardRules = new StringBuilder();
+    for (Map.Entry<ClassName, Collection<ClassName>> moduleToLazyClassKeys :
+        processedElements.asMap().entrySet()) {
+      String bindingGraphProguardName =
+          getFullyQualifiedEnclosedClassName(moduleToLazyClassKeys.getKey()) + "_LazyClassKeys.pro";
+      for (ClassName lazyClassKey : moduleToLazyClassKeys.getValue()) {
+        proguardRules.append(PROGUARD_KEEP_RULE).append(lazyClassKey).append("\n");
+      }
+      writeProguardFile(bindingGraphProguardName, proguardRules.toString(), env.getFiler());
+    }
+  }
+
+  private void writeProguardFile(String proguardFileName, String proguardRules, XFiler filer) {
+    try (OutputStream outputStream =
+            filer.writeResource(
+                Path.of("META-INF/proguard/" + proguardFileName),
+                ImmutableList.<XElement>of(),
+                XFiler.Mode.Isolating);
+        BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(outputStream, UTF_8))) {
+      writer.write(proguardRules);
+    } catch (IOException e) {
+      throw new IllegalStateException(e);
+    }
+  }
+
+  /** Returns the fully qualified class name, with _ instead of . */
+  private static String getFullyQualifiedEnclosedClassName(ClassName className) {
+    return className.packageName().replace('.', '_') + getEnclosedName(className);
+  }
+
+  public static String getEnclosedName(ClassName name) {
+    return Joiner.on('_').join(name.simpleNames());
+  }
+}

java/dagger/internal/codegen/processingstep/ProcessingStepsModule.java

diff --git a/java/dagger/internal/codegen/processingstep/ProcessingStepsModule.java b/java/dagger/internal/codegen/processingstep/ProcessingStepsModule.java
index e24dddc..f180c3d 100644
--- a/java/dagger/internal/codegen/processingstep/ProcessingStepsModule.java
+++ b/java/dagger/internal/codegen/processingstep/ProcessingStepsModule.java
@@ -39,6 +39,7 @@
       MultibindingAnnotationsProcessingStep multibindingAnnotationsProcessingStep,
       BindsInstanceProcessingStep bindsInstanceProcessingStep,
       ModuleProcessingStep moduleProcessingStep,
+      LazyClassKeyProcessingStep lazyClassKeyProcessingStep,
       ComponentProcessingStep componentProcessingStep,
       ComponentHjarProcessingStep componentHjarProcessingStep,
       BindingMethodProcessingStep bindingMethodProcessingStep,
@@ -53,9 +54,8 @@
         multibindingAnnotationsProcessingStep,
         bindsInstanceProcessingStep,
         moduleProcessingStep,
-        compilerOptions.headerCompilation()
-            ? componentHjarProcessingStep
-            : componentProcessingStep,
+        lazyClassKeyProcessingStep,
+        compilerOptions.headerCompilation() ? componentHjarProcessingStep : componentProcessingStep,
         bindingMethodProcessingStep);
   }
 

java/dagger/r8.pro

diff --git a/java/dagger/r8.pro b/java/dagger/r8.pro
index b700447..6fde1b2 100644
--- a/java/dagger/r8.pro
+++ b/java/dagger/r8.pro
@@ -1,6 +1,3 @@
 -identifiernamestring @dagger.internal.IdentifierNameString class ** {
     static java.lang.String *;
 }
--keepclassmembers,includedescriptorclasses,allowobfuscation,allowshrinking class * {
-   @dagger.internal.KeepFieldType <fields>;
-}
\ No newline at end of file