Ensure correct lens rewriting of reprocessing sets

This fixes some subtle issues that may arise due to inconsistent rewriting of reprocessing sets in the enum unboxer, the inliner, and the staticizer.

We now consistently use a LongLivedProgramMethodSet for the reprocessing sets, and maintain the invariant that all methods in these sets are always rewritten up until the same graph lens.

This ensures that we never reapply parts of a graph lens to the same method, which may lead to incorrect rewriting.

Change-Id: I6ce3fd3d52e36f3090901c82805dad8806a99def
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
index 8f6f293..60ecf33 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
@@ -56,7 +56,6 @@
 import com.android.tools.r8.ir.desugar.desugaredlibrary.DesugaredLibraryAPIConverter;
 import com.android.tools.r8.ir.desugar.itf.EmulatedInterfaceApplicationRewriter;
 import com.android.tools.r8.ir.desugar.itf.InterfaceMethodProcessorFacade;
-import com.android.tools.r8.ir.desugar.itf.InterfaceMethodRewriter;
 import com.android.tools.r8.ir.desugar.lambda.LambdaDeserializationMethodRemover;
 import com.android.tools.r8.ir.desugar.nest.D8NestBasedAccessDesugaring;
 import com.android.tools.r8.ir.optimize.AssertionsRewriter;
@@ -100,6 +99,7 @@
 import com.android.tools.r8.shaking.LibraryMethodOverrideAnalysis;
 import com.android.tools.r8.utils.Action;
 import com.android.tools.r8.utils.CfgPrinter;
+import com.android.tools.r8.utils.ConsumerUtils;
 import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.ExceptionUtils;
 import com.android.tools.r8.utils.InternalOptions;
@@ -136,7 +136,6 @@
   private final StringOptimizer stringOptimizer;
   private final StringBuilderOptimizer stringBuilderOptimizer;
   private final IdempotentFunctionCallCanonicalizer idempotentFunctionCallCanonicalizer;
-  private final InterfaceMethodRewriter interfaceMethodRewriter;
   private final ClassInliner classInliner;
   private final ClassStaticizer classStaticizer;
   private final InternalOptions options;
@@ -220,7 +219,6 @@
       // - invoke-special desugaring.
       assert options.desugarState.isOn();
       this.instructionDesugaring = CfInstructionDesugaringCollection.create(appView);
-      this.interfaceMethodRewriter = null;
       this.covariantReturnTypeAnnotationTransformer = null;
       this.dynamicTypeOptimization = null;
       this.classInliner = null;
@@ -246,10 +244,6 @@
         appView.enableWholeProgramOptimizations()
             ? CfInstructionDesugaringCollection.empty()
             : CfInstructionDesugaringCollection.create(appView);
-    this.interfaceMethodRewriter =
-        options.isInterfaceMethodDesugaringEnabled() && appView.enableWholeProgramOptimizations()
-            ? new InterfaceMethodRewriter(appView, this)
-            : null;
     this.covariantReturnTypeAnnotationTransformer =
         options.processCovariantReturnTypeAnnotations
             ? new CovariantReturnTypeAnnotationTransformer(this, appView.dexItemFactory())
@@ -346,17 +340,10 @@
         D8NestBasedAccessDesugaring::clearNestAttributes);
   }
 
-  private void staticizeClasses(
-      OptimizationFeedback feedback, ExecutorService executorService, GraphLens applied)
+  private void staticizeClasses(OptimizationFeedback feedback, ExecutorService executorService)
       throws ExecutionException {
     if (classStaticizer != null) {
-      classStaticizer.staticizeCandidates(feedback, executorService, applied);
-    }
-  }
-
-  private void collectStaticizerCandidates(DexApplication application) {
-    if (classStaticizer != null) {
-      classStaticizer.collectCandidates(application);
+      classStaticizer.staticizeCandidates(feedback, executorService);
     }
   }
 
@@ -627,7 +614,6 @@
     DexApplication application = appView.appInfo().app();
 
     computeReachabilitySensitivity(application);
-    collectStaticizerCandidates(application);
     workaroundAbstractMethodOnNonAbstractClassVerificationBug(
         executorService, simpleOptimizationFeedback);
 
@@ -639,7 +625,9 @@
 
     printPhase("Primary optimization pass");
 
-    // Setup the argument propagator for the primary optimization pass.
+    GraphLens graphLensForPrimaryOptimizationPass = appView.graphLens();
+
+    // Setup optimizations for the primary optimization pass.
     appView.withArgumentPropagator(
         argumentPropagator -> argumentPropagator.initializeCodeScanner(executorService, timing));
     appView.withCallSiteOptimizationInfoPropagator(
@@ -649,16 +637,26 @@
           optimization.abandonCallSitePropagationForPinnedMethodsAndOverrides(
               executorService, timing);
         });
+    ConsumerUtils.acceptIfNotNull(
+        enumUnboxer,
+        enumUnboxer ->
+            enumUnboxer.initializeEnumUnboxingCandidates(graphLensForPrimaryOptimizationPass));
+    ConsumerUtils.acceptIfNotNull(
+        classStaticizer,
+        classStaticizer ->
+            classStaticizer.prepareForPrimaryOptimizationPass(graphLensForPrimaryOptimizationPass));
+    ConsumerUtils.acceptIfNotNull(
+        inliner,
+        inliner -> inliner.initializeDoubleInlineCallers(graphLensForPrimaryOptimizationPass));
 
     if (fieldAccessAnalysis != null) {
       fieldAccessAnalysis.fieldAssignmentTracker().initialize();
     }
 
     // Process the application identifying outlining candidates.
-    GraphLens initialGraphLensForIR = appView.graphLens();
-    GraphLens graphLensForIR = initialGraphLensForIR;
     OptimizationFeedbackDelayed feedback = delayedOptimizationFeedback;
