src/main/java/com/android/tools/r8/repackaging/RepackagingConstraintGraph.java

// Copyright (c) 2020, 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.repackaging;

import com.android.tools.r8.graph.AppView;
import com.android.tools.r8.graph.DexClass;
import com.android.tools.r8.graph.DexDefinition;
import com.android.tools.r8.graph.DexEncodedMember;
import com.android.tools.r8.graph.DexEncodedMethod;
import com.android.tools.r8.graph.DexProgramClass;
import com.android.tools.r8.graph.MethodResolutionResult;
import com.android.tools.r8.graph.ProgramField;
import com.android.tools.r8.graph.ProgramMethod;
import com.android.tools.r8.graph.ProgramPackage;
import com.android.tools.r8.graph.ProgramPackageCollection;
import com.android.tools.r8.repackaging.Repackaging.RepackagingConfiguration;
import com.android.tools.r8.shaking.AppInfoWithLiveness;
import com.android.tools.r8.utils.InternalOptions.PackageObfuscationMode;
import com.android.tools.r8.utils.SetUtils;
import com.android.tools.r8.utils.ThreadUtils;
import com.android.tools.r8.utils.WorkList;
import java.util.IdentityHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;

/**
 * An undirected graph that contains a node for each class, field, and method in a given package.
 *
 * <p>An edge X <-> Y is added if X contains a reference to Y and Y is only accessible to X if X and
 * Y are in the same package.
 *
 * <p>Once the graph is populated, we compute the set of reachable nodes from the set of root nodes
 * that cannot be repackaged due to a -keep rule. The remaining nodes in the graph can all be
 * repackaged.
 */
public class RepackagingConstraintGraph {

  private final AppView<AppInfoWithLiveness> appView;
  private final PackageObfuscationMode packageObfuscationMode;
  private final ProgramPackage pkg;
  private final ProgramPackageCollection packages;
  private final Map<ProgramPackage, Set<DexProgramClass>> packagesWithClassesToRepackage;
  private final RepackagingConfiguration repackagingConfiguration;
  private final CrossPackageRepackagingConstraints crossPackageRepackagingConstraints;

  private final Map<DexDefinition, Node> nodes = new IdentityHashMap<>();
  private final Set<Node> pinnedNodes = ConcurrentHashMap.newKeySet();
  private final Node libraryBoundaryNode;

  public RepackagingConstraintGraph(
      AppView<AppInfoWithLiveness> appView,
      PackageObfuscationMode packageObfuscationMode,
      ProgramPackage pkg,
      ProgramPackageCollection packages,
      Map<ProgramPackage, Set<DexProgramClass>> packagesWithClassesToRepackage,
      RepackagingConfiguration repackagingConfiguration,
      CrossPackageRepackagingConstraints crossPackageRepackagingConstraints) {
    this.appView = appView;
    this.packageObfuscationMode = packageObfuscationMode;
    this.pkg = pkg;
    this.packages = packages;
    this.packagesWithClassesToRepackage = packagesWithClassesToRepackage;
    this.repackagingConfiguration = repackagingConfiguration;
    this.crossPackageRepackagingConstraints = crossPackageRepackagingConstraints;
    libraryBoundaryNode = createNode(appView.definitionFor(appView.dexItemFactory().objectType));
    pinnedNodes.add(libraryBoundaryNode);
  }

  /** Returns true if all classes in the package can be repackaged. */
  public void initializeGraph() {
    // Add all the items in the package into the graph. This way we know which items belong to the
    // package without having to extract package descriptor strings and comparing them with the
    // package descriptor.
    for (DexProgramClass clazz : pkg) {
      boolean isPinned = computeIsPinned(clazz);
      Node classNode = createNode(clazz);
      if (isPinned) {
        pinnedNodes.add(classNode);
      }
      for (DexEncodedMember<?, ?> member : clazz.members()) {
        Node memberNode = createNode(member);
        classNode.addNeighbor(memberNode);
      }
    }
  }

