blob: 675d24cc0a7b7ff170adc48721a621d625bd06e9 [file] [log] [blame]
// 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.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) {
BiMap<V, K> biMap = HashBiMap.create();
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;
}
}