Enable constructor inlining for final fields in tests

Bug: b/278964529
Change-Id: If24fe15a88c92bfc5d5fa7db11886e75fcf2b0d4
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/MethodProcessor.java b/src/main/java/com/android/tools/r8/ir/conversion/MethodProcessor.java
index 9c8b286..e26ea21 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/MethodProcessor.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/MethodProcessor.java
@@ -8,6 +8,10 @@
 
 public abstract class MethodProcessor {
 
+  public boolean hasWaves() {
+    return false;
+  }
+
   public boolean isD8MethodProcessor() {
     return false;
   }
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 d31a884..f625c8e 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
@@ -85,6 +85,11 @@
   }
 
   @Override
+  public boolean hasWaves() {
+    return true;
+  }
+
+  @Override
   public boolean isPostMethodProcessor() {
     return true;
   }
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/PrimaryMethodProcessor.java b/src/main/java/com/android/tools/r8/ir/conversion/PrimaryMethodProcessor.java
index 4e3f618..788bfcb 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/PrimaryMethodProcessor.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/PrimaryMethodProcessor.java
@@ -71,6 +71,11 @@
   }
 
   @Override
+  public boolean hasWaves() {
+    return true;
+  }
+
+  @Override
   public boolean isPrimaryMethodProcessor() {
     return true;
   }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/DefaultInliningOracle.java b/src/main/java/com/android/tools/r8/ir/optimize/DefaultInliningOracle.java
index a0bfb2d..fea6435 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/DefaultInliningOracle.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/DefaultInliningOracle.java
@@ -18,8 +18,7 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.Code;
 import com.android.tools.r8.graph.DefaultUseRegistryWithResult;
-import com.android.tools.r8.graph.DexEncodedField;
-import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.graph.DexClassAndField;
 import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexItemFactory.BoxUnboxPrimitiveMethodRoundtrip;
 import com.android.tools.r8.graph.DexMethod;