  private boolean computeIsPinned(DexProgramClass clazz) {
    if (!appView.appInfo().isRepackagingAllowed(clazz, appView)
        || !appView.options().getSyntheticItemsOptions().isRepackagingAllowed(clazz, appView)) {
      return true;
    }
    // This class is also pinned if it has a cross-package constraint saying that other classes must
    // not be repackaged, when in fact they were.
    Set<DexProgramClass> notRepackagedConstraints =
        crossPackageRepackagingConstraints.removeNotRepackagedConstraints(clazz);
    for (DexProgramClass notRepackagedConstraint : notRepackagedConstraints) {
      ProgramPackage programPackage = packages.getPackage(notRepackagedConstraint);
      Set<DexProgramClass> repackagedClassesInProgramPackage =
          packagesWithClassesToRepackage.get(programPackage);
      if (repackagedClassesInProgramPackage.contains(notRepackagedConstraint)) {
        return true;
      }
    }
    return false;
  }

  private Node createNode(DexDefinition definition) {
    Node node = new Node(definition);
    nodes.put(definition, node);
    return node;
  }

  Node getNode(DexDefinition definition) {
    if (definition.isNotProgramDefinition(appView)) {
      String packageDescriptor = definition.getContextType().getPackageDescriptor();
      if (packageDescriptor.equals(pkg.getPackageDescriptor())) {
        return libraryBoundaryNode;
      }
      return null;
    }
    return nodes.get(definition);
  }

  public void populateConstraints(ExecutorService executorService) throws ExecutionException {
    // Concurrently add references from methods to the graph.
    ThreadUtils.processItems(
        pkg::forEachMethod,
        this::registerReferencesFromMethod,
        appView.options().getThreadingModule(),
        executorService);

    // TODO(b/165783399): Evaluate if it is worth to parallelize this. The work per field and class
    //  should be little, so it may not be.
    pkg.forEachClass(this::registerReferencesFromClass);
    pkg.forEachField(this::registerReferencesFromField);
  }

  private void registerReferencesFromClass(DexProgramClass clazz) {
    RepackagingUseRegistry registry =
        new RepackagingUseRegistry(appView, this, clazz, libraryBoundaryNode);

    // Trace the references to the immediate super types.
    if (clazz.superType != null) {
      registry.registerTypeReference(clazz.getSuperType(), appView.graphLens());
    }
    clazz.interfaces.forEach(type -> registry.registerTypeReference(type, appView.graphLens()));

    // Trace the references from the class annotations.
    new RepackagingAnnotationTracer(appView, registry).trace(clazz.annotations());

    // Trace the references in the nest host and/or members.
    if (clazz.isInANest()) {
      if (clazz.isNestHost()) {
        clazz.forEachNestMember(type -> registry.registerTypeReference(type, appView.graphLens()));
      } else {
        assert clazz.isNestMember();
        registry.registerTypeReference(clazz.getNestHost(), appView.graphLens());
      }
    }

    // Trace the references to the inner and outer classes.
    clazz
        .getInnerClasses()
        .forEach(
            innerClassAttribute ->
                registry.registerInnerClassAttribute(clazz, innerClassAttribute));

    // Trace the references from the enclosing method and nest attributes.
    registry.registerEnclosingMethodAttribute(clazz.getEnclosingMethodAttribute());
    registry.registerNestHostAttribute(clazz.getNestHostClassAttribute());
    registry.registerNestMemberClassAttributes(clazz.getNestMembersClassAttributes());
  }

  private void registerReferencesFromField(ProgramField field) {
    RepackagingUseRegistry registry =
        new RepackagingUseRegistry(appView, this, field, libraryBoundaryNode);

    // Trace the type of the field.
    registry.registerTypeReference(field.getReference().getType(), appView.graphLens());

    // Trace the references in the field annotations.
    new RepackagingAnnotationTracer(appView, registry).trace(field.getDefinition().annotations());
  }

  private void registerReferencesFromMethod(ProgramMethod method) {
    DexEncodedMethod definition = method.getDefinition();
    RepackagingUseRegistry registry =
        new RepackagingUseRegistry(appView, this, method, libraryBoundaryNode);

    // Trace the type references in the method signature.
    definition
        .getProto()
        .forEachType(type -> registry.registerTypeReference(type, appView.graphLens()));

    // Check if this overrides a package-private method.
    if (method.getHolder().hasSuperType()) {
      DexClass superClass =
          appView.definitionFor(method.getHolder().getSuperType(), method.getHolder());
      if (superClass != null) {
        MethodResolutionResult resolutionResult =
            appView.appInfo().resolveMethodOnLegacy(superClass, method.getReference());
        registry.registerMemberAccess(resolutionResult);
        preserveInvalidOverridesOfPackagePrivateMethod(method, resolutionResult);
      }
    }

    // Trace the references in the method and method parameter annotations.
    RepackagingAnnotationTracer annotationTracer =
        new RepackagingAnnotationTracer(appView, registry);
    annotationTracer.trace(definition.annotations());
    annotationTracer.trace(definition.getParameterAnnotations());

    // Trace the references from the code.
    if (definition.hasCode()) {
      definition.getCode().registerCodeReferences(method, registry);
    }
  }

