Move non-null insertion to non-null-tracker

Bug: 141143236
Change-Id: I3648aab9c55c08f489540cf990c4c4b7f8dd12a6
diff --git a/src/main/java/com/android/tools/r8/ir/code/InvokeDirect.java b/src/main/java/com/android/tools/r8/ir/code/InvokeDirect.java
index 417aa63..b8178e0 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InvokeDirect.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InvokeDirect.java
@@ -120,9 +120,26 @@
   }
 
   @Override
-  public DexEncodedMethod lookupSingleTarget(
-      AppView<AppInfoWithLiveness> appView, DexType invocationContext) {
-    return appView.appInfo().lookupDirectTarget(getInvokedMethod());
+  public DexEncodedMethod lookupSingleTarget(AppView<?> appView, DexType invocationContext) {
+    DexMethod invokedMethod = getInvokedMethod();
+    if (appView.appInfo().hasLiveness()) {
+      AppInfoWithLiveness appInfo = appView.appInfo().withLiveness();
+      return appInfo.lookupDirectTarget(invokedMethod);
+    }
+    // In D8, we can treat invoke-direct instructions as having a single target if the invoke is
+    // targeting a method in the enclosing class.
+    if (invokedMethod.holder == invocationContext) {
+      DexClass clazz = appView.definitionFor(invokedMethod.holder);
+      if (clazz != null && clazz.isProgramClass()) {
+        DexEncodedMethod singleTarget = clazz.lookupDirectMethod(invokedMethod);
+        if (!singleTarget.isStatic()) {
+          return singleTarget;
+        }
+      } else {
+        assert false;
+      }
+    }
+    return null;
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/ir/code/InvokeInterface.java b/src/main/java/com/android/tools/r8/ir/code/InvokeInterface.java
index dc74f64..26f3e64 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InvokeInterface.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InvokeInterface.java
@@ -88,14 +88,17 @@
   }
 
   @Override
-  public DexEncodedMethod lookupSingleTarget(
-      AppView<AppInfoWithLiveness> appView, DexType invocationContext) {
-    AppInfoWithLiveness appInfo = appView.appInfo();
-    return appInfo.lookupSingleInterfaceTarget(
-        getInvokedMethod(),
-        invocationContext,
-        TypeAnalysis.getRefinedReceiverType(appView, this),
-        getReceiver().getExactDynamicType());
+  public DexEncodedMethod lookupSingleTarget(AppView<?> appView, DexType invocationContext) {
+    if (appView.appInfo().hasLiveness()) {
+      AppView<AppInfoWithLiveness> appViewWithLiveness = appView.withLiveness();
+      AppInfoWithLiveness appInfo = appViewWithLiveness.appInfo();
+      return appInfo.lookupSingleInterfaceTarget(
+          getInvokedMethod(),
+          invocationContext,
+          TypeAnalysis.getRefinedReceiverType(appViewWithLiveness, this),
+          getReceiver().getExactDynamicType());
+    }
+    return null;
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/ir/code/InvokeMethod.java b/src/main/java/com/android/tools/r8/ir/code/InvokeMethod.java
index ad1d6b7..c3db7ea 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InvokeMethod.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InvokeMethod.java
@@ -14,7 +14,6 @@
 import com.android.tools.r8.ir.optimize.Inliner.InlineAction;
 import com.android.tools.r8.ir.optimize.InliningOracle;
 import com.android.tools.r8.ir.regalloc.RegisterAllocator;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import java.util.Collection;
 import java.util.List;
 
@@ -59,7 +58,7 @@
   // In subclasses, e.g., invoke-virtual or invoke-super, use a narrower receiver type by using
   // receiver type and calling context---the holder of the method where the current invocation is.
   public abstract DexEncodedMethod lookupSingleTarget(
-      AppView<AppInfoWithLiveness> appView, DexType invocationContext);
+      AppView<?> appView, DexType invocationContext);
 
   public abstract Collection<DexEncodedMethod> lookupTargets(
       AppView<? extends AppInfoWithSubtyping> appView, DexType invocationContext);
diff --git a/src/main/java/com/android/tools/r8/ir/code/InvokePolymorphic.java b/src/main/java/com/android/tools/r8/ir/code/InvokePolymorphic.java
index abe39c7..c573377 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InvokePolymorphic.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InvokePolymorphic.java
@@ -19,7 +19,6 @@
 import com.android.tools.r8.ir.optimize.Inliner.InlineAction;
 import com.android.tools.r8.ir.optimize.InliningConstraints;
 import com.android.tools.r8.ir.optimize.InliningOracle;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import java.util.Collection;
 import java.util.List;
 
@@ -117,8 +116,7 @@
   }
 
   @Override
-  public DexEncodedMethod lookupSingleTarget(
-      AppView<AppInfoWithLiveness> appView, DexType invocationContext) {
+  public DexEncodedMethod lookupSingleTarget(AppView<?> appView, DexType invocationContext) {
     // TODO(herhut): Implement lookup target for invokePolymorphic.
     return null;
   }
diff --git a/src/main/java/com/android/tools/r8/ir/code/InvokeStatic.java b/src/main/java/com/android/tools/r8/ir/code/InvokeStatic.java
index f9f73f2..4f42386 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InvokeStatic.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InvokeStatic.java
@@ -9,6 +9,7 @@
 import com.android.tools.r8.code.InvokeStaticRange;
 import com.android.tools.r8.graph.AppInfoWithSubtyping;
 import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexType;
@@ -101,10 +102,26 @@
   }
 
   @Override
-  public DexEncodedMethod lookupSingleTarget(
-      AppView<AppInfoWithLiveness> appView, DexType invocationContext) {
-    DexMethod method = getInvokedMethod();
-    return appView.appInfo().lookupStaticTarget(method);
+  public DexEncodedMethod lookupSingleTarget(AppView<?> appView, DexType invocationContext) {
+    DexMethod invokedMethod = getInvokedMethod();
+    if (appView.appInfo().hasLiveness()) {
+      AppInfoWithLiveness appInfo = appView.appInfo().withLiveness();
+      return appInfo.lookupStaticTarget(invokedMethod);
+    }
+    // In D8, we can treat invoke-static instructions as having a single target if the invoke is
+    // targeting a method in the enclosing class.
+    if (invokedMethod.holder == invocationContext) {
+      DexClass clazz = appView.definitionFor(invokedMethod.holder);
+      if (clazz != null && clazz.isProgramClass()) {
+        DexEncodedMethod singleTarget = clazz.lookupDirectMethod(invokedMethod);
+        if (singleTarget.isStatic()) {
+          return singleTarget;
+        }
+      } else {
+        assert false;
+      }
+    }
+    return null;
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/ir/code/InvokeSuper.java b/src/main/java/com/android/tools/r8/ir/code/InvokeSuper.java
index b128648..e866bf4 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InvokeSuper.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InvokeSuper.java
@@ -96,16 +96,15 @@
   }
 
   @Override
-  public DexEncodedMethod lookupSingleTarget(
-      AppView<AppInfoWithLiveness> appView, DexType invocationContext) {
-    if (invocationContext == null) {
-      return null;
+  public DexEncodedMethod lookupSingleTarget(AppView<?> appView, DexType invocationContext) {
+    if (appView.appInfo().hasLiveness() && invocationContext != null) {
+      AppView<AppInfoWithLiveness> appViewWithLiveness = appView.withLiveness();
+      AppInfoWithLiveness appInfo = appViewWithLiveness.appInfo();
+      if (appInfo.isSubtype(invocationContext, getInvokedMethod().holder)) {
+        return appInfo.lookupSuperTarget(getInvokedMethod(), invocationContext);
+      }
     }
-    if (!appView.appInfo().isSubtype(invocationContext, getInvokedMethod().holder)) {
-      return null;
-    } else {
-      return appView.appInfo().lookupSuperTarget(getInvokedMethod(), invocationContext);
-    }
+    return null;
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/ir/code/InvokeVirtual.java b/src/main/java/com/android/tools/r8/ir/code/InvokeVirtual.java
index 1d1b2a8..20c8b6c 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InvokeVirtual.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InvokeVirtual.java
@@ -91,14 +91,17 @@
   }
 
   @Override
-  public DexEncodedMethod lookupSingleTarget(
-      AppView<AppInfoWithLiveness> appView, DexType invocationContext) {
-    AppInfoWithLiveness appInfo = appView.appInfo();
-    return appInfo.lookupSingleVirtualTarget(
-        getInvokedMethod(),
-        invocationContext,
-        TypeAnalysis.getRefinedReceiverType(appView, this),
-        getReceiver().getExactDynamicType());
+  public DexEncodedMethod lookupSingleTarget(AppView<?> appView, DexType invocationContext) {
+    if (appView.appInfo().hasLiveness()) {
+      AppView<AppInfoWithLiveness> appViewWithLiveness = appView.withLiveness();
+      AppInfoWithLiveness appInfo = appViewWithLiveness.appInfo();
+      return appInfo.lookupSingleVirtualTarget(
+          getInvokedMethod(),
+          invocationContext,
+          TypeAnalysis.getRefinedReceiverType(appViewWithLiveness, this),
+          getReceiver().getExactDynamicType());
+    }
+    return null;
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/MemberValuePropagation.java b/src/main/java/com/android/tools/r8/ir/optimize/MemberValuePropagation.java
index 1cbbc26..86a4de5 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/MemberValuePropagation.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/MemberValuePropagation.java
@@ -403,6 +403,8 @@
         // be updated outside the class constructor, e.g. via reflections), it is safe
         // to assume that the static-get instruction reads the value it initialized value
         // in class initializer and is never null.
+        // TODO(b/141143236): This should be subsumed entirely by the non-null propagation for
+        //  fields, and thus should be removed.
         DexClass holderDefinition = appView.definitionFor(field.holder);
         if (holderDefinition != null
             && holderDefinition.accessFlags.isFinal()
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/NonNullTracker.java b/src/main/java/com/android/tools/r8/ir/optimize/NonNullTracker.java
index 189d028..0c52b3b 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/NonNullTracker.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/NonNullTracker.java
@@ -6,10 +6,10 @@
 import static com.android.tools.r8.ir.code.DominatorTree.Assumption.MAY_HAVE_UNREACHABLE_BLOCKS;
 
 import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexEncodedField;
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.ir.analysis.type.TypeAnalysis;
 import com.android.tools.r8.ir.analysis.type.TypeLatticeElement;
@@ -24,9 +24,11 @@
 import com.android.tools.r8.ir.code.Instruction;
 import com.android.tools.r8.ir.code.InstructionIterator;
 import com.android.tools.r8.ir.code.InstructionListIterator;
+import com.android.tools.r8.ir.code.InvokeMethod;
 import com.android.tools.r8.ir.code.Phi;
 import com.android.tools.r8.ir.code.Value;
 import com.android.tools.r8.ir.optimize.info.FieldOptimizationInfo;
+import com.android.tools.r8.ir.optimize.info.MethodOptimizationInfo;
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedback;
 import com.google.common.base.Predicates;
 import com.google.common.collect.Sets;
@@ -47,6 +49,7 @@
 public class NonNullTracker {
 
   private final AppView<?> appView;
+  private final DexItemFactory dexItemFactory;
   private final Consumer<BasicBlock> splitBlockConsumer;
 
   public NonNullTracker(AppView<?> appView) {
@@ -55,6 +58,7 @@
 
   public NonNullTracker(AppView<?> appView, Consumer<BasicBlock> splitBlockConsumer) {
     this.appView = appView;
+    this.dexItemFactory = appView.dexItemFactory();
     this.splitBlockConsumer = splitBlockConsumer;
   }
 
@@ -72,72 +76,63 @@
         continue;
       }
       // Add non-null after
-      // 1) invocations that call non-overridable library methods that are known to return non null.
-      // 2) instructions that implicitly indicate receiver/array is not null.
-      // 3) parameters that are not null after the invocation.
+      // 1) instructions that implicitly indicate receiver/array is not null.
+      // 2) invocations that call non-overridable library methods that are known to return non null.
+      // 3) invocations that are guaranteed to return a non-null value.
+      // 4) parameters that are not null after the invocation.
+      // 5) field-get instructions that are guaranteed to read a non-null value.
       InstructionListIterator iterator = block.listIterator(code);
       while (iterator.hasNext()) {
         Instruction current = iterator.next();
-        if (current.isInvokeMethod()
-            && appView
-                .dexItemFactory()
-                .libraryMethodsReturningNonNull
-                .contains(current.asInvokeMethod().getInvokedMethod())) {
-          Value knownToBeNonNullValue = current.outValue();
-          // Avoid adding redundant non-null instruction.
-          // Otherwise, we will have something like:
-          // non_null_rcv <- non-null(rcv)
-          // ...
-          // another_rcv <- non-null(non_null_rcv)
-          if (knownToBeNonNullValue != null && isNullableReferenceType(knownToBeNonNullValue)) {
-            knownToBeNonNullValues.add(knownToBeNonNullValue);
-          }
-        }
+        Value outValue = current.outValue();
+
+        // Case (1), instructions that implicitly indicate receiver/array is not null.
         if (current.throwsOnNullInput()) {
           Value couldBeNonNull = current.getNonNullInput();
           if (isNullableReferenceType(couldBeNonNull)) {
             knownToBeNonNullValues.add(couldBeNonNull);
           }
         }
-        if (current.isInvokeMethod() && !current.isInvokePolymorphic()) {
-          DexEncodedMethod singleTarget = null;
-          if (appView.enableWholeProgramOptimizations()) {
-            assert appView.appInfo().hasLiveness();
-            singleTarget =
-                current
-                    .asInvokeMethod()
-                    .lookupSingleTarget(appView.withLiveness(), code.method.method.holder);
-          } else {
-            // Even in D8, invoke-{direct|static} can be resolved without liveness.
-            // Due to the incremental compilation, though, it is allowed only if the holder of the
-            // invoked method is same as that of the method we are processing now.
-            DexMethod invokedMethod = current.asInvokeMethod().getInvokedMethod();
-            if (invokedMethod.holder == code.method.method.holder) {
-              DexClass clazz = appView.definitionFor(invokedMethod.holder);
-              assert clazz != null && clazz.isProgramClass();
-              if (clazz != null) {
-                DexEncodedMethod directMethod = clazz.lookupDirectMethod(invokedMethod);
-                if (current.isInvokeDirect() && !directMethod.isStatic()) {
-                  singleTarget = directMethod;
-                } else if (current.isInvokeStatic() && directMethod.isStatic()) {
-                  singleTarget = directMethod;
-                }
-              }
+
+        if (current.isInvokeMethod()) {
+          InvokeMethod invoke = current.asInvokeMethod();
+          DexMethod invokedMethod = invoke.getInvokedMethod();
+
+          // Case (2), invocations that call non-overridable library methods that are known to
+          // return non null.
+          if (dexItemFactory.libraryMethodsReturningNonNull.contains(invokedMethod)) {
+            if (current.hasOutValue() && isNullableReferenceType(outValue)) {
+              knownToBeNonNullValues.add(outValue);
             }
           }
-          if (singleTarget != null
-              && singleTarget.getOptimizationInfo().getNonNullParamOnNormalExits() != null) {
-            BitSet facts = singleTarget.getOptimizationInfo().getNonNullParamOnNormalExits();
-            for (int i = 0; i < current.inValues().size(); i++) {
-              if (facts.get(i)) {
-                Value knownToBeNonNullValue = current.inValues().get(i);
-                if (isNullableReferenceType(knownToBeNonNullValue)) {
-                  knownToBeNonNullValues.add(knownToBeNonNullValue);
+
+          DexEncodedMethod singleTarget =
+              invoke.lookupSingleTarget(appView, code.method.method.holder);
+          if (singleTarget != null) {
+            MethodOptimizationInfo optimizationInfo = singleTarget.getOptimizationInfo();
+
+            // Case (3), invocations that are guaranteed to return a non-null value.
+            if (optimizationInfo.neverReturnsNull()) {
+              if (invoke.hasOutValue() && isNullableReferenceType(outValue)) {
+                knownToBeNonNullValues.add(outValue);
+              }
+            }
+
+            // Case (4), parameters that are not null after the invocation.
+            BitSet nonNullParamOnNormalExits = optimizationInfo.getNonNullParamOnNormalExits();
+            if (nonNullParamOnNormalExits != null) {
+              for (int i = 0; i < current.inValues().size(); i++) {
+                if (nonNullParamOnNormalExits.get(i)) {
+                  Value knownToBeNonNullValue = current.inValues().get(i);
+                  if (isNullableReferenceType(knownToBeNonNullValue)) {
+                    knownToBeNonNullValues.add(knownToBeNonNullValue);
+                  }
                 }
               }
             }
           }
         } else if (current.isFieldGet()) {
+          // Case (5), field-get instructions that are guaranteed to read a non-null value.
           FieldInstruction fieldInstruction = current.asFieldInstruction();
           DexField field = fieldInstruction.getField();
           if (field.type.isClassType()) {
@@ -145,12 +140,21 @@
             if (encodedField != null) {
               FieldOptimizationInfo optimizationInfo = encodedField.getOptimizationInfo();
               if (optimizationInfo.getDynamicType() != null
-                  && optimizationInfo.getDynamicType().isDefinitelyNotNull()) {
-                knownToBeNonNullValues.add(fieldInstruction.outValue());
+                  && optimizationInfo.getDynamicType().isDefinitelyNotNull()
+                  && isNullableReferenceType(outValue)) {
+                knownToBeNonNullValues.add(outValue);
               }
             }
           }
         }
+
+        // This is to ensure that we do not add redundant non-null instructions.
+        // Otherwise, we will have something like:
+        //   y <- assume-not-null(x)
+        //   ...
+        //   z <- assume-not-null(y)
+        assert knownToBeNonNullValues.stream().allMatch(NonNullTracker::isNullableReferenceType);
+
         if (!knownToBeNonNullValues.isEmpty()) {
           addNonNullForValues(
               code,