Enable narrowing for inputs with stack-map

Bug: 170060113
Bug: 169120386
Change-Id: I42c589b6865af0bac6b79ff837aba45efacd9a2a
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/type/TypeAnalysis.java b/src/main/java/com/android/tools/r8/ir/analysis/type/TypeAnalysis.java
index 8b63913..94b9489 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/type/TypeAnalysis.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/type/TypeAnalysis.java
@@ -64,6 +64,13 @@
     analyzeValues(values, Mode.WIDENING);
   }
 
+  public void narrowing(IRCode code) {
+    mode = Mode.NARROWING;
+    assert worklist.isEmpty();
+    code.topologicallySortedBlocks().forEach(this::analyzeBasicBlock);
+    analyze();
+  }
+
   public void narrowing(Iterable<? extends Value> values) {
     // TODO(b/125492155) Not sorting causes us to have non-deterministic behaviour. This should be
     //  removed when the bug is fixed.
diff --git a/src/main/java/com/android/tools/r8/ir/code/InstructionOrPhi.java b/src/main/java/com/android/tools/r8/ir/code/InstructionOrPhi.java
index 144bbdb..0ae5aaf 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InstructionOrPhi.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InstructionOrPhi.java
@@ -18,6 +18,10 @@
     return false;
   }
 
+  default boolean isStackMapPhi() {
+    return false;
+  }
+
   default Phi asPhi() {
     return null;
   }
diff --git a/src/main/java/com/android/tools/r8/ir/code/Phi.java b/src/main/java/com/android/tools/r8/ir/code/Phi.java
index c1de847..d7abfac 100644
--- a/src/main/java/com/android/tools/r8/ir/code/Phi.java
+++ b/src/main/java/com/android/tools/r8/ir/code/Phi.java
@@ -11,6 +11,7 @@
 import com.android.tools.r8.graph.DebugLocalInfo;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.ir.analysis.type.Nullability;
 import com.android.tools.r8.ir.analysis.type.TypeElement;
 import com.android.tools.r8.ir.code.BasicBlock.EdgeType;
 import com.android.tools.r8.ir.conversion.IRBuilder;
@@ -448,9 +449,9 @@
     return result;
   }
 
-  // TODO(b/169120386) This class is added to ensure we do not narrow or widen phi's in D8 when
-  //  having stack map information. It should be removed when we are certain to never widen or
-  //  narrowing phi's in D8.
+  // TODO(b/169120386) This class is added to ensure we do not narrow or widen the type of phi's
+  //  in D8 when having stack map information. We do allow for narrowing on nullable information.
+  //  It should be removed when we are certain to never widen or narrowing phi's in D8.
   public static class StackMapPhi extends Phi {
 
     public StackMapPhi(
@@ -478,7 +479,20 @@
     @Override
     public TypeElement computePhiType(AppView<?> appView) {
       assert !appView.enableWholeProgramOptimizations();
-      return type;
+      if (type.isPrimitiveType()) {
+        return type;
+      }
+      assert type.isReferenceType();
+      Nullability nullability = Nullability.bottom();
+      for (Value operand : getOperands()) {
+        nullability = nullability.join(operand.type.nullability());
+      }
+      return type.asReferenceType().getOrCreateVariant(nullability);
+    }
+
+    @Override
+    public boolean isStackMapPhi() {
+      return true;
     }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/IRBuilder.java b/src/main/java/com/android/tools/r8/ir/conversion/IRBuilder.java
index 6382b86..8f2c15e 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/IRBuilder.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/IRBuilder.java
@@ -723,9 +723,11 @@
       // TODO(b/169137397): We may have ended up generating StackMapPhi's before concluding
       //  having incorrect stack map types. Figure out a way to clean that up.
       new TypeAnalysis(appView).widening(ir);
+    } else {
+      assert canUseStackMapTypes() && !hasIncorrectStackMapTypes;
+      assert allPhisAreStackMapPhis(ir);
+      new TypeAnalysis(appView).narrowing(ir);
     }
-    // TODO(b/169137397): If we have canUseStackMapTypes() && !hasIncorrectStackMapTypes we should
-    //  have that all phi's are stack-map phis.
 
     // Update the IR code if collected call site optimization info has something useful.
     // While aggregation of parameter information at call sites would be more precise than static
@@ -755,6 +757,15 @@
     return !appView.enableWholeProgramOptimizations() && source.hasValidTypesFromStackMap();
   }
 
+  private boolean allPhisAreStackMapPhis(IRCode ir) {
+    ir.instructions()
+        .forEach(
+            instruction -> {
+              assert !instruction.isPhi() || instruction.isStackMapPhi();
+            });
+    return true;
+  }
+
   public void constrainType(Value value, ValueTypeConstraint constraint) {
     value.constrainType(constraint, method.getReference(), origin, appView.options().reporter);
   }
diff --git a/src/main/java/com/android/tools/r8/utils/IterableUtils.java b/src/main/java/com/android/tools/r8/utils/IterableUtils.java
index af3208f..15a8198 100644
--- a/src/main/java/com/android/tools/r8/utils/IterableUtils.java
+++ b/src/main/java/com/android/tools/r8/utils/IterableUtils.java
@@ -49,4 +49,8 @@
     iterable.forEach(result::add);
     return result;
   }
+
+  public static <T> boolean isEmpty(Iterable<T> iterable) {
+    return !iterable.iterator().hasNext();
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/PruneNotNullBranchD8Test.java b/src/test/java/com/android/tools/r8/ir/optimize/PruneNotNullBranchD8Test.java
index 870787f..ae08b92 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/PruneNotNullBranchD8Test.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/PruneNotNullBranchD8Test.java
@@ -6,6 +6,7 @@
 
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 import com.android.tools.r8.TestBase;
@@ -48,8 +49,13 @@
               assertThat(main, isPresent());
               MethodSubject mainMethod = main.mainMethod();
               assertThat(mainMethod, isPresent());
-              // TODO(b/170060113): We should be able to remove the throw.
-              assertTrue(mainMethod.streamInstructions().anyMatch(InstructionSubject::isThrow));
+              if (parameters.isDexRuntime()) {
+                assertFalse(mainMethod.streamInstructions().anyMatch(InstructionSubject::isThrow));
+              } else {
+                assert parameters.isCfRuntime();
+                // We are not going through IR when running in CF
+                assertTrue(mainMethod.streamInstructions().anyMatch(InstructionSubject::isThrow));
+              }
             });
   }