  private void preserveInvalidOverridesOfPackagePrivateMethod(
      ProgramMethod method, MethodResolutionResult resolutionResult) {
    if (!packageObfuscationMode.isRepackageClasses()) {
      // We are not merging packages.
      return;
    }

    DexEncodedMethod resolvedMethod = resolutionResult.getResolvedMethod();
    DexClass resolvedHolder = resolutionResult.getResolvedHolder();
    if (resolvedMethod == null || !resolvedMethod.getAccessFlags().isPackagePrivate()) {
      // Not a package-private method.
      return;
    }

    if (!resolvedHolder.isProgramClass()) {
      // Not a program method, so the holder class is not subject to repackaging in the first place,
      // and hence there is no merging of packages.
      return;
    }

    if (resolvedHolder.isSamePackage(method)) {
      // Not an invalid override of a package-private method.
      return;
    }

    // Make sure that the two methods do not end up in the same package, as that would change
    // semantics of virtual dispatch.
    DexProgramClass resolvedProgramHolder = resolvedHolder.asProgramClass();
    ProgramPackage resolvedPackage = packages.getPackage(resolvedProgramHolder);
    if (resolvedPackage == null) {
      // This only happens for packages that are already in the target location.
      pinnedNodes.add(getNode(method.getDefinition()));
    } else {
      Set<DexProgramClass> classesToRepackageInResolvedPackage =
          packagesWithClassesToRepackage.get(resolvedPackage);
      if (classesToRepackageInResolvedPackage == null) {
        // This package has not been processed yet. Mark the superclass as ineligible for
        // repackaging if the current class ends up being repackaged.
        crossPackageRepackagingConstraints.onlyRepackageClassIfOtherClassIsNotRepackaged(
            resolvedProgramHolder, method.getHolder());
      } else if (classesToRepackageInResolvedPackage.contains(resolvedProgramHolder)) {
        // Disallow repackaging of the current class.
        pinnedNodes.add(getNode(method.getDefinition()));
      } else if (repackagingConfiguration.isPackageInTargetLocation(resolvedPackage)) {
        assert false : "Expected resolvedPackage to be null, but was: " + resolvedPackage;
        // Disallow repackaging of the current class.
        pinnedNodes.add(getNode(method.getDefinition()));
      }
    }
  }

  public Set<DexProgramClass> computeClassesToRepackage() {
    WorkList<Node> worklist = WorkList.newIdentityWorkList(pinnedNodes);
    while (worklist.hasNext()) {
      Node pinnedNode = worklist.next();
      for (Node neighbor : pinnedNode.getNeighbors()) {
        // Mark all the immediate neighbors as ineligible for repackaging and continue the tracing
        // from the neighbors.
        worklist.addIfNotSeen(neighbor);
      }
    }
    Set<Node> pinnedNodes = worklist.getSeenSet();
    Set<DexProgramClass> classesToRepackage = new LinkedHashSet<>();
    for (DexProgramClass clazz : pkg) {
      if (!pinnedNodes.contains(getNode(clazz))) {
        classesToRepackage.add(clazz);
      }
    }
    return classesToRepackage;
  }

  static class Node {

    private final DexDefinition definitionForDebugging;
    private final Set<Node> neighbors = SetUtils.newConcurrentHashSet();

    Node(DexDefinition definitionForDebugging) {
      this.definitionForDebugging = definitionForDebugging;
    }

    public void addNeighbor(Node neighbor) {
      neighbors.add(neighbor);
      neighbor.neighbors.add(this);
    }

    public Set<Node> getNeighbors() {
      return neighbors;
    }

    @Override
    public String toString() {
      return "Node(" + definitionForDebugging.getReference().toSourceString() + ")";
    }
  }
}