Rewrite AppInfoWithLiveness using LambdaRewriterLens

Bug: 161735546
Change-Id: I5be59f19a40a1455810bd22c2de95d4e50159354

src/main/java/com/android/tools/r8/R8.java

diff --git a/src/main/java/com/android/tools/r8/R8.java b/src/main/java/com/android/tools/r8/R8.java
index 93f229c..fa6c0f3 100644
--- a/src/main/java/com/android/tools/r8/R8.java
+++ b/src/main/java/com/android/tools/r8/R8.java
@@ -702,7 +702,7 @@
                       timing)
                   .withEnumValueInfoMaps(enumValueInfoMapCollection));
           // Rerunning the enqueuer should not give rise to any method rewritings.
-          assert enqueuer.buildGraphLens(appView) == appView.graphLens();
+          assert enqueuer.buildGraphLens(appView) == null;
           appView.withGeneratedMessageLiteBuilderShrinker(
               shrinker ->
                   shrinker.rewriteDeadBuilderReferencesFromDynamicMethods(
@@ -979,7 +979,14 @@
                 options.getProguardConfiguration().getDontWarnPatterns(),
                 executorService,
                 timing));
-    appView.setGraphLens(enqueuer.buildGraphLens(appView));
+    NestedGraphLens lens = enqueuer.buildGraphLens(appView);
+    if (lens != null) {
+      appView.setGraphLens(lens);
+      appViewWithLiveness.setAppInfo(
+          appViewWithLiveness
+              .appInfo()
+              .rewrittenWithLens(appView.appInfo().app().asDirect(), lens));
+    }
     if (InternalOptions.assertionsEnabled()) {
       // Register the dead proto types. These are needed to verify that no new missing types are
       // reported and that no dead proto types are referenced in the generated application.

src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java

diff --git a/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java b/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
index 6e41b1a..3709176 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
@@ -347,9 +347,9 @@
     return knownLambdaClasses;
   }
 
-  public GraphLens buildMappingLens(AppView<?> appView) {
+  public NestedGraphLens buildMappingLens(AppView<?> appView) {
     if (originalMethodSignatures.isEmpty()) {
-      return appView.graphLens();
+      return null;
     }
     return new LambdaRewriterLens(
         originalMethodSignatures, appView.graphLens(), appView.dexItemFactory());

src/main/java/com/android/tools/r8/shaking/Enqueuer.java

diff --git a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
index 281561b..f2cdfea 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -57,7 +57,7 @@
 import com.android.tools.r8.graph.FieldAccessInfoCollectionImpl;
 import com.android.tools.r8.graph.FieldAccessInfoImpl;
 import com.android.tools.r8.graph.FieldResolutionResult;
-import com.android.tools.r8.graph.GraphLens;
+import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
 import com.android.tools.r8.graph.InnerClassAttribute;
 import com.android.tools.r8.graph.LookupLambdaTarget;
 import com.android.tools.r8.graph.LookupTarget;
@@ -2733,8 +2733,8 @@
     return appInfoWithLiveness;
   }
 
-  public GraphLens buildGraphLens(AppView<?> appView) {
-    return lambdaRewriter != null ? lambdaRewriter.buildMappingLens(appView) : appView.graphLens();
+  public NestedGraphLens buildGraphLens(AppView<?> appView) {
+    return lambdaRewriter != null ? lambdaRewriter.buildMappingLens(appView) : null;
   }
 
   private void keepClassWithRules(DexProgramClass clazz, Set<ProguardKeepRuleBase> rules) {

src/test/java/com/android/tools/r8/ir/optimize/B161735546.java

diff --git a/src/test/java/com/android/tools/r8/ir/optimize/B161735546.java b/src/test/java/com/android/tools/r8/ir/optimize/B161735546.java
index 1d880e6..ce969b7 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/B161735546.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/B161735546.java
@@ -4,10 +4,6 @@
 
 package com.android.tools.r8.ir.optimize;
 
-import static org.junit.Assert.assertTrue;
-
-import com.android.tools.r8.CompilationFailedException;
-import com.android.tools.r8.R8TestCompileResult;
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.TestParametersCollection;
@@ -31,21 +27,11 @@
 
   @Test
   public void test() throws Exception {
-    R8TestCompileResult compileResult = null;
-    try {
-      compileResult =
-          testForR8(parameters.getBackend())
-              .addInnerClasses(B161735546.class)
-              .addKeepMainRule(TestClass.class)
-              .setMinApi(parameters.getApiLevel())
-              .compile();
-    } catch (CompilationFailedException e) {
-      // TODO(b/161735546): Fix NPE.
-      assertTrue(parameters.isDexRuntime());
-      return;
-    }
-
-    compileResult
+    testForR8(parameters.getBackend())
+        .addInnerClasses(B161735546.class)
+        .addKeepMainRule(TestClass.class)
+        .setMinApi(parameters.getApiLevel())
+        .compile()
         .run(parameters.getRuntime(), TestClass.class)
         .assertSuccessWithOutputLines("1", "2", "3");
   }