Fix inadequate keep rule generating in trace references

Bug: b/397904338
Change-Id: I3d613d57f8de37a3c5670b6d843e1edc402591e7

src/main/java/com/android/tools/r8/tracereferences/Formatter.java

diff --git a/src/main/java/com/android/tools/r8/tracereferences/Formatter.java b/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
index 2f19681..e3978f3 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
@@ -10,10 +10,14 @@
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedClass;
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedField;
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedMethod;
+import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedReference;
 import com.android.tools.r8.tracereferences.internal.TraceReferencesResult;
+import com.android.tools.r8.utils.IterableUtils;
 import com.android.tools.r8.utils.ListUtils;
+import com.android.tools.r8.utils.MapUtils;
 import com.android.tools.r8.utils.StringUtils;
 import com.android.tools.r8.utils.StringUtils.BraceType;
+import com.google.common.collect.Streams;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
@@ -89,13 +93,14 @@
       Set<PackageReference> keepPackageNames,
       Map<ClassReference, Set<TracedField>> fields,
       Map<ClassReference, Set<TracedMethod>> methods) {
-    List<TracedClass> sortedTypes = new ArrayList<>(types);
-    sortedTypes.sort(Comparator.comparing(tracedClass -> tracedClass.getReference().getTypeName()));
+    List<TracedClass> sortedTypes =
+        ListUtils.sort(
+            types, Comparator.comparing(tracedClass -> tracedClass.getReference().getDescriptor()));
     for (TracedClass type : sortedTypes) {
       Set<TracedMethod> methodsForClass =
-          methods.getOrDefault(type.getReference(), Collections.emptySet());
+          MapUtils.removeOrDefault(methods, type.getReference(), Collections.emptySet());
       Set<TracedField> fieldsForClass =
-          fields.getOrDefault(type.getReference(), Collections.emptySet());
+          MapUtils.removeOrDefault(fields, type.getReference(), Collections.emptySet());
       if (type.isMissingDefinition()) {
         continue;
       }
@@ -123,6 +128,10 @@
       }
       printTypeFooter();
     }
+    assert Streams.stream(IterableUtils.flatten(fields.values()))
+        .allMatch(TracedReference::isMissingDefinition);
+    assert Streams.stream(IterableUtils.flatten(methods.values()))
+        .allMatch(TracedReference::isMissingDefinition);
     List<String> packageNamesToKeep =
         keepPackageNames.stream()
             .map(PackageReference::getPackageName)

src/main/java/com/android/tools/r8/tracereferences/UseCollector.java

diff --git a/src/main/java/com/android/tools/r8/tracereferences/UseCollector.java b/src/main/java/com/android/tools/r8/tracereferences/UseCollector.java
index 22ec0d0..6bcaf98 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/UseCollector.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/UseCollector.java
@@ -123,6 +123,12 @@
     consumer.acceptMethod(tracedMethod, diagnostics);
   }
 
+  protected void notifyPresentMethod(
+      DexClassAndMethod method, DefinitionContext referencedFrom, DexMethod reference) {
+    TracedMethodImpl tracedMethod = new TracedMethodImpl(method, referencedFrom, reference);
+    consumer.acceptMethod(tracedMethod, diagnostics);
+  }
+
   protected void notifyPackageOf(Definition definition) {
     consumer.acceptPackage(
         Reference.packageFromString(definition.getContextType().getPackageName()), diagnostics);
@@ -212,19 +218,6 @@
     }
   }
 
-  private void addSuperMethodFromTarget(
-      DexClassAndMethod method, ProgramMethod context, DefinitionContext referencedFrom) {
-    assert !method.isProgramMethod();
-    assert isTargetType(method.getHolderType());
-    // There should be no need to register the types referenced from the method signature:
-    // - The return type and the parameter types are registered when visiting the source method
-    //   that overrides this target method,
-    // - The holder type is registered from visiting the extends/implements clause of the sub
-    //   class.
-    notifyPresentMethod(method, referencedFrom);
-    ensurePackageAccessToMember(method, context.getHolder());
-  }
-
   private <R, T extends TracedReference<R, ?>> void collectMissing(
       T tracedReference, Set<R> missingCollection) {
     if (tracedReference.isMissingDefinition()) {
@@ -291,7 +284,23 @@
           if (resolvedMethod != null
               && !resolvedMethod.isProgramMethod()
               && isTargetType(resolvedMethod.getHolderType())) {
-            addSuperMethodFromTarget(resolvedMethod, method, referencedFrom);
+            // There should be no need to register the types referenced from the method signature:
+            // - The return type and the parameter types are registered when visiting the source
+            //   method that overrides this target method,
+            // - The holder type is registered from visiting the extends/implements clause of the
+            //   sub class.
+            if (resolvedMethod.getHolderType().isIdenticalTo(superType)) {
+              notifyPresentMethod(resolvedMethod, referencedFrom);
+            } else if (isTargetType(superType)) {
+              notifyPresentMethod(
+                  resolvedMethod,
+                  referencedFrom,
+                  resolvedMethod.getReference().withHolder(superType, factory));
+            } else {
+              notifyPresentMethod(resolvedMethod, referencedFrom);
+              addClass(resolvedMethod.getHolder(), referencedFrom);
+            }
+            ensurePackageAccessToMember(resolvedMethod, method.getHolder());
           }
         });
   }

src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java

diff --git a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java
index 2e315df..22f7671 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java
@@ -26,6 +26,14 @@
   }
 
   public TracedMethodImpl(
+      DexClassAndMethod method, DefinitionContext referencedFrom, DexMethod reference) {
+    this(
+        reference.asMethodReference(),
+        referencedFrom,
+        new MethodAccessFlagsImpl(method.getAccessFlags()));
+  }
+
+  public TracedMethodImpl(
       MethodReference methodReference,
       DefinitionContext referencedFrom,
       MethodAccessFlags accessFlags) {

src/test/java/com/android/tools/r8/tracereferences/TraceReferencesIndirectOverrideTest.java

diff --git a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesIndirectOverrideTest.java b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesIndirectOverrideTest.java
index cc255ed..771e4e3 100644
--- a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesIndirectOverrideTest.java
+++ b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesIndirectOverrideTest.java
@@ -4,7 +4,6 @@
 package com.android.tools.r8.tracereferences;
 
 import static org.hamcrest.CoreMatchers.containsString;
-import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 
 import com.android.tools.r8.TestBase;
@@ -57,8 +56,7 @@
                     .setOutputConsumer(testingKeepRuleConsumer)
                     .build())
             .build());
-    // TODO(b/397904338): shouldBeKept() should be kept.
-    assertThat(testingKeepRuleConsumer.get(), not(containsString("shouldBeKept")));
+    assertThat(testingKeepRuleConsumer.get(), containsString("shouldBeKept"));
   }
 
   static class A {