src/main/java/com/android/tools/r8/utils/LensUtils.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.utils;

import com.android.tools.r8.graph.DexDefinitionSupplier;
import com.android.tools.r8.graph.DexEncodedMethod;
import com.android.tools.r8.graph.DexReference;
import com.android.tools.r8.graph.lens.GraphLens;
import com.google.common.collect.Sets;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.IntFunction;

public class LensUtils {

  /**
   * Helper to lens rewrite a Map.
   *
   * @param map The Map that needs to be rewritten. The map may be mutated by this method, thus the
   *     original unrewritten map should never be used after calling this method.
   * @param mapFactory Function that given a capacity returns a Map with the given capacity.
   * @param preprocessing Function that given a (key, value) pair returns some data that is passed
   *     to the rewriting of the key using {@param keyRewriter} and to the rewriting of the value
   *     using {@param valueRewriter}. This can be used if the key and value shares some common data
   *     that needs rewriting.
   * @param keyRewriter Function that rewrites a key.
   * @param valueRewriter Function that rewrites a value.
   * @param valueMerger Function that merges two values that have the same rewritten key.
   */
  public static <K, V, P> Map<K, V> mutableRewriteMap(
      Map<K, V> map,
      IntFunction<Map<K, V>> mapFactory,
      BiFunction<K, V, P> preprocessing,
      TriFunction<K, V, P, K> keyRewriter,
      TriFunction<K, V, P, V> valueRewriter,
      TriFunction<K, V, V, V> valueMerger) {
    Collection<K> elementsToRemove = null;
    Map<K, V> rewrittenMap = null;
    for (Entry<K, V> entry : map.entrySet()) {
      K key = entry.getKey();
      V value = entry.getValue();
      P data = preprocessing.apply(key, value);
      K rewrittenKey = keyRewriter.apply(key, value, data);

      // If the key is mapped to null then remove if from `map` unless there are other changes.
      if (rewrittenKey == null) {
        if (rewrittenMap == null) {
          if (elementsToRemove == null) {
            elementsToRemove = new ArrayList<>();
          }
          elementsToRemove.add(key);
        }
        continue;
      }

      V rewrittenValue = valueRewriter.apply(key, value, data);
      if (rewrittenKey == key && rewrittenValue == value) {
        if (rewrittenMap == null) {
          // Don't write this entry until we find the first change.
          continue;
        }
      } else {
        if (rewrittenMap == null) {
          rewrittenMap = mapFactory.apply(map.size());
          MapUtils.forEachUntilExclusive(map, rewrittenMap::put, key);
          if (elementsToRemove != null) {
            assert !elementsToRemove.isEmpty();
            rewrittenMap.keySet().removeAll(elementsToRemove);
            elementsToRemove = null;
          }
        }
      }
      V existingRewrittenValue = rewrittenMap.get(rewrittenKey);
      rewrittenMap.put(
          rewrittenKey,
          existingRewrittenValue != null
              ? valueMerger.apply(rewrittenKey, rewrittenValue, existingRewrittenValue)
              : rewrittenValue);
    }
    if (rewrittenMap != null) {
      assert elementsToRemove == null;
      return MapUtils.trimCapacityIfSizeLessThan(rewrittenMap, mapFactory, map.size());
    } else {
      if (elementsToRemove != null) {
        assert !elementsToRemove.isEmpty();
        map.keySet().removeAll(elementsToRemove);
        return MapUtils.trimCapacity(map, mapFactory);
      }
      return map;
    }
  }

  public static Set<DexEncodedMethod> rewrittenWithRenamedSignature(
      Set<DexEncodedMethod> methods, DexDefinitionSupplier definitions, GraphLens lens) {
    GraphLens appliedLens = GraphLens.getIdentityLens();
    Set<DexEncodedMethod> result = Sets.newIdentityHashSet();
    for (DexEncodedMethod method : methods) {
      result.add(method.rewrittenWithLens(lens, appliedLens, definitions));
    }
    return result;
  }

  public static <T extends DexReference> void rewriteAndApplyIfNotPrimitiveType(
      GraphLens graphLens, T reference, Consumer<T> rewrittenConsumer) {
    T rewrittenReference = graphLens.rewriteReference(reference);
    // Enum unboxing can change a class type to int which leads to errors going forward since
    // the root set should not have primitive types.
    if (rewrittenReference.isDexType() && rewrittenReference.asDexType().isPrimitiveType()) {
      return;
    }
    rewrittenConsumer.accept(rewrittenReference);
  }
}