src/main/java/com/android/tools/r8/optimize/BridgeHoistingToSharedSyntheticSuperClass.java

// Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.
package com.android.tools.r8.optimize;

import static com.android.tools.r8.ir.optimize.info.OptimizationFeedback.getSimpleFeedback;
import static com.android.tools.r8.utils.internal.MapUtils.ignoreKey;

import com.android.tools.r8.contexts.CompilationContext.MainThreadContext;
import com.android.tools.r8.graph.AppView;
import com.android.tools.r8.graph.Code;
import com.android.tools.r8.graph.DexClass;
import com.android.tools.r8.graph.DexEncodedMethod;
import com.android.tools.r8.graph.DexItemFactory;
import com.android.tools.r8.graph.DexMethod;
import com.android.tools.r8.graph.DexMethodSignature;
import com.android.tools.r8.graph.DexProgramClass;
import com.android.tools.r8.graph.DexType;
import com.android.tools.r8.graph.MethodAccessFlags;
import com.android.tools.r8.graph.MethodResolutionResult;
import com.android.tools.r8.graph.ProgramMethod;
import com.android.tools.r8.graph.bytecodemetadata.BytecodeMetadataProvider;
import com.android.tools.r8.ir.code.IRCode;
import com.android.tools.r8.ir.code.InvokeStatic;
import com.android.tools.r8.ir.code.InvokeVirtual;
import com.android.tools.r8.ir.code.Value;
import com.android.tools.r8.ir.conversion.IRFinalizer;
import com.android.tools.r8.ir.optimize.DeadCodeRemover;
import com.android.tools.r8.ir.optimize.info.bridge.BridgeAnalyzer;
import com.android.tools.r8.ir.optimize.info.bridge.BridgeInfo;
import com.android.tools.r8.ir.optimize.info.bridge.StaticBridgeExcludingReceiverInfo;
import com.android.tools.r8.ir.optimize.info.bridge.VirtualBridgeInfo;
import com.android.tools.r8.optimize.bridgehoisting.BridgeHoisting;
import com.android.tools.r8.profile.rewriting.ProfileCollectionAdditions;
import com.android.tools.r8.shaking.AppInfoWithLiveness;
import com.android.tools.r8.synthesis.SyntheticItems;
import com.android.tools.r8.utils.InternalOptions;
import com.android.tools.r8.utils.InternalOptions.TestingOptions;
import com.android.tools.r8.utils.collections.DexMethodSignatureMap;
import com.android.tools.r8.utils.internal.ListUtils;
import com.android.tools.r8.utils.internal.OptionalBool;
import com.android.tools.r8.utils.internal.SetUtils;
import com.android.tools.r8.utils.timing.Timing;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.function.BiConsumer;

public class BridgeHoistingToSharedSyntheticSuperClass {

  private final AppView<AppInfoWithLiveness> appView;
  private final DexItemFactory factory;

  BridgeHoistingToSharedSyntheticSuperClass(AppView<AppInfoWithLiveness> appView) {
    this.appView = appView;
    this.factory = appView.dexItemFactory();
  }

  public static void run(
      AppView<AppInfoWithLiveness> appView, ExecutorService executorService, Timing timing)
      throws ExecutionException {
    InternalOptions options = appView.options();
    if (!options.isOptimizing() || !options.isShrinking()) {
      return;
    }
    if (!appView.options().canInitNewInstanceUsingSuperclassConstructor()) {
      // TODO(b/309575527): Extend to all runtimes.
      return;
    }
    TestingOptions testingOptions = options.getTestingOptions();
    if (!testingOptions.enableBridgeHoistingToSharedSyntheticSuperclass) {
      return;
    }
    timing.time(
        "BridgeHoistingToSharedSyntheticSuperClass",
        () ->
            new BridgeHoistingToSharedSyntheticSuperClass(appView)
                .internalRun(executorService, timing));
  }

  private void internalRun(ExecutorService executorService, Timing timing)
      throws ExecutionException {
    Collection<Group> groups = createInitialGroups(appView);
    ProfileCollectionAdditions profileCollectionAdditions =
        ProfileCollectionAdditions.create(appView);
    groups = refineGroups(groups, profileCollectionAdditions);
    if (!groups.isEmpty()) {
      rewriteApplication(groups);
      commitPendingSyntheticClasses(timing);
      updateArtProfiles(groups, profileCollectionAdditions);
      new BridgeHoisting(appView).run(executorService, timing);
    }
    appView.getTypeElementFactory().clearTypeElementsCache();
  }

