diff --git a/src/main/java/com/android/tools/r8/graph/AppView.java b/src/main/java/com/android/tools/r8/graph/AppView.java
index 72cddc2..0225c1d 100644
--- a/src/main/java/com/android/tools/r8/graph/AppView.java
+++ b/src/main/java/com/android/tools/r8/graph/AppView.java
@@ -97,7 +97,7 @@
   private InitializedClassesInInstanceMethods initializedClassesInInstanceMethods;
   private HorizontallyMergedClasses horizontallyMergedClasses = HorizontallyMergedClasses.empty();
   private VerticallyMergedClasses verticallyMergedClasses;
-  private EnumDataMap unboxedEnums = EnumDataMap.empty();
+  private EnumDataMap unboxedEnums = null;
   // TODO(b/169115389): Remove
   private Set<DexMethod> cfByteCodePassThrough = ImmutableSet.of();
   private Map<DexType, DexValueString> sourceDebugExtensions = new IdentityHashMap<>();
@@ -575,12 +575,16 @@
     testing().verticallyMergedClassesConsumer.accept(dexItemFactory(), verticallyMergedClasses);
   }
 
+  public boolean hasUnboxedEnums() {
+    return unboxedEnums != null;
+  }
+
   public EnumDataMap unboxedEnums() {
-    return unboxedEnums;
+    return hasUnboxedEnums() ? unboxedEnums : EnumDataMap.empty();
   }
 
   public void setUnboxedEnums(EnumDataMap unboxedEnums) {
-    assert this.unboxedEnums.isEmpty();
+    assert !hasUnboxedEnums();
     this.unboxedEnums = unboxedEnums;
     testing().unboxedEnumsConsumer.accept(dexItemFactory(), unboxedEnums);
   }
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/MethodOptimizationFeedback.java b/src/main/java/com/android/tools/r8/ir/conversion/MethodOptimizationFeedback.java
index 13b7140..53bb277 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/MethodOptimizationFeedback.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/MethodOptimizationFeedback.java
@@ -66,6 +66,8 @@
   void setEnumUnboxerMethodClassification(
       ProgramMethod method, EnumUnboxerMethodClassification enumUnboxerMethodClassification);
 
+  void unsetEnumUnboxerMethodClassification(ProgramMethod method);
+
   void setInstanceInitializerInfoCollection(
       DexEncodedMethod method, InstanceInitializerInfoCollection instanceInitializerInfoCollection);
 
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxerImpl.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxerImpl.java
index 46eec15..e4229c2 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxerImpl.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxerImpl.java
@@ -10,15 +10,18 @@
 import static com.android.tools.r8.ir.code.Opcodes.ARRAY_PUT;
 import static com.android.tools.r8.ir.code.Opcodes.ASSUME;
 import static com.android.tools.r8.ir.code.Opcodes.CHECK_CAST;
+import static com.android.tools.r8.ir.code.Opcodes.CONST_CLASS;
 import static com.android.tools.r8.ir.code.Opcodes.IF;
 import static com.android.tools.r8.ir.code.Opcodes.INSTANCE_GET;
 import static com.android.tools.r8.ir.code.Opcodes.INSTANCE_PUT;
+import static com.android.tools.r8.ir.code.Opcodes.INVOKE_CUSTOM;
 import static com.android.tools.r8.ir.code.Opcodes.INVOKE_DIRECT;
 import static com.android.tools.r8.ir.code.Opcodes.INVOKE_INTERFACE;
 import static com.android.tools.r8.ir.code.Opcodes.INVOKE_STATIC;
 import static com.android.tools.r8.ir.code.Opcodes.INVOKE_SUPER;
 import static com.android.tools.r8.ir.code.Opcodes.INVOKE_VIRTUAL;
 import static com.android.tools.r8.ir.code.Opcodes.RETURN;
+import static com.android.tools.r8.ir.code.Opcodes.STATIC_GET;
 import static com.android.tools.r8.ir.code.Opcodes.STATIC_PUT;
 import static com.android.tools.r8.utils.MapUtils.ignoreKey;
 
@@ -42,7 +45,6 @@
 import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.MethodResolutionResult;
 import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.graph.PrunedItems;
 import com.android.tools.r8.graph.UseRegistry;
 import com.android.tools.r8.ir.analysis.fieldvalueanalysis.StaticFieldValues;
 import com.android.tools.r8.ir.analysis.fieldvalueanalysis.StaticFieldValues.EnumStaticFieldValues;