-    PostMethodProcessor.Builder postMethodProcessorBuilder = new PostMethodProcessor.Builder();
+    PostMethodProcessor.Builder postMethodProcessorBuilder =
+        new PostMethodProcessor.Builder(graphLensForPrimaryOptimizationPass);
     {
       timing.begin("Build primary method processor");
       PrimaryMethodProcessor primaryMethodProcessor =
@@ -669,6 +667,7 @@
       if (outliner != null) {
         outliner.createOutlineMethodIdentifierGenerator();
       }
+      assert appView.graphLens() == graphLensForPrimaryOptimizationPass;
       primaryMethodProcessor.forEachMethod(
           (method, methodProcessingContext) ->
               processDesugaredMethod(
@@ -677,8 +676,8 @@
           this::waveDone,
           timing,
           executorService);
+      assert appView.graphLens() == graphLensForPrimaryOptimizationPass;
       timing.end();
-      assert graphLensForIR == appView.graphLens();
     }
 
     // The field access info collection is not maintained during IR processing.
@@ -696,42 +695,54 @@
     // Commit synthetics from the primary optimization pass.
     commitPendingSyntheticItemsR8(appView);
 
+    // Post processing:
+    //   1) Second pass for methods whose collected call site information become more precise.
+    //   2) Second inlining pass for dealing with double inline callers.
+    printPhase("Post optimization pass");
+
     // Analyze the data collected by the argument propagator, use the analysis result to update
     // the parameter optimization infos, and rewrite the application.
     appView.withArgumentPropagator(
         argumentPropagator ->
             argumentPropagator.tearDownCodeScanner(
                 postMethodProcessorBuilder, executorService, timing));
+    appView.withCallSiteOptimizationInfoPropagator(
+        callSiteOptimizationInfoPropagator ->
+            callSiteOptimizationInfoPropagator.enqueueMethodsForReprocessing(
+                postMethodProcessorBuilder));
 
     if (libraryMethodOverrideAnalysis != null) {
       libraryMethodOverrideAnalysis.finish();
     }
 
-    // Post processing:
-    //   1) Second pass for methods whose collected call site information become more precise.
-    //   2) Second inlining pass for dealing with double inline callers.
-    printPhase("Post optimization pass");
-    appView.withCallSiteOptimizationInfoPropagator(
-        optimization ->
-            postMethodProcessorBuilder.put(appView.callSiteOptimizationInfoPropagator()));
-    if (inliner != null) {
-      postMethodProcessorBuilder.put(inliner);
-    }
+    ConsumerUtils.acceptIfNotNull(
+        inliner, inliner -> inliner.enqueueMethodsForReprocessing(postMethodProcessorBuilder));
+
     if (!options.debug) {
       new TrivialFieldAccessReprocessor(appView.withLiveness(), postMethodProcessorBuilder)
           .run(executorService, feedback, timing);
     }
+
     if (enumUnboxer != null) {
       enumUnboxer.unboxEnums(this, postMethodProcessorBuilder, executorService, feedback);
     } else {
       appView.setUnboxedEnums(EnumDataMap.empty());
     }
+
+    GraphLens graphLensForSecondaryOptimizationPass = appView.graphLens();
+
+    ConsumerUtils.acceptIfNotNull(
+        classStaticizer,
+        classStaticizer ->
+            classStaticizer.prepareForSecondaryOptimizationPass(
+                graphLensForSecondaryOptimizationPass));
+
     timing.begin("IR conversion phase 2");
-    graphLensForIR = appView.graphLens();
     PostMethodProcessor postMethodProcessor =
         postMethodProcessorBuilder.build(appView, executorService, timing);
     if (postMethodProcessor != null) {
       assert !options.debug;
+      assert appView.graphLens() == graphLensForSecondaryOptimizationPass;
       postMethodProcessor.forEachMethod(
           (method, methodProcessingContext) ->
               processDesugaredMethod(
@@ -739,7 +750,7 @@
           feedback,
           executorService);
       feedback.updateVisibleOptimizationInfo();
-      assert graphLensForIR == appView.graphLens();
+      assert appView.graphLens() == graphLensForSecondaryOptimizationPass;
     }
     timing.end();
 
@@ -756,7 +767,7 @@
     if (!options.isGeneratingClassFiles()) {
       printPhase("Class staticizer post processing");
       // TODO(b/127694949): Adapt to PostOptimization.
-      staticizeClasses(feedback, executorService, initialGraphLensForIR);
+      staticizeClasses(feedback, executorService);
       feedback.updateVisibleOptimizationInfo();
       // The class staticizer lens shall not be applied through lens code rewriting or it breaks
       // the lambda merger.
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/PostMethodProcessor.java b/src/main/java/com/android/tools/r8/ir/conversion/PostMethodProcessor.java
index a8e5508..36c8dc6 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/PostMethodProcessor.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/PostMethodProcessor.java
@@ -55,15 +55,22 @@
 
   public static class Builder {
 
-    private final LongLivedProgramMethodSetBuilder<?> methodsToReprocessBuilder =
-        LongLivedProgramMethodSetBuilder.createForIdentitySet();
+    private final LongLivedProgramMethodSetBuilder<ProgramMethodSet> methodsToReprocessBuilder;
 
-    Builder() {}
+    Builder(GraphLens graphLensForPrimaryOptimizationPass) {
+      this.methodsToReprocessBuilder =
+          LongLivedProgramMethodSetBuilder.createForIdentitySet(
+              graphLensForPrimaryOptimizationPass);
+    }
 
     public void add(ProgramMethod method) {
       methodsToReprocessBuilder.add(method);
     }
 
+    public LongLivedProgramMethodSetBuilder<ProgramMethodSet> getMethodsToReprocessBuilder() {
+      return methodsToReprocessBuilder;
+    }
+
     public void put(ProgramMethodSet methodsToRevisit) {
       methodsToRevisit.forEach(this::add);
     }
@@ -72,15 +79,11 @@
       put(postOptimization.methodsToRevisit());
     }
 
-    public void removePrunedMethods(Iterable<DexMethod> prunedMethod) {
-      methodsToReprocessBuilder.removeAll(prunedMethod);
-    }
-
     // Some optimizations may change methods, creating new instances of the encoded methods with a
     // new signature. The compiler needs to update the set of methods that must be reprocessed
     // according to the graph lens.
-    public void rewrittenWithLens(AppView<AppInfoWithLiveness> appView, GraphLens applied) {
-      methodsToReprocessBuilder.rewrittenWithLens(appView, applied);
+    public void rewrittenWithLens(AppView<AppInfoWithLiveness> appView) {
+      methodsToReprocessBuilder.rewrittenWithLens(appView);
     }
 
     PostMethodProcessor build(
@@ -103,8 +106,7 @@
         // Nothing to revisit.
         return null;
       }
-      ProgramMethodSet methodsToReprocess =
-          methodsToReprocessBuilder.build(appView, appView.graphLens());
+      ProgramMethodSet methodsToReprocess = methodsToReprocessBuilder.build(appView);
       CallGraph callGraph =
           new PartialCallGraphBuilder(appView, methodsToReprocess).build(executorService, timing);
       return new PostMethodProcessor(appView, callGraph);
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java b/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
index ce0646d..d39f509 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
@@ -11,6 +11,7 @@
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexMethodHandle;
 import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.LookupResult;
 import com.android.tools.r8.graph.MethodResolutionResult.SingleResolutionResult;
 import com.android.tools.r8.graph.ProgramMethod;
@@ -19,7 +20,7 @@
 import com.android.tools.r8.ir.code.InvokeCustom;
 import com.android.tools.r8.ir.code.InvokeMethod;
 import com.android.tools.r8.ir.code.InvokeMethodWithReceiver;
-import com.android.tools.r8.ir.conversion.PostOptimization;
+import com.android.tools.r8.ir.conversion.PostMethodProcessor;
 import com.android.tools.r8.ir.optimize.info.CallSiteOptimizationInfo;
 import com.android.tools.r8.ir.optimize.info.ConcreteCallSiteOptimizationInfo;
 import com.android.tools.r8.logging.Log;
@@ -31,6 +32,7 @@
 import com.android.tools.r8.utils.LazyBox;
 import com.android.tools.r8.utils.ThreadUtils;
 import com.android.tools.r8.utils.Timing;
+import com.android.tools.r8.utils.collections.LongLivedProgramMethodSetBuilder;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
 import com.google.common.collect.Sets;
 import java.util.Set;
@@ -38,7 +40,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.function.Consumer;
 
-public class CallSiteOptimizationInfoPropagator implements PostOptimization {
+public class CallSiteOptimizationInfoPropagator {
 
   // TODO(b/139246447): should we revisit new targets over and over again?
   //   Maybe piggy-back on MethodProcessor's wave/batch processing?
@@ -61,16 +63,17 @@
   private final AppView<AppInfoWithLiveness> appView;
   private final InternalOptions options;
   private final CallSiteOptimizationOptions optimizationOptions;
-  private ProgramMethodSet revisitedMethods = null;
   private Mode mode = Mode.COLLECT;
 
+  private ProgramMethodSet revisitedMethodsForTesting = null;
+
   public CallSiteOptimizationInfoPropagator(AppView<AppInfoWithLiveness> appView) {
     assert appView.enableWholeProgramOptimizations();
     this.appView = appView;
     this.options = appView.options();
     this.optimizationOptions = appView.options().callSiteOptimizationOptions();
     if (Log.isLoggingEnabledFor(CallSiteOptimizationInfoPropagator.class)) {
-      revisitedMethods = ProgramMethodSet.create();
+      revisitedMethodsForTesting = ProgramMethodSet.create();
     }
   }
 
@@ -80,9 +83,9 @@
 
   public void logResults() {
     assert Log.ENABLED;
-    if (revisitedMethods != null) {
-      Log.info(getClass(), "# of methods to revisit: %s", revisitedMethods.size());
-      for (ProgramMethod m : revisitedMethods) {
+    if (revisitedMethodsForTesting != null) {
+      Log.info(getClass(), "# of methods to revisit: %s", revisitedMethodsForTesting.size());
+      for (ProgramMethod m : revisitedMethodsForTesting) {
         Log.info(
             getClass(),
             "%s: %s",
@@ -379,10 +382,19 @@
     return callSiteOptimizationInfo;
   }
 
-  @Override
-  public ProgramMethodSet methodsToRevisit() {
+  public void enqueueMethodsForReprocessing(
+      PostMethodProcessor.Builder postMethodProcessorBuilder) {
+    postMethodProcessorBuilder
+        .getMethodsToReprocessBuilder()
+        .rewrittenWithLens(appView.graphLens())
+        .merge(methodsToRevisit());
+  }
+
+  private LongLivedProgramMethodSetBuilder<ProgramMethodSet> methodsToRevisit() {
     mode = Mode.REVISIT;
-    ProgramMethodSet targetsToRevisit = ProgramMethodSet.create();
+    GraphLens currentGraphLens = appView.graphLens();
+    LongLivedProgramMethodSetBuilder<ProgramMethodSet> builder =
+        LongLivedProgramMethodSetBuilder.createForIdentitySet(currentGraphLens);
     for (DexProgramClass clazz : appView.appInfo().classes()) {
       clazz.forEachProgramMethodMatching(
           definition -> {
@@ -398,14 +410,14 @@
             return callSiteOptimizationInfo.hasUsefulOptimizationInfo(appView, definition);
           },
           method -> {
-            targetsToRevisit.add(method);
+            builder.add(method, currentGraphLens);
             appView.options().testing.callSiteOptimizationInfoInspector.accept(method);
           });
     }
-    if (revisitedMethods != null) {
-      revisitedMethods.addAll(targetsToRevisit);
+    if (revisitedMethodsForTesting != null) {
+      revisitedMethodsForTesting.addAll(builder.build(appView));
     }
-    return targetsToRevisit;
+    return builder;
   }
 
   private synchronized boolean verifyAllProgramDispatchTargetsHaveBeenAbandoned(
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/Inliner.java b/src/main/java/com/android/tools/r8/ir/optimize/Inliner.java
index fdc9208..0ddc915 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/Inliner.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/Inliner.java
@@ -48,7 +48,7 @@
 import com.android.tools.r8.ir.code.Value;
 import com.android.tools.r8.ir.conversion.LensCodeRewriter;
 import com.android.tools.r8.ir.conversion.MethodProcessor;
-import com.android.tools.r8.ir.conversion.PostOptimization;
+import com.android.tools.r8.ir.conversion.PostMethodProcessor;
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedback;
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedbackIgnore;
 import com.android.tools.r8.ir.optimize.inliner.DefaultInliningReasonStrategy;
@@ -64,6 +64,7 @@
 import com.android.tools.r8.utils.ListUtils;
 import com.android.tools.r8.utils.SetUtils;
 import com.android.tools.r8.utils.Timing;
+import com.android.tools.r8.utils.collections.LongLivedProgramMethodSetBuilder;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -77,7 +78,7 @@
 import java.util.Map;
 import java.util.Set;
 
-public class Inliner implements PostOptimization {
+public class Inliner {
 
   protected final AppView<AppInfoWithLiveness> appView;
   private final Set<DexMethod> extraNeverInlineMethods;
@@ -86,7 +87,7 @@
 
   // State for inlining methods which are known to be called twice.
   private boolean applyDoubleInlining = false;
-  private final ProgramMethodSet doubleInlineCallers = ProgramMethodSet.create();
+  private LongLivedProgramMethodSetBuilder<ProgramMethodSet> doubleInlineCallers;
   private final ProgramMethodSet doubleInlineSelectedTargets = ProgramMethodSet.create();
   private final Map<DexEncodedMethod, ProgramMethod> doubleInlineeCandidates =
       new IdentityHashMap<>();
@@ -226,9 +227,10 @@
 
     if (doubleInlineeCandidates.containsKey(target.getDefinition())) {
       // Both calls can be inlined.
+      GraphLens currentGraphLens = appView.graphLens();
       ProgramMethod doubleInlineeCandidate = doubleInlineeCandidates.get(target.getDefinition());
-      doubleInlineCallers.add(doubleInlineeCandidate);
-      doubleInlineCallers.add(method);
+      doubleInlineCallers.add(doubleInlineeCandidate, currentGraphLens);
+      doubleInlineCallers.add(method, currentGraphLens);
       doubleInlineSelectedTargets.add(target);
     } else {
       // First call can be inlined.
@@ -236,10 +238,20 @@
     }
   }
 
-  @Override
-  public ProgramMethodSet methodsToRevisit() {
+  public void initializeDoubleInlineCallers(GraphLens graphLensForPrimaryOptimizationPass) {
+    assert appView.graphLens() == graphLensForPrimaryOptimizationPass;
+    doubleInlineCallers =
+        LongLivedProgramMethodSetBuilder.createForIdentitySet(graphLensForPrimaryOptimizationPass);
+  }
+
+  public void enqueueMethodsForReprocessing(
+      PostMethodProcessor.Builder postMethodProcessorBuilder) {
+    // The double inline callers are always rewritten up until the graph lens of the primary
+    // optimization pass, so we can safely merge them into the methods to reprocess (which may be
+    // rewritten with a newer graph lens).
+    postMethodProcessorBuilder.getMethodsToReprocessBuilder().merge(doubleInlineCallers);
+    doubleInlineCallers = null;
     applyDoubleInlining = true;
-    return doubleInlineCallers;
   }
 
   /**
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxer.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxer.java
index 2e74bb2..705ce93 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxer.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxer.java
@@ -26,6 +26,7 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexCallSite;
 import com.android.tools.r8.graph.DexClass;
+import com.android.tools.r8.graph.DexClassAndField;
 import com.android.tools.r8.graph.DexClassAndMethod;
 import com.android.tools.r8.graph.DexEncodedField;
 import com.android.tools.r8.graph.DexEncodedMember;
@@ -101,6 +102,7 @@
 import com.android.tools.r8.utils.Reporter;
 import com.android.tools.r8.utils.StringDiagnostic;
 import com.android.tools.r8.utils.collections.ImmutableInt2ReferenceSortedMap;
+import com.android.tools.r8.utils.collections.LongLivedProgramMethodSetBuilder;
 import com.android.tools.r8.utils.collections.ProgramMethodMap;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
 import com.google.common.collect.HashMultiset;
@@ -135,18 +137,20 @@
   private final DexItemFactory factory;
   // Map the enum candidates with their dependencies, i.e., the methods to reprocess for the given
   // enum if the optimization eventually decides to unbox it.
-  private final EnumUnboxingCandidateInfoCollection enumUnboxingCandidatesInfo;
+  private EnumUnboxingCandidateInfoCollection enumUnboxingCandidatesInfo;
   private final Set<DexProgramClass> candidatesToRemoveInWave = Sets.newConcurrentHashSet();
   private final Map<DexType, EnumStaticFieldValues> staticFieldValuesMap =
       new ConcurrentHashMap<>();
-  private final ProgramMethodSet methodsDependingOnLibraryModelisation =
-      ProgramMethodSet.createConcurrent();
+
+  // Methods depending on library modelisation need to be reprocessed so they are peephole
+  // optimized.
+  private LongLivedProgramMethodSetBuilder<ProgramMethodSet> methodsDependingOnLibraryModelisation;
 
   // Map from checkNotNull() methods to the enums that use the given method.
   private final ProgramMethodMap<Set<DexProgramClass>> checkNotNullMethods =
       ProgramMethodMap.createConcurrent();
 
-  private final DexEncodedField ordinalField;
+  private final DexClassAndField ordinalField;
 
   private EnumUnboxingRewriter enumUnboxerRewriter;
 
@@ -164,21 +168,18 @@
       debugLogs = null;
     }
     assert !appView.options().debug;
-    enumUnboxingCandidatesInfo = new EnumUnboxingCandidateAnalysis(appView, this).findCandidates();
-
     ordinalField =
-        appView.appInfo().resolveField(factory.enumMembers.ordinalField).getResolvedField();
-    if (ordinalField == null) {
-      // This can happen when compiling for non standard libraries, in that case, this effectively
-      // disables the enum unboxer.
-      enumUnboxingCandidatesInfo.clear();
-    }
+        appView.appInfo().resolveField(factory.enumMembers.ordinalField).getResolutionPair();
   }
 
   public static int ordinalToUnboxedInt(int ordinal) {
     return ordinal + 1;
   }
 
+  public DexClassAndField getOrdinalField() {
+    return ordinalField;
+  }
+
   public void updateEnumUnboxingCandidatesInfo() {
     for (DexProgramClass candidate : candidatesToRemoveInWave) {
       enumUnboxingCandidatesInfo.removeCandidate(candidate);
@@ -203,7 +204,7 @@
   }
 
   private void markMethodDependsOnLibraryModelisation(ProgramMethod method) {
-    methodsDependingOnLibraryModelisation.add(method);
+    methodsDependingOnLibraryModelisation.add(method, appView.graphLens());
   }
 
   private DexProgramClass getEnumUnboxingCandidateOrNull(TypeElement lattice) {
@@ -298,7 +299,7 @@
         enumUnboxingCandidatesInfo.addMethodDependency(eligibleEnum, code.context());
       }
     }
-    if (methodsDependingOnLibraryModelisation.contains(code.context())) {
+    if (methodsDependingOnLibraryModelisation.contains(code.context(), appView.graphLens())) {
       conversionOptions.disablePeepholeOptimizations();
     }
   }
@@ -555,6 +556,16 @@
     return result;
   }
 
+  public void initializeEnumUnboxingCandidates(GraphLens graphLensForPrimaryOptimizationPass) {
+    assert enumUnboxingCandidatesInfo == null;
+    enumUnboxingCandidatesInfo =
+        new EnumUnboxingCandidateAnalysis(appView, this)
+            .findCandidates(graphLensForPrimaryOptimizationPass);
+    methodsDependingOnLibraryModelisation =
+        LongLivedProgramMethodSetBuilder.createConcurrentForIdentitySet(
+            graphLensForPrimaryOptimizationPass);
+  }
+
   public void unboxEnums(
       IRConverter converter,
       PostMethodProcessor.Builder postBuilder,
@@ -577,7 +588,8 @@
     ImmutableSet<DexType> enumsToUnbox = enumUnboxingCandidatesInfo.candidates();
     ImmutableSet<DexProgramClass> enumClassesToUnbox =
         enumUnboxingCandidatesInfo.candidateClasses();
-    ProgramMethodSet dependencies = enumUnboxingCandidatesInfo.allMethodDependencies();
+    LongLivedProgramMethodSetBuilder<ProgramMethodSet> dependencies =
+        enumUnboxingCandidatesInfo.allMethodDependencies();
     enumUnboxingCandidatesInfo.clear();
     // Update keep info on any of the enum methods of the removed classes.
     updateKeepInfo(enumsToUnbox);
@@ -587,11 +599,32 @@
             .synthesizeEnumUnboxingUtilityClasses(enumClassesToUnbox, enumDataMap)
             .build(converter, executorService);
 
+    // Fixup the application.
     EnumUnboxingTreeFixer.Result treeFixerResult =
         new EnumUnboxingTreeFixer(
                 appView, checkNotNullMethods, enumDataMap, enumClassesToUnbox, utilityClasses)
             .fixupTypeReferences(converter, executorService);
     EnumUnboxingLens enumUnboxingLens = treeFixerResult.getLens();
+    appView.setUnboxedEnums(enumDataMap);
+
+    // Update the graph lens.
+    appView.rewriteWithLens(enumUnboxingLens);
+
+    // Enqueue the (lens rewritten) methods that require reprocessing.
+    //
+    // Note that the reprocessing set must be rewritten to the new enum unboxing lens before pruning
+    // the builders with the methods removed by the tree fixer (since these methods references are
+    // already fully lens rewritten).
+    postBuilder
+        .getMethodsToReprocessBuilder()
+        .rewrittenWithLens(appView)
+        .merge(dependencies)
+        .merge(methodsDependingOnLibraryModelisation)
+        .removeAll(treeFixerResult.getPrunedItems().getRemovedMethods());
+    methodsDependingOnLibraryModelisation.clear();
+
+    updateOptimizationInfos(executorService, feedback, treeFixerResult.getPrunedItems());
+
     enumUnboxerRewriter =
         new EnumUnboxingRewriter(
             appView,
@@ -600,17 +633,6 @@
             enumUnboxingLens,
             enumDataMap,
             utilityClasses);
-    appView.setUnboxedEnums(enumDataMap);
-    GraphLens previousLens = appView.graphLens();
-    appView.rewriteWithLens(enumUnboxingLens);
-    updateOptimizationInfos(executorService, feedback, treeFixerResult.getPrunedItems());
-    postBuilder.put(dependencies);
-    // Methods depending on library modelisation need to be reprocessed so they are peephole
-    // optimized.
-    postBuilder.put(methodsDependingOnLibraryModelisation);
-    methodsDependingOnLibraryModelisation.clear();
-    postBuilder.removePrunedMethods(treeFixerResult.getPrunedItems().getRemovedMethods());
-    postBuilder.rewrittenWithLens(appView, previousLens);
   }
 
   private void updateOptimizationInfos(
@@ -842,7 +864,7 @@
   }
 
   private OptionalInt getOrdinal(ObjectState state) {
-    AbstractValue field = state.getAbstractFieldValue(ordinalField);
+    AbstractValue field = state.getAbstractFieldValue(getOrdinalField().getDefinition());
     if (field.isSingleNumberValue()) {
       return OptionalInt.of(field.asSingleNumberValue().getIntValue());
     }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingCandidateAnalysis.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingCandidateAnalysis.java
index a675735..4225447 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingCandidateAnalysis.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingCandidateAnalysis.java
@@ -11,6 +11,7 @@
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexProto;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.ir.optimize.enums.eligibility.Reason;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.shaking.KeepInfoCollection;
@@ -35,10 +36,16 @@
     factory = appView.dexItemFactory();
   }
 
-  EnumUnboxingCandidateInfoCollection findCandidates() {
+  EnumUnboxingCandidateInfoCollection findCandidates(
+      GraphLens graphLensForPrimaryOptimizationPass) {
+    if (enumUnboxer.getOrdinalField() == null || enumUnboxer.getOrdinalField().isProgramField()) {
+      // This can happen when compiling for non standard libraries, in that case, this effectively
+      // disables the enum unboxer.
+      return enumToUnboxCandidates;
+    }
     for (DexProgramClass clazz : appView.appInfo().classes()) {
       if (isEnumUnboxingCandidate(clazz)) {
-        enumToUnboxCandidates.addCandidate(clazz);
+        enumToUnboxCandidates.addCandidate(appView, clazz, graphLensForPrimaryOptimizationPass);
       }
     }
     removeEnumsInAnnotations();
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingCandidateInfoCollection.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingCandidateInfoCollection.java
index 7d85b8d..386bcb0 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingCandidateInfoCollection.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingCandidateInfoCollection.java
@@ -4,13 +4,18 @@
 
 package com.android.tools.r8.ir.optimize.enums;
 
+import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.collections.LongLivedProgramMethodSetBuilder;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
@@ -21,9 +26,14 @@
 
   private final Map<DexType, EnumUnboxingCandidateInfo> enumTypeToInfo = new ConcurrentHashMap<>();
 
-  public void addCandidate(DexProgramClass enumClass) {
+  public void addCandidate(
+      AppView<AppInfoWithLiveness> appView,
+      DexProgramClass enumClass,
+      GraphLens graphLensForPrimaryOptimizationPass) {
     assert !enumTypeToInfo.containsKey(enumClass.type);
-    enumTypeToInfo.put(enumClass.type, new EnumUnboxingCandidateInfo(enumClass));
+    enumTypeToInfo.put(
+        enumClass.type,
+        new EnumUnboxingCandidateInfo(appView, enumClass, graphLensForPrimaryOptimizationPass));
   }
 
   public void removeCandidate(DexProgramClass enumClass) {
@@ -62,10 +72,13 @@
     return info.enumClass;
   }
 
-  public ProgramMethodSet allMethodDependencies() {
-    ProgramMethodSet allMethodDependencies = ProgramMethodSet.create();
-    for (EnumUnboxingCandidateInfo info : enumTypeToInfo.values()) {
-      allMethodDependencies.addAll(info.methodDependencies);
+  public LongLivedProgramMethodSetBuilder<ProgramMethodSet> allMethodDependencies() {
+    Iterator<EnumUnboxingCandidateInfo> candidateInfoIterator = enumTypeToInfo.values().iterator();
+    assert candidateInfoIterator.hasNext();
+    LongLivedProgramMethodSetBuilder<ProgramMethodSet> allMethodDependencies =
+        candidateInfoIterator.next().methodDependencies;
+    while (candidateInfoIterator.hasNext()) {
+      allMethodDependencies.merge(candidateInfoIterator.next().methodDependencies);
     }
     return allMethodDependencies;
   }
@@ -111,20 +124,27 @@
   private static class EnumUnboxingCandidateInfo {
 
     private final DexProgramClass enumClass;
-    private final ProgramMethodSet methodDependencies = ProgramMethodSet.createConcurrent();
+    private final LongLivedProgramMethodSetBuilder<ProgramMethodSet> methodDependencies;
     private final Set<DexField> requiredInstanceFieldData = Sets.newConcurrentHashSet();
 
-    public EnumUnboxingCandidateInfo(DexProgramClass enumClass) {
+    public EnumUnboxingCandidateInfo(
+        AppView<AppInfoWithLiveness> appView,
+        DexProgramClass enumClass,
+        GraphLens graphLensForPrimaryOptimizationPass) {
       assert enumClass != null;
+      assert appView.graphLens() == graphLensForPrimaryOptimizationPass;
       this.enumClass = enumClass;
+      this.methodDependencies =
+          LongLivedProgramMethodSetBuilder.createConcurrentForIdentitySet(
+              graphLensForPrimaryOptimizationPass);
     }
 
     public DexProgramClass getEnumClass() {
       return enumClass;
     }
 
-    public void addMethodDependency(ProgramMethod programMethod) {
-      methodDependencies.add(programMethod);
+    public void addMethodDependency(ProgramMethod method) {
+      methodDependencies.add(method);
     }
 
     public void addRequiredInstanceFieldData(DexField field) {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/staticizer/ClassStaticizer.java b/src/main/java/com/android/tools/r8/ir/optimize/staticizer/ClassStaticizer.java
index 1f6c00a..0fbb615 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/staticizer/ClassStaticizer.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/staticizer/ClassStaticizer.java
@@ -7,7 +7,6 @@
 import static com.android.tools.r8.graph.DexProgramClass.asProgramClassOrNull;
 
 import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.DexApplication;
 import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexClassAndField;
 import com.android.tools.r8.graph.DexEncodedField;
@@ -58,6 +57,8 @@
   private final DexItemFactory factory;
   private final IRConverter converter;
 
+  private GraphLens graphLensForOptimizationPass;
+
   // Represents a staticizing candidate with all information
   // needed for staticizing.
   final class CandidateInfo {
@@ -113,14 +114,36 @@
     this.converter = converter;
   }
 
+  public void prepareForPrimaryOptimizationPass(GraphLens graphLensForPrimaryOptimizationPass) {
+    collectCandidates();
+    this.graphLensForOptimizationPass = graphLensForPrimaryOptimizationPass;
+  }
+
+  public void prepareForSecondaryOptimizationPass(GraphLens graphLensForSecondaryOptimizationPass) {
+    // Rewrite all the referenced from sets such that they are all rewritten up until the lens of
+    // the second optimization pass. This is needed to ensure all elements in the referenced from
+    // sets are rewritten up until the same graph lens, in case any referenced from sets are
+    // extended during the secondary optimization pass.
+    assert appView.graphLens() == graphLensForSecondaryOptimizationPass;
+    referencedFrom
+        .values()
+        .forEach(
+            referencedFromBuilder ->
+                referencedFromBuilder.rewrittenWithLens(graphLensForSecondaryOptimizationPass));
+    this.graphLensForOptimizationPass = graphLensForSecondaryOptimizationPass;
+  }
+
   // Before doing any usage-based analysis we collect a set of classes that can be
   // candidates for staticizing. This analysis is very simple, but minimizes the
   // set of eligible classes staticizer tracks and thus time and memory it needs.
-  public final void collectCandidates(DexApplication app) {
+  public final void collectCandidates() {
     Set<DexType> notEligible = Sets.newIdentityHashSet();
     Map<DexType, DexEncodedField> singletonFields = new HashMap<>();
 
-    app.classes()
+    assert !appView.getSyntheticItems().hasPendingSyntheticClasses();
+    appView
+        .appInfo()
+        .classes()
         .forEach(
             cls -> {
               // We only consider classes eligible for staticizing if there is just
@@ -169,20 +192,26 @@
             });
 
     // Finalize the set of the candidates.
-    app.classes().forEach(cls -> {
-      DexType type = cls.type;
-      if (!notEligible.contains(type)) {
-        DexEncodedField field = singletonFields.get(type);
-        if (field != null && // Singleton field found
-            !field.accessFlags.isVolatile() && // Don't remove volatile fields.
-            !isPinned(cls, field)) { // Don't remove pinned objects.
-          assert field.accessFlags.isStatic();
-          // Note: we don't check that the field is final, since we will analyze
-          //       later how and where it is initialized.
-          new CandidateInfo(cls, field); // will self-register
-        }
-      }
-    });
+    appView
+        .appInfo()
+        .classes()
+        .forEach(
+            cls -> {
+              DexType type = cls.type;
+              if (!notEligible.contains(type)) {
+                DexEncodedField field = singletonFields.get(type);
+                if (field != null
+                    && // Singleton field found
+                    !field.accessFlags.isVolatile()
+                    && // Don't remove volatile fields.
+                    !isPinned(cls, field)) { // Don't remove pinned objects.
+                  assert field.accessFlags.isStatic();
+                  // Note: we don't check that the field is final, since we will analyze
+                  //       later how and where it is initialized.
+                  new CandidateInfo(cls, field); // will self-register
+                }
+              }
+            });
   }
 
   private void markNotEligible(DexType type, Set<DexType> notEligible) {
@@ -374,9 +403,13 @@
   }
 
   private void addReferencedFrom(CandidateInfo info, ProgramMethod context) {
+    GraphLens currentGraphLens = appView.graphLens();
+    assert currentGraphLens == graphLensForOptimizationPass;
     LongLivedProgramMethodSetBuilder<?> builder =
         referencedFrom.computeIfAbsent(
-            info, ignore -> LongLivedProgramMethodSetBuilder.createConcurrentForIdentitySet());
+            info,
+            ignore ->
+                LongLivedProgramMethodSetBuilder.createConcurrentForIdentitySet(currentGraphLens));
     builder.add(context);
   }
 
@@ -689,9 +722,8 @@
   //  3. Rewrite methods referencing staticized members, also remove instance creation
   //
   public final void staticizeCandidates(
-      OptimizationFeedback feedback, ExecutorService executorService, GraphLens applied)
-      throws ExecutionException {
-    new StaticizingProcessor(appView, this, converter, applied).run(feedback, executorService);
+      OptimizationFeedback feedback, ExecutorService executorService) throws ExecutionException {
+    new StaticizingProcessor(appView, this, converter).run(feedback, executorService);
   }
 
   private class CallSiteReferencesInvalidator extends UseRegistry {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/staticizer/StaticizingProcessor.java b/src/main/java/com/android/tools/r8/ir/optimize/staticizer/StaticizingProcessor.java
index 3e59aa8..2fbaa44 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/staticizer/StaticizingProcessor.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/staticizer/StaticizingProcessor.java
@@ -19,7 +19,6 @@
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.analysis.type.TypeElement;
 import com.android.tools.r8.ir.code.IRCode;
@@ -87,17 +86,14 @@
   private final Map<DexField, CandidateInfo> singletonFields = new IdentityHashMap<>();
   private final Map<DexMethod, CandidateInfo> singletonGetters = new IdentityHashMap<>();
   private final Map<DexType, DexType> candidateToHostMapping = new IdentityHashMap<>();
-  private final GraphLens applied;
 
   StaticizingProcessor(
       AppView<AppInfoWithLiveness> appView,
       ClassStaticizer classStaticizer,
-      IRConverter converter,
-      GraphLens applied) {
+      IRConverter converter) {
     this.appView = appView;
     this.classStaticizer = classStaticizer;
     this.converter = converter;
-    this.applied = applied;
   }
 
   final void run(OptimizationFeedback feedback, ExecutorService executorService)
@@ -238,7 +234,7 @@
         LongLivedProgramMethodSetBuilder<?> referencedFromBuilder =
             classStaticizer.referencedFrom.remove(info);
         assert referencedFromBuilder != null;
-        referencedFrom = referencedFromBuilder.build(appView, applied);
+        referencedFrom = referencedFromBuilder.build(appView);
         materializedReferencedFromCollections.put(info, referencedFrom);
       } else {
         referencedFrom = ProgramMethodSet.empty();
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
index 2761296..61f52a9 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
@@ -9,6 +9,7 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.ImmediateProgramSubtypingInfo;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.code.IRCode;
@@ -23,6 +24,8 @@
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.ThreadUtils;
 import com.android.tools.r8.utils.Timing;
+import com.android.tools.r8.utils.collections.LongLivedProgramMethodSetBuilder;
+import com.android.tools.r8.utils.collections.ProgramMethodSet;
 import java.util.Collection;
 import java.util.List;
 import java.util.Set;
@@ -198,6 +201,14 @@
    * postMethodProcessorBuilder}.
    */
   private void enqueueMethodsForProcessing(PostMethodProcessor.Builder postMethodProcessorBuilder) {
+    // Before adding any methods to the reprocessing set, make sure that the set is rewritten up
+    // until the point of the current graph lens, such that all elements in the set are rewritten
+    // up until the same lens.
+    GraphLens currentGraphLens = appView.graphLens();
+    LongLivedProgramMethodSetBuilder<ProgramMethodSet> methodsToReprocessBuilder =
+        postMethodProcessorBuilder
+            .getMethodsToReprocessBuilder()
+            .rewrittenWithLens(currentGraphLens);
     for (DexProgramClass clazz : appView.appInfo().classes()) {
       clazz.forEachProgramMethodMatching(
           DexEncodedMethod::hasCode,
@@ -206,7 +217,7 @@
                 method.getDefinition().getCallSiteOptimizationInfo();
             if (callSiteOptimizationInfo.isConcreteCallSiteOptimizationInfo()
                 && !appView.appInfo().isNeverReprocessMethod(method.getReference())) {
-              postMethodProcessorBuilder.add(method);
+              methodsToReprocessBuilder.add(method, currentGraphLens);
               appView.testing().callSiteOptimizationInfoInspector.accept(method);
             }
           });
diff --git a/src/main/java/com/android/tools/r8/utils/ConsumerUtils.java b/src/main/java/com/android/tools/r8/utils/ConsumerUtils.java
index 7259af1..4f652cf 100644
--- a/src/main/java/com/android/tools/r8/utils/ConsumerUtils.java
+++ b/src/main/java/com/android/tools/r8/utils/ConsumerUtils.java
@@ -19,6 +19,12 @@
     return curry(function).apply(arg);
   }
 
+  public static <T> void acceptIfNotNull(T element, Consumer<T> consumer) {
+    if (element != null) {
+      consumer.accept(element);
+    }
+  }
+
   public static <T> Consumer<T> acceptIfNotSeen(Consumer<T> consumer, Set<T> seen) {
     return element -> {
       if (seen.add(element)) {
diff --git a/src/main/java/com/android/tools/r8/utils/SetUtils.java b/src/main/java/com/android/tools/r8/utils/SetUtils.java
index bd97bd7..04b399d 100644
--- a/src/main/java/com/android/tools/r8/utils/SetUtils.java
+++ b/src/main/java/com/android/tools/r8/utils/SetUtils.java
@@ -10,6 +10,7 @@
 import java.util.HashSet;
 import java.util.IdentityHashMap;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 
 public class SetUtils {
@@ -23,6 +24,10 @@
     return false;
   }
 
+  public static <T> Set<T> newConcurrentHashSet(int capacity) {
+    return Collections.newSetFromMap(new ConcurrentHashMap<>(capacity));
+  }
+
   public static <T> HashSet<T> newHashSet(T element) {
     HashSet<T> result = new HashSet<>(1);
     result.add(element);
diff --git a/src/main/java/com/android/tools/r8/utils/collections/LongLivedProgramMethodSetBuilder.java b/src/main/java/com/android/tools/r8/utils/collections/LongLivedProgramMethodSetBuilder.java
index 5152dde..4002425 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/LongLivedProgramMethodSetBuilder.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/LongLivedProgramMethodSetBuilder.java
@@ -10,70 +10,133 @@
 import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.google.common.collect.Sets;
+import com.android.tools.r8.utils.SetUtils;
 import java.util.Set;
 import java.util.function.IntFunction;
 
 public class LongLivedProgramMethodSetBuilder<T extends ProgramMethodSet> {
 
+  // Factory for creating the final ProgramMethodSet.
   private final IntFunction<T> factory;
-  private final Set<DexMethod> methods;
 
-  private LongLivedProgramMethodSetBuilder(IntFunction<T> factory, Set<DexMethod> methods) {
+  // Factory for creating a Set<DexMethod>.
+  private final IntFunction<Set<DexMethod>> factoryForBuilder;
+
+  // The graph lens that this collection has been rewritten up until.
+  private GraphLens appliedGraphLens;
+
+  // The methods in this collection.
+  private Set<DexMethod> methods;
+
+  private LongLivedProgramMethodSetBuilder(
+      GraphLens currentGraphLens,
+      IntFunction<T> factory,
+      IntFunction<Set<DexMethod>> factoryBuilder) {
+    this.appliedGraphLens = currentGraphLens;
     this.factory = factory;
-    this.methods = methods;
+    this.factoryForBuilder = factoryBuilder;
+    this.methods = factoryBuilder.apply(2);
   }
 
-  public static LongLivedProgramMethodSetBuilder<?> createForIdentitySet() {
+  public static LongLivedProgramMethodSetBuilder<ProgramMethodSet> createForIdentitySet(
+      GraphLens currentGraphLens) {
     return new LongLivedProgramMethodSetBuilder<>(
-        ProgramMethodSet::create, Sets.newIdentityHashSet());
+        currentGraphLens, ProgramMethodSet::create, SetUtils::newIdentityHashSet);
   }
 
-  public static LongLivedProgramMethodSetBuilder<SortedProgramMethodSet> createForSortedSet() {
+  public static LongLivedProgramMethodSetBuilder<ProgramMethodSet> createConcurrentForIdentitySet(
+      GraphLens currentGraphLens) {
     return new LongLivedProgramMethodSetBuilder<>(
-        ignore -> SortedProgramMethodSet.create(), Sets.newIdentityHashSet());
+        currentGraphLens, ProgramMethodSet::create, SetUtils::newConcurrentHashSet);
   }
 
-  public static LongLivedProgramMethodSetBuilder<?> createConcurrentForIdentitySet() {
-    return new LongLivedProgramMethodSetBuilder<>(
-        ignore -> ProgramMethodSet.create(), Sets.newConcurrentHashSet());
-  }
-
+  @Deprecated
   public void add(ProgramMethod method) {
     methods.add(method.getReference());
   }
 
+  public void add(ProgramMethod method, GraphLens currentGraphLens) {
+    // All methods in a long lived program method set should be rewritten up until the same graph
+    // lens.
+    assert verifyIsRewrittenWithLens(currentGraphLens);
+    methods.add(method.getReference());
+  }
+
+  @Deprecated
   public void addAll(Iterable<ProgramMethod> methods) {
     methods.forEach(this::add);
   }
 
+  public void clear() {
+    methods.clear();
+  }
+
+  public boolean contains(ProgramMethod method, GraphLens currentGraphLens) {
+    // We can only query a long lived program method set that is fully lens rewritten.
+    assert verifyIsRewrittenWithLens(currentGraphLens);
+    return methods.contains(method.getReference());
+  }
+
+  public boolean isRewrittenWithLens(GraphLens graphLens) {
+    return appliedGraphLens == graphLens;
+  }
+
+  public LongLivedProgramMethodSetBuilder<T> merge(LongLivedProgramMethodSetBuilder<T> builder) {
+    // Check that the two builders are rewritten up until the same lens (if not we could rewrite the
+    // methods in the given builder up until the applied graph lens of this builder, but it must be
+    // such that this builder has the same or a newer graph lens than the given builder).
+    if (isRewrittenWithLens(builder.appliedGraphLens)) {
+      methods.addAll(builder.methods);
+    } else {
+      // Rewrite the methods in the given builder up until the applied graph lens of this builder.
+      // Note that this builder must have a newer graph lens than the given builder.
+      assert verifyIsRewrittenWithNewerLens(builder.appliedGraphLens);
+      for (DexMethod method : builder.methods) {
+        methods.add(appliedGraphLens.getRenamedMethodSignature(method, builder.appliedGraphLens));
+      }
+    }
+    return this;
+  }
+
   public void remove(DexMethod method) {
     methods.remove(method);
   }
 
-  public void removeAll(Iterable<DexMethod> methods) {
+  public LongLivedProgramMethodSetBuilder<T> removeAll(Iterable<DexMethod> methods) {
     methods.forEach(this::remove);
+    return this;
   }
 
-  public void rewrittenWithLens(AppView<AppInfoWithLiveness> appView, GraphLens applied) {
-    Set<DexMethod> newMethods = Sets.newIdentityHashSet();
-    for (DexMethod method : methods) {
-      newMethods.add(appView.graphLens().getRenamedMethodSignature(method, applied));
+  public LongLivedProgramMethodSetBuilder<T> rewrittenWithLens(
+      AppView<AppInfoWithLiveness> appView) {
+    return rewrittenWithLens(appView.graphLens());
+  }
+
+  public LongLivedProgramMethodSetBuilder<T> rewrittenWithLens(GraphLens newGraphLens) {
+    // Check if the graph lens has changed (otherwise lens rewriting is not needed).
+    if (newGraphLens == appliedGraphLens) {
+      return this;
     }
-    methods.clear();
-    methods.addAll(newMethods);
+
+    // Rewrite the backing.
+    Set<DexMethod> newMethods = factoryForBuilder.apply(methods.size());
+    for (DexMethod method : methods) {
+      newMethods.add(newGraphLens.getRenamedMethodSignature(method, appliedGraphLens));
+    }
+    methods = newMethods;
+
+    // Record that this collection is now rewritten up until the given graph lens.
+    appliedGraphLens = newGraphLens;
+    return this;
   }
 
   public T build(AppView<AppInfoWithLiveness> appView) {
-    return build(appView, null);
-  }
-
-  public T build(AppView<AppInfoWithLiveness> appView, GraphLens appliedGraphLens) {
     T result = factory.apply(methods.size());
-    for (DexMethod oldMethod : methods) {
-      DexMethod method = appView.graphLens().getRenamedMethodSignature(oldMethod, appliedGraphLens);
-      DexProgramClass holder = appView.definitionForHolder(method).asProgramClass();
-      result.createAndAdd(holder, holder.lookupMethod(method));
+    for (DexMethod method : methods) {
+      DexMethod rewrittenMethod =
+          appView.graphLens().getRenamedMethodSignature(method, appliedGraphLens);
+      DexProgramClass holder = appView.definitionForHolder(rewrittenMethod).asProgramClass();
+      result.createAndAdd(holder, holder.lookupMethod(rewrittenMethod));
     }
     return result;
   }
@@ -81,4 +144,18 @@
   public boolean isEmpty() {
     return methods.isEmpty();
   }
+
+  public boolean verifyIsRewrittenWithLens(GraphLens graphLens) {
+    assert isRewrittenWithLens(graphLens);
+    return true;
+  }
+
+  public boolean verifyIsRewrittenWithNewerLens(GraphLens graphLens) {
+    assert appliedGraphLens != graphLens;
+    assert appliedGraphLens.isNonIdentityLens();
+    assert graphLens.isIdentityLens()
+        || appliedGraphLens.asNonIdentityLens().findPrevious(previous -> previous == graphLens)
+            != null;
+    return true;
+  }
 }