Towards unifying AssumeDynamicType and AssumeNotNull

This extends each Assume instruction to have both a DynamicTypeAssumption and NonNullAssumption.

It is still the case that each Assume instruction in the IR represents an AssumeDynamicType instruction or an AssumeNonNull instruction. A follow up should remove redundant Assume instructions by exploiting that a single Assume instruction can now hold multiple assumptions.

Change-Id: I92048bb46b4d1950ffdae07750fac3d9e760fc1a
diff --git a/src/main/java/com/android/tools/r8/ir/code/Assume.java b/src/main/java/com/android/tools/r8/ir/code/Assume.java
index c6012bd..117af7b 100644
--- a/src/main/java/com/android/tools/r8/ir/code/Assume.java
+++ b/src/main/java/com/android/tools/r8/ir/code/Assume.java
@@ -12,50 +12,60 @@
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.analysis.type.ClassTypeElement;
 import com.android.tools.r8.ir.analysis.type.TypeElement;
-import com.android.tools.r8.ir.code.Assume.Assumption;
 import com.android.tools.r8.ir.conversion.CfBuilder;
 import com.android.tools.r8.ir.conversion.DexBuilder;
 import com.android.tools.r8.ir.optimize.Inliner.ConstraintWithTarget;
 import com.android.tools.r8.ir.optimize.InliningConstraints;
+import com.android.tools.r8.utils.BooleanUtils;
 import java.util.Objects;
 import java.util.Set;
 