@@ -69,7 +71,6 @@
 import com.android.tools.r8.ir.code.InvokeStatic;
 import com.android.tools.r8.ir.code.InvokeVirtual;
 import com.android.tools.r8.ir.code.MemberType;
-import com.android.tools.r8.ir.code.Opcodes;
 import com.android.tools.r8.ir.code.Phi;
 import com.android.tools.r8.ir.code.Return;
 import com.android.tools.r8.ir.code.Value;
@@ -102,6 +103,8 @@
 import com.android.tools.r8.utils.Reporter;
 import com.android.tools.r8.utils.StringDiagnostic;
 import com.android.tools.r8.utils.collections.ImmutableInt2ReferenceSortedMap;
+import com.android.tools.r8.utils.collections.LongLivedClassSetBuilder;
+import com.android.tools.r8.utils.collections.LongLivedProgramMethodMapBuilder;
 import com.android.tools.r8.utils.collections.LongLivedProgramMethodSetBuilder;
 import com.android.tools.r8.utils.collections.ProgramMethodMap;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
@@ -147,8 +150,8 @@
   private LongLivedProgramMethodSetBuilder<ProgramMethodSet> methodsDependingOnLibraryModelisation;
 
   // Map from checkNotNull() methods to the enums that use the given method.
-  private final ProgramMethodMap<Set<DexProgramClass>> checkNotNullMethods =
-      ProgramMethodMap.createConcurrent();
+  private LongLivedProgramMethodMapBuilder<LongLivedClassSetBuilder<DexProgramClass>>
+      checkNotNullMethodsBuilder;
 
   private final DexClassAndField ordinalField;
 