  /** Returns the set of (non-singleton) groups that have the same superclass. */
  private Collection<Group> createInitialGroups(AppView<AppInfoWithLiveness> appView) {
    Map<DexClass, Group> groups = new LinkedHashMap<>();
    for (DexProgramClass clazz : appView.appInfo().classesWithDeterministicOrder()) {
      if (!clazz.hasSuperType()) {
        continue;
      }
      DexClass superclass = appView.definitionFor(clazz.getSuperType());
      if (superclass != null) {
        groups.computeIfAbsent(superclass, ignoreKey(Group::new)).addClass(clazz);
      }
    }
    groups.values().removeIf(Group::isSingleton);
    return groups.values();
  }

  private Collection<Group> refineGroups(
      Collection<Group> groups, ProfileCollectionAdditions profileCollectionAdditions) {
    Collection<Group> newGroups = new ArrayList<>();
    for (Group group : groups) {
      Iterables.addAll(newGroups, refineGroup(group, profileCollectionAdditions));
    }
    return newGroups;
  }

  /**
   * Splits the group into a collection of smaller groups that should receive a shared superclass.
   *
   * <p>For each class, this creates a specification of the bridges (a mapping from bridge method
   * signatures to their bridge implementation). Two classes are selected for getting a shared
   * synthetic super class if the bridge specification of one is a subset of the other (i.e., a
   * subset of the bridges can be shared and there are no bridges with the same signature that have
   * different behavior).
   */
  private Iterable<Group> refineGroup(
      Group group, ProfileCollectionAdditions profileCollectionAdditions) {
    List<Group> newGroups = new ArrayList<>();
    for (DexProgramClass clazz : group) {
      BridgeSpecification bridgeSpecification =
          getBridgeSpecification(clazz, profileCollectionAdditions);
      if (bridgeSpecification.isEmpty()) {
        continue;
      }
      Group targetGroup = getGroupForClass(newGroups, clazz, bridgeSpecification);
      if (targetGroup == null) {
        newGroups.add(new Group(clazz, bridgeSpecification));
      }
    }
    // Only introduce a shared super class for non-singleton groups that do not already have a
    // shared superclass in the first place.
    return Iterables.filter(
        newGroups, newGroup -> !newGroup.isSingleton() && newGroup.size() < group.size());
  }

  // TODO(b/309575527): Avoid building IR for all methods.
  private BridgeSpecification getBridgeSpecification(
      DexProgramClass clazz, ProfileCollectionAdditions profileCollectionAdditions) {
    BridgeSpecification bridgeSpecification = new BridgeSpecification();
    List<DexEncodedMethod> pendingMethods = new ArrayList<>();
    clazz.forEachProgramVirtualMethodMatching(
        DexEncodedMethod::hasCode,
        method -> {
          IRCode code = method.buildIR(appView);
          BridgeInfo bridgeInfo =
              BridgeAnalyzer.analyzeMethod(appView, method.getDefinition(), code);
          if (bridgeInfo == null
              || bridgeInfo.getInvokedMethod().getProto().isIdenticalTo(method.getProto())) {
            return;
          }
          if (bridgeInfo.isStaticBridgeExcludingReceiverInfo()) {
            tryMaterializeSpecializedOverloadOnLambdaClass(
                clazz,
                method,
                code,
                bridgeInfo.asStaticBridgeExcludingReceiverInfo(),
                bridgeSpecification,
                pendingMethods,
                profileCollectionAdditions);
          } else if (bridgeInfo.isVirtualBridgeInfo()) {
            VirtualBridgeInfo virtualBridgeInfo = bridgeInfo.asVirtualBridgeInfo();
            boolean isInvokedMethodPresentOnSuper =
                appView
                    .appInfo()
                    .resolveMethodOnClass(
                        clazz.getSuperType(), virtualBridgeInfo.getInvokedMethod())
                    .isSingleResolution();
            if (isInvokedMethodPresentOnSuper) {
              // No need to insert a method on a synthetic super class in this case.
              return;
            }
            bridgeSpecification.addBridge(method, virtualBridgeInfo);
            getSimpleFeedback().setBridgeInfo(method, virtualBridgeInfo);
          }
        });

    // Commit the synthesized methods, if any.
    clazz.addVirtualMethods(pendingMethods);

    return bridgeSpecification;
  }

