Apply initial graph lens to lookup in staticizer

Bug: 181493713
Change-Id: I79c64757a28f066a99f165d699c91d4107d6aab2
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 72f9ac7..8e5c11c 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
@@ -361,10 +361,11 @@
         D8NestBasedAccessDesugaring::reportDesugarDependencies);
   }
 
-  private void staticizeClasses(OptimizationFeedback feedback, ExecutorService executorService)
+  private void staticizeClasses(
+      OptimizationFeedback feedback, ExecutorService executorService, GraphLens applied)
       throws ExecutionException {
     if (classStaticizer != null) {
-      classStaticizer.staticizeCandidates(feedback, executorService);
+      classStaticizer.staticizeCandidates(feedback, executorService, applied);
     }
   }
 
@@ -723,7 +724,7 @@
     if (!options.isGeneratingClassFiles()) {
       printPhase("Class staticizer post processing");
       // TODO(b/127694949): Adapt to PostOptimization.
-      staticizeClasses(feedback, executorService);
+      staticizeClasses(feedback, executorService, initialGraphLensForIR);
       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/optimize/staticizer/ClassStaticizer.java b/src/main/java/com/android/tools/r8/ir/optimize/staticizer/ClassStaticizer.java
index d5ef586..5a2f63a 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
@@ -16,6 +16,7 @@
 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.graph.ResolutionResult;
 import com.android.tools.r8.graph.UseRegistry;
@@ -686,8 +687,9 @@
   //  3. Rewrite methods referencing staticized members, also remove instance creation
   //
   public final void staticizeCandidates(
-      OptimizationFeedback feedback, ExecutorService executorService) throws ExecutionException {
-    new StaticizingProcessor(appView, this, converter).run(feedback, executorService);
+      OptimizationFeedback feedback, ExecutorService executorService, GraphLens applied)
+      throws ExecutionException {
+    new StaticizingProcessor(appView, this, converter, applied).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 5f54ab6..27806d5 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,6 +19,7 @@
 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;
@@ -85,14 +86,17 @@
   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) {
+      IRConverter converter,
+      GraphLens applied) {
     this.appView = appView;
     this.classStaticizer = classStaticizer;
     this.converter = converter;
+    this.applied = applied;
   }
 
   final void run(OptimizationFeedback feedback, ExecutorService executorService)
@@ -233,7 +237,7 @@
         LongLivedProgramMethodSetBuilder<?> referencedFromBuilder =
             classStaticizer.referencedFrom.remove(info);
         assert referencedFromBuilder != null;
-        referencedFrom = referencedFromBuilder.build(appView);
+        referencedFrom = referencedFromBuilder.build(appView, applied);
         materializedReferencedFromCollections.put(info, referencedFrom);
       } else {
         referencedFrom = ProgramMethodSet.empty();
diff --git a/src/test/java/com/android/tools/r8/classmerging/horizontal/MergedVirtualMethodStaticizerTest.java b/src/test/java/com/android/tools/r8/classmerging/horizontal/MergedVirtualMethodStaticizerTest.java
index 290ccbb..e7eda20 100644
--- a/src/test/java/com/android/tools/r8/classmerging/horizontal/MergedVirtualMethodStaticizerTest.java
+++ b/src/test/java/com/android/tools/r8/classmerging/horizontal/MergedVirtualMethodStaticizerTest.java
@@ -30,11 +30,7 @@
             enableHorizontalClassMerging,
             inspector -> inspector.assertMergedInto(Program.B.class, Program.A.class))
         .run(parameters.getRuntime(), Program.Main.class)
-        .applyIf(
-            parameters.isDexRuntime() && enableHorizontalClassMerging,
-            // TODO(b/181493713): Should not fail.
-            result -> result.assertFailureWithErrorThatThrows(IncompatibleClassChangeError.class),
-            result -> result.assertSuccessWithOutputLines("A::foo", "Staticized::foo", "B::foo"));
+        .assertSuccessWithOutputLines("A::foo", "Staticized::foo", "B::foo");
   }
 
   public static class Program {