@@ -261,21 +264,21 @@
           }
         }
         switch (instruction.opcode()) {
-          case Opcodes.CONST_CLASS:
+          case CONST_CLASS:
             analyzeConstClass(instruction.asConstClass(), eligibleEnums, code.context());
             break;
-          case Opcodes.CHECK_CAST:
+          case CHECK_CAST:
             analyzeCheckCast(instruction.asCheckCast(), eligibleEnums);
             break;
-          case Opcodes.INVOKE_CUSTOM:
+          case INVOKE_CUSTOM:
             analyzeInvokeCustom(instruction.asInvokeCustom(), eligibleEnums);
             break;
           case INVOKE_STATIC:
             analyzeInvokeStatic(instruction.asInvokeStatic(), eligibleEnums, code.context());
             break;
-          case Opcodes.STATIC_GET:
-          case Opcodes.INSTANCE_GET:
-          case Opcodes.STATIC_PUT:
+          case STATIC_GET:
+          case INSTANCE_GET:
+          case STATIC_PUT:
           case INSTANCE_PUT:
             analyzeFieldInstruction(
                 instruction.asFieldInstruction(), eligibleEnums, code.context());
@@ -560,9 +563,18 @@
 
   @Override
   public void prepareForPrimaryOptimizationPass(GraphLens graphLensForPrimaryOptimizationPass) {
+    assert appView.graphLens() == graphLensForPrimaryOptimizationPass;
+    initializeCheckNotNullMethods(graphLensForPrimaryOptimizationPass);
     initializeEnumUnboxingCandidates(graphLensForPrimaryOptimizationPass);
   }
 
+  private void initializeCheckNotNullMethods(GraphLens graphLensForPrimaryOptimizationPass) {
+    assert checkNotNullMethodsBuilder == null;
+    checkNotNullMethodsBuilder =
+        LongLivedProgramMethodMapBuilder.createConcurrentBuilderForNonConcurrentMap(
+            graphLensForPrimaryOptimizationPass);
+  }
+
   private void initializeEnumUnboxingCandidates(GraphLens graphLensForPrimaryOptimizationPass) {
     assert enumUnboxingCandidatesInfo == null;
     enumUnboxingCandidatesInfo =
@@ -577,10 +589,12 @@
   public void unboxEnums(
       AppView<AppInfoWithLiveness> appView,
       IRConverter converter,
-      Builder postBuilder,
+      Builder postMethodProcessorBuilder,
       ExecutorService executorService,
       OptimizationFeedbackDelayed feedback)
       throws ExecutionException {
+    assert feedback.noUpdatesLeft();
+
     assert candidatesToRemoveInWave.isEmpty();
     EnumDataMap enumDataMap = finishAnalysis();
     assert candidatesToRemoveInWave.isEmpty();
@@ -588,9 +602,10 @@
     // At this point the enum unboxing candidates are no longer candidates, they will all be
     // unboxed. We extract the now immutable enums to unbox information and clear the candidate
     // info.
+    appView.setUnboxedEnums(enumDataMap);
+
     if (enumUnboxingCandidatesInfo.isEmpty()) {
       assert enumDataMap.isEmpty();
-      appView.setUnboxedEnums(enumDataMap);
       return;
     }
 
@@ -609,12 +624,15 @@
             .build(converter, executorService);
 
     // Fixup the application.
+    ProgramMethodMap<Set<DexProgramClass>> checkNotNullMethods =
+        checkNotNullMethodsBuilder
+            .rewrittenWithLens(appView, (enumClasses, appliedGraphLens) -> enumClasses)
+            .build(appView, builder -> builder.build(appView));
     EnumUnboxingTreeFixer.Result treeFixerResult =
         new EnumUnboxingTreeFixer(
                 appView, checkNotNullMethods, enumDataMap, enumClassesToUnbox, utilityClasses)
             .fixupTypeReferences(converter, executorService);
     EnumUnboxingLens enumUnboxingLens = treeFixerResult.getLens();
-    appView.setUnboxedEnums(enumDataMap);
 
     // Update the graph lens.
     appView.rewriteWithLens(enumUnboxingLens);
@@ -624,7 +642,7 @@
     // Note that the reprocessing set must be rewritten to the new enum unboxing lens before pruning
     // the builders with the methods removed by the tree fixer (since these methods references are
     // already fully lens rewritten).
-    postBuilder
+    postMethodProcessorBuilder
         .getMethodsToReprocessBuilder()
         .rewrittenWithLens(appView)
         .merge(dependencies)
@@ -632,7 +650,7 @@
         .removeAll(treeFixerResult.getPrunedItems().getRemovedMethods());
     methodsDependingOnLibraryModelisation.clear();
 
-    updateOptimizationInfos(executorService, feedback, treeFixerResult.getPrunedItems());
+    updateOptimizationInfos(executorService, feedback, treeFixerResult);
 
     enumUnboxerRewriter =
         new EnumUnboxingRewriter(
@@ -647,7 +665,7 @@
   private void updateOptimizationInfos(
       ExecutorService executorService,
       OptimizationFeedbackDelayed feedback,
-      PrunedItems prunedItems)
+      EnumUnboxingTreeFixer.Result treeFixerResult)
       throws ExecutionException {
     feedback.fixupOptimizationInfos(
         appView,
@@ -666,7 +684,16 @@
             optimizationInfo
                 .fixupClassTypeReferences(appView, appView.graphLens())
                 .fixupAbstractReturnValue(appView, appView.graphLens())
-                .fixupInstanceInitializerInfo(appView, appView.graphLens(), prunedItems);
+                .fixupInstanceInitializerInfo(
+                    appView, appView.graphLens(), treeFixerResult.getPrunedItems());
+
+            // Clear the enum unboxer method classification for check-not-null methods (these
+            // classifications are transferred to the synthesized check-not-zero methods by now).
+            if (!treeFixerResult
+                .getCheckNotNullToCheckNotZeroMapping()
+                .containsValue(method.getReference())) {
+              optimizationInfo.unsetEnumUnboxerMethodClassification();
+            }
           }
         });
   }
@@ -1316,10 +1343,16 @@
             classification.asCheckNotNullClassification();
         if (checkNotNullClassification.isUseEligibleForUnboxing(
             invoke.asInvokeStatic(), enumValue)) {
-          checkNotNullMethods
+          GraphLens graphLens = appView.graphLens();
+          checkNotNullMethodsBuilder
               .computeIfAbsent(
-                  singleTarget.asProgramMethod(), ignoreKey(Sets::newConcurrentHashSet))
-              .add(enumClass);
+                  singleTarget.asProgramMethod(),
+                  ignoreKey(
+                      () ->
+                          LongLivedClassSetBuilder.createConcurrentBuilderForIdentitySet(
+                              graphLens)),
+                  graphLens)
+              .add(enumClass, graphLens);
           return Reason.ELIGIBLE;
         }
       }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingRewriter.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingRewriter.java
