Add support for LIR rewriting using custom lens

In order to implement library desugaring as a LIR-to-LIR transformation using a graph lens, it is desirable not to commit a library desugaring graph lens to the graph lens chain, since the graph lens only represents a rewriting of the code (i.e., it is not required for building the mapping file).

The LirLensCodeRewriter currently applies the lenses starting from the current (unapplied) graph lens to the code lens. This CL adds support for applying an uncommitted graph lens.

Change-Id: I514207174be6e82aed8dacba4717a7a0edf975e9
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java b/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
index 208c233..d68abb1 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
@@ -189,7 +189,17 @@
 
   /** Replace type appearances, invoke targets and field accesses with actual definitions. */
   public void rewrite(IRCode code, ProgramMethod method, MethodProcessor methodProcessor) {
-    Deque<GraphLensInterval> unappliedLenses = getUnappliedLenses(method);
+    GraphLens codeLens = method.getDefinition().getCode().getCodeLens(appView);
+    rewrite(code, method, methodProcessor, appView.graphLens(), codeLens);
+  }
+
+  public void rewrite(
+      IRCode code,
+      ProgramMethod method,
+      MethodProcessor methodProcessor,
+      GraphLens graphLens,
+      GraphLens codeLens) {
+    Deque<GraphLensInterval> unappliedLenses = getUnappliedLenses(method, graphLens, codeLens);
     DexMethod originalMethodReference =
         appView.graphLens().getOriginalMethodSignature(method.getReference());
     while (!unappliedLenses.isEmpty()) {
@@ -227,7 +237,9 @@
     rewriteArguments(
         code, originalMethodReference, prototypeChanges, affectedPhis, unusedArguments);
     if (graphLens.hasCustomLensCodeRewriter()) {
-      assert graphLens.getPrevious() == codeLens;
+      assert graphLens.getPrevious() == codeLens
+          || (graphLens.getPrevious().isMemberRebindingIdentityLens()
+              && graphLens.getPrevious().asMemberRebindingIdentityLens().getPrevious() == codeLens);
       CustomLensCodeRewriter customLensCodeRewriter = graphLens.getCustomLensCodeRewriter();
       affectedPhis.addAll(
           customLensCodeRewriter.rewriteCode(code, methodProcessor, prototypeChanges, graphLens));
@@ -1103,17 +1115,18 @@
     affectedValues.narrowingWithAssumeRemoval(appView, code);
   }
 
-  private Deque<GraphLensInterval> getUnappliedLenses(ProgramMethod method) {
+  private Deque<GraphLensInterval> getUnappliedLenses(
+      ProgramMethod method, GraphLens graphLens, GraphLens codeLens) {
     Deque<GraphLensInterval> unappliedLenses = new ArrayDeque<>(8);
-    GraphLens codeLens = method.getDefinition().getCode().getCodeLens(appView);
-    GraphLens currentLens = appView.graphLens();
+    GraphLens currentLens = graphLens;
     DexMethod currentMethod = method.getReference();
     while (currentLens != codeLens) {
       assert currentLens.isNonIdentityLens();
       NonIdentityGraphLens currentNonIdentityLens = currentLens.asNonIdentityLens();
       NonIdentityGraphLens fromInclusiveLens = currentNonIdentityLens;
-      if (!currentNonIdentityLens.hasCustomLensCodeRewriter()) {
-        GraphLens fromInclusiveLensPredecessor = fromInclusiveLens.getPrevious();
+      GraphLens fromInclusiveLensPredecessor = fromInclusiveLens.getPrevious();
+      if (!currentNonIdentityLens.hasCustomLensCodeRewriter()
+          || fromInclusiveLensPredecessor.isMemberRebindingIdentityLens()) {
         while (fromInclusiveLensPredecessor.isNonIdentityLens()
             && !fromInclusiveLensPredecessor.hasCustomLensCodeRewriter()
             && fromInclusiveLensPredecessor != codeLens) {
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/LirConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/LirConverter.java
index a92bbe3..70e6ff9 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/LirConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/LirConverter.java
@@ -142,7 +142,7 @@
         clazz ->
             clazz.forEachProgramMethodMatching(
                 m -> m.hasCode() && m.getCode().isLirCode(),
-                m -> rewriteLirMethodWithLens(m, appView, rewriterUtils)),
+                m -> rewriteLirMethodWithLens(m, appView, appView.graphLens(), rewriterUtils)),
         appView.options().getThreadingModule(),
         executorService);
 
@@ -151,9 +151,13 @@
   }
 
   public static void rewriteLirMethodWithLens(
-      ProgramMethod method, AppView<?> appView, LensCodeRewriterUtils rewriterUtils) {
+      ProgramMethod method,
+      AppView<?> appView,
+      GraphLens graphLens,
+      LensCodeRewriterUtils rewriterUtils) {
     LirCode<Integer> lirCode = method.getDefinition().getCode().asLirCode();
-    LirCode<Integer> rewrittenLirCode = lirCode.rewriteWithLens(method, appView, rewriterUtils);
+    LirCode<Integer> rewrittenLirCode =
+        lirCode.rewriteWithLens(method, appView, graphLens, rewriterUtils);
     if (ObjectUtils.notIdentical(lirCode, rewrittenLirCode)) {
       method.setCode(rewrittenLirCode, appView);
     }
diff --git a/src/main/java/com/android/tools/r8/lightir/LirCode.java b/src/main/java/com/android/tools/r8/lightir/LirCode.java
index d4284af..db84383 100644
--- a/src/main/java/com/android/tools/r8/lightir/LirCode.java
+++ b/src/main/java/com/android/tools/r8/lightir/LirCode.java
@@ -833,15 +833,17 @@
   }
 
   public LirCode<EV> rewriteWithLens(
-      ProgramMethod context, AppView<?> appView, LensCodeRewriterUtils rewriterUtils) {
-    GraphLens graphLens = appView.graphLens();
+      ProgramMethod context,
+      AppView<?> appView,
+      GraphLens graphLens,
+      LensCodeRewriterUtils rewriterUtils) {
     assert graphLens.isNonIdentityLens();
     if (graphLens.isMemberRebindingIdentityLens()) {
       return this;
     }
 
     LirLensCodeRewriter<EV> rewriter =
-        new LirLensCodeRewriter<>(appView, this, context, rewriterUtils);
+        new LirLensCodeRewriter<>(appView, this, context, graphLens, rewriterUtils);
     return rewriter.rewrite();
   }
 
diff --git a/src/main/java/com/android/tools/r8/lightir/LirLensCodeRewriter.java b/src/main/java/com/android/tools/r8/lightir/LirLensCodeRewriter.java
index d289215..eb4ec46 100644
--- a/src/main/java/com/android/tools/r8/lightir/LirLensCodeRewriter.java
+++ b/src/main/java/com/android/tools/r8/lightir/LirLensCodeRewriter.java
@@ -71,12 +71,16 @@
   private boolean hasNonTrivialRewritings = false;
 
   public LirLensCodeRewriter(
-      AppView<?> appView, LirCode<EV> code, ProgramMethod context, LensCodeRewriterUtils helper) {
+      AppView<?> appView,
+      LirCode<EV> code,
+      ProgramMethod context,
+      GraphLens graphLens,
+      LensCodeRewriterUtils helper) {
     super(code);
     this.appView = appView;
     this.context = context;
     this.contextReference = context.getReference();
-    this.graphLens = appView.graphLens();
+    this.graphLens = graphLens;
     this.codeLens = context.getDefinition().getCode().getCodeLens(appView);
     this.helper = helper;
     NonIdentityGraphLens nonStartupInStartupOutlinerLens =
@@ -451,7 +455,8 @@
                 .setFinalizeAfterLensCodeRewriter());
     // MethodProcessor argument is only used by unboxing lenses.
     MethodProcessor methodProcessor = null;
-    new LensCodeRewriter(appViewWithClassHierarchy).rewrite(code, context, methodProcessor);
+    new LensCodeRewriter(appViewWithClassHierarchy)
+        .rewrite(code, context, methodProcessor, graphLens, codeLens);
     IRToLirFinalizer finalizer = new IRToLirFinalizer(appView);
     LirCode<?> rewritten =
         finalizer.finalizeCode(code, BytecodeMetadataProvider.empty(), Timing.empty());
diff --git a/src/main/java/com/android/tools/r8/partial/R8PartialApplicationWriter.java b/src/main/java/com/android/tools/r8/partial/R8PartialApplicationWriter.java
index 4ad128d..e21b159 100644
--- a/src/main/java/com/android/tools/r8/partial/R8PartialApplicationWriter.java
+++ b/src/main/java/com/android/tools/r8/partial/R8PartialApplicationWriter.java
@@ -79,7 +79,8 @@
     } else if (code.isDexCode()) {
       rewriteDexCodeWithLens(code.asDexCode(), method);
     } else {
-      LirConverter.rewriteLirMethodWithLens(method, appView, LensCodeRewriterUtils.empty());
+      LirConverter.rewriteLirMethodWithLens(
+          method, appView, appView.graphLens(), LensCodeRewriterUtils.empty());
     }
   }