Improve pruning of art profiles

Change-Id: Ifa7d4ac4c35fabb87214a75b65b0401f95f561c5
diff --git a/src/main/java/com/android/tools/r8/profile/art/ArtProfile.java b/src/main/java/com/android/tools/r8/profile/art/ArtProfile.java
index 784adbf..3eee55a 100644
--- a/src/main/java/com/android/tools/r8/profile/art/ArtProfile.java
+++ b/src/main/java/com/android/tools/r8/profile/art/ArtProfile.java
@@ -176,21 +176,8 @@
   }
 
   public ArtProfile withoutPrunedItems(PrunedItems prunedItems) {
-    return transform(
-        (classRule, classRuleBuilderFactory) -> {
-          if (!prunedItems.isRemoved(classRule.getType())) {
-            classRuleBuilderFactory.accept(classRule.getType());
-          }
-        },
-        (methodRule, classRuleBuilderFactory, methodRuleBuilderFactory) -> {
-          if (!prunedItems.isRemoved(methodRule.getMethod())) {
-            methodRuleBuilderFactory
-                .apply(methodRule.getMethod())
-                .acceptMethodRuleInfoBuilder(
-                    methodRuleInfoBuilder ->
-                        methodRuleInfoBuilder.merge(methodRule.getMethodRuleInfo()));
-          }
-        });
+    rules.keySet().removeIf(prunedItems::isRemoved);
+    return this;
   }
 
   private ArtProfile transform(
diff --git a/src/main/java/com/android/tools/r8/profile/art/NonEmptyArtProfileCollection.java b/src/main/java/com/android/tools/r8/profile/art/NonEmptyArtProfileCollection.java
index 51132d8..bcf3221 100644
--- a/src/main/java/com/android/tools/r8/profile/art/NonEmptyArtProfileCollection.java
+++ b/src/main/java/com/android/tools/r8/profile/art/NonEmptyArtProfileCollection.java
@@ -101,10 +101,9 @@
   @Override
   public NonEmptyArtProfileCollection withoutPrunedItems(PrunedItems prunedItems, Timing timing) {
     timing.begin("Prune NonEmptyArtProfileCollection");
-    NonEmptyArtProfileCollection result =
-        map(artProfile -> artProfile.withoutPrunedItems(prunedItems));
+    forEach(artProfile -> artProfile.withoutPrunedItems(prunedItems));
     timing.end();
-    return result;
+    return this;
   }
 
   private NonEmptyArtProfileCollection map(Function<ArtProfile, ArtProfile> fn) {