index 02d8cc1..5fc743e 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingRewriter.java
@@ -398,26 +398,38 @@
       return;
     }
 
-    if (singleTarget.isProgramMethod()) {
-      EnumUnboxerMethodClassification classification =
-          singleTarget.getOptimizationInfo().getEnumUnboxerMethodClassification();
-      if (classification.isCheckNotNullClassification()) {
-        CheckNotNullEnumUnboxerMethodClassification checkNotNullClassification =
-            classification.asCheckNotNullClassification();
-        Value argument = invoke.getArgument(checkNotNullClassification.getArgumentIndex());
-        DexType enumType = getEnumTypeOrNull(argument, convertedEnums);
-        if (enumType != null) {
-          InvokeStatic replacement =
-              InvokeStatic.builder()
-                  .setMethod(checkNotNullToCheckNotZeroMapping.get(singleTarget.getReference()))
-                  .setArguments(invoke.arguments())
-                  .setPosition(invoke.getPosition())
-                  .build();
-          instructionIterator.replaceCurrentInstruction(replacement);
-          convertedEnums.put(replacement, enumType);
+    if (singleTarget.isProgramMethod()
+        && checkNotNullToCheckNotZeroMapping.containsKey(singleTarget.getReference())) {
+      DexMethod checkNotZeroMethodReference =
+          checkNotNullToCheckNotZeroMapping.get(singleTarget.getReference());
+      ProgramMethod checkNotZeroMethod =
+          appView
+              .appInfo()
+              .resolveMethodOnClass(checkNotZeroMethodReference)
+              .getResolvedProgramMethod();
+      if (checkNotZeroMethod != null) {
+        EnumUnboxerMethodClassification classification =
+            checkNotZeroMethod.getOptimizationInfo().getEnumUnboxerMethodClassification();
+        if (classification.isCheckNotNullClassification()) {
+          CheckNotNullEnumUnboxerMethodClassification checkNotNullClassification =
+              classification.asCheckNotNullClassification();
+          Value argument = invoke.getArgument(checkNotNullClassification.getArgumentIndex());
+          DexType enumType = getEnumTypeOrNull(argument, convertedEnums);
+          if (enumType != null) {
+            InvokeStatic replacement =
+                InvokeStatic.builder()
+                    .setMethod(checkNotZeroMethod)
+                    .setArguments(invoke.arguments())
+                    .setPosition(invoke.getPosition())
+                    .build();
+            instructionIterator.replaceCurrentInstruction(replacement);
+            convertedEnums.put(replacement, enumType);
+          }
+        } else {
+          assert false;
         }
       } else {
-        assert !checkNotNullToCheckNotZeroMapping.containsKey(singleTarget.getReference());
+        assert false;
       }
     }
   }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingTreeFixer.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingTreeFixer.java
index c7bcc37..8a97aef 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingTreeFixer.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingTreeFixer.java
@@ -54,10 +54,11 @@
 import com.android.tools.r8.utils.ThreadUtils;
 import com.android.tools.r8.utils.Timing;
 import com.android.tools.r8.utils.collections.ProgramMethodMap;
+import com.google.common.collect.BiMap;
+import com.google.common.collect.HashBiMap;
 import com.google.common.collect.Sets;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.IdentityHashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.ListIterator;
@@ -130,16 +131,16 @@
     }
 
     // Create mapping from checkNotNull() to checkNotZero() methods.
-    Map<DexMethod, DexMethod> checkNotNullToCheckNotZeroMapping =
+    BiMap<DexMethod, DexMethod> checkNotNullToCheckNotZeroMapping =
         duplicateCheckNotNullMethods(converter, executorService);
 
     return new Result(
         checkNotNullToCheckNotZeroMapping, lensBuilder.build(appView), prunedItemsBuilder.build());
   }
 
