Support for class merging of Java lambdas

Change-Id: I1a2ca49e38bb98e319f21583bcd3e41676ebc85a
diff --git a/src/main/java/com/android/tools/r8/R8.java b/src/main/java/com/android/tools/r8/R8.java
index 9158d3b..c54ab80 100644
--- a/src/main/java/com/android/tools/r8/R8.java
+++ b/src/main/java/com/android/tools/r8/R8.java
@@ -513,7 +513,7 @@
                 PrunedItems.builder()
                     .setPrunedApp(appView.appInfo().app())
                     .addRemovedClasses(appView.horizontallyMergedClasses().getSources())
-                    .addNoLongerSyntheticItems(appView.horizontallyMergedClasses().getTargets())
+                    .addNoLongerSyntheticItems(appView.horizontallyMergedClasses().getSources())
                     .build());
           }
           timing.end();
diff --git a/src/main/java/com/android/tools/r8/horizontalclassmerging/CanOnlyMergeIntoClassPolicy.java b/src/main/java/com/android/tools/r8/horizontalclassmerging/CanOnlyMergeIntoClassPolicy.java
deleted file mode 100644
index fa9873a..0000000
--- a/src/main/java/com/android/tools/r8/horizontalclassmerging/CanOnlyMergeIntoClassPolicy.java
+++ /dev/null
@@ -1,17 +0,0 @@
-// 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.graph.DexProgramClass;
-
-public abstract class CanOnlyMergeIntoClassPolicy extends SingleClassPolicy {
-  public abstract boolean canOnlyMergeInto(DexProgramClass clazz);
-
-  @Override
-  public boolean canMerge(DexProgramClass program) {
-    // TODO(b/165577835): Allow merging of classes that must be the target of their merge group.
-    return !canOnlyMergeInto(program);
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/horizontalclassmerging/HorizontalClassMerger.java b/src/main/java/com/android/tools/r8/horizontalclassmerging/HorizontalClassMerger.java
index d97f12f..5c86a3a 100644
--- a/src/main/java/com/android/tools/r8/horizontalclassmerging/HorizontalClassMerger.java
+++ b/src/main/java/com/android/tools/r8/horizontalclassmerging/HorizontalClassMerger.java
@@ -11,7 +11,6 @@
 import com.android.tools.r8.horizontalclassmerging.policies.CheckAbstractClasses;
 import com.android.tools.r8.horizontalclassmerging.policies.DontInlinePolicy;
 import com.android.tools.r8.horizontalclassmerging.policies.DontMergeSynchronizedClasses;
-import com.android.tools.r8.horizontalclassmerging.policies.IgnoreSynthetics;
 import com.android.tools.r8.horizontalclassmerging.policies.LimitGroups;
 import com.android.tools.r8.horizontalclassmerging.policies.MinimizeFieldCasts;
 import com.android.tools.r8.horizontalclassmerging.policies.NoAnnotations;
@@ -37,6 +36,7 @@
 import com.android.tools.r8.horizontalclassmerging.policies.SameInstanceFields;
 import com.android.tools.r8.horizontalclassmerging.policies.SameNestHost;
 import com.android.tools.r8.horizontalclassmerging.policies.SameParentClass;
+import com.android.tools.r8.horizontalclassmerging.policies.SyntheticItemsPolicy;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.shaking.FieldAccessInfoCollectionModifier;
 import com.android.tools.r8.shaking.KeepInfoCollection;
@@ -57,7 +57,6 @@
     assert appView.options().enableInlining;
   }
 
-  // TODO(b/165577835): replace Collection<DexProgramClass> with MergeGroup
   public HorizontalClassMergerResult run(
       DirectMappedDexApplication.Builder appBuilder,
       RuntimeTypeCheckInfo runtimeTypeCheckInfo) {
@@ -128,7 +127,7 @@
         new NoAnnotations(),
         new NoEnums(appView),
         new CheckAbstractClasses(appView),
-        new IgnoreSynthetics(appView),
+        new SyntheticItemsPolicy(appView),
         new NoClassesOrMembersWithAnnotations(appView),
         new NoInnerClasses(),
         new NoClassInitializerWithObservableSideEffects(),
@@ -145,7 +144,7 @@
         new PreventMergeIntoDifferentMainDexGroups(appView),
         new AllInstantiatedOrUninstantiated(appView),
         new SameParentClass(),
-        new SameNestHost(),
+        new SameNestHost(appView),
         new PreserveMethodCharacteristics(appView),
         new SameFeatureSplit(appView),
         new RespectPackageBoundaries(appView),
diff --git a/src/main/java/com/android/tools/r8/horizontalclassmerging/MultiClassSameReferencePolicy.java b/src/main/java/com/android/tools/r8/horizontalclassmerging/MultiClassSameReferencePolicy.java
index 3cea5125..79db17b 100644
--- a/src/main/java/com/android/tools/r8/horizontalclassmerging/MultiClassSameReferencePolicy.java
+++ b/src/main/java/com/android/tools/r8/horizontalclassmerging/MultiClassSameReferencePolicy.java
@@ -15,11 +15,18 @@
   public final Collection<MergeGroup> apply(MergeGroup group) {
     Map<T, MergeGroup> groups = new LinkedHashMap<>();
     for (DexProgramClass clazz : group) {
-      groups.computeIfAbsent(getMergeKey(clazz), ignore -> new MergeGroup()).add(clazz);
+      T mergeKey = getMergeKey(clazz);
+      if (mergeKey != null) {
+        groups.computeIfAbsent(mergeKey, ignore -> new MergeGroup()).add(clazz);
+      }
     }
     removeTrivialGroups(groups.values());
     return groups.values();
   }
 
   public abstract T getMergeKey(DexProgramClass clazz);
+
+  protected final T ineligibleForClassInlining() {
+    return null;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/IgnoreSynthetics.java b/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/IgnoreSynthetics.java
deleted file mode 100644
index e45e2d3..0000000
--- a/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/IgnoreSynthetics.java
+++ /dev/null
@@ -1,28 +0,0 @@
-// 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.policies;
-
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.DexProgramClass;
-import com.android.tools.r8.horizontalclassmerging.SingleClassPolicy;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-
-public class IgnoreSynthetics extends SingleClassPolicy {
-
-  private final AppView<AppInfoWithLiveness> appView;
-
-  public IgnoreSynthetics(AppView<AppInfoWithLiveness> appView) {
-    this.appView = appView;
-  }
-
-  @Override
-  public boolean canMerge(DexProgramClass program) {
-    if (appView.getSyntheticItems().isSyntheticClass(program)) {
-      return appView.options().horizontalClassMergerOptions().isJavaLambdaMergingEnabled()
-          && appView.getSyntheticItems().isLegacySyntheticClass(program);
-    }
-    return true;
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/SameNestHost.java b/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/SameNestHost.java
index 6619878..88d2984 100644
--- a/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/SameNestHost.java
+++ b/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/SameNestHost.java
@@ -4,13 +4,23 @@
 
 package com.android.tools.r8.horizontalclassmerging.policies;
 
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.horizontalclassmerging.MultiClassSameReferencePolicy;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
 public class SameNestHost extends MultiClassSameReferencePolicy<DexType> {
+
+  private final DexItemFactory dexItemFactory;
+
+  public SameNestHost(AppView<AppInfoWithLiveness> appView) {
+    this.dexItemFactory = appView.dexItemFactory();
+  }
+
   @Override
   public DexType getMergeKey(DexProgramClass clazz) {
-    return clazz.getNestHost();
+    return clazz.isInANest() ? clazz.getNestHost() : dexItemFactory.objectType;
   }
 }
diff --git a/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/SyntheticItemsPolicy.java b/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/SyntheticItemsPolicy.java
new file mode 100644
index 0000000..ffc7b34
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/horizontalclassmerging/policies/SyntheticItemsPolicy.java
@@ -0,0 +1,51 @@
+// 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.policies;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.horizontalclassmerging.MultiClassSameReferencePolicy;
+import com.android.tools.r8.horizontalclassmerging.policies.SyntheticItemsPolicy.ClassKind;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.synthesis.SyntheticItems;
+import com.android.tools.r8.synthesis.SyntheticNaming.SyntheticKind;
+
+public class SyntheticItemsPolicy extends MultiClassSameReferencePolicy<ClassKind> {
+
+  enum ClassKind {
+    SYNTHETIC,
+    NOT_SYNTHETIC
+  }
+
+  private final AppView<AppInfoWithLiveness> appView;
+
+  public SyntheticItemsPolicy(AppView<AppInfoWithLiveness> appView) {
+    this.appView = appView;
+  }
+
+  @Override
+  public ClassKind getMergeKey(DexProgramClass clazz) {
+    SyntheticItems syntheticItems = appView.getSyntheticItems();
+
+    // Allow merging non-synthetics with non-synthetics.
+    if (!syntheticItems.isSyntheticClass(clazz)) {
+      return ClassKind.NOT_SYNTHETIC;
+    }
+
+    // Do not allow merging synthetics that are not lambdas.
+    if (!syntheticItems.isNonLegacySynthetic(clazz)
+        || syntheticItems.getNonLegacySyntheticKind(clazz) != SyntheticKind.LAMBDA) {
+      return ineligibleForClassInlining();
+    }
+
+    // Allow merging Java lambdas with Java lambdas.
+    if (appView.options().horizontalClassMergerOptions().isJavaLambdaMergingEnabled()) {
+      return ClassKind.SYNTHETIC;
+    }
+
+    // Java lambda merging is disabled.
+    return ineligibleForClassInlining();
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/synthesis/SyntheticItems.java b/src/main/java/com/android/tools/r8/synthesis/SyntheticItems.java
index 9922011..60a9543 100644
--- a/src/main/java/com/android/tools/r8/synthesis/SyntheticItems.java
+++ b/src/main/java/com/android/tools/r8/synthesis/SyntheticItems.java
@@ -212,6 +212,22 @@
     return isCommittedSynthetic(type) || isPendingSynthetic(type);
   }
 
+  public SyntheticKind getNonLegacySyntheticKind(DexProgramClass clazz) {
+    assert isNonLegacySynthetic(clazz);
+    SyntheticReference<?, ?, ?> reference = committed.getNonLegacyItem(clazz.getType());
+    if (reference == null) {
+      SyntheticDefinition<?, ?, ?> definition = pending.nonLegacyDefinitions.get(clazz.getType());
+      if (definition != null) {
+        reference = definition.toReference();
+      }
+    }
+    if (reference != null) {
+      return reference.getKind();
+    }
+    assert false;
+    return null;
+  }
+
   public boolean isSyntheticClass(DexType type) {
     return isLegacySyntheticClass(type) || isNonLegacySynthetic(type);
   }
diff --git a/src/main/java/com/android/tools/r8/synthesis/SyntheticProgramClassReference.java b/src/main/java/com/android/tools/r8/synthesis/SyntheticProgramClassReference.java
index 5a17b04..807af94 100644
--- a/src/main/java/com/android/tools/r8/synthesis/SyntheticProgramClassReference.java
+++ b/src/main/java/com/android/tools/r8/synthesis/SyntheticProgramClassReference.java
@@ -42,9 +42,6 @@
     // If the reference has been non-trivially rewritten the compiler has changed it and it can no
     // longer be considered a synthetic. The context may or may not have changed.
     if (type != rewritten && !lens.isSimpleRenaming(type, rewritten)) {
-      // If the referenced item is rewritten, it should be moved to another holder as the
-      // synthetic holder is no longer part of the synthetic collection.
-      assert SyntheticNaming.verifyNotInternalSynthetic(rewritten);
       return null;
     }
     if (rewrittenContext == getContext() && rewritten == type) {
diff --git a/src/main/java/com/android/tools/r8/utils/InternalOptions.java b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
index ea310aa..59d3b10 100644
--- a/src/main/java/com/android/tools/r8/utils/InternalOptions.java
+++ b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
@@ -1151,7 +1151,6 @@
 
     public boolean enable = true;
     public boolean enableConstructorMerging = true;
-    // TODO(b/174809311): Update or remove the option and its tests after new lambdas synthetics.
     public boolean enableJavaLambdaMerging = false;
 
     public int syntheticArgumentCount = 3;
diff --git a/src/test/java/com/android/tools/r8/classmerging/horizontal/JavaLambdaMergingTest.java b/src/test/java/com/android/tools/r8/classmerging/horizontal/JavaLambdaMergingTest.java
index 63c1f18..394c381 100644
--- a/src/test/java/com/android/tools/r8/classmerging/horizontal/JavaLambdaMergingTest.java
+++ b/src/test/java/com/android/tools/r8/classmerging/horizontal/JavaLambdaMergingTest.java
@@ -15,7 +15,6 @@
 import com.android.tools.r8.utils.codeinspector.VerticallyMergedClassesInspector;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.junit.Ignore;
 import org.junit.Test;
 
 public class JavaLambdaMergingTest extends HorizontalClassMergingTestBase {
@@ -25,7 +24,6 @@
   }
 
   @Test
-  @Ignore("b/174809311): Test does not work with hygienic lambdas. Rewrite or remove")
   public void test() throws Exception {
     testForR8(parameters.getBackend())
         .addInnerClasses(getClass())
@@ -59,7 +57,7 @@
   }
 
   private static boolean isLambda(DexType type) {
-    return SyntheticItemsTestUtils.isExternalLambda(
+    return SyntheticItemsTestUtils.isInternalLambda(
         Reference.classFromDescriptor(type.toDescriptorString()));
   }