-public class Assume<An extends Assumption> extends Instruction {
+public class Assume extends Instruction {
 
   private static final String ERROR_MESSAGE =
       "Expected Assume instructions to be removed after IR processing.";
 
-  private final An assumption;
+  private final DynamicTypeAssumption dynamicTypeAssumption;
+  private final NonNullAssumption nonNullAssumption;
   private final Instruction origin;
 
-  private Assume(An assumption, Value dest, Value src, Instruction origin, AppView<?> appView) {
+  private Assume(
+      DynamicTypeAssumption dynamicTypeAssumption,
+      NonNullAssumption nonNullAssumption,
+      Value dest,
+      Value src,
+      Instruction origin,
+      AppView<?> appView) {
     super(dest, src);
-    assert assumption != null;
-    assert assumption.verifyCorrectnessOfValues(dest, src, appView);
+    assert dynamicTypeAssumption != null || nonNullAssumption != null;
+    assert BooleanUtils.intValue(dynamicTypeAssumption != null)
+            + BooleanUtils.intValue(nonNullAssumption != null)
+        == 1;
+    assert dynamicTypeAssumption == null
+        || dynamicTypeAssumption.verifyCorrectnessOfValues(dest, src, appView);
+    assert nonNullAssumption == null
+        || nonNullAssumption.verifyCorrectnessOfValues(dest, src, appView);
     assert dest != null;
-    this.assumption = assumption;
+    this.dynamicTypeAssumption = dynamicTypeAssumption;
+    this.nonNullAssumption = nonNullAssumption;
     this.origin = origin;
   }
 
-  public static Assume<NoAssumption> createAssumeNoneInstruction(
+  public static Assume createAssumeNonNullInstruction(
       Value dest, Value src, Instruction origin, AppView<?> appView) {
-    return new Assume<>(NoAssumption.get(), dest, src, origin, appView);
+    return new Assume(null, NonNullAssumption.get(), dest, src, origin, appView);
   }
 
-  public static Assume<NonNullAssumption> createAssumeNonNullInstruction(
-      Value dest, Value src, Instruction origin, AppView<?> appView) {
-    return new Assume<>(NonNullAssumption.get(), dest, src, origin, appView);
-  }
-
-  public static Assume<DynamicTypeAssumption> createAssumeDynamicTypeInstruction(
+  public static Assume createAssumeDynamicTypeInstruction(
       TypeElement dynamicUpperBoundType,
       ClassTypeElement dynamicLowerBoundType,
       Value dest,
       Value src,
       Instruction origin,
       AppView<?> appView) {
-    return new Assume<>(
+    return new Assume(
         new DynamicTypeAssumption(dynamicUpperBoundType, dynamicLowerBoundType),
+        null,
         dest,
         src,
         origin,
@@ -69,7 +79,7 @@
 
   public boolean verifyInstructionIsNeeded(AppView<?> appView) {
     if (isAssumeDynamicType()) {
-      assert assumption.verifyCorrectnessOfValues(outValue(), src(), appView);
+      assert dynamicTypeAssumption.verifyCorrectnessOfValues(outValue(), src(), appView);
     }
     return true;
   }
@@ -79,8 +89,12 @@
     return visitor.visit(this);
   }
 
-  public An getAssumption() {
-    return assumption;
+  public DynamicTypeAssumption getDynamicTypeAssumption() {
+    return dynamicTypeAssumption;
+  }
+
+  public NonNullAssumption getNonNullAssumption() {
+    return nonNullAssumption;
   }
 
   public Value src() {
@@ -98,9 +112,6 @@
 
   @Override
   public String getInstructionName() {
-    if (isAssumeNone()) {
-      return "AssumeNone";
-    }
     if (isAssumeDynamicType()) {
       return "AssumeDynamicType";
     }
@@ -116,51 +127,18 @@
   }
 
   @Override
-  public Assume<An> asAssume() {
+  public Assume asAssume() {
     return this;
   }
 
   @Override
-  public boolean isAssumeNone() {
-    return assumption.isAssumeNone();
-  }
-
-  @Override
-  public Assume<NoAssumption> asAssumeNone() {
-    assert isAssumeNone();
-    @SuppressWarnings("unchecked")
-    Assume<NoAssumption> self = (Assume<NoAssumption>) this;
-    return self;
-  }
-
-  @Override
   public boolean isAssumeDynamicType() {
-    return assumption.isAssumeDynamicType();
-  }
-
-  @Override
-  public Assume<DynamicTypeAssumption> asAssumeDynamicType() {
-    assert isAssumeDynamicType();
-    @SuppressWarnings("unchecked")
-    Assume<DynamicTypeAssumption> self = (Assume<DynamicTypeAssumption>) this;
-    return self;
+    return dynamicTypeAssumption != null;
   }
 
   @Override
   public boolean isAssumeNonNull() {
-    return assumption.isAssumeNonNull();
-  }
-
-  @Override
-  public Assume<NonNullAssumption> asAssumeNonNull() {
-    assert isAssumeNonNull();
-    @SuppressWarnings("unchecked")
-    Assume<NonNullAssumption> self = (Assume<NonNullAssumption>) this;
-    return self;
-  }
-
-  public boolean mayAffectStaticType() {
-    return isAssumeNonNull();
+    return nonNullAssumption != null;
   }
 
   @Override
@@ -171,12 +149,8 @@
     if (outType.isPrimitiveType()) {
       return false;
     }
-    if (assumption.isAssumeNone()) {
-      // The main purpose of AssumeNone is to test local alias tracking.
-      return true;
-    }
-    if (assumption.isAssumeDynamicType()) {
-      outType = asAssumeDynamicType().assumption.getDynamicUpperBoundType();
+    if (isAssumeDynamicType()) {
+      outType = dynamicTypeAssumption.getDynamicUpperBoundType();
     }
     if (appView.appInfo().hasLiveness()) {
       if (outType.isClassType()
@@ -233,8 +207,9 @@
     if (!other.isAssume()) {
       return false;
     }
-    Assume<?> assumeInstruction = other.asAssume();
-    return assumption.equals(assumeInstruction.assumption);
+    Assume assumeInstruction = other.asAssume();
+    return Objects.equals(dynamicTypeAssumption, assumeInstruction.dynamicTypeAssumption)
+        && Objects.equals(nonNullAssumption, assumeInstruction.nonNullAssumption);
   }
 
   @Override
@@ -245,10 +220,10 @@
 
   @Override
   public TypeElement evaluate(AppView<?> appView) {
-    if (assumption.isAssumeNone() || assumption.isAssumeDynamicType()) {
+    if (isAssumeDynamicType()) {
       return src().getType();
     }
-    if (assumption.isAssumeNonNull()) {
+    if (isAssumeNonNull()) {
       assert src().getType().isReferenceType();
       return src().getType().asReferenceType().asMeetWithNotNull();
     }
@@ -281,7 +256,7 @@
 
     TypeElement inType = src().getType();
     TypeElement outType = getOutType();
-    if (isAssumeNone() || isAssumeDynamicType()) {
+    if (isAssumeDynamicType()) {
       assert inType.isReferenceType() : inType;
       assert outType.equals(inType)
           : "At " + this + System.lineSeparator() + outType + " != " + inType;
@@ -303,63 +278,22 @@
     //     are still valid.
     String originString =
         origin.hasBlock() ? " (origin: `" + origin.toString() + "`)" : " (obsolete origin)";
-    if (isAssumeNone() || isAssumeNonNull()) {
+    if (isAssumeNonNull()) {
       return super.toString() + originString;
     }
     if (isAssumeDynamicType()) {
-      DynamicTypeAssumption assumption = asAssumeDynamicType().getAssumption();
       return super.toString()
           + "; upper bound: "
-          + assumption.dynamicUpperBoundType
-          + (assumption.dynamicLowerBoundType != null
-              ? "; lower bound: " + assumption.dynamicLowerBoundType
+          + dynamicTypeAssumption.dynamicUpperBoundType
+          + (dynamicTypeAssumption.dynamicLowerBoundType != null
+              ? "; lower bound: " + dynamicTypeAssumption.dynamicLowerBoundType
               : "")
           + originString;
     }
     return super.toString();
   }
 
-  abstract static class Assumption {
-
-    public boolean isAssumeNone() {
-      return false;
-    }
-
-    public boolean isAssumeDynamicType() {
-      return false;
-    }
-
-    public boolean isAssumeNonNull() {
-      return false;
-    }
-
-    public boolean verifyCorrectnessOfValues(Value dest, Value src, AppView<?> appView) {
-      return true;
-    }
-  }
-
-  public static class NoAssumption extends Assumption {
-    private static final NoAssumption instance = new NoAssumption();
-
-    private NoAssumption() {}
-
-    static NoAssumption get() {
-      return instance;
-    }
-
-    @Override
-    public boolean isAssumeNone() {
-      return true;
-    }
-
-    @Override
-    public boolean verifyCorrectnessOfValues(Value dest, Value src, AppView<?> appView) {
-      assert dest.getType() == src.getType();
-      return true;
-    }
-  }
-
-  public static class DynamicTypeAssumption extends Assumption {
+  public static class DynamicTypeAssumption {
 
     private final TypeElement dynamicUpperBoundType;
     private final ClassTypeElement dynamicLowerBoundType;
@@ -378,12 +312,6 @@
       return dynamicLowerBoundType;
     }
 
-    @Override
-    public boolean isAssumeDynamicType() {
-      return true;
-    }
-
-    @Override
     public boolean verifyCorrectnessOfValues(Value dest, Value src, AppView<?> appView) {
       assert dynamicUpperBoundType.lessThanOrEqualUpToNullability(src.getType(), appView);
       return true;
@@ -408,7 +336,7 @@
     }
   }
 
-  public static class NonNullAssumption extends Assumption {
+  public static class NonNullAssumption {
 
     private static final NonNullAssumption instance = new NonNullAssumption();
 
@@ -418,12 +346,6 @@
       return instance;
     }
 
-    @Override
-    public boolean isAssumeNonNull() {
-      return true;
-    }
-
-    @Override
     public boolean verifyCorrectnessOfValues(Value dest, Value src, AppView<?> appView) {
       assert !src.isNeverNull();
       return true;
diff --git a/src/main/java/com/android/tools/r8/ir/code/DefaultInstructionVisitor.java b/src/main/java/com/android/tools/r8/ir/code/DefaultInstructionVisitor.java
index fdfbe61..f8793ab 100644
--- a/src/main/java/com/android/tools/r8/ir/code/DefaultInstructionVisitor.java
+++ b/src/main/java/com/android/tools/r8/ir/code/DefaultInstructionVisitor.java
@@ -35,7 +35,7 @@
   }
 
   @Override
-  public T visit(Assume<?> instruction) {
+  public T visit(Assume instruction) {
     return null;
   }
 
diff --git a/src/main/java/com/android/tools/r8/ir/code/IRCode.java b/src/main/java/com/android/tools/r8/ir/code/IRCode.java
index 4a6a082..a0af5e2 100644
--- a/src/main/java/com/android/tools/r8/ir/code/IRCode.java
+++ b/src/main/java/com/android/tools/r8/ir/code/IRCode.java
@@ -620,7 +620,7 @@
     for (BasicBlock block : blocks) {
       for (Instruction instruction : block.getInstructions()) {
         if (instruction.isAssumeDynamicType()) {
-          assert instruction.asAssumeDynamicType().verifyInstructionIsNeeded(appView);
+          assert instruction.asAssume().verifyInstructionIsNeeded(appView);
         }
       }
     }
diff --git a/src/main/java/com/android/tools/r8/ir/code/Instruction.java b/src/main/java/com/android/tools/r8/ir/code/Instruction.java
index 6d298df..1400455 100644
--- a/src/main/java/com/android/tools/r8/ir/code/Instruction.java
+++ b/src/main/java/com/android/tools/r8/ir/code/Instruction.java
@@ -22,9 +22,6 @@
 import com.android.tools.r8.ir.analysis.type.TypeElement;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.ir.analysis.value.UnknownValue;
-import com.android.tools.r8.ir.code.Assume.DynamicTypeAssumption;
-import com.android.tools.r8.ir.code.Assume.NoAssumption;
-import com.android.tools.r8.ir.code.Assume.NonNullAssumption;
 import com.android.tools.r8.ir.conversion.CfBuilder;
 import com.android.tools.r8.ir.conversion.DexBuilder;
 import com.android.tools.r8.ir.optimize.Inliner.ConstraintWithTarget;
@@ -716,15 +713,7 @@
     return false;
   }
 
-  public Assume<?> asAssume() {
-    return null;
-  }
-
-  public boolean isAssumeNone() {
-    return false;
-  }
-
-  public Assume<NoAssumption> asAssumeNone() {
+  public Assume asAssume() {
     return null;
   }
 
@@ -732,18 +721,10 @@
     return false;
   }
 
-  public Assume<DynamicTypeAssumption> asAssumeDynamicType() {
-    return null;
-  }
-
   public boolean isAssumeNonNull() {
     return false;
   }
 
-  public Assume<NonNullAssumption> asAssumeNonNull() {
-    return null;
-  }
-
   public boolean isBinop() {
     return false;
   }
diff --git a/src/main/java/com/android/tools/r8/ir/code/InstructionVisitor.java b/src/main/java/com/android/tools/r8/ir/code/InstructionVisitor.java
index 4bb29bb..a374439 100644
--- a/src/main/java/com/android/tools/r8/ir/code/InstructionVisitor.java
+++ b/src/main/java/com/android/tools/r8/ir/code/InstructionVisitor.java
@@ -24,7 +24,7 @@
 
   T visit(ArrayPut instruction);
 
-  T visit(Assume<?> instruction);
+  T visit(Assume instruction);
 
   T visit(CheckCast instruction);
 
diff --git a/src/main/java/com/android/tools/r8/ir/code/Value.java b/src/main/java/com/android/tools/r8/ir/code/Value.java
index b79a063..53af0fc 100644
--- a/src/main/java/com/android/tools/r8/ir/code/Value.java
+++ b/src/main/java/com/android/tools/r8/ir/code/Value.java
@@ -1173,16 +1173,15 @@
       return root.getDynamicUpperBoundType(appView);
     }
 
-    // Try to find an alias of the receiver, which is defined by an instruction of the type
-    // Assume<DynamicTypeAssumption>.
+    // Try to find an alias of the receiver, which is defined by an instruction of the type Assume.
     Value aliasedValue =
         getSpecificAliasedValue(value -> !value.isPhi() && value.definition.isAssumeDynamicType());
     TypeElement lattice;
     if (aliasedValue != null) {
-      // If there is an alias of the receiver, which is defined by an Assume<DynamicTypeAssumption>
-      // instruction, then use the dynamic type as the refined receiver type.
+      // If there is an alias of the receiver, which is defined by an Assume instruction that
+      // carries a dynamic type, then use the dynamic type as the refined receiver type.
       lattice =
-          aliasedValue.definition.asAssumeDynamicType().getAssumption().getDynamicUpperBoundType();
+          aliasedValue.definition.asAssume().getDynamicTypeAssumption().getDynamicUpperBoundType();
 
       // For precision, verify that the dynamic type is at least as precise as the static type.
       assert lattice.lessThanOrEqualUpToNullability(type, appView) : type + " < " + lattice;
@@ -1228,12 +1227,11 @@
       return null;
     }
 
-    // Try to find an alias of the receiver, which is defined by an instruction of the type
-    // Assume<DynamicTypeAssumption>.
+    // Try to find an alias of the receiver, which is defined by an instruction of the type Assume.
     Value aliasedValue = getSpecificAliasedValue(value -> value.definition.isAssumeDynamicType());
     if (aliasedValue != null) {
       ClassTypeElement lattice =
-          aliasedValue.definition.asAssumeDynamicType().getAssumption().getDynamicLowerBoundType();
+          aliasedValue.definition.asAssume().getDynamicTypeAssumption().getDynamicLowerBoundType();
       return lattice != null && type.isDefinitelyNotNull() && lattice.isNullable()
           ? lattice.asMeetWithNotNull()
           : lattice;
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/AssumeDynamicTypeRemover.java b/src/main/java/com/android/tools/r8/ir/optimize/AssumeDynamicTypeRemover.java
index 9102f53..7f45db3 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/AssumeDynamicTypeRemover.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/AssumeDynamicTypeRemover.java
@@ -7,7 +7,6 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.ir.analysis.type.TypeAnalysis;
 import com.android.tools.r8.ir.code.Assume;
-import com.android.tools.r8.ir.code.Assume.DynamicTypeAssumption;
 import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.Instruction;
@@ -34,8 +33,7 @@
   private final IRCode code;
 
   private final Set<Value> affectedValues = Sets.newIdentityHashSet();
-  private final Set<Assume<DynamicTypeAssumption>> assumeDynamicTypeInstructionsToRemove =
-      Sets.newIdentityHashSet();
+  private final Set<Assume> assumeDynamicTypeInstructionsToRemove = Sets.newIdentityHashSet();
 
   private boolean mayHaveIntroducedTrivialPhi = false;
 
@@ -48,21 +46,20 @@
     return mayHaveIntroducedTrivialPhi;
   }
 
-  public void markForRemoval(Assume<DynamicTypeAssumption> assumeDynamicTypeInstruction) {
+  public void markForRemoval(Assume assumeDynamicTypeInstruction) {
     assumeDynamicTypeInstructionsToRemove.add(assumeDynamicTypeInstruction);
   }
 
   public void markUsersForRemoval(Value value) {
     for (Instruction user : value.aliasedUsers()) {
       if (user.isAssumeDynamicType()) {
-        markForRemoval(user.asAssumeDynamicType());
+        markForRemoval(user.asAssume());
       }
     }
   }
 
   public void removeIfMarked(
-      Assume<DynamicTypeAssumption> assumeDynamicTypeInstruction,
-      InstructionListIterator instructionIterator) {
+      Assume assumeDynamicTypeInstruction, InstructionListIterator instructionIterator) {
     if (assumeDynamicTypeInstructionsToRemove.remove(assumeDynamicTypeInstruction)) {
       Value inValue = assumeDynamicTypeInstruction.src();
       Value outValue = assumeDynamicTypeInstruction.outValue();
@@ -91,7 +88,7 @@
         while (instructionIterator.hasNext()) {
           Instruction instruction = instructionIterator.next();
           if (instruction.isAssumeDynamicType()) {
-            removeIfMarked(instruction.asAssumeDynamicType(), instructionIterator);
+            removeIfMarked(instruction.asAssume(), instructionIterator);
           }
         }
       }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java b/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
index d05d64e..0e956c0 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/CallSiteOptimizationInfoPropagator.java
@@ -19,8 +19,6 @@
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.ir.analysis.value.SingleValue;
 import com.android.tools.r8.ir.code.Assume;
-import com.android.tools.r8.ir.code.Assume.DynamicTypeAssumption;
-import com.android.tools.r8.ir.code.Assume.NonNullAssumption;
 import com.android.tools.r8.ir.code.ConstNumber;
 import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.Instruction;
@@ -309,7 +307,7 @@
       return;
     }
     Set<Value> affectedValues = Sets.newIdentityHashSet();
-    List<Assume<?>> assumeInstructions = new LinkedList<>();
+    List<Assume> assumeInstructions = new LinkedList<>();
     List<Instruction> constants = new LinkedList<>();
     int argumentsSeen = 0;
     InstructionListIterator iterator = code.entryBlock().listIterator(code);
@@ -356,7 +354,7 @@
         specializedArg = code.createValue(originalArg.getType());
         affectedValues.addAll(originalArg.affectedValues());
         originalArg.replaceUsers(specializedArg);
-        Assume<DynamicTypeAssumption> assumeType =
+        Assume assumeType =
             Assume.createAssumeDynamicTypeInstruction(
                 dynamicUpperBoundType, null, specializedArg, originalArg, instr, appView);
         assumeType.setPosition(instr.getPosition());
@@ -372,7 +370,7 @@
               code.createValue(specializedArg.getType().asReferenceType().asMeetWithNotNull());
           affectedValues.addAll(specializedArg.affectedValues());
           specializedArg.replaceUsers(nonNullArg);
-          Assume<NonNullAssumption> assumeNotNull =
+          Assume assumeNotNull =
               Assume.createAssumeNonNullInstruction(nonNullArg, specializedArg, instr, appView);
           assumeNotNull.setPosition(instr.getPosition());
           assumeInstructions.add(assumeNotNull);
@@ -387,7 +385,7 @@
             + code.method().method.getArity()
             + ", static: "
             + code.method().isStatic();
-    // After packed Argument instructions, add Assume<?> and constant instructions.
+    // After packed Argument instructions, add Assume and constant instructions.
     assert !iterator.peekPrevious().isArgument();
     iterator.previous();
     assert iterator.peekPrevious().isArgument();
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/CodeRewriter.java b/src/main/java/com/android/tools/r8/ir/optimize/CodeRewriter.java
index 52844be..172a6b3 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/CodeRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/CodeRewriter.java
@@ -176,8 +176,8 @@
 
   public static void removeAssumeInstructions(AppView<?> appView, IRCode code) {
     // We need to update the types of all values whose definitions depend on a non-null value.
-    // This is needed to preserve soundness of the types after the Assume<NonNullAssumption>
-    // instructions have been removed.
+    // This is needed to preserve soundness of the types after the Assume instructions have been
+    // removed.
     //
     // As an example, consider a check-cast instruction on the form "z = (T) y". If y used to be
     // defined by a NonNull instruction, then the type analysis could have used this information
@@ -201,7 +201,7 @@
       //
       //   x.foo()
       if (instruction.isAssume()) {
-        Assume<?> assumeInstruction = instruction.asAssume();
+        Assume assumeInstruction = instruction.asAssume();
         Value src = assumeInstruction.src();
         Value dest = assumeInstruction.outValue();
 
@@ -1507,8 +1507,8 @@
           TypeElement dynamicType =
               aliasedValue
                   .definition
-                  .asAssumeDynamicType()
-                  .getAssumption()
+                  .asAssume()
+                  .getDynamicTypeAssumption()
                   .getDynamicUpperBoundType();
           if (dynamicType.isDefinitelyNull()) {
             result = InstanceOfResult.FALSE;
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/Devirtualizer.java b/src/main/java/com/android/tools/r8/ir/optimize/Devirtualizer.java
index 1f91212..7752d6c 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/Devirtualizer.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/Devirtualizer.java
@@ -13,7 +13,6 @@
 import com.android.tools.r8.ir.analysis.type.TypeAnalysis;
 import com.android.tools.r8.ir.analysis.type.TypeElement;
 import com.android.tools.r8.ir.code.Assume;
-import com.android.tools.r8.ir.code.Assume.NonNullAssumption;
 import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.CheckCast;
 import com.android.tools.r8.ir.code.DominatorTree;
@@ -78,7 +77,7 @@
         // ...
         // non_null_rcv <- non-null rcv_c  // <- Update the input rcv to the non-null, too.
         if (current.isAssumeNonNull()) {
-          Assume<NonNullAssumption> nonNull = current.asAssumeNonNull();
+          Assume nonNull = current.asAssume();
           Instruction origin = nonNull.origin();
           if (origin.isInvokeInterface()
               && !origin.asInvokeInterface().getReceiver().hasLocalInfo()
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/DynamicTypeOptimization.java b/src/main/java/com/android/tools/r8/ir/optimize/DynamicTypeOptimization.java
index 4aabc5b..6e2d9c7 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/DynamicTypeOptimization.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/DynamicTypeOptimization.java
@@ -15,7 +15,6 @@
 import com.android.tools.r8.ir.analysis.type.ClassTypeElement;
 import com.android.tools.r8.ir.analysis.type.TypeElement;
 import com.android.tools.r8.ir.code.Assume;
-import com.android.tools.r8.ir.code.Assume.DynamicTypeAssumption;
 import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.BasicBlockIterator;
 import com.android.tools.r8.ir.code.IRCode;
@@ -143,7 +142,7 @@
       outValue.replaceUsers(specializedOutValue);
 
       // Insert AssumeDynamicType instruction.
-      Assume<DynamicTypeAssumption> assumeInstruction =
+      Assume assumeInstruction =
           Assume.createAssumeDynamicTypeInstruction(
               dynamicUpperBoundType,
               dynamicLowerBoundType,
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 fe6cbb4..bf37ac8 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
@@ -1084,7 +1084,7 @@
             blockIterator.next();
           }
         } else if (current.isAssumeDynamicType()) {
-          assumeDynamicTypeRemover.removeIfMarked(current.asAssumeDynamicType(), iterator);
+          assumeDynamicTypeRemover.removeIfMarked(current.asAssume(), iterator);
         }
       }
     }
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 e08aeae..6918398 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
@@ -13,7 +13,6 @@
 import com.android.tools.r8.ir.analysis.type.TypeAnalysis;
 import com.android.tools.r8.ir.analysis.type.TypeElement;
 import com.android.tools.r8.ir.code.Assume;
-import com.android.tools.r8.ir.code.Assume.NonNullAssumption;
 import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.BasicBlockIterator;
 import com.android.tools.r8.ir.code.DominatorTree;
@@ -390,7 +389,7 @@
           }
           affectedValues.addAll(newValue.affectedValues());
 
-          Assume<NonNullAssumption> assumeInstruction =
+          Assume assumeInstruction =
               Assume.createAssumeNonNullInstruction(newValue, nonNullValue, instruction, appView);
           assumeInstruction.setPosition(instruction.getPosition());
           if (insertionBlock != block) {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/MethodOptimizationInfoCollector.java b/src/main/java/com/android/tools/r8/ir/optimize/info/MethodOptimizationInfoCollector.java
index a1a6638..cac5801 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/MethodOptimizationInfoCollector.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/MethodOptimizationInfoCollector.java
@@ -1157,7 +1157,7 @@
       nullCheckedBlocks.clear();
       for (Instruction user : argument.uniqueUsers()) {
         if (user.isAssumeNonNull()) {
-          nullCheckedBlocks.add(user.asAssumeNonNull().getBlock());
+          nullCheckedBlocks.add(user.asAssume().getBlock());
         }
         if (user.isIf()
             && user.asIf().isZeroTest()
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/library/EnumMethodOptimizer.java b/src/main/java/com/android/tools/r8/ir/optimize/library/EnumMethodOptimizer.java
index fa0e683..cdc1793 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/library/EnumMethodOptimizer.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/library/EnumMethodOptimizer.java
@@ -12,7 +12,6 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.ir.analysis.type.TypeElement;
 import com.android.tools.r8.ir.code.Assume;
-import com.android.tools.r8.ir.code.Assume.DynamicTypeAssumption;
 import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.InstructionListIterator;
 import com.android.tools.r8.ir.code.InvokeMethod;
@@ -69,7 +68,7 @@
     outValue.replaceUsers(specializedOutValue);
 
     // Insert AssumeDynamicType instruction.
-    Assume<DynamicTypeAssumption> assumeInstruction =
+    Assume assumeInstruction =
         Assume.createAssumeDynamicTypeInstruction(
             dynamicUpperBoundType, null, specializedOutValue, outValue, invoke, appView);
     assumeInstruction.setPosition(appView.options().debug ? invoke.getPosition() : Position.none());
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizer.java b/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizer.java
index 53122b0..a01be35 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizer.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizer.java
@@ -767,7 +767,7 @@
         }
         // If there are aliasing instructions, they should be removed before new-instance.
         if (instr.isAssume() && buildersToRemove.contains(instr.outValue().getAliasedValue())) {
-          Assume<?> assumeInstruction = instr.asAssume();
+          Assume assumeInstruction = instr.asAssume();
           Value src = assumeInstruction.src();
           Value dest = assumeInstruction.outValue();
           dest.replaceUsers(src);
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/NonNullTrackerTest.java b/src/test/java/com/android/tools/r8/ir/optimize/NonNullTrackerTest.java
index 06305cf..f48d061 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/NonNullTrackerTest.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/NonNullTrackerTest.java
@@ -69,8 +69,7 @@
             || (prev.isIf() && prev.asIf().isZeroTest())
             || !curr.getBlock().getPredecessors().contains(prev.getBlock()));
         // Make sure non-null is used or inserted for arguments.
-        assertTrue(
-            curr.outValue().numberOfAllUsers() > 0 || curr.asAssumeNonNull().src().isArgument());
+        assertTrue(curr.outValue().numberOfAllUsers() > 0 || curr.asAssume().src().isArgument());
         count++;
       }
     }