@@ -726,20 +725,19 @@
       } else if (instruction.isInstancePut()) {
         // Final fields may not be initialized outside of a constructor in the enclosing class.
         InstancePut instancePut = instruction.asInstancePut();
-        DexField field = instancePut.getField();
-        DexEncodedField target = appView.appInfo().lookupInstanceTarget(field);
+        DexClassAndField target =
+            instancePut.resolveField(appView, singleTarget).getResolutionPair();
         if (target == null) {
           whyAreYouNotInliningReporter.reportUnsafeConstructorInliningDueToMissingFieldAssignment(
               instancePut);
           return false;
         }
-        if (target.isFinal()) {
-          // TODO(b/278964529): We should unset the final flag for good measure. Find a way to do
-          //  this in a thread safe manner.
+        if (target.getAccessFlags().isFinal()) {
           if (inlinerOptions.enableConstructorInliningWithFinalFields
-              && options.canAssignFinalInstanceFieldOutsideConstructor()
-              && options.canUseJavaLangVarHandleStoreStoreFence(appView)) {
-            actionBuilder.setShouldEnsureStoreStoreBarrier();
+              && options.canUseJavaLangVarHandleStoreStoreFence(appView)
+              && methodProcessor.hasWaves()
+              && target.isProgramField()) {
+            actionBuilder.setShouldEnsureStoreStoreFence(target.asProgramField());
           } else {
             whyAreYouNotInliningReporter.reportUnsafeConstructorInliningDueToFinalFieldAssignment(
                 instancePut);
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 b58048e..af554ca 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
@@ -21,6 +21,7 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.MethodResolutionResult.SingleResolutionResult;
 import com.android.tools.r8.graph.NestMemberClassAttribute;
+import com.android.tools.r8.graph.ProgramField;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.analysis.ClassInitializationAnalysis;
 import com.android.tools.r8.ir.analysis.proto.ProtoInliningReasonStrategy;
@@ -69,6 +70,7 @@
 import com.android.tools.r8.utils.ListUtils;
 import com.android.tools.r8.utils.Timing;
 import com.android.tools.r8.utils.collections.LongLivedProgramMethodSetBuilder;
+import com.android.tools.r8.utils.collections.ProgramFieldSet;
 import com.android.tools.r8.utils.collections.ProgramMethodMap;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
 import com.google.common.collect.ImmutableList;
@@ -84,6 +86,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -100,6 +103,10 @@
 
   private final MultiCallerInliner multiCallerInliner;
 
+  // The set of final fields that have been assigned by a constructor that has been inlined in the
+  // current wave.
+  private final ProgramFieldSet inlinedFinalFieldsInWave = ProgramFieldSet.createConcurrent();
+
   // The set of methods that have been single caller inlined in the current wave. These need to be
   // pruned when the wave ends.
   private final Map<DexProgramClass, ProgramMethodMap<ProgramMethod>>
@@ -109,6 +116,8 @@
 
   private final AvailableApiExceptions availableApiExceptions;
 
+  private final AtomicBoolean waveDoneScheduled = new AtomicBoolean();
+
   public Inliner(AppView<AppInfoWithLiveness> appView) {
     this(appView, null);
   }
@@ -565,7 +574,7 @@
     final Reason reason;
 
     private boolean shouldEnsureStaticInitialization;
-    private boolean shouldEnsureStoreStoreBarrier;
+    private ProgramFieldSet shouldEnsureStoreStoreFenceCauses;
 
     private DexProgramClass downcastClass;
 
@@ -596,8 +605,8 @@
       shouldEnsureStaticInitialization = true;
     }
 
-    void setShouldEnsureStoreStoreBarrier() {
-      shouldEnsureStoreStoreBarrier = true;
+    void setShouldEnsureStoreStoreFenceCauses(ProgramFieldSet shouldEnsureStoreStoreFenceCauses) {
+      this.shouldEnsureStoreStoreFenceCauses = shouldEnsureStoreStoreFenceCauses;
     }
 
     boolean mustBeInlined() {
@@ -626,9 +635,9 @@
             failingBlock -> synthesizeInitClass(code, failingBlock));
       }
 
-      // Insert a store-store barrier at the end of the method.
-      if (shouldEnsureStoreStoreBarrier) {
-        synthesizeStoreStoreBarrier(appView, code);
+      // Insert a store-store fence at the end of the method.
+      if (shouldEnsureStoreStoreFenceCauses != null) {
+        synthesizeStoreStoreFence(appView, code);
       }
 
       // Insert a null check if this is needed to preserve the implicit null check for the receiver.
@@ -825,7 +834,7 @@
       }
     }
 
-    private void synthesizeStoreStoreBarrier(AppView<AppInfoWithLiveness> appView, IRCode code) {
+    private void synthesizeStoreStoreFence(AppView<AppInfoWithLiveness> appView, IRCode code) {
       assert target.getDefinition().isInstanceInitializer();
       BasicBlockIterator blockIterator = code.listIterator();
       while (blockIterator.hasNext()) {
@@ -833,10 +842,23 @@
         if (!block.isReturnBlock()) {
           continue;
         }
-        // It is an invariant that return blocks do not have catch handlers.
-        assert !block.hasCatchHandlers();
-        InstructionListIterator instructionIterator =
-            block.listIterator(block.getInstructions().size() - 1);
+        InstructionListIterator instructionIterator;
+        if (block.hasCatchHandlers()) {
+          instructionIterator = block.listIterator();
+          Instruction throwingInstruction =
+              instructionIterator.nextUntil(Instruction::instructionTypeCanThrow);
+          if (throwingInstruction != null) {
+            instructionIterator =
+                instructionIterator
+                    .splitCopyCatchHandlers(code, blockIterator, appView.options())
+                    .listIterator();
+          } else {
+            Instruction previousInstruction = instructionIterator.previous();
+            assert previousInstruction.isReturn();
+          }
+        } else {
+          instructionIterator = block.listIterator(block.getInstructions().size() - 1);
+        }
         DexMethod storeStoreFenceMethod =
             appView.dexItemFactory().javaLangInvokeVarHandleMembers.storeStoreFence;
         assert appView.definitionFor(storeStoreFenceMethod) != null;
@@ -867,7 +889,7 @@
       private InvokeMethod invoke;
       private Reason reason;
       private boolean shouldEnsureStaticInitialization;
-      private boolean shouldEnsureStoreStoreBarrier;
+      private ProgramFieldSet shouldEnsureStoreStoreFenceCauses;
       private ProgramMethod target;
 
       Builder setDowncastClass(DexProgramClass downcastClass) {
@@ -890,8 +912,11 @@
         return this;
       }
 
-      Builder setShouldEnsureStoreStoreBarrier() {
-        this.shouldEnsureStoreStoreBarrier = true;
+      Builder setShouldEnsureStoreStoreFence(ProgramField finalFieldCause) {
+        if (shouldEnsureStoreStoreFenceCauses == null) {
+          shouldEnsureStoreStoreFenceCauses = ProgramFieldSet.create();
+        }
+        shouldEnsureStoreStoreFenceCauses.add(finalFieldCause);
         return this;
       }
 
@@ -908,8 +933,8 @@
         if (shouldEnsureStaticInitialization) {
           action.setShouldEnsureStaticInitialization();
         }
-        if (shouldEnsureStoreStoreBarrier) {
-          action.setShouldEnsureStoreStoreBarrier();
+        if (shouldEnsureStoreStoreFenceCauses != null) {
+          action.setShouldEnsureStoreStoreFenceCauses(shouldEnsureStoreStoreFenceCauses);
         }
         return action;
       }
@@ -1149,6 +1174,13 @@
           iterator.inlineInvoke(
               appView, code, inlinee, blockIterator, blocksToRemove, action.getDowncastClass());
 
+          if (action.shouldEnsureStoreStoreFenceCauses != null) {
+            assert !action.shouldEnsureStoreStoreFenceCauses.isEmpty();
+            assert converter.isInWave();
+            inlinedFinalFieldsInWave.addAll(action.shouldEnsureStoreStoreFenceCauses);
+            scheduleWaveDone();
+          }
+
           if (methodProcessor.getCallSiteInformation().hasSingleCallSite(singleTarget, context)) {
             feedback.markInlinedIntoSingleCallSite(singleTargetMethod);
             appView.withArgumentPropagator(
@@ -1157,9 +1189,7 @@
                         singleTarget, context, methodProcessor));
             if (!(methodProcessor instanceof OneTimeMethodProcessor)) {
               assert converter.isInWave();
-              if (singleCallerInlinedMethodsInWave.isEmpty()) {
-                converter.addWaveDoneAction(this::onWaveDone);
-              }
+              scheduleWaveDone();
               singleCallerInlinedMethodsInWave
                   .computeIfAbsent(
                       singleTarget.getHolder(), ignoreKey(ProgramMethodMap::createConcurrent))
@@ -1339,6 +1369,7 @@
   }
 
   private void onWaveDone() {
+    inlinedFinalFieldsInWave.forEach(field -> field.getAccessFlags().demoteFromFinal());
     singleCallerInlinedMethodsInWave.forEach(
         (clazz, singleCallerInlinedMethodsForClass) -> {
           // Convert and remove virtual single caller inlined methods to abstract or throw null.
@@ -1384,7 +1415,15 @@
                 });
           }
         });
+    inlinedFinalFieldsInWave.clear();
     singleCallerInlinedMethodsInWave.clear();
+    waveDoneScheduled.set(false);
+  }
+
+  private void scheduleWaveDone() {
+    if (!waveDoneScheduled.getAndSet(true)) {
+      converter.addWaveDoneAction(this::onWaveDone);
+    }
   }
 
   public void onLastWaveDone(
diff --git a/src/main/java/com/android/tools/r8/utils/InternalOptions.java b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
index 6709446..db25078 100644
--- a/src/main/java/com/android/tools/r8/utils/InternalOptions.java
+++ b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
@@ -2681,11 +2681,6 @@
     return !appView.enableWholeProgramOptimizations();
   }
 
-  public boolean canAssignFinalInstanceFieldOutsideConstructor() {
-    // ART does not check this property.
-    return isGeneratingDex();
-  }
-
   /**
    * Dex2Oat issues a warning for abstract methods on non-abstract classes, so we never allow this.
    *
@@ -2887,7 +2882,7 @@
   }
 
   public boolean canUseJavaLangVarHandleStoreStoreFence(DexDefinitionSupplier definitions) {
-    if (isGeneratingDex() && hasMinApi(AndroidApiLevel.P)) {
+    if (isGeneratingDex() && hasMinApi(AndroidApiLevel.T)) {
       DexItemFactory factory = definitions.dexItemFactory();
       return definitions.hasDefinitionFor(factory.javaLangInvokeVarHandleMembers.storeStoreFence);
     }
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/inliner/InlineConstructorWithFinalFieldsTest.java b/src/test/java/com/android/tools/r8/ir/optimize/inliner/InlineConstructorWithFinalFieldsTest.java
index e5cca20..6e8e2f9 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/inliner/InlineConstructorWithFinalFieldsTest.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/inliner/InlineConstructorWithFinalFieldsTest.java
@@ -5,9 +5,10 @@
 
 import static com.android.tools.r8.utils.codeinspector.CodeMatchers.invokesMethod;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isAbsentIf;
+import static com.android.tools.r8.utils.codeinspector.Matchers.isFinal;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
+import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.Assert.assertFalse;
 
 import com.android.tools.r8.AlwaysInline;
 import com.android.tools.r8.TestBase;
@@ -15,9 +16,9 @@
 import com.android.tools.r8.TestParametersCollection;
 import com.android.tools.r8.ToolHelper;
 import com.android.tools.r8.references.Reference;
-import com.android.tools.r8.utils.AndroidApiLevel;
 import com.android.tools.r8.utils.MethodReferenceUtils;
 import com.android.tools.r8.utils.codeinspector.ClassSubject;
+import com.android.tools.r8.utils.codeinspector.FieldSubject;
 import com.android.tools.r8.utils.codeinspector.MethodSubject;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -45,11 +46,6 @@
             parameters.isDexRuntime(),
             testBuilder -> testBuilder.addLibraryFiles(ToolHelper.getMostRecentAndroidJar()))
         .addKeepMainRule(Main.class)
-        .addOptionsModification(
-            options -> {
-              assertFalse(options.inlinerOptions().enableConstructorInliningWithFinalFields);
-              options.inlinerOptions().enableConstructorInliningWithFinalFields = true;
-            })
         .enableAlwaysInliningAnnotations()
         .setMinApi(parameters)
         .compile()
@@ -58,17 +54,24 @@
               ClassSubject mainClassSubject = inspector.clazz(Main.class);
               assertThat(mainClassSubject, isPresent());
 
+              FieldSubject xFieldSubject = mainClassSubject.uniqueFieldWithOriginalName("x");
+              assertThat(xFieldSubject, isPresent());
+
+              FieldSubject yFieldSubject = mainClassSubject.uniqueFieldWithOriginalName("y");
+              assertThat(yFieldSubject, isPresent());
+
               MethodSubject initMethodSubject = mainClassSubject.init("int", "int");
               assertThat(
                   initMethodSubject,
-                  isAbsentIf(
-                      parameters.isDexRuntime()
-                          && parameters.getApiLevel().isGreaterThanOrEqualTo(AndroidApiLevel.P)));
+                  isAbsentIf(parameters.canUseJavaLangInvokeVarHandleStoreStoreFence()));
 
               MethodSubject mainMethodSubject = mainClassSubject.mainMethod();
               assertThat(mainMethodSubject, isPresent());
+
               if (initMethodSubject.isPresent()) {
                 assertThat(mainMethodSubject, invokesMethod(initMethodSubject));
+                assertThat(xFieldSubject, isFinal());
+                assertThat(yFieldSubject, isFinal());
               } else {
                 assertThat(
                     mainMethodSubject,
@@ -78,6 +81,8 @@
                     invokesMethod(
                         Reference.methodFromDescriptor(
                             "Ljava/lang/invoke/VarHandle;", "storeStoreFence", "()V")));
+                assertThat(xFieldSubject, not(isFinal()));
+                assertThat(yFieldSubject, not(isFinal()));
               }
             })
         .run(parameters.getRuntime(), Main.class, "20", "22")
diff --git a/src/test/testbase/java/com/android/tools/r8/TestCompilerBuilder.java b/src/test/testbase/java/com/android/tools/r8/TestCompilerBuilder.java
index 2a9697a..f4f0034 100644
--- a/src/test/testbase/java/com/android/tools/r8/TestCompilerBuilder.java
+++ b/src/test/testbase/java/com/android/tools/r8/TestCompilerBuilder.java
@@ -67,6 +67,7 @@
         options.testing.listIterationRewritingEnabled = true;
         options.horizontalClassMergerOptions().enable();
         options.horizontalClassMergerOptions().setEnableInterfaceMerging();
+        options.inlinerOptions().enableConstructorInliningWithFinalFields = true;
         options
             .getCfCodeAnalysisOptions()
             .setAllowUnreachableCfBlocks(false)
diff --git a/src/test/testbase/java/com/android/tools/r8/TestParameters.java b/src/test/testbase/java/com/android/tools/r8/TestParameters.java
index fc90c26..ab3077c 100644
--- a/src/test/testbase/java/com/android/tools/r8/TestParameters.java
+++ b/src/test/testbase/java/com/android/tools/r8/TestParameters.java
@@ -85,6 +85,10 @@
         .isGreaterThanOrEqualTo(TestBase.apiLevelWithDefaultInterfaceMethodsSupport());
   }
 
+  public boolean canUseJavaLangInvokeVarHandleStoreStoreFence() {
+    return isDexRuntime() && getApiLevel().isGreaterThanOrEqualTo(AndroidApiLevel.T);
+  }
+
   public boolean canUseNativeDexPC() {
     assert isCfRuntime() || isDexRuntime();
     return isDexRuntime() && getDexRuntimeVersion().isNewerThanOrEqual(DexVm.Version.V8_1_0);