  /**
   * For a lambda class with a main method `Object apply(Object)` that targets a static javac
   * synthetic method `Integer lambda$0(Integer)`, try to:
   *
   * <p>1. Insert a `Integer apply(Integer)` method that calls the `Integer lambda$0(Integer)`
   * method.
   *
   * <p>2. Update the `Object apply(Object)` method to call the newly added overload.
   */
  private void tryMaterializeSpecializedOverloadOnLambdaClass(
      DexProgramClass clazz,
      ProgramMethod method,
      IRCode code,
      StaticBridgeExcludingReceiverInfo bridgeInfo,
      BridgeSpecification bridgeSpecification,
      List<DexEncodedMethod> pendingMethods,
      ProfileCollectionAdditions profileCollectionAdditions) {
    SyntheticItems syntheticItems = appView.getSyntheticItems();
    if (!syntheticItems.isSynthetic(clazz)
        || !syntheticItems.hasKindThatMatches(clazz, (kind, n) -> kind.equals(n.LAMBDA))) {
      return;
    }

    // For a lambda main method such as `Object apply(Object)` that targets a method
    // `Integer lambda$0(Integer)`, we want to introduce a method `Integer apply(Integer)` on the
    // lambda class. First check that this method signature does not already exist in the hierarchy.
    DexMethod bridgeTarget = bridgeInfo.getInvokedMethod();
    DexMethod bridgeMethodReference =
        method.getReference().withProto(bridgeTarget.getProto(), factory);
    MethodResolutionResult resolutionResult =
        appView.appInfo().resolveMethodOn(clazz, bridgeMethodReference);
    if (resolutionResult.isSingleResolution()) {
      return;
    }

    // Synthesize a bridge method on the lambda class with the target signature.
    // TODO(b/309575527): Consider only materializing this method later if this actually leads to
    //  any sharing.
    DexEncodedMethod bridgeMethod =
        method.getDefinition().toTypeSubstitutedMethodAsInlining(bridgeMethodReference, factory);
    pendingMethods.add(bridgeMethod);
    profileCollectionAdditions.addMethodIfContextIsInProfile(
        bridgeMethod.asProgramMethod(clazz), method);

    // Update the current method to call the bridge instead.
    InvokeStatic invoke =
        code.<InvokeStatic>instructions(
                i ->
                    i.isInvokeStatic()
                        && i.asInvokeStatic().getInvokedMethod().isIdenticalTo(bridgeTarget))
            .iterator()
            .next();
    invoke.replace(
        InvokeVirtual.builder()
            .setArguments(
                ImmutableList.<Value>builder()
                    .add(code.getThis())
                    .addAll(invoke.arguments())
                    .build())
            .setIsInterface(clazz.isInterface())
            .setMethod(bridgeMethodReference)
            .setFreshOutValue(code, bridgeMethodReference.getReturnType().toTypeElement(appView))
            .setPosition(invoke)
            .build());

    // Commit the code.
    IRFinalizer<?> finalizer =
        code.getConversionOptions().getFinalizer(new DeadCodeRemover(appView), appView);
    Code newCode = finalizer.finalizeCode(code, BytecodeMetadataProvider.empty(), Timing.empty());
    method.setCode(newCode, appView);

    // Now the lambda main method is a virtual bridge to the newly synthesized method on the lambda.
    VirtualBridgeInfo virtualBridgeInfo = new VirtualBridgeInfo(bridgeMethodReference);
    bridgeSpecification.addBridge(method, virtualBridgeInfo);
    getSimpleFeedback().setBridgeInfo(method, virtualBridgeInfo);
  }

  private Group getGroupForClass(
      Collection<Group> groups, DexProgramClass clazz, BridgeSpecification bridgeSpecification) {
    for (Group group : groups) {
      if (bridgeSpecification.lessThanOrEquals(group.getBridgeSpecification())) {
        group.addClass(clazz);
        return group;
      } else if (group.getBridgeSpecification().lessThanOrEquals(bridgeSpecification)) {
        group.addClass(clazz);
        group.setBridgeSpecification(bridgeSpecification);
        return group;
      }
    }
    return null;
  }

  private void rewriteApplication(Collection<Group> groups) {
    MainThreadContext mainThreadContext =
        appView.createProcessorContext().createMainThreadContext();
    for (Group group : groups) {
      DexProgramClass representative = ListUtils.first(group.getClasses());
      Set<DexType> interfaces = SetUtils.newIdentityHashSet(representative.getInterfaces());
      for (DexProgramClass clazz : Iterables.skip(group.getClasses(), 1)) {
        interfaces.removeIf(type -> !clazz.getInterfaces().contains(type));
      }
      DexProgramClass syntheticSuperclass =
          appView
              .getSyntheticItems()
              .createClass(
                  kinds -> kinds.SHARED_SUPER_CLASS,
                  mainThreadContext.createUniqueContext(representative),
                  appView,
                  classBuilder -> {
                    classBuilder
                        .setAbstract()
                        .setSuperType(representative.getSuperType())
                        .setInterfaces(ListUtils.sort(interfaces, Comparator.naturalOrder()));
                    group
                        .getBridgeSpecification()
                        .forEach(
                            (bridge, target) ->
                                classBuilder.addMethod(
                                    methodBuilder ->
                                        methodBuilder
                                            .setAccessFlags(
                                                MethodAccessFlags.builder()
                                                    .setAbstract()
                                                    .setPublic()
                                                    .build())
                                            // TODO(b/309575527): Set correct api level.
                                            .setApiLevelForDefinition(appView.computedMinApiLevel())
                                            // TODO(b/309575527): Set correct library override info.
                                            .setIsLibraryMethodOverride(OptionalBool.FALSE)
                                            .setName(target.getName())
                                            .setProto(target.getProto())));
                  });

      // Fixup class hierarchy.
      for (DexProgramClass clazz : group) {
        clazz.setSuperType(syntheticSuperclass.getType());
        clazz.setInterfaces(clazz.getInterfaces().removeIf(interfaces::contains));
      }

      // Fixup instantiated hierarchy.
      Map<DexType, Set<DexClass>> instantiatedHierarchy =
          appView.appInfo().getMutableObjectAllocationInfoCollection().getInstantiatedHierarchy();
      for (DexProgramClass clazz : group) {
        for (DexType oldSupertype : syntheticSuperclass.allImmediateSupertypes()) {
          Set<DexClass> instantiatedSubclasses = instantiatedHierarchy.get(oldSupertype);
          if (instantiatedSubclasses != null && instantiatedSubclasses.remove(clazz)) {
            instantiatedSubclasses.add(syntheticSuperclass);
            instantiatedHierarchy
                .computeIfAbsent(syntheticSuperclass.getType(), ignoreKey(Sets::newIdentityHashSet))
                .add(clazz);
          }
        }
      }
    }
  }

