Always insert Assume instructions after inlining

Change-Id: Ibf5f983df64bd327d5f5f4ab30c27e91e5f014bd
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 6b7c389..c7e511a 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
@@ -47,7 +47,6 @@
 import java.util.List;
 import java.util.ListIterator;
 import java.util.Set;
-import java.util.function.Consumer;
 import java.util.function.Predicate;
 
 public final class DefaultInliningOracle implements InliningOracle, InliningStrategy {
@@ -692,48 +691,6 @@
     instructionAllowance -= Inliner.numberOfInstructions(inlinee.code);
   }
 
-  @Override
-  public void updateTypeInformationIfNeeded(
-      IRCode inlinee, ListIterator<BasicBlock> blockIterator, BasicBlock block) {
-    boolean assumersEnabled =
-        appView.options().enableNonNullTracking
-            || appView.options().enableDynamicTypeOptimization
-            || appView.options().testing.forceAssumeNoneInsertion;
-    if (assumersEnabled) {
-      BasicBlock state = IteratorUtils.peekNext(blockIterator);
-
-      Set<BasicBlock> inlineeBlocks = Sets.newIdentityHashSet();
-      inlineeBlocks.addAll(inlinee.blocks);
-
-      // Introduce aliases only to the inlinee blocks.
-      if (appView.options().testing.forceAssumeNoneInsertion) {
-        insertAssumeInstructionsToInlinee(
-            new AliasIntroducer(appView), code, block, blockIterator, inlineeBlocks);
-      }
-
-      // Add non-null IRs only to the inlinee blocks.
-      if (appView.options().enableNonNullTracking) {
-        Consumer<BasicBlock> splitBlockConsumer = inlineeBlocks::add;
-        Assumer nonNullTracker = new NonNullTracker(appView, splitBlockConsumer);
-        insertAssumeInstructionsToInlinee(
-            nonNullTracker, code, block, blockIterator, inlineeBlocks);
-      }
-
-      // Add dynamic type assumptions only to the inlinee blocks.
-      if (appView.options().enableDynamicTypeOptimization) {
-        insertAssumeInstructionsToInlinee(
-            new DynamicTypeOptimization(appView), code, block, blockIterator, inlineeBlocks);
-      }
-
-      // Restore the old state of the iterator.
-      while (blockIterator.hasPrevious() && blockIterator.previous() != state) {
-        // Do nothing.
-      }
-      assert IteratorUtils.peekNext(blockIterator) == state;
-    }
-    // TODO(b/72693244): need a test where refined env in inlinee affects the caller.
-  }
-
   private void insertAssumeInstructionsToInlinee(
       Assumer assumer,
       IRCode code,
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/ForcedInliningOracle.java b/src/main/java/com/android/tools/r8/ir/optimize/ForcedInliningOracle.java
index c791772..15281a5 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/ForcedInliningOracle.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/ForcedInliningOracle.java
@@ -17,7 +17,6 @@
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedback;
 import com.android.tools.r8.ir.optimize.inliner.WhyAreYouNotInliningReporter;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import java.util.ListIterator;
 import java.util.Map;
 
 final class ForcedInliningOracle implements InliningOracle, InliningStrategy {
@@ -92,10 +91,6 @@
   }
 
   @Override
-  public void updateTypeInformationIfNeeded(
-      IRCode inlinee, ListIterator<BasicBlock> blockIterator, BasicBlock block) {}
-
-  @Override
   public boolean allowInliningOfInvokeInInlinee(
       InlineAction action,
       int inliningDepth,
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 ab3274d..ed35f73 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
@@ -65,6 +65,7 @@
 import java.util.ListIterator;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Consumer;
 
 public class Inliner implements PostOptimization {
 
@@ -969,7 +970,7 @@
           }
 
           classInitializationAnalysis.notifyCodeHasChanged();
-          strategy.updateTypeInformationIfNeeded(inlinee.code, blockIterator, block);
+          applyAssumersToInlinee(code, inlinee.code, blockIterator, block);
 
           // The synthetic and bridge flags are maintained only if the inlinee has also these flags.
           if (context.accessFlags.isBridge() && !inlinee.code.method.accessFlags.isBridge()) {
@@ -1040,6 +1041,62 @@
     return null;
   }
 
+  private void applyAssumersToInlinee(
+      IRCode code, IRCode inlinee, ListIterator<BasicBlock> blockIterator, BasicBlock block) {
+    boolean assumersEnabled =
+        appView.options().enableNonNullTracking
+            || appView.options().enableDynamicTypeOptimization
+            || appView.options().testing.forceAssumeNoneInsertion;
+    if (assumersEnabled) {
+      BasicBlock state = IteratorUtils.peekNext(blockIterator);
+
+      Set<BasicBlock> inlineeBlocks = Sets.newIdentityHashSet();
+      inlineeBlocks.addAll(inlinee.blocks);
+
+      // Introduce aliases only to the inlinee blocks.
+      if (appView.options().testing.forceAssumeNoneInsertion) {
+        applyAssumerToInlinee(
+            new AliasIntroducer(appView), code, block, blockIterator, inlineeBlocks);
+      }
+
+      // Add non-null IRs only to the inlinee blocks.
+      if (appView.options().enableNonNullTracking) {
+        Consumer<BasicBlock> splitBlockConsumer = inlineeBlocks::add;
+        Assumer nonNullTracker = new NonNullTracker(appView, splitBlockConsumer);
+        applyAssumerToInlinee(nonNullTracker, code, block, blockIterator, inlineeBlocks);
+      }
+
+      // Add dynamic type assumptions only to the inlinee blocks.
+      if (appView.options().enableDynamicTypeOptimization) {
+        applyAssumerToInlinee(
+            new DynamicTypeOptimization(appView), code, block, blockIterator, inlineeBlocks);
+      }
+
+      // Restore the old state of the iterator.
+      while (blockIterator.hasPrevious() && blockIterator.previous() != state) {
+        // Do nothing.
+      }
+      assert IteratorUtils.peekNext(blockIterator) == state;
+    }
+    // TODO(b/72693244): need a test where refined env in inlinee affects the caller.
+  }
+
+  private void applyAssumerToInlinee(
+      Assumer assumer,
+      IRCode code,
+      BasicBlock block,
+      ListIterator<BasicBlock> blockIterator,
+      Set<BasicBlock> inlineeBlocks) {
+    // Move the cursor back to where the first inlinee block was added.
+    while (blockIterator.hasPrevious() && blockIterator.previous() != block) {
+      // Do nothing.
+    }
+    assert IteratorUtils.peekNext(blockIterator) == block;
+
+    assumer.insertAssumeInstructionsInBlocks(code, blockIterator, inlineeBlocks::contains);
+    assert !blockIterator.hasNext();
+  }
+
   public static boolean verifyNoMethodsInlinedDueToSingleCallSite(AppView<?> appView) {
     for (DexProgramClass clazz : appView.appInfo().classes()) {
       for (DexEncodedMethod method : clazz.methods()) {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/InliningStrategy.java b/src/main/java/com/android/tools/r8/ir/optimize/InliningStrategy.java
index e0b0a2c..0f22b14 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/InliningStrategy.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/InliningStrategy.java
@@ -13,7 +13,6 @@
 import com.android.tools.r8.ir.optimize.Inliner.InlineeWithReason;
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedback;
 import com.android.tools.r8.ir.optimize.inliner.WhyAreYouNotInliningReporter;
-import java.util.ListIterator;
 
 interface InliningStrategy {
 
@@ -47,8 +46,5 @@
   void ensureMethodProcessed(
       DexEncodedMethod target, IRCode inlinee, OptimizationFeedback feedback);
 
-  void updateTypeInformationIfNeeded(
-      IRCode inlinee, ListIterator<BasicBlock> blockIterator, BasicBlock block);
-
   DexType getReceiverTypeIfKnown(InvokeMethod invoke);
 }