| // 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.horizontalclassmerging; |
| |
| import com.android.tools.r8.utils.collections.BidirectionalOneToOneHashMap; |
| import com.google.common.collect.BiMap; |
| import com.google.common.collect.HashBiMap; |
| import java.util.HashMap; |
| import java.util.HashSet; |
| import java.util.IdentityHashMap; |
| import java.util.Map; |
| import java.util.Map.Entry; |
| import java.util.Set; |
| import java.util.function.Function; |
| |
| /** |
| * This mapping class is used to track method mappings for horizontal class merging. Essentially it |
| * is a bidirectional many to one map, but with support for having unidirectional mappings and with |
| * support for remapping the values to new values using {@link ManyToOneMap#remap(BiMap, Function, |
| * Function)}. It also supports generating an inverse mapping {@link ManyToOneInverseMap} that can |
| * be used by the graph lens using {@link ManyToOneMap#inverse(Function)}. The inverse map is a |
| * bidirectional one to one map with additional non-bidirectional representative entries. |
| */ |
| public class ManyToOneMap<K, V> { |
| private final Map<K, V> forwardMap = new IdentityHashMap<>(); |
| private final Map<V, Set<K>> inverseMap = new IdentityHashMap<>(); |
| private final Map<V, K> representativeMap = new IdentityHashMap<>(); |
| |
| public Map<K, V> getForwardMap() { |
| return forwardMap; |
| } |
| |
| public Set<K> lookupReverse(V to) { |
| return inverseMap.get(to); |
| } |
| |
| public V put(K from, V to) { |
| return forwardMap.put(from, to); |
| } |
| |
| public void putInverse(K from, V to) { |
| inverseMap.computeIfAbsent(to, ignore -> new HashSet<>()).add(from); |
| } |
| |
| public K setRepresentative(K from, V to) { |
| putInverse(from, to); |
| return representativeMap.put(to, from); |
| } |
| |
| public ManyToOneInverseMap<K, V> inverse(Function<Set<K>, K> pickRepresentative) { |
| BidirectionalOneToOneHashMap<V, K> biMap = new BidirectionalOneToOneHashMap<>(); |
| Map<V, K> extraMap = new HashMap<>(); |
| for (Entry<V, Set<K>> entry : inverseMap.entrySet()) { |
| K representative = representativeMap.get(entry.getKey()); |
| if (entry.getValue().size() == 1) { |
| K singleton = entry.getValue().iterator().next(); |
| assert representative == null || singleton == representative; |
| if (representative == null) { |
| biMap.put(entry.getKey(), singleton); |
| } else { |
| extraMap.put(entry.getKey(), singleton); |
| } |
| } else { |
| if (representative == null) { |
| representative = pickRepresentative.apply(entry.getValue()); |
| } else { |
| assert representative == entry.getKey() || entry.getValue().contains(representative); |
| } |
| extraMap.put(entry.getKey(), representative); |
| } |
| } |
| |
| return new ManyToOneInverseMap<>(biMap, extraMap); |
| } |
| |
| public <NewV> ManyToOneMap<K, NewV> remap( |
| BiMap<V, NewV> biMap, Function<V, NewV> notInBiMap, Function<V, K> notInForwardMap) { |
| ManyToOneMap<K, NewV> newMap = new ManyToOneMap<>(); |
| |
| // All entries that should be remapped and are already in the forward and/or inverse mappings |
| // should only be remapped in the directions they are already mapped in. |
| BiMap<V, NewV> biMapCopy = HashBiMap.create(biMap); |
| for (Entry<V, Set<K>> entry : inverseMap.entrySet()) { |
| NewV to = biMapCopy.remove(entry.getKey()); |
| if (to == null) { |
| to = biMap.getOrDefault(entry.getKey(), notInBiMap.apply(entry.getKey())); |
| } |
| newMap.inverseMap.put(to, entry.getValue()); |
| } |
| for (Entry<K, V> entry : forwardMap.entrySet()) { |
| NewV newTo = biMapCopy.remove(entry.getValue()); |
| if (newTo == null) { |
| newTo = biMap.getOrDefault(entry.getValue(), notInBiMap.apply(entry.getValue())); |
| } |
| newMap.forwardMap.put(entry.getKey(), newTo); |
| } |
| |
| // All new entries should be mapped in both directions. |
| for (Entry<V, NewV> entry : biMapCopy.entrySet()) { |
| newMap.forwardMap.put(notInForwardMap.apply(entry.getKey()), entry.getValue()); |
| newMap |
| .inverseMap |
| .computeIfAbsent(entry.getValue(), ignore -> new HashSet<>()) |
| .add(notInForwardMap.apply(entry.getKey())); |
| } |
| |
| // Representatives are always in the inverse mapping, so they should always be remapped as new |
| // representatives. |
| for (Entry<V, K> entry : representativeMap.entrySet()) { |
| NewV newTo = biMap.get(entry.getKey()); |
| if (newTo == null) { |
| newTo = notInBiMap.apply(entry.getKey()); |
| } |
| newMap.representativeMap.put(newTo, entry.getValue()); |
| } |
| |
| return newMap; |
| } |
| } |