Support custom implementations of emulated core interface methods

RELNOTES: None.

PiperOrigin-RevId: 188760099

3 files changed, 58 insertions, 32 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibraryInvocationRewriter.java b/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibraryInvocationRewriter.java
index 0e0610f48d..77db915f4b 100644
--- a/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibraryInvocationRewriter.java
+++ b/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibraryInvocationRewriter.java
@@ -52,19 +52,8 @@ public class CoreLibraryInvocationRewriter extends ClassVisitor {
     public void visitMethodInsn(int opcode, String owner, String name, String desc, boolean itf) {
       Class<?> coreInterface =
           support.getCoreInterfaceRewritingTarget(opcode, owner, name, desc, itf);
-      String newOwner = support.getMoveTarget(owner, name);
 
-      if (newOwner != null) {
-        checkState(coreInterface == null,
-            "Can't move and use companion: %s.%s : %s", owner, name, desc);
-        if (opcode != Opcodes.INVOKESTATIC) {
-          // assuming a static method
-          desc = InterfaceDesugaring.companionDefaultMethodDescriptor(owner, desc);
-          opcode = Opcodes.INVOKESTATIC;
-        }
-        owner = newOwner;
-        itf = false; // assuming a class
-      } else if (coreInterface != null) {
+      if (coreInterface != null) {
         String coreInterfaceName = coreInterface.getName().replace('.', '/');
         name =
             InterfaceDesugaring.normalizeInterfaceMethodName(
@@ -84,6 +73,17 @@ public class CoreLibraryInvocationRewriter extends ClassVisitor {
 
         opcode = Opcodes.INVOKESTATIC;
         itf = false;
+      } else {
+        String newOwner = support.getMoveTarget(owner, name);
+        if (newOwner != null) {
+          if (opcode != Opcodes.INVOKESTATIC) {
+            // assuming a static method
+            desc = InterfaceDesugaring.companionDefaultMethodDescriptor(owner, desc);
+            opcode = Opcodes.INVOKESTATIC;
+          }
+          owner = newOwner;
+          itf = false; // assuming a class
+        }
       }
       super.visitMethodInsn(opcode, owner, name, desc, itf);
     }
diff --git a/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibrarySupport.java b/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibrarySupport.java
index b90222bd69..da23c1268a 100644
--- a/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibrarySupport.java
+++ b/src/tools/android/java/com/google/devtools/build/android/desugar/CoreLibrarySupport.java
@@ -14,7 +14,9 @@
 package com.google.devtools.build.android.desugar;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
+import static java.util.stream.Stream.concat;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Splitter;
@@ -265,10 +267,13 @@ class CoreLibrarySupport {
       Class<?> root = group
           .stream()
           .map(EmulatedMethod::owner)
-          .max(DefaultMethodClassFixer.InterfaceComparator.INSTANCE)
+          .max(DefaultMethodClassFixer.SubtypeComparator.INSTANCE)
           .get();
       checkState(group.stream().map(m -> m.owner()).allMatch(o -> root.isAssignableFrom(o)),
           "Not a single unique method: %s", group);
+      String methodName = group.stream().findAny().get().name();
+
+      ImmutableList<Class<?>> customOverrides = findCustomOverrides(root, methodName);
 
       for (EmulatedMethod methodDefinition : group) {
         Class<?> owner = methodDefinition.owner();
@@ -288,20 +293,39 @@ class CoreLibrarySupport {
 
         // Types to check for before calling methodDefinition's companion, sub- before super-types
         ImmutableList<Class<?>> typechecks =
-            group
-                .stream()
-                .map(EmulatedMethod::owner)
+            concat(group.stream().map(EmulatedMethod::owner), customOverrides.stream())
                 .filter(o -> o != owner && owner.isAssignableFrom(o))
-                .distinct()  // should already be but just in case
-                .sorted(DefaultMethodClassFixer.InterfaceComparator.INSTANCE)
+                .distinct() // should already be but just in case
+                .sorted(DefaultMethodClassFixer.SubtypeComparator.INSTANCE)
                 .collect(ImmutableList.toImmutableList());
         makeDispatchHelperMethod(dispatchHelper, methodDefinition, typechecks);
       }
     }
   }
 
+  private ImmutableList<Class<?>> findCustomOverrides(Class<?> root, String methodName) {
+    ImmutableList.Builder<Class<?>> customOverrides = ImmutableList.builder();
+    for (ImmutableMap.Entry<String, String> move : memberMoves.entrySet()) {
+      // move.getKey is a string <owner>#<name> which we validated in the constructor.
+      // We need to take the string apart here to compare owner and name separately.
+      if (!methodName.equals(move.getKey().substring(move.getKey().indexOf('#') + 1))) {
+        continue;
+      }
+      Class<?> target =
+          loadFromInternal(
+              rewriter.getPrefix() + move.getKey().substring(0, move.getKey().indexOf('#')));
+      if (!root.isAssignableFrom(target)) {
+        continue;
+      }
+      checkState(!target.isInterface(), "can't move emulated interface method: %s", move);
+      customOverrides.add(target);
+    }
+    return customOverrides.build();
+  }
+
   private void makeDispatchHelperMethod(
       ClassVisitor helper, EmulatedMethod method, ImmutableList<Class<?>> typechecks) {
+    checkArgument(method.owner().isInterface());
     String owner = method.owner().getName().replace('.', '/');
     Type methodType = Type.getMethodType(method.descriptor());
     String companionDesc =
@@ -341,24 +365,27 @@ class CoreLibrarySupport {
       dispatchMethod.visitFrame(Opcodes.F_SAME, 0, EMPTY_FRAME, 0, EMPTY_FRAME);
     }
 
-    // Next, check for emulated subtypes and call their companion methods
+    // Next, check for subtypes with specialized implementations and call them
     for (Class<?> tested : typechecks) {
-      checkState(tested.isInterface(), "Dispatch emulation not supported for classes: %s", tested);
       Label fallthrough = new Label();
-      String emulatedInterface = tested.getName().replace('.', '/');
+      String testedName = tested.getName().replace('.', '/');
+      // In case of a class this must be a member move; for interfaces use the companion.
+      String target =
+          tested.isInterface()
+              ? InterfaceDesugaring.getCompanionClassName(testedName)
+              : checkNotNull(memberMoves.get(rewriter.unprefix(testedName) + '#' + method.name()));
       dispatchMethod.visitVarInsn(Opcodes.ALOAD, 0);  // load "receiver"
-      dispatchMethod.visitTypeInsn(Opcodes.INSTANCEOF, emulatedInterface);
+      dispatchMethod.visitTypeInsn(Opcodes.INSTANCEOF, testedName);
       dispatchMethod.visitJumpInsn(Opcodes.IFEQ, fallthrough);
       dispatchMethod.visitVarInsn(Opcodes.ALOAD, 0);  // load "receiver"
-      dispatchMethod.visitTypeInsn(Opcodes.CHECKCAST, emulatedInterface);  // make verifier happy
+      dispatchMethod.visitTypeInsn(Opcodes.CHECKCAST, testedName);  // make verifier happy
 
       visitLoadArgs(dispatchMethod, methodType, 1 /* receiver already loaded above */);
       dispatchMethod.visitMethodInsn(
           Opcodes.INVOKESTATIC,
-          InterfaceDesugaring.getCompanionClassName(emulatedInterface),
+          target,
           method.name(),
-          InterfaceDesugaring.companionDefaultMethodDescriptor(
-              emulatedInterface, method.descriptor()),
+          InterfaceDesugaring.companionDefaultMethodDescriptor(testedName, method.descriptor()),
           /*itf=*/ false);
       dispatchMethod.visitInsn(methodType.getReturnType().getOpcode(Opcodes.IRETURN));
 
@@ -400,7 +427,7 @@ class CoreLibrarySupport {
     return collectImplementedInterfaces(clazz, new LinkedHashSet<>())
         .stream()
         // search more subtypes before supertypes
-        .sorted(DefaultMethodClassFixer.InterfaceComparator.INSTANCE)
+        .sorted(DefaultMethodClassFixer.SubtypeComparator.INSTANCE)
         .map(itf -> findMethod(itf, name, desc))
         .filter(Objects::nonNull)
         .findFirst()
diff --git a/src/tools/android/java/com/google/devtools/build/android/desugar/DefaultMethodClassFixer.java b/src/tools/android/java/com/google/devtools/build/android/desugar/DefaultMethodClassFixer.java
index 292e14204d..853ed092ad 100644
--- a/src/tools/android/java/com/google/devtools/build/android/desugar/DefaultMethodClassFixer.java
+++ b/src/tools/android/java/com/google/devtools/build/android/desugar/DefaultMethodClassFixer.java
@@ -203,7 +203,7 @@ public class DefaultMethodClassFixer extends ClassVisitor {
   }
 
   private void stubMissingDefaultAndBridgeMethods() {
-    TreeSet<Class<?>> allInterfaces = new TreeSet<>(InterfaceComparator.INSTANCE);
+    TreeSet<Class<?>> allInterfaces = new TreeSet<>(SubtypeComparator.INSTANCE);
     for (String direct : directInterfaces) {
       // Loading ensures all transitively implemented interfaces can be loaded, which is necessary
       // to produce correct default method stubs in all cases.  We could do without classloading but
@@ -647,18 +647,17 @@ public class DefaultMethodClassFixer extends ClassVisitor {
     }
   }
 
-  /** Comparator for interfaces that compares by whether interfaces extend one another. */
-  enum InterfaceComparator implements Comparator<Class<?>> {
+  /** Comparator for classes and interfaces that compares by whether subtyping relationship. */
+  enum  SubtypeComparator implements Comparator<Class<?>> {
     /** Orders subtypes before supertypes and breaks ties lexicographically. */
     INSTANCE;
 
     @Override
     public int compare(Class<?> o1, Class<?> o2) {
-      checkArgument(o1.isInterface());
-      checkArgument(o2.isInterface());
       if (o1 == o2) {
         return 0;
       }
+      // order subtypes before supertypes
       if (o1.isAssignableFrom(o2)) { // o1 is supertype of o2
         return 1; // we want o1 to come after o2
       }