  private void commitPendingSyntheticClasses(Timing timing) {
    assert appView.getSyntheticItems().hasPendingSyntheticClasses();
    appView.rebuildAppInfo(timing);
  }

  private void updateArtProfiles(
      Collection<Group> groups, ProfileCollectionAdditions profileCollectionAdditions) {
    if (profileCollectionAdditions.isNop()) {
      return;
    }
    for (Group group : groups) {
      for (DexProgramClass clazz : group) {
        profileCollectionAdditions.applyIfContextIsInProfile(
            clazz, additionsBuilder -> additionsBuilder.addClassRule(clazz.getSuperType()));
        group
            .getBridgeSpecification()
            .forEach(
                (bridge, target) -> {
                  DexEncodedMethod targetMethod = clazz.getMethodCollection().getMethod(target);
                  if (targetMethod != null) {
                    profileCollectionAdditions.applyIfContextIsInProfile(
                        targetMethod.getReference(),
                        additionsBuilder ->
                            additionsBuilder.addMethodRule(
                                target.withHolder(clazz.getSuperType(), appView.dexItemFactory())));
                  }
                });
      }
    }
    profileCollectionAdditions.commit(appView);
  }

  private static class Group implements Iterable<DexProgramClass> {

    private final List<DexProgramClass> classes;
    private BridgeSpecification bridgeSpecification;

    public Group() {
      this.classes = new ArrayList<>();
      this.bridgeSpecification = null;
    }

    public Group(DexProgramClass clazz, BridgeSpecification bridgeSpecification) {
      this.classes = ListUtils.newArrayList(clazz);
      this.bridgeSpecification = bridgeSpecification;
    }

    void addClass(DexProgramClass clazz) {
      classes.add(clazz);
    }

    BridgeSpecification getBridgeSpecification() {
      return bridgeSpecification;
    }

    List<DexProgramClass> getClasses() {
      return classes;
    }

    void setBridgeSpecification(BridgeSpecification bridgeSpecification) {
      this.bridgeSpecification = bridgeSpecification;
    }

    boolean isSingleton() {
      return size() == 1;
    }

    @Override
    public Iterator<DexProgramClass> iterator() {
      return classes.iterator();
    }

    public int size() {
      return classes.size();
    }
  }

  private static class BridgeSpecification {

    private final DexMethodSignatureMap<DexMethodSignature> bridges =
        DexMethodSignatureMap.create();

    void addBridge(ProgramMethod method, VirtualBridgeInfo bridgeInfo) {
      bridges.put(method, bridgeInfo.getInvokedMethod().getSignature());
    }

    boolean containsBridgeWithTarget(DexMethodSignature method, DexMethodSignature target) {
      return target.equals(bridges.get(method));
    }

    void forEach(BiConsumer<? super DexMethodSignature, ? super DexMethodSignature> consumer) {
      bridges.forEach(consumer);
    }

    boolean isEmpty() {
      return bridges.isEmpty();
    }

    boolean lessThanOrEquals(BridgeSpecification bridgeSpecification) {
      if (size() > bridgeSpecification.size()) {
        return false;
      }
      for (Entry<DexMethodSignature, DexMethodSignature> entry : bridges.entrySet()) {
        DexMethodSignature method = entry.getKey();
        DexMethodSignature target = entry.getValue();
        if (!bridgeSpecification.containsBridgeWithTarget(method, target)) {
          return false;
        }
      }
      return true;
    }

    int size() {
      return bridges.size();
    }
  }
}