-  private Map<DexMethod, DexMethod> duplicateCheckNotNullMethods(
+  private BiMap<DexMethod, DexMethod> duplicateCheckNotNullMethods(
       IRConverter converter, ExecutorService executorService) throws ExecutionException {
-    Map<DexMethod, DexMethod> checkNotNullToCheckNotZeroMapping = new IdentityHashMap<>();
+    BiMap<DexMethod, DexMethod> checkNotNullToCheckNotZeroMapping = HashBiMap.create();
     ProcessorContext processorContext = appView.createProcessorContext();
     OneTimeMethodProcessor.Builder methodProcessorBuilder =
         OneTimeMethodProcessor.builder(processorContext);
@@ -186,6 +187,11 @@
                               .setApiLevelForDefinition(minApiLevelIfEnabledOrUnknown(appView))
                               .setApiLevelForCode(minApiLevelIfEnabledOrUnknown(appView))
                               .setCode(method -> new CheckNotZeroCode(checkNotNullMethod))
+                              .setOptimizationInfo(
+                                  checkNotNullMethod
+                                      .getOptimizationInfo()
+                                      .asMutableMethodOptimizationInfo()
+                                      .mutableCopy())
                               .setProto(newProto));
           checkNotNullToCheckNotZeroMapping.put(
               checkNotNullMethod.getReference(), checkNotZeroMethod.getReference());
@@ -603,12 +609,12 @@
 
   public static class Result {
 
-    private final Map<DexMethod, DexMethod> checkNotNullToCheckNotZeroMapping;
+    private final BiMap<DexMethod, DexMethod> checkNotNullToCheckNotZeroMapping;
     private final EnumUnboxingLens lens;
     private final PrunedItems prunedItems;
 
     Result(
-        Map<DexMethod, DexMethod> checkNotNullToCheckNotZeroMapping,
+        BiMap<DexMethod, DexMethod> checkNotNullToCheckNotZeroMapping,
         EnumUnboxingLens lens,
         PrunedItems prunedItems) {
       this.checkNotNullToCheckNotZeroMapping = checkNotNullToCheckNotZeroMapping;
@@ -616,7 +622,7 @@
       this.prunedItems = prunedItems;
     }
 
-    Map<DexMethod, DexMethod> getCheckNotNullToCheckNotZeroMapping() {
+    BiMap<DexMethod, DexMethod> getCheckNotNullToCheckNotZeroMapping() {
       return checkNotNullToCheckNotZeroMapping;
     }
 
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/classification/EnumUnboxerMethodClassificationAnalysis.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/classification/EnumUnboxerMethodClassificationAnalysis.java
index f01abfc..3d8bccc 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/classification/EnumUnboxerMethodClassificationAnalysis.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/classification/EnumUnboxerMethodClassificationAnalysis.java
@@ -46,13 +46,25 @@
     }
 
     // Look for an argument with a single if-zero user.
+    EnumUnboxerMethodClassification currentClassification =
+        method.getOptimizationInfo().getEnumUnboxerMethodClassification();
     DexItemFactory dexItemFactory = appView.dexItemFactory();
     InstructionIterator entryIterator = code.entryBlock().iterator();
     for (int index = 0; index < method.getParameters().size(); index++) {
       Argument argument = entryIterator.next().asArgument();
       DexType parameter = method.getParameter(index);
-      if (parameter != dexItemFactory.objectType) {
-        continue;
+      // Before enum unboxing, we classify methods with `object != null` as check-not-null methods.
+      // After enum unboxing, we check correctness of the classification for check-not-zero methods.
+      if (appView.hasUnboxedEnums()) {
+        if (parameter != dexItemFactory.intType
+            || !currentClassification.isCheckNotNullClassification()
+            || currentClassification.asCheckNotNullClassification().getArgumentIndex() != index) {
+          continue;
+        }
+      } else {
+        if (parameter != dexItemFactory.objectType) {
+          continue;
+        }
       }
 
       if (onlyHasCheckNotNullUsers(argument, methodProcessor)) {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/code/CheckNotZeroCode.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/code/CheckNotZeroCode.java
index 38ad5cb..c2cbec9 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/code/CheckNotZeroCode.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/code/CheckNotZeroCode.java
@@ -49,7 +49,7 @@
     // Start iterating at the argument instruction for the checked argument.
     IteratorUtils.skip(
         instructionIterator,
-        checkNotNullMethod
+        checkNotZeroMethod
             .getOptimizationInfo()
             .getEnumUnboxerMethodClassification()
             .asCheckNotNullClassification()
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 6abfb55..db7c3d0 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
@@ -808,9 +808,33 @@
       IRCode code,
       OptimizationFeedback feedback,
       MethodProcessor methodProcessor) {
-    EnumUnboxerMethodClassification enumUnboxerMethodClassification =
-        EnumUnboxerMethodClassificationAnalysis.analyze(appView, method, code, methodProcessor);
-    feedback.setEnumUnboxerMethodClassification(method, enumUnboxerMethodClassification);
+    if (appView.hasUnboxedEnums()) {
+      if (appView.unboxedEnums().isEmpty()) {
+        feedback.unsetEnumUnboxerMethodClassification(method);
+      } else {
+        assert verifyEnumUnboxerMethodClassificationCorrect(method, code, methodProcessor);
+      }
+    } else {
+      EnumUnboxerMethodClassification enumUnboxerMethodClassification =
+          EnumUnboxerMethodClassificationAnalysis.analyze(appView, method, code, methodProcessor);
+      feedback.setEnumUnboxerMethodClassification(method, enumUnboxerMethodClassification);
+    }
+  }
+
+  private boolean verifyEnumUnboxerMethodClassificationCorrect(
+      ProgramMethod method, IRCode code, MethodProcessor methodProcessor) {
+    EnumUnboxerMethodClassification existingClassification =
+        method.getOptimizationInfo().getEnumUnboxerMethodClassification();
+    if (existingClassification.isCheckNotNullClassification()) {
+      EnumUnboxerMethodClassification computedClassification =
+          EnumUnboxerMethodClassificationAnalysis.analyze(appView, method, code, methodProcessor);
+      assert computedClassification.isCheckNotNullClassification();
+      assert computedClassification.asCheckNotNullClassification().getArgumentIndex()
+          == existingClassification.asCheckNotNullClassification().getArgumentIndex();
+    } else {
+      assert existingClassification.isUnknownClassification();
+    }
+    return true;
   }
 
   private void computeSimpleInliningConstraint(
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java b/src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java
index c4a08f4..d8bbe6b 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java
@@ -280,6 +280,10 @@
     this.enumUnboxerMethodClassification = enumUnboxerMethodClassification;
   }
 
+  public void unsetEnumUnboxerMethodClassification() {
+    this.enumUnboxerMethodClassification = EnumUnboxerMethodClassification.unknown();
+  }
+
   public MutableMethodOptimizationInfo fixupEnumUnboxerMethodClassification(
       MethodOptimizationInfoFixer fixer) {
     enumUnboxerMethodClassification =
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackDelayed.java b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackDelayed.java
index c9390d7..067407d 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackDelayed.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackDelayed.java
@@ -272,6 +272,11 @@
   }
 
   @Override
+  public void unsetEnumUnboxerMethodClassification(ProgramMethod method) {
+    getMethodOptimizationInfoForUpdating(method).unsetEnumUnboxerMethodClassification();
+  }
+
+  @Override
   public synchronized void setInstanceInitializerInfoCollection(
       DexEncodedMethod method,
       InstanceInitializerInfoCollection instanceInitializerInfoCollection) {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackIgnore.java b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackIgnore.java
index 623d8d2..1894532 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackIgnore.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackIgnore.java
@@ -122,6 +122,9 @@
       ProgramMethod method, EnumUnboxerMethodClassification enumUnboxerMethodClassification) {}
 
   @Override
+  public void unsetEnumUnboxerMethodClassification(ProgramMethod method) {}
+
+  @Override
   public void setInstanceInitializerInfoCollection(
       DexEncodedMethod method,
       InstanceInitializerInfoCollection instanceInitializerInfoCollection) {}
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackSimple.java b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackSimple.java
index 47f1b8e..d5859c1 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackSimple.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackSimple.java
@@ -194,6 +194,16 @@
   }
 
   @Override
+  public void unsetEnumUnboxerMethodClassification(ProgramMethod method) {
+    if (method.getOptimizationInfo().isMutableOptimizationInfo()) {
+      method
+          .getOptimizationInfo()
+          .asMutableMethodOptimizationInfo()
+          .unsetEnumUnboxerMethodClassification();
+    }
+  }
+
+  @Override
   public void setInstanceInitializerInfoCollection(
       DexEncodedMethod method,
       InstanceInitializerInfoCollection instanceInitializerInfoCollection) {
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
index 76845cc..a7a68ff 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
@@ -394,6 +394,13 @@
           || method.getDefinition().isInstanceInitializer()) {
         return ArgumentInfoCollection.empty();
       }
+      // TODO(b/199864962): Allow parameter removal from check-not-null classified methods.
+      if (method
+          .getOptimizationInfo()
+          .getEnumUnboxerMethodClassification()
+          .isCheckNotNullClassification()) {
+        return ArgumentInfoCollection.empty();
+      }
       return computeRemovableParametersFromMethod(method);
     }
 
diff --git a/src/test/java/com/android/tools/r8/optimize/argumentpropagation/CheckNotZeroMethodWithArgumentRemovalTest.java b/src/test/java/com/android/tools/r8/optimize/argumentpropagation/CheckNotZeroMethodWithArgumentRemovalTest.java
new file mode 100644
index 0000000..be25553
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/optimize/argumentpropagation/CheckNotZeroMethodWithArgumentRemovalTest.java
@@ -0,0 +1,79 @@
+// Copyright (c) 2021, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.optimize.argumentpropagation;
+
+import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertEquals;
+
+import com.android.tools.r8.NeverInline;
+import com.android.tools.r8.TestBase;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.TestParametersCollection;
+import com.android.tools.r8.utils.codeinspector.ClassSubject;
+import com.android.tools.r8.utils.codeinspector.MethodSubject;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class CheckNotZeroMethodWithArgumentRemovalTest extends TestBase {
+
+  @Parameter(0)
+  public TestParameters parameters;
+
+  @Parameters(name = "{0}")
+  public static TestParametersCollection parameters() {
+    return getTestParameters().withAllRuntimesAndApiLevels().build();
+  }
+
+  @Test
+  public void test() throws Exception {
+    testForR8(parameters.getBackend())
+        .addInnerClasses(getClass())
+        .addKeepMainRule(Main.class)
+        .addEnumUnboxingInspector(inspector -> inspector.assertUnboxed(MyEnum.class))
+        .enableInliningAnnotations()
+        // TODO(b/173398086): uniqueMethodWithName() does not work with argument removal.
+        .noMinification()
+        .setMinApi(parameters.getApiLevel())
+        .compile()
+        .inspect(
+            inspector -> {
+              ClassSubject mainClassSubject = inspector.clazz(Main.class);
+              assertThat(mainClassSubject, isPresent());
+
+              MethodSubject checkNotNullSubject =
+                  mainClassSubject.uniqueMethodWithName("checkNotNull");
+              assertThat(checkNotNullSubject, isPresent());
+              // TODO(b/199864962): Allow parameter removal from check-not-null classified methods.
+              assertEquals(2, checkNotNullSubject.getProgramMethod().getReference().getArity());
+            })
+        .run(parameters.getRuntime(), Main.class)
+        .assertSuccessWithEmptyOutput();
+  }
+
+  static class Main {
+
+    public static void main(String[] args) {
+      checkNotNull(System.currentTimeMillis() > 0 ? MyEnum.A : null, "x");
+      checkNotNull(System.currentTimeMillis() > 0 ? new Object() : null, "x");
+    }
+
+    @NeverInline
+    static void checkNotNull(Object o, String name) {
+      if (o == null) {
+        throw new NullPointerException("Expected not null, but " + name + " was null");
+      }
+    }
+  }
+
+  enum MyEnum {
+    A,
+    B
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/optimize/argumentpropagation/MixedArgumentRemovalAndEnumUnboxingTest.java b/src/test/java/com/android/tools/r8/optimize/argumentpropagation/MixedArgumentRemovalAndEnumUnboxingTest.java
index ad4c08a..58f2423 100644
--- a/src/test/java/com/android/tools/r8/optimize/argumentpropagation/MixedArgumentRemovalAndEnumUnboxingTest.java
+++ b/src/test/java/com/android/tools/r8/optimize/argumentpropagation/MixedArgumentRemovalAndEnumUnboxingTest.java
@@ -84,7 +84,7 @@
     }
 
     @NeverInline
-    static void test(Object alwaysNull, MyEnum alwaysA, Object alsoAlwaysNull, MyEnum alwaysB) {
+    static void test(Main alwaysNull, MyEnum alwaysA, Main alsoAlwaysNull, MyEnum alwaysB) {
       if (alwaysNull == null && alsoAlwaysNull == null) {
         System.out.println(alwaysA.name());
         System.out.println(alwaysB.name());