diff --git a/.gitignore b/.gitignore
index 27d89f9..ad943a3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -160,6 +160,8 @@
 third_party/kotlin/kotlin-compiler-1.8.0
 third_party/kotlin/kotlin-compiler-1.9.21.tar.gz
 third_party/kotlin/kotlin-compiler-1.9.21
+third_party/kotlin/kotlin-compiler-2.0.20.tar.gz
+third_party/kotlin/kotlin-compiler-2.0.20
 third_party/kotlin/kotlin-compiler-dev.tar.gz
 third_party/kotlin/kotlin-compiler-dev
 third_party/kotlinx-coroutines-1.3.6.tar.gz
diff --git a/d8_r8/commonBuildSrc/src/main/kotlin/DependenciesPlugin.kt b/d8_r8/commonBuildSrc/src/main/kotlin/DependenciesPlugin.kt
index b54edc3..1154694 100644
--- a/d8_r8/commonBuildSrc/src/main/kotlin/DependenciesPlugin.kt
+++ b/d8_r8/commonBuildSrc/src/main/kotlin/DependenciesPlugin.kt
@@ -870,6 +870,7 @@
     "kotlin-compiler-1.7.0",
     "kotlin-compiler-1.8.0",
     "kotlin-compiler-1.9.21",
+    "kotlin-compiler-2.0.20",
     "kotlin-compiler-dev")
     .map { ThirdPartyDependency(
       it,
diff --git a/src/main/java/com/android/tools/r8/DexSegments.java b/src/main/java/com/android/tools/r8/DexSegments.java
index fafce61..81b904c 100644
--- a/src/main/java/com/android/tools/r8/DexSegments.java
+++ b/src/main/java/com/android/tools/r8/DexSegments.java
@@ -140,38 +140,30 @@
     }
   }
 
-  public static Map<Integer, SegmentInfo> run(Command command)
+  public static Map<Integer, SegmentInfo> runForTesting(Command command)
+      throws IOException, ResourceException {
+    return run(command);
+  }
+
+  public static Int2ReferenceMap<SegmentInfo> run(Command command)
       throws IOException, ResourceException {
     if (command.isPrintHelp()) {
       System.out.println(Command.USAGE_MESSAGE);
       return null;
     }
-    AndroidApp app = command.getInputApp();
+    return run(command.getInputApp());
+  }
 
+  public static Map<Integer, SegmentInfo> runForTesting(AndroidApp app)
+      throws IOException, ResourceException {
+    return run(app);
+  }
+
+  public static Int2ReferenceMap<SegmentInfo> run(AndroidApp app)
+      throws IOException, ResourceException {
     Int2ReferenceMap<SegmentInfo> result = new Int2ReferenceLinkedOpenHashMap<>();
-    // Fill the results with all benchmark items otherwise golem may report missing benchmarks.
-    int[] benchmarks =
-        new int[] {
-          Constants.TYPE_ENCODED_ARRAY_ITEM,
-          Constants.TYPE_HEADER_ITEM,
-          Constants.TYPE_DEBUG_INFO_ITEM,
-          Constants.TYPE_FIELD_ID_ITEM,
-          Constants.TYPE_ANNOTATION_SET_REF_LIST,
-          Constants.TYPE_STRING_ID_ITEM,
-          Constants.TYPE_MAP_LIST,
-          Constants.TYPE_PROTO_ID_ITEM,
-          Constants.TYPE_METHOD_ID_ITEM,
-          Constants.TYPE_TYPE_ID_ITEM,
-          Constants.TYPE_STRING_DATA_ITEM,
-          Constants.TYPE_CLASS_DATA_ITEM,
-          Constants.TYPE_TYPE_LIST,
-          Constants.TYPE_ANNOTATIONS_DIRECTORY_ITEM,
-          Constants.TYPE_ANNOTATION_ITEM,
-          Constants.TYPE_ANNOTATION_SET_ITEM,
-          Constants.TYPE_CLASS_DEF_ITEM
-        };
-    for (int benchmark : benchmarks) {
-      result.computeIfAbsent(benchmark, (key) -> new SegmentInfo());
+    for (int benchmark : DexSection.getConstants()) {
+      result.put(benchmark, new SegmentInfo());
     }
     try (Closer closer = Closer.create()) {
       for (ProgramResource resource : app.computeAllProgramResources()) {
@@ -179,7 +171,8 @@
           for (DexSection dexSection :
               DexParser.parseMapFrom(
                   closer.register(resource.getByteStream()), resource.getOrigin())) {
-            SegmentInfo info = result.computeIfAbsent(dexSection.type, (key) -> new SegmentInfo());
+            assert result.containsKey(dexSection.type) : dexSection.typeName();
+            SegmentInfo info = result.get(dexSection.type);
             info.increment(dexSection.length, dexSection.size());
           }
         }
diff --git a/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeAnnotationTransformer.java b/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeAnnotationTransformer.java
new file mode 100644
index 0000000..47357cd
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeAnnotationTransformer.java
@@ -0,0 +1,295 @@
+// Copyright (c) 2024, 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.desugar.covariantreturntype;
+
+import com.android.tools.r8.errors.CompilationError;
+import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexAnnotation;
+import com.android.tools.r8.graph.DexAnnotationElement;
+import com.android.tools.r8.graph.DexEncodedAnnotation;
+import com.android.tools.r8.graph.DexEncodedMethod;
+import com.android.tools.r8.graph.DexItemFactory;
+import com.android.tools.r8.graph.DexMethod;
+import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.DexString;
+import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.DexValue;
+import com.android.tools.r8.graph.DexValue.DexValueAnnotation;
+import com.android.tools.r8.graph.DexValue.DexValueArray;
+import com.android.tools.r8.graph.DexValue.DexValueType;
+import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.ir.conversion.IRConverter;
+import com.android.tools.r8.ir.conversion.MethodConversionOptions;
+import com.android.tools.r8.ir.conversion.MethodProcessorEventConsumer;
+import com.android.tools.r8.ir.synthetic.ForwardMethodBuilder;
+import com.android.tools.r8.utils.ForEachable;
+import com.android.tools.r8.utils.ListUtils;
+import com.android.tools.r8.utils.ThreadUtils;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.function.Consumer;
+
+// Responsible for processing the annotations dalvik.annotation.codegen.CovariantReturnType and
+// dalvik.annotation.codegen.CovariantReturnType$CovariantReturnTypes.
+//
+// Consider the following class:
+//   public class B extends A {
+//     @CovariantReturnType(returnType = B.class, presentAfter = 25)
+//     @Override
+//     public A m(...) { ... return new B(); }
+//   }
+//
+// The annotation is used to indicate that the compiler should insert a synthetic method that is
+// equivalent to method m, but has return type B instead of A. Thus, for this example, this
+// component is responsible for inserting the following method in class B (in addition to the
+// existing method m):
+//   public B m(...) { A result = "invoke B.m(...)A;"; return (B) result; }
+//
+// Note that a method may be annotated with more than one CovariantReturnType annotation. In this
+// case there will be a CovariantReturnType$CovariantReturnTypes annotation on the method that wraps
+// several CovariantReturnType annotations. In this case, a new method is synthesized for each of
+// the contained CovariantReturnType annotations.
+public final class CovariantReturnTypeAnnotationTransformer {
+
+  private final AppView<?> appView;
+  private final IRConverter converter;
+  private final DexItemFactory factory;
+  private final CovariantReturnTypeReferences references;
+
+  public CovariantReturnTypeAnnotationTransformer(
+      AppView<? extends AppInfoWithClassHierarchy> appView) {
+    this(appView, null);
+  }
+
+  private CovariantReturnTypeAnnotationTransformer(AppView<?> appView, IRConverter converter) {
+    this.appView = appView;
+    this.converter = converter;
+    this.factory = appView.dexItemFactory();
+    this.references = new CovariantReturnTypeReferences(factory);
+  }
+
+  public CovariantReturnTypeReferences getReferences() {
+    return references;
+  }
+
+  public static void runIfNecessary(
+      AppView<?> appView,
+      IRConverter converter,
+      CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer,
+      ExecutorService executorService)
+      throws ExecutionException {
+    if (shouldRun(appView)) {
+      new CovariantReturnTypeAnnotationTransformer(appView, converter)
+          .run(eventConsumer, executorService);
+    }
+  }
+
+  public static boolean shouldRun(AppView<?> appView) {
+    if (!appView.options().processCovariantReturnTypeAnnotations) {
+      return false;
+    }
+    assert !appView.options().isDesugaredLibraryCompilation();
+    DexItemFactory factory = appView.dexItemFactory();
+    DexString covariantReturnTypeDescriptor =
+        factory.createString(CovariantReturnTypeReferences.COVARIANT_RETURN_TYPE_DESCRIPTOR);
+    return factory.lookupType(covariantReturnTypeDescriptor) != null;
+  }
+
+  private void run(
+      CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer,
+      ExecutorService executorService)
+      throws ExecutionException {
+    List<ProgramMethod> covariantReturnTypeMethods = new ArrayList<>();
+    for (DexProgramClass clazz : appView.appInfo().classes()) {
+      List<ProgramMethod> newCovariantReturnTypeMethods =
+          processClass(clazz, clazz::forEachProgramVirtualMethod, eventConsumer);
+      covariantReturnTypeMethods.addAll(newCovariantReturnTypeMethods);
+    }
+    // Convert methods to DEX.
+    converter.optimizeSynthesizedMethods(
+        covariantReturnTypeMethods,
+        MethodProcessorEventConsumer.empty(),
+        MethodConversionOptions.forD8(appView),
+        executorService);
+  }
+
+  public void processMethods(
+      Map<DexProgramClass, List<ProgramMethod>> methodsToProcess,
+      CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer,
+      ExecutorService executorService)
+      throws ExecutionException {
+    ThreadUtils.processMap(
+        methodsToProcess,
+        (clazz, methods) -> {
+          List<ProgramMethod> sortedMethods =
+              ListUtils.destructiveSort(methods, Comparator.comparing(ProgramMethod::getReference));
+          processClass(clazz, sortedMethods::forEach, eventConsumer);
+        },
+        appView.options().getThreadingModule(),
+        executorService);
+  }
+
+  // Processes all the dalvik.annotation.codegen.CovariantReturnType and dalvik.annotation.codegen.
+  // CovariantReturnTypes annotations in the given DexClass. Adds the newly constructed, synthetic
+  // methods to the list covariantReturnTypeMethods.
+  private List<ProgramMethod> processClass(
+      DexProgramClass clazz,
+      ForEachable<ProgramMethod> methodsToProcess,
+      CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer) {
+    List<ProgramMethod> covariantReturnTypeMethods = new ArrayList<>();
+    methodsToProcess.forEach(
+        method ->
+            processMethod(
+                method,
+                covariantReturnTypeMethod -> {
+                  covariantReturnTypeMethods.add(covariantReturnTypeMethod);
+                  eventConsumer.acceptCovariantReturnTypeBridgeMethod(
+                      covariantReturnTypeMethod, method);
+                }));
+    clazz.getMethodCollection().addVirtualClassMethods(covariantReturnTypeMethods);
+    return covariantReturnTypeMethods;
+  }
+
+  // Processes all the dalvik.annotation.codegen.CovariantReturnType and dalvik.annotation.Co-
+  // variantReturnTypes annotations on the given method. Adds the newly constructed, synthetic
+  // methods to the list covariantReturnTypeMethods.
+  private void processMethod(ProgramMethod method, Consumer<ProgramMethod> consumer) {
+    for (DexType covariantReturnType : clearCovariantReturnTypeAnnotations(method)) {
+      consumer.accept(buildCovariantReturnTypeMethod(method, covariantReturnType));
+    }
+  }
+
+  // Builds a synthetic method that invokes the given method, casts the result to
+  // covariantReturnType, and then returns the result. The newly created method will have return
+  // type covariantReturnType.
+  //
+  // Note: any "synchronized" or "strictfp" modifier could be dropped safely.
+  private ProgramMethod buildCovariantReturnTypeMethod(
+      ProgramMethod method, DexType covariantReturnType) {
+    DexMethod covariantReturnTypeMethodReference =
+        method.getReference().withReturnType(covariantReturnType, factory);
+    failIfPresent(method.getHolder(), covariantReturnTypeMethodReference);
+    DexEncodedMethod definition =
+        DexEncodedMethod.syntheticBuilder()
+            .setMethod(covariantReturnTypeMethodReference)
+            .setAccessFlags(
+                method.getAccessFlags().copy().setBridge().setSynthetic().unsetAbstract())
+            .setGenericSignature(method.getDefinition().getGenericSignature())
+            .setAnnotations(method.getAnnotations())
+            .setParameterAnnotations(method.getParameterAnnotations())
+            .setCode(
+                ForwardMethodBuilder.builder(factory)
+                    .setNonStaticSource(covariantReturnTypeMethodReference)
+                    .setVirtualTarget(method.getReference(), method.getHolder().isInterface())
+                    .setCastResult()
+                    .buildCf())
+            .setApiLevelForDefinition(method.getDefinition().getApiLevelForDefinition())
+            .setApiLevelForCode(method.getDefinition().getApiLevelForCode())
+            .build();
+    return new ProgramMethod(method.getHolder(), definition);
+  }
+
+  private void failIfPresent(DexProgramClass clazz, DexMethod covariantReturnTypeMethodReference) {
+    // It is a compilation error if the class already has a method with a signature similar to one
+    // of the methods in covariantReturnTypeMethods.
+    if (clazz.lookupMethod(covariantReturnTypeMethodReference) != null) {
+      throw appView
+          .reporter()
+          .fatalError(
+              String.format(
+                  "Cannot process CovariantReturnType annotation: Class %s already "
+                      + "has a method \"%s\"",
+                  clazz.getTypeName(), covariantReturnTypeMethodReference.toSourceString()));
+    }
+  }
+
+  // Returns the set of covariant return types for method.
+  //
+  // If the method is:
+  //   @dalvik.annotation.codegen.CovariantReturnType(returnType=SubOfFoo, presentAfter=25)
+  //   @dalvik.annotation.codegen.CovariantReturnType(returnType=SubOfSubOfFoo, presentAfter=28)
+  //   @Override
+  //   public Foo foo() { ... return new SubOfSubOfFoo(); }
+  // then this method returns the set { SubOfFoo, SubOfSubOfFoo }.
+  private Set<DexType> clearCovariantReturnTypeAnnotations(ProgramMethod method) {
+    Set<DexType> covariantReturnTypes = new LinkedHashSet<>();
+    for (DexAnnotation annotation : method.getAnnotations().getAnnotations()) {
+      if (references.isOneOfCovariantReturnTypeAnnotations(annotation.getAnnotationType())) {
+        getCovariantReturnTypesFromAnnotation(
+            method, annotation.getAnnotation(), covariantReturnTypes);
+      }
+    }
+    if (!covariantReturnTypes.isEmpty()) {
+      method
+          .getDefinition()
+          .setAnnotations(
+              method
+                  .getAnnotations()
+                  .removeIf(
+                      annotation ->
+                          references.isOneOfCovariantReturnTypeAnnotations(
+                              annotation.getAnnotationType())));
+    }
+    return covariantReturnTypes;
+  }
+
+  private void getCovariantReturnTypesFromAnnotation(
+      ProgramMethod method, DexEncodedAnnotation annotation, Set<DexType> covariantReturnTypes) {
+    boolean hasPresentAfterElement = false;
+    for (DexAnnotationElement element : annotation.elements) {
+      DexString name = element.getName();
+      if (references.isCovariantReturnTypeAnnotation(annotation.getType())) {
+        if (name.isIdenticalTo(references.returnTypeName)) {
+          DexValueType dexValueType = element.getValue().asDexValueType();
+          if (dexValueType == null) {
+            throw new CompilationError(
+                String.format(
+                    "Expected element \"returnType\" of CovariantReturnType annotation to "
+                        + "reference a type (method: \"%s\", was: %s)",
+                    method.toSourceString(), element.value.getClass().getCanonicalName()));
+          }
+          covariantReturnTypes.add(dexValueType.getValue());
+        } else if (name.isIdenticalTo(references.presentAfterName)) {
+          hasPresentAfterElement = true;
+        }
+      } else {
+        assert references.isCovariantReturnTypesAnnotation(annotation.getType());
+        if (name.isIdenticalTo(references.valueName)) {
+          DexValueArray array = element.getValue().asDexValueArray();
+          if (array == null) {
+            throw new CompilationError(
+                String.format(
+                    "Expected element \"value\" of CovariantReturnTypes annotation to "
+                        + "be an array (method: \"%s\", was: %s)",
+                    method.toSourceString(), element.getValue().getClass().getCanonicalName()));
+          }
+
+          // Handle the inner dalvik.annotation.codegen.CovariantReturnType annotations recursively.
+          for (DexValue value : array.getValues()) {
+            assert value.isDexValueAnnotation();
+            DexValueAnnotation innerAnnotation = value.asDexValueAnnotation();
+            getCovariantReturnTypesFromAnnotation(
+                method, innerAnnotation.getValue(), covariantReturnTypes);
+          }
+        }
+      }
+    }
+
+    if (references.isCovariantReturnTypeAnnotation(annotation.getType())
+        && !hasPresentAfterElement) {
+      throw new CompilationError(
+          String.format(
+              "CovariantReturnType annotation for method \"%s\" is missing mandatory element "
+                  + "\"presentAfter\" (class %s)",
+              method.toSourceString(), method.getHolder().getType()));
+    }
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/CovariantReturnTypeAnnotationTransformerEventConsumer.java b/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeAnnotationTransformerEventConsumer.java
similarity index 93%
rename from src/main/java/com/android/tools/r8/ir/desugar/CovariantReturnTypeAnnotationTransformerEventConsumer.java
rename to src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeAnnotationTransformerEventConsumer.java
index f5c97fe..a0ff741 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/CovariantReturnTypeAnnotationTransformerEventConsumer.java
+++ b/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeAnnotationTransformerEventConsumer.java
@@ -1,8 +1,7 @@
-// Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
+// Copyright (c) 2024, 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.ir.desugar;
+package com.android.tools.r8.desugar.covariantreturntype;
 
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.profile.rewriting.ProfileCollectionAdditions;
diff --git a/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeEnqueuerExtension.java b/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeEnqueuerExtension.java
new file mode 100644
index 0000000..4e8573f
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeEnqueuerExtension.java
@@ -0,0 +1,116 @@
+// Copyright (c) 2024, 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.desugar.covariantreturntype;
+
+import static com.android.tools.r8.utils.MapUtils.ignoreKey;
+
+import com.android.tools.r8.Diagnostic;
+import com.android.tools.r8.errors.NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic;
+import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexAnnotationSet;
+import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.ProgramDefinition;
+import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.graph.analysis.EnqueuerAnalysis;
+import com.android.tools.r8.shaking.Enqueuer;
+import com.android.tools.r8.shaking.EnqueuerWorklist;
+import com.android.tools.r8.shaking.KeepInfo;
+import com.android.tools.r8.shaking.KeepMethodInfo;
+import com.android.tools.r8.shaking.MinimumKeepInfoCollection;
+import com.android.tools.r8.utils.InternalOptions;
+import com.android.tools.r8.utils.Timing;
+import com.android.tools.r8.utils.collections.ProgramMethodMap;
+import java.util.ArrayList;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+
+public class CovariantReturnTypeEnqueuerExtension extends EnqueuerAnalysis {
+
+  private final AppView<? extends AppInfoWithClassHierarchy> appView;
+  private final CovariantReturnTypeAnnotationTransformer transformer;
+
+  private final Map<DexProgramClass, List<ProgramMethod>> pendingCovariantReturnTypeDesugaring =
+      new IdentityHashMap<>();
+
+  public CovariantReturnTypeEnqueuerExtension(
+      AppView<? extends AppInfoWithClassHierarchy> appView) {
+    this.appView = appView;
+    this.transformer = new CovariantReturnTypeAnnotationTransformer(appView);
+  }
+
+  public static void register(
+      AppView<? extends AppInfoWithClassHierarchy> appView, Enqueuer enqueuer) {
+    if (enqueuer.getMode().isInitialTreeShaking()
+        && CovariantReturnTypeAnnotationTransformer.shouldRun(appView)) {
+      enqueuer.registerAnalysis(new CovariantReturnTypeEnqueuerExtension(appView));
+    }
+  }
+
+  @Override
+  public void processNewlyLiveMethod(
+      ProgramMethod method,
+      ProgramDefinition context,
+      Enqueuer enqueuer,
+      EnqueuerWorklist worklist) {
+    if (hasCovariantReturnTypeAnnotation(method)) {
+      pendingCovariantReturnTypeDesugaring
+          .computeIfAbsent(method.getHolder(), ignoreKey(ArrayList::new))
+          .add(method);
+    }
+  }
+
+  private boolean hasCovariantReturnTypeAnnotation(ProgramMethod method) {
+    CovariantReturnTypeReferences references = transformer.getReferences();
+    DexAnnotationSet annotations = method.getAnnotations();
+    return annotations.hasAnnotation(
+        annotation ->
+            references.isOneOfCovariantReturnTypeAnnotations(annotation.getAnnotationType()));
+  }
+
+  @Override
+  public void notifyFixpoint(
+      Enqueuer enqueuer, EnqueuerWorklist worklist, ExecutorService executorService, Timing timing)
+      throws ExecutionException {
+    if (pendingCovariantReturnTypeDesugaring.isEmpty()) {
+      return;
+    }
+    ProgramMethodMap<Diagnostic> errors = ProgramMethodMap.createConcurrent();
+    transformer.processMethods(
+        pendingCovariantReturnTypeDesugaring,
+        (bridge, target) -> {
+          KeepMethodInfo.Joiner bridgeKeepInfo =
+              getKeepInfoForCovariantReturnTypeBridge(target, errors);
+          enqueuer.getKeepInfo().registerCompilerSynthesizedMethod(bridge);
+          enqueuer.applyMinimumKeepInfoWhenLiveOrTargeted(bridge, bridgeKeepInfo);
+          enqueuer.getProfileCollectionAdditions().addMethodIfContextIsInProfile(bridge, target);
+        },
+        executorService);
+    errors.forEachValue(appView.reporter()::error);
+    pendingCovariantReturnTypeDesugaring.clear();
+  }
+
+  private KeepMethodInfo.Joiner getKeepInfoForCovariantReturnTypeBridge(
+      ProgramMethod target, ProgramMethodMap<Diagnostic> errors) {
+    KeepInfo.Joiner<?, ?, ?> targetKeepInfo =
+        appView
+            .rootSet()
+            .getDependentMinimumKeepInfo()
+            .getUnconditionalMinimumKeepInfoOrDefault(MinimumKeepInfoCollection.empty())
+            .getOrDefault(target.getReference(), null);
+    if (targetKeepInfo == null) {
+      targetKeepInfo = KeepMethodInfo.newEmptyJoiner();
+    }
+    InternalOptions options = appView.options();
+    if ((options.isMinifying() && targetKeepInfo.isMinificationAllowed())
+        || (options.isOptimizing() && targetKeepInfo.isOptimizationAllowed())
+        || (options.isShrinking() && targetKeepInfo.isShrinkingAllowed())) {
+      errors.computeIfAbsent(target, NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic::new);
+    }
+    return targetKeepInfo.asMethodJoiner();
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeReferences.java b/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeReferences.java
new file mode 100644
index 0000000..fa26148
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/desugar/covariantreturntype/CovariantReturnTypeReferences.java
@@ -0,0 +1,42 @@
+// Copyright (c) 2024, 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.desugar.covariantreturntype;
+
+import com.android.tools.r8.graph.DexItemFactory;
+import com.android.tools.r8.graph.DexString;
+import com.android.tools.r8.graph.DexType;
+
+public class CovariantReturnTypeReferences {
+
+  static String COVARIANT_RETURN_TYPE_DESCRIPTOR =
+      "Ldalvik/annotation/codegen/CovariantReturnType;";
+
+  final DexType annotationCovariantReturnType;
+  final DexType annotationCovariantReturnTypes;
+
+  final DexString presentAfterName;
+  final DexString returnTypeName;
+  final DexString valueName;
+
+  CovariantReturnTypeReferences(DexItemFactory factory) {
+    this.annotationCovariantReturnType = factory.createType(COVARIANT_RETURN_TYPE_DESCRIPTOR);
+    this.annotationCovariantReturnTypes =
+        factory.createType("Ldalvik/annotation/codegen/CovariantReturnType$CovariantReturnTypes;");
+    this.presentAfterName = factory.createString("presentAfter");
+    this.returnTypeName = factory.createString("returnType");
+    this.valueName = factory.createString("value");
+  }
+
+  boolean isCovariantReturnTypeAnnotation(DexType type) {
+    return type.isIdenticalTo(annotationCovariantReturnType);
+  }
+
+  boolean isCovariantReturnTypesAnnotation(DexType type) {
+    return type.isIdenticalTo(annotationCovariantReturnTypes);
+  }
+
+  boolean isOneOfCovariantReturnTypeAnnotations(DexType type) {
+    return isCovariantReturnTypeAnnotation(type) || isCovariantReturnTypesAnnotation(type);
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/dex/DexSection.java b/src/main/java/com/android/tools/r8/dex/DexSection.java
index 39a01b2..79b20bc 100644
--- a/src/main/java/com/android/tools/r8/dex/DexSection.java
+++ b/src/main/java/com/android/tools/r8/dex/DexSection.java
@@ -20,6 +20,29 @@
     this.end = -1;
   }
 
+  public static int[] getConstants() {
+    return new int[] {
+      Constants.TYPE_ENCODED_ARRAY_ITEM,
+      Constants.TYPE_HEADER_ITEM,
+      Constants.TYPE_DEBUG_INFO_ITEM,
+      Constants.TYPE_FIELD_ID_ITEM,
+      Constants.TYPE_ANNOTATION_SET_REF_LIST,
+      Constants.TYPE_STRING_ID_ITEM,
+      Constants.TYPE_MAP_LIST,
+      Constants.TYPE_PROTO_ID_ITEM,
+      Constants.TYPE_METHOD_ID_ITEM,
+      Constants.TYPE_TYPE_ID_ITEM,
+      Constants.TYPE_STRING_DATA_ITEM,
+      Constants.TYPE_CLASS_DATA_ITEM,
+      Constants.TYPE_TYPE_LIST,
+      Constants.TYPE_ANNOTATIONS_DIRECTORY_ITEM,
+      Constants.TYPE_ANNOTATION_ITEM,
+      Constants.TYPE_ANNOTATION_SET_ITEM,
+      Constants.TYPE_CLASS_DEF_ITEM,
+      Constants.TYPE_CODE_ITEM
+    };
+  }
+
   void setEnd(int end) {
     this.end = end;
   }
diff --git a/src/main/java/com/android/tools/r8/errors/NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic.java b/src/main/java/com/android/tools/r8/errors/NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic.java
new file mode 100644
index 0000000..a737aa4
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/errors/NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic.java
@@ -0,0 +1,43 @@
+// Copyright (c) 2024, 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.errors;
+
+import com.android.tools.r8.Diagnostic;
+import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.keepanno.annotations.KeepForApi;
+import com.android.tools.r8.origin.Origin;
+import com.android.tools.r8.position.MethodPosition;
+import com.android.tools.r8.position.Position;
+import com.android.tools.r8.references.MethodReference;
+import com.android.tools.r8.utils.MethodReferenceUtils;
+
+@KeepForApi
+public class NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic implements Diagnostic {
+
+  private final Origin origin;
+  private final MethodReference method;
+  private final MethodPosition position;
+
+  public NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic(ProgramMethod method) {
+    this.origin = method.getOrigin();
+    this.method = method.getMethodReference();
+    this.position = MethodPosition.create(method);
+  }
+
+  @Override
+  public Origin getOrigin() {
+    return origin;
+  }
+
+  @Override
+  public Position getPosition() {
+    return position;
+  }
+
+  @Override
+  public String getDiagnosticMessage() {
+    return "Methods with @CovariantReturnType annotations should be kept, but was not: "
+        + MethodReferenceUtils.toSourceString(method);
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/graph/AccessFlags.java b/src/main/java/com/android/tools/r8/graph/AccessFlags.java
index 4151107..58011d6 100644
--- a/src/main/java/com/android/tools/r8/graph/AccessFlags.java
+++ b/src/main/java/com/android/tools/r8/graph/AccessFlags.java
@@ -226,8 +226,9 @@
     return isSet(Constants.ACC_SYNTHETIC);
   }
 
-  public void setSynthetic() {
+  public T setSynthetic() {
     set(Constants.ACC_SYNTHETIC);
+    return self();
   }
 
   public T unsetSynthetic() {
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 b742d92..9f3483e 100644
--- a/src/main/java/com/android/tools/r8/graph/AppView.java
+++ b/src/main/java/com/android/tools/r8/graph/AppView.java
@@ -534,7 +534,6 @@
   }
 
   public ComposeReferences getComposeReferences() {
-    assert options().getJetpackComposeOptions().isAnyOptimizationsEnabled();
     if (composeReferences == null) {
       composeReferences = new ComposeReferences(dexItemFactory());
     }
diff --git a/src/main/java/com/android/tools/r8/graph/DexAnnotation.java b/src/main/java/com/android/tools/r8/graph/DexAnnotation.java
index ec6b70a..cee4f28 100644
--- a/src/main/java/com/android/tools/r8/graph/DexAnnotation.java
+++ b/src/main/java/com/android/tools/r8/graph/DexAnnotation.java
@@ -16,7 +16,6 @@
 import com.android.tools.r8.graph.DexValue.DexValueString;
 import com.android.tools.r8.graph.DexValue.DexValueType;
 import com.android.tools.r8.graph.GenericSignature.FieldTypeSignature;
-import com.android.tools.r8.ir.desugar.CovariantReturnTypeAnnotationTransformer;
 import com.android.tools.r8.origin.Origin;
 import com.android.tools.r8.synthesis.SyntheticItems;
 import com.android.tools.r8.synthesis.SyntheticNaming.SyntheticKind;
@@ -79,6 +78,10 @@
     this.annotation = annotation;
   }
 
+  public DexEncodedAnnotation getAnnotation() {
+    return annotation;
+  }
+
   public boolean isTypeAnnotation() {
     return false;
   }
@@ -136,22 +139,21 @@
     mixedItems.add(this);
   }
 
-  @SuppressWarnings("ReferenceEquality")
-  public static boolean retainCompileTimeAnnotation(DexType annotation, InternalOptions options) {
+  public static boolean retainCompileTimeAnnotation(
+      DexType annotationType, InternalOptions options) {
     if (options.retainCompileTimeAnnotations) {
       return true;
     }
-    if (annotation == options.itemFactory.annotationSynthesizedClass
-        || annotation
-            .getDescriptor()
-            .startsWith(options.itemFactory.dalvikAnnotationOptimizationPrefix)) {
+    DexItemFactory factory = options.itemFactory;
+    if (annotationType.isIdenticalTo(factory.annotationSynthesizedClass)) {
       return true;
     }
-    if (options.processCovariantReturnTypeAnnotations) {
-      // @CovariantReturnType annotations are processed by CovariantReturnTypeAnnotationTransformer,
-      // they thus need to be read here and will then be removed as part of the processing.
-      return CovariantReturnTypeAnnotationTransformer.isCovariantReturnTypeAnnotation(
-          annotation, options.itemFactory);
+    DexString descriptor = annotationType.getDescriptor();
+    if (descriptor.startsWith(factory.dalvikAnnotationPrefix)) {
+      if (descriptor.startsWith(factory.dalvikAnnotationCodegenCovariantReturnTypePrefix)) {
+        return options.processCovariantReturnTypeAnnotations;
+      }
+      return descriptor.startsWith(factory.dalvikAnnotationOptimizationPrefix);
     }
     return false;
   }
diff --git a/src/main/java/com/android/tools/r8/graph/DexAnnotationSet.java b/src/main/java/com/android/tools/r8/graph/DexAnnotationSet.java
index 6dab6e8..c76fb95 100644
--- a/src/main/java/com/android/tools/r8/graph/DexAnnotationSet.java
+++ b/src/main/java/com/android/tools/r8/graph/DexAnnotationSet.java
@@ -156,9 +156,17 @@
     return getFirstMatching(type) != null;
   }
 
+  public boolean hasAnnotation(Predicate<DexAnnotation> predicate) {
+    return getFirstMatching(predicate) != null;
+  }
+
   public DexAnnotation getFirstMatching(DexType type) {
+    return getFirstMatching(annotation -> annotation.getAnnotationType().isIdenticalTo(type));
+  }
+
+  public DexAnnotation getFirstMatching(Predicate<DexAnnotation> predicate) {
     for (DexAnnotation annotation : annotations) {
-      if (annotation.getAnnotationType().isIdenticalTo(type)) {
+      if (predicate.test(annotation)) {
         return annotation;
       }
     }
diff --git a/src/main/java/com/android/tools/r8/graph/DexEncodedAnnotation.java b/src/main/java/com/android/tools/r8/graph/DexEncodedAnnotation.java
index 143dbe4..884b892 100644
--- a/src/main/java/com/android/tools/r8/graph/DexEncodedAnnotation.java
+++ b/src/main/java/com/android/tools/r8/graph/DexEncodedAnnotation.java
@@ -41,6 +41,10 @@
     return DexEncodedAnnotation::specify;
   }
 
+  public DexType getType() {
+    return type;
+  }
+
   public void collectIndexedItems(AppView<?> appView, IndexedItemCollection indexedItems) {
     type.collectIndexedItems(appView, indexedItems);
     for (DexAnnotationElement element : elements) {
diff --git a/src/main/java/com/android/tools/r8/graph/DexItemFactory.java b/src/main/java/com/android/tools/r8/graph/DexItemFactory.java
index 3c056f8..8f6be1b 100644
--- a/src/main/java/com/android/tools/r8/graph/DexItemFactory.java
+++ b/src/main/java/com/android/tools/r8/graph/DexItemFactory.java
@@ -81,8 +81,6 @@
       "Lcom/android/tools/r8/DesugarMethodHandlesLookup;";
   public static final String methodHandlesLookupDescriptorString =
       "Ljava/lang/invoke/MethodHandles$Lookup;";
-  public static final String dalvikAnnotationOptimizationPrefixString =
-      "Ldalvik/annotation/optimization/";
   public static final String androidUtilSparseArrayDescriptorString = "Landroid/util/SparseArray;";
   public static final String androidContentResTypedArrayDescriptorString =
       "Landroid/content/res/TypedArray;";
@@ -380,8 +378,11 @@
   public final DexString apiLevelString = createString("apiLevel");
 
   // Prefix for runtime affecting yet potential class-retained annotations.
+  public final DexString dalvikAnnotationPrefix = createString("Ldalvik/annotation/");
+  public final DexString dalvikAnnotationCodegenCovariantReturnTypePrefix =
+      createString("Ldalvik/annotation/codegen/CovariantReturnType");
   public final DexString dalvikAnnotationOptimizationPrefix =
-      createString(dalvikAnnotationOptimizationPrefixString);
+      createString("Ldalvik/annotation/optimization/");
 
   // Method names used on VarHandle.
   public final DexString getString = createString("get");
@@ -811,11 +812,6 @@
   public final DexType annotationThrows = createStaticallyKnownType("Ldalvik/annotation/Throws;");
   public final DexType annotationSynthesizedClass =
       createStaticallyKnownType("Lcom/android/tools/r8/annotations/SynthesizedClassV2;");
-  public final DexType annotationCovariantReturnType =
-      createStaticallyKnownType("Ldalvik/annotation/codegen/CovariantReturnType;");
-  public final DexType annotationCovariantReturnTypes =
-      createStaticallyKnownType(
-          "Ldalvik/annotation/codegen/CovariantReturnType$CovariantReturnTypes;");
 
   public final String annotationReachabilitySensitiveDesc =
       "Ldalvik/annotation/optimization/ReachabilitySensitive;";
diff --git a/src/main/java/com/android/tools/r8/graph/DexMethod.java b/src/main/java/com/android/tools/r8/graph/DexMethod.java
index ea4be31..4abdb68 100644
--- a/src/main/java/com/android/tools/r8/graph/DexMethod.java
+++ b/src/main/java/com/android/tools/r8/graph/DexMethod.java
@@ -360,4 +360,8 @@
   public DexMethod withProto(DexProto proto, DexItemFactory dexItemFactory) {
     return dexItemFactory.createMethod(holder, proto, name);
   }
+
+  public DexMethod withReturnType(DexType returnType, DexItemFactory dexItemFactory) {
+    return withProto(dexItemFactory.createProto(returnType, getParameters()), dexItemFactory);
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/graph/DexProgramClass.java b/src/main/java/com/android/tools/r8/graph/DexProgramClass.java
index 5b0aec4..e890a6a 100644
--- a/src/main/java/com/android/tools/r8/graph/DexProgramClass.java
+++ b/src/main/java/com/android/tools/r8/graph/DexProgramClass.java
@@ -755,11 +755,6 @@
     methodCollection.addMethod(method);
   }
 
-  public void replaceVirtualMethod(
-      DexMethod virtualMethod, Function<DexEncodedMethod, DexEncodedMethod> replacement) {
-    methodCollection.replaceVirtualMethod(virtualMethod, replacement);
-  }
-
   public void addExtraInterfaces(List<ClassTypeSignature> extraInterfaces, DexItemFactory factory) {
     if (extraInterfaces.isEmpty()) {
       return;
diff --git a/src/main/java/com/android/tools/r8/graph/MethodAccessFlags.java b/src/main/java/com/android/tools/r8/graph/MethodAccessFlags.java
index 8ede6f7..742b628 100644
--- a/src/main/java/com/android/tools/r8/graph/MethodAccessFlags.java
+++ b/src/main/java/com/android/tools/r8/graph/MethodAccessFlags.java
@@ -152,8 +152,9 @@
     return isSet(Constants.ACC_BRIDGE);
   }
 
-  public void setBridge() {
+  public MethodAccessFlags setBridge() {
     set(Constants.ACC_BRIDGE);
+    return this;
   }
 
   public void unsetBridge() {
@@ -204,8 +205,9 @@
     promote(Constants.ACC_ABSTRACT);
   }
 
-  public void unsetAbstract() {
+  public MethodAccessFlags unsetAbstract() {
     unset(Constants.ACC_ABSTRACT);
+    return this;
   }
 
   public boolean isStrict() {
diff --git a/src/main/java/com/android/tools/r8/graph/MethodArrayBacking.java b/src/main/java/com/android/tools/r8/graph/MethodArrayBacking.java
index ab17919..7daef60 100644
--- a/src/main/java/com/android/tools/r8/graph/MethodArrayBacking.java
+++ b/src/main/java/com/android/tools/r8/graph/MethodArrayBacking.java
@@ -184,12 +184,12 @@
   }
 
   @Override
-  void addVirtualMethods(Collection<DexEncodedMethod> methods) {
+  <T> void addVirtualMethods(Collection<T> methods, Function<? super T, DexEncodedMethod> fn) {
     DexEncodedMethod[] newMethods = new DexEncodedMethod[virtualMethods.length + methods.size()];
     System.arraycopy(virtualMethods, 0, newMethods, 0, virtualMethods.length);
     int i = virtualMethods.length;
-    for (DexEncodedMethod method : methods) {
-      newMethods[i] = method;
+    for (T method : methods) {
+      newMethods[i] = fn.apply(method);
       i++;
     }
     virtualMethods = newMethods;
diff --git a/src/main/java/com/android/tools/r8/graph/MethodCollection.java b/src/main/java/com/android/tools/r8/graph/MethodCollection.java
index 0f10f5a..17b7e8b 100644
--- a/src/main/java/com/android/tools/r8/graph/MethodCollection.java
+++ b/src/main/java/com/android/tools/r8/graph/MethodCollection.java
@@ -5,6 +5,9 @@
 import com.android.tools.r8.utils.InternalOptions;
 import com.android.tools.r8.utils.IterableUtils;
 import com.android.tools.r8.utils.TraversalContinuation;
+import com.google.common.base.Function;
+import com.google.common.base.Functions;
+import com.google.common.collect.Iterables;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -12,7 +15,6 @@
 import java.util.List;
 import java.util.Set;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.function.Predicate;
 
 public class MethodCollection {
@@ -354,9 +356,21 @@
   }
 
   public void addVirtualMethods(Collection<DexEncodedMethod> methods) {
-    assert verifyCorrectnessOfMethodHolders(methods);
+    internalAddVirtualMethods(methods, Functions.identity());
+  }
+
+  public void addVirtualClassMethods(Collection<? extends DexClassAndMethod> methods) {
+    internalAddVirtualMethods(methods, DexClassAndMethod::getDefinition);
+  }
+
+  private <T> void internalAddVirtualMethods(
+      Collection<T> methods, Function<? super T, DexEncodedMethod> fn) {
+    if (methods.isEmpty()) {
+      return;
+    }
+    assert verifyCorrectnessOfMethodHolders(Iterables.transform(methods, fn));
     resetVirtualMethodCaches();
-    backing.addVirtualMethods(methods);
+    backing.addVirtualMethods(methods, fn);
   }
 
   public void clearVirtualMethods() {
diff --git a/src/main/java/com/android/tools/r8/graph/MethodCollectionBacking.java b/src/main/java/com/android/tools/r8/graph/MethodCollectionBacking.java
index 18a72a3..9de91a2 100644
--- a/src/main/java/com/android/tools/r8/graph/MethodCollectionBacking.java
+++ b/src/main/java/com/android/tools/r8/graph/MethodCollectionBacking.java
@@ -6,6 +6,7 @@
 import static com.google.common.base.Predicates.alwaysTrue;
 
 import com.android.tools.r8.utils.TraversalContinuation;
+import com.google.common.base.Functions;
 import java.util.Collection;
 import java.util.Set;
 import java.util.function.Consumer;
@@ -99,7 +100,12 @@
 
   abstract void addDirectMethods(Collection<DexEncodedMethod> methods);
 
-  abstract void addVirtualMethods(Collection<DexEncodedMethod> methods);
+  final void addVirtualMethods(Collection<DexEncodedMethod> methods) {
+    addVirtualMethods(methods, Functions.identity());
+  }
+
+  abstract <T> void addVirtualMethods(
+      Collection<T> methods, Function<? super T, DexEncodedMethod> fn);
 
   // Removal methods.
 
diff --git a/src/main/java/com/android/tools/r8/graph/MethodCollectionConcurrencyChecked.java b/src/main/java/com/android/tools/r8/graph/MethodCollectionConcurrencyChecked.java
index 78cfebe8..fa3a0a9 100644
--- a/src/main/java/com/android/tools/r8/graph/MethodCollectionConcurrencyChecked.java
+++ b/src/main/java/com/android/tools/r8/graph/MethodCollectionConcurrencyChecked.java
@@ -4,11 +4,11 @@
 package com.android.tools.r8.graph;
 
 import com.android.tools.r8.utils.TraversalContinuation;
+import com.google.common.base.Function;
 import java.util.Collection;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.function.Predicate;
 
 public class MethodCollectionConcurrencyChecked extends MethodCollection {
diff --git a/src/main/java/com/android/tools/r8/graph/MethodMapBacking.java b/src/main/java/com/android/tools/r8/graph/MethodMapBacking.java
index fa553d9..f878e46 100644
--- a/src/main/java/com/android/tools/r8/graph/MethodMapBacking.java
+++ b/src/main/java/com/android/tools/r8/graph/MethodMapBacking.java
@@ -198,9 +198,9 @@
   }
 
   @Override
-  void addVirtualMethods(Collection<DexEncodedMethod> methods) {
-    for (DexEncodedMethod method : methods) {
-      addVirtualMethod(method);
+  <T> void addVirtualMethods(Collection<T> methods, Function<? super T, DexEncodedMethod> fn) {
+    for (T method : methods) {
+      addVirtualMethod(fn.apply(method));
     }
   }
 
diff --git a/src/main/java/com/android/tools/r8/graph/analysis/EnqueuerAnalysis.java b/src/main/java/com/android/tools/r8/graph/analysis/EnqueuerAnalysis.java
index d587b61..791b3b4 100644
--- a/src/main/java/com/android/tools/r8/graph/analysis/EnqueuerAnalysis.java
+++ b/src/main/java/com/android/tools/r8/graph/analysis/EnqueuerAnalysis.java
@@ -15,6 +15,8 @@
 import com.android.tools.r8.shaking.Enqueuer;
 import com.android.tools.r8.shaking.EnqueuerWorklist;
 import com.android.tools.r8.utils.Timing;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
 
 public abstract class EnqueuerAnalysis {
 
@@ -57,7 +59,9 @@
    * Called when the Enqueuer reaches a fixpoint. This may happen multiple times, since each
    * analysis may enqueue items into the worklist upon the fixpoint using {@param worklist}.
    */
-  public void notifyFixpoint(Enqueuer enqueuer, EnqueuerWorklist worklist, Timing timing) {}
+  public void notifyFixpoint(
+      Enqueuer enqueuer, EnqueuerWorklist worklist, ExecutorService executorService, Timing timing)
+      throws ExecutionException {}
 
   /**
    * Called when the Enqueuer has reached the final fixpoint. Each analysis may use this callback to
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysis.java b/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysis.java
index 3d61f0e..097cdc8 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysis.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysis.java
@@ -7,6 +7,7 @@
 import com.android.tools.r8.ir.analysis.framework.intraprocedural.IntraproceduralDataflowAnalysis;
 import com.android.tools.r8.ir.analysis.path.state.PathConstraintAnalysisState;
 import com.android.tools.r8.ir.code.IRCode;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldValueFactory;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameterFactory;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
@@ -41,12 +42,13 @@
   public PathConstraintAnalysis(
       AppView<AppInfoWithLiveness> appView,
       IRCode code,
+      FieldValueFactory fieldValueFactory,
       MethodParameterFactory methodParameterFactory) {
     super(
         appView,
         PathConstraintAnalysisState.bottom(),
         code,
         new PathConstraintAnalysisTransferFunction(
-            appView.abstractValueFactory(), code.context(), methodParameterFactory));
+            appView, code, code.context(), fieldValueFactory, methodParameterFactory));
   }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysisTransferFunction.java b/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysisTransferFunction.java
index c61df2a..ed423d7 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysisTransferFunction.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysisTransferFunction.java
@@ -3,31 +3,37 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.ir.analysis.path;
 
+import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.analysis.framework.intraprocedural.AbstractTransferFunction;
 import com.android.tools.r8.ir.analysis.framework.intraprocedural.TransferFunctionResult;
 import com.android.tools.r8.ir.analysis.path.state.ConcretePathConstraintAnalysisState;
 import com.android.tools.r8.ir.analysis.path.state.PathConstraintAnalysisState;
-import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
 import com.android.tools.r8.ir.code.BasicBlock;
+import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.If;
 import com.android.tools.r8.ir.code.Instruction;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldValueFactory;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameterFactory;
-import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeBuilder;
 import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
+import com.android.tools.r8.optimize.argumentpropagation.computation.DefaultComputationTreeBuilder;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
 public class PathConstraintAnalysisTransferFunction
     implements AbstractTransferFunction<BasicBlock, Instruction, PathConstraintAnalysisState> {
 
-  private final ComputationTreeBuilder computationTreeBuilder;
+  private final DefaultComputationTreeBuilder computationTreeBuilder;
 
   PathConstraintAnalysisTransferFunction(
-      AbstractValueFactory abstractValueFactory,
+      AppView<AppInfoWithLiveness> appView,
+      IRCode code,
       ProgramMethod method,
+      FieldValueFactory fieldValueFactory,
       MethodParameterFactory methodParameterFactory) {
     computationTreeBuilder =
-        new ComputationTreeBuilder(abstractValueFactory, method, methodParameterFactory);
+        new DefaultComputationTreeBuilder(
+            appView, code, method, fieldValueFactory, methodParameterFactory);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintSupplier.java b/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintSupplier.java
index cb81605..649b2bb 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintSupplier.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/path/PathConstraintSupplier.java
@@ -5,16 +5,21 @@
 
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.ir.analysis.framework.intraprocedural.DataflowAnalysisResult.SuccessfulDataflowAnalysisResult;
+import com.android.tools.r8.ir.analysis.path.state.ConcretePathConstraintAnalysisState;
 import com.android.tools.r8.ir.analysis.path.state.PathConstraintAnalysisState;
+import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.IRCode;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldValueFactory;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameterFactory;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
 public class PathConstraintSupplier {
 
   private final AppView<AppInfoWithLiveness> appView;
   private final IRCode code;
+  private final FieldValueFactory fieldValueFactory;
   private final MethodParameterFactory methodParameterFactory;
 
   private SuccessfulDataflowAnalysisResult<BasicBlock, PathConstraintAnalysisState>
@@ -23,16 +28,32 @@
   public PathConstraintSupplier(
       AppView<AppInfoWithLiveness> appView,
       IRCode code,
+      FieldValueFactory fieldValueFactory,
       MethodParameterFactory methodParameterFactory) {
     this.appView = appView;
     this.code = code;
+    this.fieldValueFactory = fieldValueFactory;
     this.methodParameterFactory = methodParameterFactory;
   }
 
+  public ComputationTreeNode getDifferentiatingPathConstraint(
+      BasicBlock block, BasicBlock otherBlock) {
+    ConcretePathConstraintAnalysisState state = getPathConstraint(block).asConcreteState();
+    if (state == null) {
+      return AbstractValue.unknown();
+    }
+    ConcretePathConstraintAnalysisState otherState =
+        getPathConstraint(otherBlock).asConcreteState();
+    if (otherState == null) {
+      return AbstractValue.unknown();
+    }
+    return state.getDifferentiatingPathConstraint(otherState);
+  }
+
   public PathConstraintAnalysisState getPathConstraint(BasicBlock block) {
     if (pathConstraintAnalysisResult == null) {
       PathConstraintAnalysis analysis =
-          new PathConstraintAnalysis(appView, code, methodParameterFactory);
+          new PathConstraintAnalysis(appView, code, fieldValueFactory, methodParameterFactory);
       pathConstraintAnalysisResult = analysis.run(code.entryBlock()).asSuccessfulAnalysisResult();
       assert pathConstraintAnalysisResult != null;
     }
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/path/state/ConcretePathConstraintAnalysisState.java b/src/main/java/com/android/tools/r8/ir/analysis/path/state/ConcretePathConstraintAnalysisState.java
index 7bcb650..28bdcf1 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/path/state/ConcretePathConstraintAnalysisState.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/path/state/ConcretePathConstraintAnalysisState.java
@@ -4,6 +4,7 @@
 package com.android.tools.r8.ir.analysis.path.state;
 
 import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
 import java.util.Collections;
 import java.util.HashMap;
@@ -62,26 +63,22 @@
   @Override
   public PathConstraintAnalysisState add(ComputationTreeNode pathConstraint, boolean negate) {
     PathConstraintKind previousKind = pathConstraints.get(pathConstraint);
+    PathConstraintKind newKind;
     if (previousKind != null) {
-      if (previousKind == PathConstraintKind.DISABLED) {
-        // There is a loop.
+      // There is a loop.
+      newKind = PathConstraintKind.DISABLED;
+      if (previousKind == newKind) {
         return this;
       }
-      if (previousKind == PathConstraintKind.get(negate)) {
-        // This branch is dominated by a previous if-condition that has the same branch condition,
-        // e.g., if (x) { if (x) { ...
-        return this;
-      }
-      // This branch is dominated by a previous if-condition that has the negated branch condition,
-      // e.g., if (x) { if (!x) { ...
-      return bottom();
+    } else {
+      newKind = PathConstraintKind.get(negate);
     }
     // No jumps can dominate the entry of their own block, so when adding the condition of a jump
     // this cannot currently be active.
     Map<ComputationTreeNode, PathConstraintKind> newPathConstraints =
         new HashMap<>(pathConstraints.size() + 1);
     newPathConstraints.putAll(pathConstraints);
-    newPathConstraints.put(pathConstraint, PathConstraintKind.get(negate));
+    newPathConstraints.put(pathConstraint, newKind);
     return new ConcretePathConstraintAnalysisState(newPathConstraints);
   }
 
@@ -153,7 +150,7 @@
         return pathConstraint;
       }
     }
-    return null;
+    return AbstractValue.unknown();
   }
 
   public ConcretePathConstraintAnalysisState join(ConcretePathConstraintAnalysisState other) {
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/path/state/PathConstraintKind.java b/src/main/java/com/android/tools/r8/ir/analysis/path/state/PathConstraintKind.java
index 4e97d27..6891163 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/path/state/PathConstraintKind.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/path/state/PathConstraintKind.java
@@ -24,6 +24,9 @@
   }
 
   PathConstraintKind join(PathConstraintKind other) {
+    if (other == null) {
+      return this;
+    }
     return this == other ? this : DISABLED;
   }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedMessageLiteBuilderShrinker.java b/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedMessageLiteBuilderShrinker.java
index 04c491e..f3299a3 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedMessageLiteBuilderShrinker.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedMessageLiteBuilderShrinker.java
@@ -135,7 +135,11 @@
     return new EnqueuerAnalysis() {
       @Override
       @SuppressWarnings("ReferenceEquality")
-      public void notifyFixpoint(Enqueuer enqueuer, EnqueuerWorklist worklist, Timing timing) {
+      public void notifyFixpoint(
+          Enqueuer enqueuer,
+          EnqueuerWorklist worklist,
+          ExecutorService executorService,
+          Timing timing) {
         builders.forEach(
             (builder, dynamicMethod) -> {
               if (seen.add(builder)) {
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java b/src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java
index c9a6021..0ce0317 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java
@@ -54,6 +54,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutorService;
 import java.util.function.Predicate;
 
 // TODO(b/112437944): Handle cycles in the graph + add a test that fails with the current
@@ -179,7 +180,11 @@
   }
 
   @Override
-  public void notifyFixpoint(Enqueuer enqueuer, EnqueuerWorklist worklist, Timing timing) {
+  public void notifyFixpoint(
+      Enqueuer enqueuer,
+      EnqueuerWorklist worklist,
+      ExecutorService executorService,
+      Timing timing) {
     timing.begin("[Proto] Extend fixpoint");
     populateExtensionGraph(enqueuer);
 
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/value/AbstractValue.java b/src/main/java/com/android/tools/r8/ir/analysis/value/AbstractValue.java
index 6350b1e..9fa7cf7 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/value/AbstractValue.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/value/AbstractValue.java
@@ -9,10 +9,15 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.lens.GraphLens;
 import com.android.tools.r8.ir.analysis.value.objectstate.ObjectState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlow;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlowComparator;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlowKind;
 import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import java.util.function.IntFunction;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.function.Function;
 
 public abstract class AbstractValue implements ComputationTreeNode {
 
@@ -26,15 +31,37 @@
 
   @Override
   public AbstractValue evaluate(
-      IntFunction<AbstractValue> argumentAssignment, AbstractValueFactory abstractValueFactory) {
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
     return this;
   }
 
   @Override
-  public MethodParameter getSingleOpenVariable() {
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    // Abstract values do not contain any open variables.
+    return TraversalContinuation.doContinue();
+  }
+
+  @Override
+  public InFlowKind getKind() {
+    throw new Unreachable();
+  }
+
+  @Override
+  public BaseInFlow getSingleOpenVariable() {
     return null;
   }
 
+  @Override
+  public int internalCompareToSameKind(InFlow inFlow, InFlowComparator comparator) {
+    throw new Unreachable();
+  }
+
+  @Override
+  public final boolean isComputationLeaf() {
+    return true;
+  }
+
   public abstract boolean isNonTrivial();
 
   public boolean isSingleBoolean() {
diff --git a/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java b/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
index b55931c..59fde05 100644
--- a/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
+++ b/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
@@ -370,6 +370,10 @@
     return predecessors.get(0);
   }
 
+  public BasicBlock getPredecessor(int i) {
+    return getPredecessors().get(i);
+  }
+
   public List<BasicBlock> getPredecessors() {
     return ListUtils.unmodifiableForTesting(predecessors);
   }
diff --git a/src/main/java/com/android/tools/r8/ir/code/IfType.java b/src/main/java/com/android/tools/r8/ir/code/IfType.java
index 5babdae..8a0d96b 100644
--- a/src/main/java/com/android/tools/r8/ir/code/IfType.java
+++ b/src/main/java/com/android/tools/r8/ir/code/IfType.java
@@ -4,63 +4,88 @@
 package com.android.tools.r8.ir.code;
 
 import com.android.tools.r8.errors.Unreachable;
+import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
-import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
-import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
 public enum IfType {
   EQ {
     @Override
-    public AbstractValue evaluate(
-        SingleNumberValue operand, AbstractValueFactory abstractValueFactory) {
-      return abstractValueFactory.createSingleBooleanValue(operand.getValue() == 0);
+    public boolean evaluate(int operand) {
+      return operand == 0;
+    }
+
+    @Override
+    public String getSymbol() {
+      return "==";
     }
   },
   GE {
     @Override
-    public AbstractValue evaluate(
-        SingleNumberValue operand, AbstractValueFactory abstractValueFactory) {
-      return abstractValueFactory.createSingleBooleanValue(operand.getValue() >= 0);
+    public boolean evaluate(int operand) {
+      return operand >= 0;
+    }
+
+    @Override
+    public String getSymbol() {
+      return ">=";
     }
   },
   GT {
     @Override
-    public AbstractValue evaluate(
-        SingleNumberValue operand, AbstractValueFactory abstractValueFactory) {
-      return abstractValueFactory.createSingleBooleanValue(operand.getValue() > 0);
+    public boolean evaluate(int operand) {
+      return operand > 0;
+    }
+
+    @Override
+    public String getSymbol() {
+      return ">";
     }
   },
   LE {
     @Override
-    public AbstractValue evaluate(
-        SingleNumberValue operand, AbstractValueFactory abstractValueFactory) {
-      return abstractValueFactory.createSingleBooleanValue(operand.getValue() <= 0);
+    public boolean evaluate(int operand) {
+      return operand <= 0;
+    }
+
+    @Override
+    public String getSymbol() {
+      return "<=";
     }
   },
   LT {
     @Override
-    public AbstractValue evaluate(
-        SingleNumberValue operand, AbstractValueFactory abstractValueFactory) {
-      return abstractValueFactory.createSingleBooleanValue(operand.getValue() < 0);
+    public boolean evaluate(int operand) {
+      return operand < 0;
+    }
+
+    @Override
+    public String getSymbol() {
+      return "<";
     }
   },
   NE {
     @Override
-    public AbstractValue evaluate(
-        SingleNumberValue operand, AbstractValueFactory abstractValueFactory) {
-      return abstractValueFactory.createSingleBooleanValue(operand.getValue() != 0);
+    public boolean evaluate(int operand) {
+      return operand != 0;
+    }
+
+    @Override
+    public String getSymbol() {
+      return "!=";
     }
   };
 
-  public AbstractValue evaluate(AbstractValue operand, AbstractValueFactory abstractValueFactory) {
+  public AbstractValue evaluate(AbstractValue operand, AppView<AppInfoWithLiveness> appView) {
     if (operand.isSingleNumberValue()) {
-      return evaluate(operand.asSingleNumberValue(), abstractValueFactory);
+      int operandValue = operand.asSingleNumberValue().getIntValue();
+      boolean result = evaluate(operandValue);
+      return appView.abstractValueFactory().createSingleBooleanValue(result);
     }
     return AbstractValue.unknown();
   }
 
-  public abstract AbstractValue evaluate(
-      SingleNumberValue operand, AbstractValueFactory abstractValueFactory);
+  public abstract boolean evaluate(int operand);
 
   public boolean isEqualsOrNotEquals() {
     return this == EQ || this == NE;
@@ -103,4 +128,6 @@
         throw new Unreachable("Unknown if condition type.");
     }
   }
+
+  public abstract String getSymbol();
 }
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 0ae5aaf..34d9483 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
@@ -4,27 +4,4 @@
 
 package com.android.tools.r8.ir.code;
 
-public interface InstructionOrPhi {
-
-  default boolean isInstruction() {
-    return false;
-  }
-
-  default Instruction asInstruction() {
-    return null;
-  }
-
-  default boolean isPhi() {
-    return false;
-  }
-
-  default boolean isStackMapPhi() {
-    return false;
-  }
-
-  default Phi asPhi() {
-    return null;
-  }
-
-  BasicBlock getBlock();
-}
+public interface InstructionOrPhi extends InstructionOrValue {}
diff --git a/src/main/java/com/android/tools/r8/ir/code/InstructionOrValue.java b/src/main/java/com/android/tools/r8/ir/code/InstructionOrValue.java
new file mode 100644
index 0000000..b45fd92
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/ir/code/InstructionOrValue.java
@@ -0,0 +1,33 @@
+// Copyright (c) 2024, 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.ir.code;
+
+public interface InstructionOrValue {
+
+  default boolean isInstruction() {
+    return false;
+  }
+
+  default Instruction asInstruction() {
+    return null;
+  }
+
+  default boolean isPhi() {
+    return false;
+  }
+
+  default boolean isStackMapPhi() {
+    return false;
+  }
+
+  default Phi asPhi() {
+    return null;
+  }
+
+  default Value asValue() {
+    return null;
+  }
+
+  BasicBlock getBlock();
+}
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 fbccb7c..d08946f 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
@@ -51,7 +51,7 @@
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 
-public class Value implements Comparable<Value> {
+public class Value implements Comparable<Value>, InstructionOrValue {
 
   public void constrainType(
       ValueTypeConstraint constraint, ProgramMethod method, Reporter reporter) {
@@ -371,8 +371,17 @@
   }
 
   public Instruction singleUniqueUser() {
-    assert ImmutableSet.copyOf(users).size() == 1;
-    return users.getFirst();
+    assert hasSingleUniqueUser();
+    return firstUser();
+  }
+
+  public boolean hasSingleUniquePhiUser() {
+    return uniquePhiUsers().size() == 1;
+  }
+
+  public Phi singleUniquePhiUser() {
+    assert hasSingleUniquePhiUser();
+    return firstPhiUser();
   }
 
   public Set<Instruction> aliasedUsers() {
@@ -397,6 +406,11 @@
     }
   }
 
+  public Instruction firstUser() {
+    assert !users.isEmpty();
+    return users.getFirst();
+  }
+
   public Phi firstPhiUser() {
     assert !phiUsers.isEmpty();
     return phiUsers.getFirst();
@@ -775,6 +789,11 @@
   }
 
   @Override
+  public Value asValue() {
+    return this;
+  }
+
+  @Override
   public int compareTo(Value value) {
     return Integer.compare(this.number, value.number);
   }
@@ -929,14 +948,6 @@
     return predicate.test(definition);
   }
 
-  public boolean isPhi() {
-    return false;
-  }
-
-  public Phi asPhi() {
-    return null;
-  }
-
   /**
    * Returns whether this value is known to never be <code>null</code>.
    */
@@ -1106,6 +1117,7 @@
     return definition.hasBlock();
   }
 
+  @Override
   public BasicBlock getBlock() {
     return definition.getBlock();
   }
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
index 13b059d..fd87cbc 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
@@ -41,7 +41,6 @@
 import com.android.tools.r8.ir.conversion.passes.ThrowCatchOptimizer;
 import com.android.tools.r8.ir.conversion.passes.TrivialPhiSimplifier;
 import com.android.tools.r8.ir.desugar.CfInstructionDesugaringCollection;
-import com.android.tools.r8.ir.desugar.CovariantReturnTypeAnnotationTransformer;
 import com.android.tools.r8.ir.optimize.AssertionErrorTwoArgsConstructorRewriter;
 import com.android.tools.r8.ir.optimize.AssertionsRewriter;
 import com.android.tools.r8.ir.optimize.AssumeInserter;
@@ -120,7 +119,6 @@
   protected final Inliner inliner;
   protected final IdentifierNameStringMarker identifierNameStringMarker;
   private final Devirtualizer devirtualizer;
-  protected final CovariantReturnTypeAnnotationTransformer covariantReturnTypeAnnotationTransformer;
   private final TypeChecker typeChecker;
   protected EnumUnboxer enumUnboxer;
   protected final NumberUnboxer numberUnboxer;
@@ -190,7 +188,6 @@
       assert options.desugarState.isOn();
       this.instructionDesugaring =
           CfInstructionDesugaringCollection.create(appView, appView.apiLevelCompute());
-      this.covariantReturnTypeAnnotationTransformer = null;
       this.dynamicTypeOptimization = null;
       this.classInliner = null;
       this.fieldAccessAnalysis = null;
@@ -213,10 +210,6 @@
         appView.enableWholeProgramOptimizations()
             ? CfInstructionDesugaringCollection.empty()
             : CfInstructionDesugaringCollection.create(appView, appView.apiLevelCompute());
-    this.covariantReturnTypeAnnotationTransformer =
-        options.processCovariantReturnTypeAnnotations
-            ? new CovariantReturnTypeAnnotationTransformer(appView, this)
-            : null;
     removeVerificationErrorForUnknownReturnedValues =
         (appView.options().apiModelingOptions().enableLibraryApiModeling
                 && appView.options().canHaveVerifyErrorForUnknownUnusedReturnValue())
@@ -871,10 +864,6 @@
     appView.withArgumentPropagator(
         argumentPropagator -> argumentPropagator.scan(method, code, methodProcessor, timing));
 
-    if (methodProcessor.isComposeMethodProcessor()) {
-      methodProcessor.asComposeMethodProcessor().scan(method, code, timing);
-    }
-
     if (methodProcessor.isPrimaryMethodProcessor()) {
       enumUnboxer.analyzeEnums(code, methodProcessor);
       numberUnboxer.analyze(code);
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java b/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
index 72e0051..e8b022d 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
@@ -35,7 +35,6 @@
 import static com.android.tools.r8.utils.ObjectUtils.getBooleanOrElse;
 
 import com.android.tools.r8.errors.CompilationError;
-import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.graph.AccessControl;
 import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
 import com.android.tools.r8.graph.AppView;
@@ -75,6 +74,7 @@
 import com.android.tools.r8.ir.code.ConstClass;
 import com.android.tools.r8.ir.code.ConstMethodHandle;
 import com.android.tools.r8.ir.code.ConstMethodType;
+import com.android.tools.r8.ir.code.ConstNumber;
 import com.android.tools.r8.ir.code.DexItemBasedConstString;
 import com.android.tools.r8.ir.code.FieldGet;
 import com.android.tools.r8.ir.code.FieldInstruction;
@@ -874,7 +874,7 @@
       }
     }
     if (mayHaveUnreachableBlocks) {
-      code.removeUnreachableBlocks(affectedValues, prunedValue -> affectedPhis.remove(prunedValue));
+      code.removeUnreachableBlocks(affectedValues, affectedPhis::remove);
     }
     affectedValues.narrowingWithAssumeRemoval(appView, code);
     if (!affectedPhis.isEmpty()) {
@@ -884,7 +884,7 @@
     nullCheckInserter.processWorklist();
     code.removeAllDeadAndTrivialPhis();
     code.removeRedundantBlocks();
-    removeUnusedArguments(method, code, unusedArguments);
+    removeUnusedArguments(code, unusedArguments);
 
     // Finalize cast and null check insertion.
     interfaceTypeToClassTypeRewriterHelper.processWorklist();
@@ -1075,16 +1075,31 @@
     return replacement;
   }
 
-  private void removeUnusedArguments(
-      ProgramMethod method, IRCode code, Set<UnusedArgument> unusedArguments) {
+  private void removeUnusedArguments(IRCode code, Set<UnusedArgument> unusedArguments) {
+    AffectedValues affectedValues = new AffectedValues();
     for (UnusedArgument unusedArgument : unusedArguments) {
+      InstructionListIterator instructionIterator =
+          unusedArgument.getBlock().listIterator(code, unusedArgument);
       if (unusedArgument.outValue().hasAnyUsers()) {
-        throw new Unreachable("Unused argument with users in " + method.toSourceString());
+        // This is an unused argument with a default value. The unused argument is an operand of the
+        // phi. This use is eliminated after constant propagation + branch pruning. We eliminate the
+        // UnusedArgument instruction early by replacing it with a const 0.
+        assert unusedArgument.outValue().hasPhiUsers();
+        assert !unusedArgument.outValue().hasUsers();
+        assert !unusedArgument.outValue().hasDebugUsers();
+        TypeElement type = unusedArgument.outValue().getType();
+        instructionIterator.replaceCurrentInstruction(
+            ConstNumber.builder()
+                .setFreshOutValue(code, type.isReferenceType() ? TypeElement.getNull() : type)
+                .setPosition(Position.none())
+                .setValue(0)
+                .build(),
+            affectedValues);
+      } else {
+        instructionIterator.removeOrReplaceByDebugLocalRead();
       }
-      InstructionListIterator instructionIterator = unusedArgument.getBlock().listIterator(code);
-      instructionIterator.nextUntil(instruction -> instruction == unusedArgument);
-      instructionIterator.removeOrReplaceByDebugLocalRead();
     }
+    affectedValues.narrowingWithAssumeRemoval(appView, code);
   }
 
   private Deque<GraphLensInterval> getUnappliedLenses(ProgramMethod method) {
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/MethodProcessor.java b/src/main/java/com/android/tools/r8/ir/conversion/MethodProcessor.java
index 9eb24e1..9c8b286 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/MethodProcessor.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/MethodProcessor.java
@@ -5,18 +5,9 @@
 
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.conversion.callgraph.CallSiteInformation;
-import com.android.tools.r8.optimize.compose.ComposeMethodProcessor;
 
 public abstract class MethodProcessor {
 
-  public boolean isComposeMethodProcessor() {
-    return false;
-  }
-
-  public ComposeMethodProcessor asComposeMethodProcessor() {
-    return null;
-  }
-
   public boolean isD8MethodProcessor() {
     return false;
   }
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/PrimaryD8L8IRConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/PrimaryD8L8IRConverter.java
index eb16aef..93b681a 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/PrimaryD8L8IRConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/PrimaryD8L8IRConverter.java
@@ -8,6 +8,8 @@
 import static com.android.tools.r8.ir.desugar.lambda.D8LambdaDesugaring.rewriteEnclosingLambdaMethodAttributes;
 
 import com.android.tools.r8.contexts.CompilationContext.MethodProcessingContext;
+import com.android.tools.r8.desugar.covariantreturntype.CovariantReturnTypeAnnotationTransformer;
+import com.android.tools.r8.desugar.covariantreturntype.CovariantReturnTypeAnnotationTransformerEventConsumer;
 import com.android.tools.r8.errors.CompilationError;
 import com.android.tools.r8.graph.AppInfo;
 import com.android.tools.r8.graph.AppView;
@@ -22,7 +24,6 @@
 import com.android.tools.r8.ir.desugar.CfInstructionDesugaringEventConsumer;
 import com.android.tools.r8.ir.desugar.CfPostProcessingDesugaringCollection;
 import com.android.tools.r8.ir.desugar.CfPostProcessingDesugaringEventConsumer;
-import com.android.tools.r8.ir.desugar.CovariantReturnTypeAnnotationTransformerEventConsumer;
 import com.android.tools.r8.ir.desugar.ProgramAdditions;
 import com.android.tools.r8.ir.desugar.desugaredlibrary.apiconversion.DesugaredLibraryAPIConverter;
 import com.android.tools.r8.ir.desugar.itf.EmulatedInterfaceApplicationRewriter;
@@ -77,6 +78,8 @@
 
     application = commitPendingSyntheticItems(appView, application);
 
+    processCovariantReturnTypeAnnotations(profileCollectionAdditions, executorService);
+
     // Build a new application with jumbo string info,
     Builder<?> builder = application.builder();
 
@@ -85,8 +88,6 @@
       new L8InnerOuterAttributeEraser(appView).run();
     }
 
-    processCovariantReturnTypeAnnotations(builder, profileCollectionAdditions);
-
     timing.end();
 
     application = builder.build();
@@ -272,6 +273,7 @@
               appView.appInfo().getMainDexInfo()));
       application = appView.appInfo().app();
     }
+    assert application == appView.appInfo().app();
     return application;
   }
 
@@ -338,14 +340,13 @@
     programAdditions.apply(threadingModule, executorService);
   }
 
-  @SuppressWarnings("BadImport")
   private void processCovariantReturnTypeAnnotations(
-      Builder<?> builder, ProfileCollectionAdditions profileCollectionAdditions) {
-    if (covariantReturnTypeAnnotationTransformer != null) {
-      covariantReturnTypeAnnotationTransformer.process(
-          builder,
-          CovariantReturnTypeAnnotationTransformerEventConsumer.create(profileCollectionAdditions));
-    }
+      ProfileCollectionAdditions profileCollectionAdditions, ExecutorService executorService)
+      throws ExecutionException {
+    CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer =
+        CovariantReturnTypeAnnotationTransformerEventConsumer.create(profileCollectionAdditions);
+    CovariantReturnTypeAnnotationTransformer.runIfNecessary(
+        appView, this, eventConsumer, executorService);
   }
 
   Timing rewriteNonDesugaredCode(
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/PrimaryR8IRConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/PrimaryR8IRConverter.java
index 3b3ba7a..d25c1b5 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/PrimaryR8IRConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/PrimaryR8IRConverter.java
@@ -15,7 +15,6 @@
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedbackDelayed;
 import com.android.tools.r8.naming.IdentifierMinifier;
 import com.android.tools.r8.optimize.argumentpropagation.ArgumentPropagator;
-import com.android.tools.r8.optimize.compose.ComposableOptimizationPass;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.Timing;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
@@ -224,8 +223,6 @@
       identifierNameStringMarker.decoupleIdentifierNameStringsInFields(executorService);
     }
 
-    ComposableOptimizationPass.run(appView, this, executorService, timing);
-
     // Assure that no more optimization feedback left after post processing.
     assert feedback.noUpdatesLeft();
     return appView.appInfo().app();
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/CovariantReturnTypeAnnotationTransformer.java b/src/main/java/com/android/tools/r8/ir/desugar/CovariantReturnTypeAnnotationTransformer.java
deleted file mode 100644
index f84950a..0000000
--- a/src/main/java/com/android/tools/r8/ir/desugar/CovariantReturnTypeAnnotationTransformer.java
+++ /dev/null
@@ -1,309 +0,0 @@
-// Copyright (c) 2018, 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.ir.desugar;
-
-import com.android.tools.r8.errors.CompilationError;
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.DexAnnotation;
-import com.android.tools.r8.graph.DexAnnotationElement;
-import com.android.tools.r8.graph.DexApplication;
-import com.android.tools.r8.graph.DexClass;
-import com.android.tools.r8.graph.DexEncodedAnnotation;
-import com.android.tools.r8.graph.DexEncodedMethod;
-import com.android.tools.r8.graph.DexItemFactory;
-import com.android.tools.r8.graph.DexMethod;
-import com.android.tools.r8.graph.DexProgramClass;
-import com.android.tools.r8.graph.DexProto;
-import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.DexValue;
-import com.android.tools.r8.graph.DexValue.DexValueAnnotation;
-import com.android.tools.r8.graph.DexValue.DexValueArray;
-import com.android.tools.r8.graph.DexValue.DexValueType;
-import com.android.tools.r8.graph.MethodAccessFlags;
-import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.ir.conversion.IRConverter;
-import com.android.tools.r8.ir.conversion.MethodConversionOptions;
-import com.android.tools.r8.ir.conversion.MethodProcessorEventConsumer;
-import com.android.tools.r8.ir.synthetic.ForwardMethodBuilder;
-import com.google.common.base.Predicates;
-import java.util.HashSet;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Set;
-
-// Responsible for processing the annotations dalvik.annotation.codegen.CovariantReturnType and
-// dalvik.annotation.codegen.CovariantReturnType$CovariantReturnTypes.
-//
-// Consider the following class:
-//   public class B extends A {
-//     @CovariantReturnType(returnType = B.class, presentAfter = 25)
-//     @Override
-//     public A m(...) { ... return new B(); }
-//   }
-//
-// The annotation is used to indicate that the compiler should insert a synthetic method that is
-// equivalent to method m, but has return type B instead of A. Thus, for this example, this
-// component is responsible for inserting the following method in class B (in addition to the
-// existing method m):
-//   public B m(...) { A result = "invoke B.m(...)A;"; return (B) result; }
-//
-// Note that a method may be annotated with more than one CovariantReturnType annotation. In this
-// case there will be a CovariantReturnType$CovariantReturnTypes annotation on the method that wraps
-// several CovariantReturnType annotations. In this case, a new method is synthesized for each of
-// the contained CovariantReturnType annotations.
-public final class CovariantReturnTypeAnnotationTransformer {
-
-  private final IRConverter converter;
-  private final MethodProcessorEventConsumer methodProcessorEventConsumer =
-      MethodProcessorEventConsumer.empty();
-  private final DexItemFactory factory;
-
-  public CovariantReturnTypeAnnotationTransformer(AppView<?> appView, IRConverter converter) {
-    this.converter = converter;
-    this.factory = appView.dexItemFactory();
-  }
-
-  // TODO(b/270398965): Replace LinkedList.
-  @SuppressWarnings("JdkObsolete")
-  public void process(
-      DexApplication.Builder<?> builder,
-      CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer) {
-    // List of methods that should be added to the next class.
-    List<DexEncodedMethod> methodsWithCovariantReturnTypeAnnotation = new LinkedList<>();
-    List<DexEncodedMethod> covariantReturnTypeMethods = new LinkedList<>();
-    for (DexProgramClass clazz : builder.getProgramClasses()) {
-      // Construct the methods that should be added to clazz.
-      buildCovariantReturnTypeMethodsForClass(
-          clazz,
-          methodsWithCovariantReturnTypeAnnotation,
-          covariantReturnTypeMethods,
-          eventConsumer);
-      if (covariantReturnTypeMethods.isEmpty()) {
-        continue;
-      }
-      updateClass(clazz, methodsWithCovariantReturnTypeAnnotation, covariantReturnTypeMethods);
-      // Reset lists for the next class that will have a CovariantReturnType or
-      // CovariantReturnType$CovariantReturnTypes annotation.
-      methodsWithCovariantReturnTypeAnnotation.clear();
-      covariantReturnTypeMethods.clear();
-    }
-  }
-
-  private void updateClass(
-      DexClass clazz,
-      List<DexEncodedMethod> methodsWithCovariantReturnTypeAnnotation,
-      List<DexEncodedMethod> covariantReturnTypeMethods) {
-    // It is a compilation error if the class already has a method with a signature similar to one
-    // of the methods in covariantReturnTypeMethods.
-    for (DexEncodedMethod syntheticMethod : covariantReturnTypeMethods) {
-      if (hasVirtualMethodWithSignature(clazz, syntheticMethod)) {
-        throw new CompilationError(
-            String.format(
-                "Cannot process CovariantReturnType annotation: Class %s already "
-                    + "has a method \"%s\"",
-                clazz.getType(), syntheticMethod.toSourceString()));
-      }
-    }
-    // Remove the CovariantReturnType annotations.
-    for (DexEncodedMethod method : methodsWithCovariantReturnTypeAnnotation) {
-      method.setAnnotations(
-          method.annotations().keepIf(x -> !isCovariantReturnTypeAnnotation(x.annotation)));
-    }
-    // Add the newly constructed methods to the class.
-    clazz.addVirtualMethods(covariantReturnTypeMethods);
-  }
-
-  // Processes all the dalvik.annotation.codegen.CovariantReturnType and dalvik.annotation.codegen.
-  // CovariantReturnTypes annotations in the given DexClass. Adds the newly constructed, synthetic
-  // methods to the list covariantReturnTypeMethods.
-  private void buildCovariantReturnTypeMethodsForClass(
-      DexProgramClass clazz,
-      List<DexEncodedMethod> methodsWithCovariantReturnTypeAnnotation,
-      List<DexEncodedMethod> covariantReturnTypeMethods,
-      CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer) {
-    clazz.forEachProgramVirtualMethod(
-        method -> {
-          if (methodHasCovariantReturnTypeAnnotation(method.getDefinition())) {
-            methodsWithCovariantReturnTypeAnnotation.add(method.getDefinition());
-            buildCovariantReturnTypeMethodsForMethod(
-                method, covariantReturnTypeMethods, eventConsumer);
-          }
-        });
-  }
-
-  private boolean methodHasCovariantReturnTypeAnnotation(DexEncodedMethod method) {
-    for (DexAnnotation annotation : method.annotations().annotations) {
-      if (isCovariantReturnTypeAnnotation(annotation.annotation)) {
-        return true;
-      }
-    }
-    return false;
-  }
-
-  // Processes all the dalvik.annotation.codegen.CovariantReturnType and dalvik.annotation.Co-
-  // variantReturnTypes annotations on the given method. Adds the newly constructed, synthetic
-  // methods to the list covariantReturnTypeMethods.
-  private void buildCovariantReturnTypeMethodsForMethod(
-      ProgramMethod method,
-      List<DexEncodedMethod> covariantReturnTypeMethods,
-      CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer) {
-    assert methodHasCovariantReturnTypeAnnotation(method.getDefinition());
-    for (DexType covariantReturnType : getCovariantReturnTypes(method)) {
-      DexEncodedMethod covariantReturnTypeMethod =
-          buildCovariantReturnTypeMethod(method, covariantReturnType, eventConsumer);
-      covariantReturnTypeMethods.add(covariantReturnTypeMethod);
-    }
-  }
-
-  // Builds a synthetic method that invokes the given method, casts the result to
-  // covariantReturnType, and then returns the result. The newly created method will have return
-  // type covariantReturnType.
-  //
-  // Note: any "synchronized" or "strictfp" modifier could be dropped safely.
-  private DexEncodedMethod buildCovariantReturnTypeMethod(
-      ProgramMethod method,
-      DexType covariantReturnType,
-      CovariantReturnTypeAnnotationTransformerEventConsumer eventConsumer) {
-    DexProgramClass methodHolder = method.getHolder();
-    DexMethod methodReference = method.getReference();
-    DexEncodedMethod methodDefinition = method.getDefinition();
-    DexProto newProto = factory.createProto(covariantReturnType, methodReference.proto.parameters);
-    MethodAccessFlags newAccessFlags = methodDefinition.accessFlags.copy();
-    newAccessFlags.setBridge();
-    newAccessFlags.setSynthetic();
-    newAccessFlags.unsetAbstract(); // Synthetic bridge has code, so never abstract.
-    DexMethod newMethod =
-        factory.createMethod(methodHolder.getType(), newProto, methodReference.getName());
-    ForwardMethodBuilder forwardMethodBuilder =
-        ForwardMethodBuilder.builder(factory)
-            .setNonStaticSource(newMethod)
-            .setVirtualTarget(methodReference, methodHolder.isInterface())
-            .setCastResult();
-    DexEncodedMethod newVirtualMethod =
-        DexEncodedMethod.syntheticBuilder()
-            .setMethod(newMethod)
-            .setAccessFlags(newAccessFlags)
-            .setGenericSignature(methodDefinition.getGenericSignature())
-            .setAnnotations(
-                methodDefinition
-                    .annotations()
-                    .keepIf(x -> !isCovariantReturnTypeAnnotation(x.annotation)))
-            .setParameterAnnotations(
-                methodDefinition.parameterAnnotationsList.keepIf(Predicates.alwaysTrue()))
-            .setCode(forwardMethodBuilder.buildCf())
-            .setApiLevelForDefinition(methodDefinition.getApiLevelForDefinition())
-            .setApiLevelForCode(methodDefinition.getApiLevelForCode())
-            .build();
-    // Optimize to generate DexCode instead of CfCode.
-    ProgramMethod programMethod = new ProgramMethod(methodHolder, newVirtualMethod);
-    converter.optimizeSynthesizedMethod(
-        programMethod,
-        methodProcessorEventConsumer,
-        MethodConversionOptions.forD8(converter.appView));
-    eventConsumer.acceptCovariantReturnTypeBridgeMethod(programMethod, method);
-    return newVirtualMethod;
-  }
-
-  // Returns the set of covariant return types for method.
-  //
-  // If the method is:
-  //   @dalvik.annotation.codegen.CovariantReturnType(returnType=SubOfFoo, presentAfter=25)
-  //   @dalvik.annotation.codegen.CovariantReturnType(returnType=SubOfSubOfFoo, presentAfter=28)
-  //   @Override
-  //   public Foo foo() { ... return new SubOfSubOfFoo(); }
-  // then this method returns the set { SubOfFoo, SubOfSubOfFoo }.
-  private Set<DexType> getCovariantReturnTypes(ProgramMethod method) {
-    Set<DexType> covariantReturnTypes = new HashSet<>();
-    for (DexAnnotation annotation : method.getDefinition().annotations().annotations) {
-      if (isCovariantReturnTypeAnnotation(annotation.annotation)) {
-        getCovariantReturnTypesFromAnnotation(
-            method.getHolder(),
-            method.getDefinition(),
-            annotation.annotation,
-            covariantReturnTypes);
-      }
-    }
-    return covariantReturnTypes;
-  }
-
-  @SuppressWarnings("ReferenceEquality")
-  private void getCovariantReturnTypesFromAnnotation(
-      DexClass clazz,
-      DexEncodedMethod method,
-      DexEncodedAnnotation annotation,
-      Set<DexType> covariantReturnTypes) {
-    assert isCovariantReturnTypeAnnotation(annotation);
-    boolean hasPresentAfterElement = false;
-    for (DexAnnotationElement element : annotation.elements) {
-      String name = element.name.toString();
-      if (annotation.type == factory.annotationCovariantReturnType) {
-        if (name.equals("returnType")) {
-          DexValueType dexValueType = element.value.asDexValueType();
-          if (dexValueType == null) {
-            throw new CompilationError(
-                String.format(
-                    "Expected element \"returnType\" of CovariantReturnType annotation to "
-                        + "reference a type (method: \"%s\", was: %s)",
-                    method.toSourceString(), element.value.getClass().getCanonicalName()));
-          }
-          covariantReturnTypes.add(dexValueType.value);
-        } else if (name.equals("presentAfter")) {
-          hasPresentAfterElement = true;
-        }
-      } else {
-        if (name.equals("value")) {
-          DexValueArray array = element.value.asDexValueArray();
-          if (array == null) {
-            throw new CompilationError(
-                String.format(
-                    "Expected element \"value\" of CovariantReturnTypes annotation to "
-                        + "be an array (method: \"%s\", was: %s)",
-                    method.toSourceString(), element.value.getClass().getCanonicalName()));
-          }
-
-          // Handle the inner dalvik.annotation.codegen.CovariantReturnType annotations recursively.
-          for (DexValue value : array.getValues()) {
-            assert value.isDexValueAnnotation();
-            DexValueAnnotation innerAnnotation = value.asDexValueAnnotation();
-            getCovariantReturnTypesFromAnnotation(
-                clazz, method, innerAnnotation.value, covariantReturnTypes);
-          }
-        }
-      }
-    }
-
-    if (annotation.type == factory.annotationCovariantReturnType && !hasPresentAfterElement) {
-      throw new CompilationError(
-          String.format(
-              "CovariantReturnType annotation for method \"%s\" is missing mandatory element "
-                  + "\"presentAfter\" (class %s)",
-              clazz.getType(), method.toSourceString()));
-    }
-  }
-
-  public boolean isCovariantReturnTypeAnnotation(DexEncodedAnnotation annotation) {
-    return isCovariantReturnTypeAnnotation(annotation, factory);
-  }
-
-  public static boolean isCovariantReturnTypeAnnotation(
-      DexEncodedAnnotation annotation, DexItemFactory factory) {
-    return isCovariantReturnTypeAnnotation(annotation.type, factory);
-  }
-
-  @SuppressWarnings("ReferenceEquality")
-  public static boolean isCovariantReturnTypeAnnotation(DexType type, DexItemFactory factory) {
-    return type == factory.annotationCovariantReturnType
-        || type == factory.annotationCovariantReturnTypes;
-  }
-
-  private static boolean hasVirtualMethodWithSignature(DexClass clazz, DexEncodedMethod method) {
-    for (DexEncodedMethod existingMethod : clazz.virtualMethods()) {
-      if (existingMethod.getReference().equals(method.getReference())) {
-        return true;
-      }
-    }
-    return false;
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java b/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java
index 520e2d8..85215e5 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java
@@ -324,7 +324,8 @@
     }
 
     // Check that the service is not kept.
-    if (appView().appInfo().isPinnedWithDefinitionLookup(serviceType)) {
+    if (!options.allowServiceLoaderRewritingPinnedTypes
+        && appView().appInfo().isPinnedWithDefinitionLookup(serviceType)) {
       report(code.context(), serviceType, "The service loader type is kept");
       return null;
     }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/CallSiteOptimizationInfo.java b/src/main/java/com/android/tools/r8/ir/optimize/info/CallSiteOptimizationInfo.java
index b651b1f..41a3bc2 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/CallSiteOptimizationInfo.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/CallSiteOptimizationInfo.java
@@ -6,6 +6,7 @@
 import com.android.tools.r8.ir.analysis.type.DynamicType;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.ir.analysis.value.UnknownValue;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
 
 // A flat lattice structure: TOP and a lattice element that holds accumulated argument info.
 public abstract class CallSiteOptimizationInfo {
@@ -32,5 +33,9 @@
     return UnknownValue.getInstance();
   }
 
+  public final AbstractValue getAbstractArgumentValue(MethodParameter methodParameter) {
+    return getAbstractArgumentValue(methodParameter.getIndex());
+  }
+
   // TODO(b/139249918): propagate classes that are guaranteed to be initialized.
 }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/DefaultMethodOptimizationInfo.java b/src/main/java/com/android/tools/r8/ir/optimize/info/DefaultMethodOptimizationInfo.java
index 8b2f63d..b348230 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/DefaultMethodOptimizationInfo.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/DefaultMethodOptimizationInfo.java
@@ -9,7 +9,6 @@
 import com.android.tools.r8.ir.analysis.inlining.SimpleInliningConstraint;
 import com.android.tools.r8.ir.analysis.type.DynamicType;
 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.InvokeDirect;
 import com.android.tools.r8.ir.code.InvokeMethod;
 import com.android.tools.r8.ir.optimize.classinliner.constraint.ClassInlinerMethodConstraint;
@@ -33,7 +32,6 @@
   static final Set<DexType> UNKNOWN_INITIALIZED_CLASSES_ON_NORMAL_EXIT = ImmutableSet.of();
   static final int UNKNOWN_RETURNED_ARGUMENT = -1;
   static final boolean UNKNOWN_NEVER_RETURNS_NORMALLY = false;
-  static final AbstractValue UNKNOWN_ABSTRACT_RETURN_VALUE = UnknownValue.getInstance();
   static final boolean UNKNOWN_INITIALIZER_ENABLING_JAVA_ASSERTIONS = false;
   static final boolean UNKNOWN_MAY_HAVE_SIDE_EFFECTS = true;
   static final boolean UNKNOWN_RETURN_VALUE_ONLY_DEPENDS_ON_ARGUMENTS = false;
@@ -138,12 +136,12 @@
 
   @Override
   public AbstractFunction getAbstractFunction() {
-    return AbstractFunction.unknown();
+    return AbstractValue.unknown();
   }
 
   @Override
   public AbstractValue getAbstractReturnValue() {
-    return UNKNOWN_ABSTRACT_RETURN_VALUE;
+    return AbstractValue.unknown();
   }
 
   @Override
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 fbf3767..ab2fc60 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
@@ -104,7 +104,7 @@
 import com.android.tools.r8.kotlin.Kotlin.Intrinsics;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
 import com.android.tools.r8.optimize.compose.ComposeUtils;
-import com.android.tools.r8.optimize.compose.UpdateChangedFlagsAbstractFunction;
+import com.android.tools.r8.optimize.compose.ComputationTreeUnopUpdateChangedFlagsNode;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.BooleanUtils;
 import com.android.tools.r8.utils.InternalOptions;
@@ -755,10 +755,9 @@
   private void computeAbstractFunction(IRCode code, OptimizationFeedback feedback) {
     if (ComposeUtils.isUpdateChangedFlags(code, appView.dexItemFactory())) {
       MethodParameter methodParameter = new MethodParameter(code.context(), 0);
-      UpdateChangedFlagsAbstractFunction updateChangedFlagsAbstractFunction =
-          new UpdateChangedFlagsAbstractFunction(methodParameter);
-      feedback.setAbstractFunction(
-          code.context().getDefinition(), updateChangedFlagsAbstractFunction);
+      ComputationTreeUnopUpdateChangedFlagsNode node =
+          new ComputationTreeUnopUpdateChangedFlagsNode(methodParameter);
+      feedback.setAbstractFunction(code.context().getDefinition(), node);
     }
   }
 
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/MethodResolutionOptimizationInfo.java b/src/main/java/com/android/tools/r8/ir/optimize/info/MethodResolutionOptimizationInfo.java
index 6c59889..727328a 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/MethodResolutionOptimizationInfo.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/MethodResolutionOptimizationInfo.java
@@ -77,7 +77,7 @@
 
   @Override
   public AbstractFunction getAbstractFunction() {
-    return AbstractFunction.unknown();
+    return AbstractValue.unknown();
   }
 
   @Override
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 efe5cb6..f8c1f07 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
@@ -44,9 +44,8 @@
   private Set<DexType> initializedClassesOnNormalExit =
       DefaultMethodOptimizationInfo.UNKNOWN_INITIALIZED_CLASSES_ON_NORMAL_EXIT;
   private int returnedArgument = DefaultMethodOptimizationInfo.UNKNOWN_RETURNED_ARGUMENT;
-  private AbstractFunction abstractFunction = AbstractFunction.unknown();
-  private AbstractValue abstractReturnValue =
-      DefaultMethodOptimizationInfo.UNKNOWN_ABSTRACT_RETURN_VALUE;
+  private AbstractFunction abstractFunction = AbstractValue.unknown();
+  private AbstractValue abstractReturnValue = AbstractValue.unknown();
   private ClassInlinerMethodConstraint classInlinerConstraint =
       ClassInlinerMethodConstraint.alwaysFalse();
   private boolean convertCheckNotNull = false;
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
index 2273dbb..c059f0a 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
@@ -130,11 +130,15 @@
       AbstractValueSupplier abstractValueSupplier =
           value -> value.getAbstractValue(appView, method);
       PathConstraintSupplier pathConstraintSupplier =
-          new PathConstraintSupplier(appView, code, codeScanner.getMethodParameterFactory());
+          new PathConstraintSupplier(
+              appView,
+              code,
+              codeScanner.getFieldValueFactory(),
+              codeScanner.getMethodParameterFactory());
       codeScanner.scan(method, code, abstractValueSupplier, pathConstraintSupplier, timing);
 
       assert effectivelyUnusedArgumentsAnalysis != null;
-      effectivelyUnusedArgumentsAnalysis.scan(method, code);
+      effectivelyUnusedArgumentsAnalysis.scan(method, code, pathConstraintSupplier);
 
       assert reprocessingCriteriaCollection != null;
       reprocessingCriteriaCollection.analyzeArgumentUses(method, code);
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScanner.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScanner.java
index ba31341..99759a9 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScanner.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScanner.java
@@ -13,6 +13,7 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.MethodResolutionResult.SingleResolutionResult;
 import com.android.tools.r8.graph.ProgramField;
+import com.android.tools.r8.graph.ProgramMember;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.analysis.path.PathConstraintSupplier;
 import com.android.tools.r8.ir.analysis.path.state.ConcretePathConstraintAnalysisState;
@@ -26,7 +27,6 @@
 import com.android.tools.r8.ir.code.AbstractValueSupplier;
 import com.android.tools.r8.ir.code.AliasedValueConfiguration;
 import com.android.tools.r8.ir.code.AssumeAndCheckCastAliasedValueConfiguration;
-import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.CheckCast;
 import com.android.tools.r8.ir.code.FieldGet;
 import com.android.tools.r8.ir.code.FieldPut;
@@ -38,7 +38,6 @@
 import com.android.tools.r8.ir.code.Phi;
 import com.android.tools.r8.ir.code.Position.SourcePosition;
 import com.android.tools.r8.ir.code.Value;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.AbstractFunction;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.CastAbstractFunction;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteArrayTypeValueState;
@@ -66,6 +65,7 @@
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.StateCloner;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.UnknownMethodState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ValueState;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComposableComputationTreeBuilder;
 import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
 import com.android.tools.r8.optimize.argumentpropagation.reprocessingcriteria.ArgumentPropagatorReprocessingCriteriaCollection;
 import com.android.tools.r8.optimize.argumentpropagation.reprocessingcriteria.MethodReprocessingCriteria;
@@ -77,10 +77,9 @@
 import com.android.tools.r8.utils.DeterminismChecker.LineCallback;
 import com.android.tools.r8.utils.ListUtils;
 import com.android.tools.r8.utils.Timing;
+import com.android.tools.r8.utils.TraversalUtils;
 import com.android.tools.r8.utils.structural.StructuralItem;
 import com.google.common.collect.Sets;
-import it.unimi.dsi.fastutil.objects.Object2IntMap;
-import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -105,8 +104,6 @@
 
   private final AppView<AppInfoWithLiveness> appView;
 
-  private final ArgumentPropagatorCodeScannerModeling modeling;
-
   private final FieldValueFactory fieldValueFactory = new FieldValueFactory();
 
   final MethodParameterFactory methodParameterFactory = new MethodParameterFactory();
@@ -145,7 +142,6 @@
       AppView<AppInfoWithLiveness> appView,
       ArgumentPropagatorReprocessingCriteriaCollection reprocessingCriteriaCollection) {
     this.appView = appView;
-    this.modeling = new ArgumentPropagatorCodeScannerModeling(appView);
     this.reprocessingCriteriaCollection = reprocessingCriteriaCollection;
   }
 
@@ -161,6 +157,10 @@
     return fieldStates;
   }
 
+  public FieldValueFactory getFieldValueFactory() {
+    return fieldValueFactory;
+  }
+
   public MethodParameterFactory getMethodParameterFactory() {
     return methodParameterFactory;
   }
@@ -242,8 +242,6 @@
     protected final ProgramMethod context;
     private final PathConstraintSupplier pathConstraintSupplier;
 
-    private Object2IntMap<Phi> phiNumbering = null;
-
     protected CodeScanner(
         AbstractValueSupplier abstractValueSupplier,
         IRCode code,
@@ -320,6 +318,7 @@
       NonEmptyValueState inFlowState =
           computeInFlowState(
               field.getType(),
+              field,
               value,
               initialValue,
               context,
@@ -356,7 +355,7 @@
         MethodParameter inParameter =
             methodParameterFactory.create(
                 context, valueRoot.getDefinition().asArgument().getIndex());
-        if (!widenBaseInFlow(staticType, inParameter, context).isUnknownAbstractFunction()) {
+        if (!widenBaseInFlow(staticType, inParameter, context).isUnknown()) {
           return inParameter;
         }
       } else if (valueRoot.isDefinedByInstructionSatisfying(Instruction::isFieldGet)) {
@@ -364,7 +363,7 @@
         ProgramField field = fieldGet.resolveField(appView, context).getProgramField();
         if (field != null) {
           FieldValue fieldValue = fieldValueFactory.create(field);
-          if (!widenBaseInFlow(staticType, fieldValue, context).isUnknownAbstractFunction()) {
+          if (!widenBaseInFlow(staticType, fieldValue, context).isUnknown()) {
             return fieldValue;
           }
         }
@@ -380,6 +379,7 @@
     // TODO(b/302281503): Canonicalize computed in flow.
     private InFlow computeInFlow(
         DexType staticType,
+        ProgramMember<?, ?> target,
         Value value,
         Value initialValue,
         ProgramMethod context,
@@ -411,16 +411,28 @@
         return castBaseInFlow(
             widenBaseInFlow(staticType, fieldValueFactory.create(field), context), value);
       } else if (value.isPhi()) {
+        // TODO(b/302281503): Replace IfThenElseAbstractFunction by ComputationTreeNode (?).
         return computeIfThenElseAbstractFunction(value.asPhi(), valueStateSupplier);
+      } else if (target != null && appView.getComposeReferences().isComposable(target)) {
+        ComputationTreeNode node =
+            new ComposableComputationTreeBuilder(
+                    appView,
+                    code,
+                    code.context(),
+                    fieldValueFactory,
+                    methodParameterFactory,
+                    pathConstraintSupplier)
+                .getOrBuildComputationTree(value);
+        if (!node.isComputationLeaf() && TraversalUtils.hasNext(node::traverseBaseInFlow)) {
+          recordComputationTreePosition(node, value);
+          return node;
+        }
       }
       return null;
     }
 
     private IfThenElseAbstractFunction computeIfThenElseAbstractFunction(
         Phi phi, Function<Value, NonEmptyValueState> valueStateSupplier) {
-      if (!appView.testing().enableIfThenElseAbstractFunction) {
-        return null;
-      }
       if (phi.getOperands().size() != 2 || !phi.hasOperandThatMatches(Value::isArgument)) {
         return null;
       }
@@ -438,11 +450,14 @@
       ComputationTreeNode condition =
           leftPredecessorPathConstraint.getDifferentiatingPathConstraint(
               rightPredecessorPathConstraint);
-      if (condition == null || condition.getSingleOpenVariable() == null) {
+      if (condition.getSingleOpenVariable() == null) {
         return null;
       }
       NonEmptyValueState leftValue = valueStateSupplier.apply(phi.getOperand(0));
       NonEmptyValueState rightValue = valueStateSupplier.apply(phi.getOperand(1));
+      if (leftValue.isUnknown() && rightValue.isUnknown()) {
+        return null;
+      }
       IfThenElseAbstractFunction result =
           leftPredecessorPathConstraint.isNegated(condition)
               ? new IfThenElseAbstractFunction(condition, rightValue, leftValue)
@@ -451,30 +466,27 @@
       return result;
     }
 
+    private void recordComputationTreePosition(ComputationTreeNode computation, Value value) {
+      inFlowComparatorBuilder.addComputationTreePosition(
+          computation,
+          SourcePosition.builder()
+              .setMethod(code.context().getReference())
+              .setLine(value.getNumber())
+              .build());
+    }
+
     private void recordIfThenElsePosition(
         IfThenElseAbstractFunction ifThenElseAbstractFunction, Phi phi) {
       inFlowComparatorBuilder.addIfThenElsePosition(
           ifThenElseAbstractFunction,
           SourcePosition.builder()
               .setMethod(code.context().getReference())
-              .setLine(getOrCreatePhiNumbering().getInt(phi))
+              .setLine(phi.getNumber())
               .build());
     }
 
-    private Object2IntMap<Phi> getOrCreatePhiNumbering() {
-      if (phiNumbering == null) {
-        phiNumbering = new Object2IntOpenHashMap<>();
-        for (BasicBlock block : code.getBlocks()) {
-          for (Phi phi : block.getPhis()) {
-            phiNumbering.put(phi, phiNumbering.size());
-          }
-        }
-      }
-      return phiNumbering;
-    }
-
     private InFlow castBaseInFlow(InFlow inFlow, Value value) {
-      if (inFlow.isUnknownAbstractFunction()) {
+      if (inFlow.isUnknown()) {
         return inFlow;
       }
       assert inFlow.isBaseInFlow();
@@ -489,12 +501,12 @@
     private InFlow widenBaseInFlow(DexType staticType, BaseInFlow inFlow, ProgramMethod context) {
       if (inFlow.isFieldValue()) {
         if (isFieldValueAlreadyUnknown(staticType, inFlow.asFieldValue().getField())) {
-          return AbstractFunction.unknown();
+          return AbstractValue.unknown();
         }
       } else {
         assert inFlow.isMethodParameter();
         if (isMethodParameterAlreadyUnknown(staticType, inFlow.asMethodParameter(), context)) {
-          return AbstractFunction.unknown();
+          return AbstractValue.unknown();
         }
       }
       return inFlow;
@@ -502,19 +514,22 @@
 
     private NonEmptyValueState computeInFlowState(
         DexType staticType,
+        ProgramMember<?, ?> target,
         Value value,
         Value initialValue,
         ProgramMethod context,
         Function<Value, NonEmptyValueState> valueStateSupplier) {
       assert value == initialValue || initialValue.getAliasedValue().isPhi();
-      InFlow inFlow = computeInFlow(staticType, value, initialValue, context, valueStateSupplier);
+      InFlow inFlow =
+          computeInFlow(staticType, target, value, initialValue, context, valueStateSupplier);
       if (inFlow == null) {
         return null;
       }
-      if (inFlow.isUnknownAbstractFunction()) {
+      if (inFlow.isUnknown()) {
         return ValueState.unknown();
       }
       assert inFlow.isBaseInFlow()
+          || inFlow.isAbstractComputation()
           || inFlow.isCastAbstractFunction()
           || inFlow.isIfThenElseAbstractFunction()
           || inFlow.isInstanceFieldReadAbstractFunction();
@@ -860,13 +875,8 @@
         Value value,
         Value initialValue,
         ConcreteMonomorphicMethodStateOrBottom existingMethodState) {
+      assert invoke.isInvokeStatic() || argumentIndex > 0;
       assert value == initialValue || initialValue.getAliasedValue().isPhi();
-      NonEmptyValueState modeledState =
-          modeling.modelParameterStateForArgumentToFunction(
-              invoke, singleTarget, argumentIndex, value, context);
-      if (modeledState != null) {
-        return modeledState;
-      }
 
       // Don't compute a state for this parameter if the stored state is already unknown.
       if (existingMethodState.isMonomorphic()
@@ -885,6 +895,7 @@
       NonEmptyValueState inFlowState =
           computeInFlowState(
               parameterType,
+              singleTarget,
               value,
               initialValue,
               context,
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScannerModeling.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScannerModeling.java
deleted file mode 100644
index 0a021d2..0000000
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScannerModeling.java
+++ /dev/null
@@ -1,40 +0,0 @@
-// Copyright (c) 2023, 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 com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.ir.code.InvokeMethod;
-import com.android.tools.r8.ir.code.Value;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.NonEmptyValueState;
-import com.android.tools.r8.optimize.compose.ArgumentPropagatorComposeModeling;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-
-public class ArgumentPropagatorCodeScannerModeling {
-
-  private final ArgumentPropagatorComposeModeling composeModeling;
-
-  ArgumentPropagatorCodeScannerModeling(AppView<AppInfoWithLiveness> appView) {
-    this.composeModeling =
-        appView
-                .options()
-                .getJetpackComposeOptions()
-                .isModelingChangedArgumentsToComposableFunctions()
-            ? new ArgumentPropagatorComposeModeling(appView)
-            : null;
-  }
-
-  NonEmptyValueState modelParameterStateForArgumentToFunction(
-      InvokeMethod invoke,
-      ProgramMethod singleTarget,
-      int argumentIndex,
-      Value argument,
-      ProgramMethod context) {
-    if (composeModeling != null) {
-      return composeModeling.modelParameterStateForChangedOrDefaultArgumentToComposableFunction(
-          invoke, singleTarget, argumentIndex, argument, context);
-    }
-    return null;
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/AbstractFunction.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/AbstractFunction.java
index 6830918..7ffaddd 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/AbstractFunction.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/AbstractFunction.java
@@ -4,7 +4,9 @@
 package com.android.tools.r8.optimize.argumentpropagation.codescanner;
 
 import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.TraversalContinuation;
 
 public interface AbstractFunction extends InFlow {
 
@@ -12,17 +14,14 @@
     return IdentityAbstractFunction.get();
   }
 
-  static UnknownAbstractFunction unknown() {
-    return UnknownAbstractFunction.get();
-  }
-
   /**
-   * Applies the current abstract function to its declared inputs (in {@link #getBaseInFlow()}).
+   * Applies the current abstract function to its declared inputs (from {@link
+   * #traverseBaseInFlow}).
    *
    * <p>It is guaranteed by the caller that the given {@param state} is the abstract state for the
    * field or parameter that caused this function to be reevaluated. If this abstract function takes
    * a single input, then {@param state} is guaranteed to be the state for the node returned by
-   * {@link #getBaseInFlow()}, and {@param flowGraphStateProvider} should never be used.
+   * {@link #traverseBaseInFlow}, and {@param flowGraphStateProvider} should never be used.
    *
    * <p>Abstract functions that depend on multiple inputs can lookup the state for each input in
    * {@param flowGraphStateProvider}. Attempting to lookup the state of a non-declared input is an
@@ -31,16 +30,8 @@
   ValueState apply(
       AppView<AppInfoWithLiveness> appView,
       FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState inState);
-
-  /** Returns true if the given {@param inFlow} is a declared input of this abstract function. */
-  boolean verifyContainsBaseInFlow(BaseInFlow inFlow);
-
-  /**
-   * Returns the program field or parameter graph nodes that this function depends on. Upon any
-   * change to the abstract state of any of these nodes this abstract function must be re-evaluated.
-   */
-  Iterable<BaseInFlow> getBaseInFlow();
+      ConcreteValueState inState,
+      DexType outStaticType);
 
   default boolean usesFlowGraphStateProvider() {
     return false;
@@ -59,4 +50,16 @@
   default boolean isIdentity() {
     return false;
   }
+
+  default boolean isUpdateChangedFlags() {
+    return false;
+  }
+
+  /** Verifies that {@param stoppingCriterion} is a declared input of this abstract function. */
+  default boolean verifyContainsBaseInFlow(BaseInFlow stoppingCriterion) {
+    assert traverseBaseInFlow(
+            inFlow -> TraversalContinuation.breakIf(inFlow.equals(stoppingCriterion)))
+        .shouldBreak();
+    return true;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomArrayTypeValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomArrayTypeValueState.java
index 19f1a3d..645d3bc 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomArrayTypeValueState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomArrayTypeValueState.java
@@ -34,6 +34,10 @@
     if (inState.isUnknown()) {
       return inState;
     }
+    if (inState.isUnused()) {
+      assert inState.identical(unusedArrayTypeState());
+      return inState;
+    }
     assert inState.isConcrete();
     assert inState.asConcrete().isReferenceState();
     ConcreteReferenceTypeValueState concreteState = inState.asConcrete().asReferenceState();
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomClassTypeValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomClassTypeValueState.java
index 6eca904..1019811 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomClassTypeValueState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomClassTypeValueState.java
@@ -36,6 +36,10 @@
     if (inState.isUnknown()) {
       return inState;
     }
+    if (inState.isUnused()) {
+      assert inState.identical(unusedClassTypeState());
+      return inState;
+    }
     assert inState.isConcrete();
     assert inState.asConcrete().isReferenceState();
     ConcreteReferenceTypeValueState concreteState = inState.asConcrete().asReferenceState();
@@ -60,7 +64,6 @@
     if (outStaticType != null) {
       return WideningUtils.widenDynamicNonReceiverType(appView, joinedDynamicType, outStaticType);
     } else {
-      assert !joinedDynamicType.isUnknown();
       return joinedDynamicType;
     }
   }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomPrimitiveTypeValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomPrimitiveTypeValueState.java
index b429da0..08f7a1d 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomPrimitiveTypeValueState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomPrimitiveTypeValueState.java
@@ -34,6 +34,10 @@
     if (inState.isUnknown()) {
       return inState;
     }
+    if (inState.isUnused()) {
+      assert inState.identical(unusedPrimitiveTypeState());
+      return inState;
+    }
     assert inState.isPrimitiveState();
     return cloner.mutableCopy(inState);
   }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/CastAbstractFunction.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/CastAbstractFunction.java
index 2586f8d..c6d7ba6 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/CastAbstractFunction.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/CastAbstractFunction.java
@@ -6,8 +6,9 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import java.util.Collections;
+import com.android.tools.r8.utils.TraversalContinuation;
 import java.util.Objects;
+import java.util.function.Function;
 
 public class CastAbstractFunction implements AbstractFunction {
 
@@ -23,19 +24,15 @@
   public ValueState apply(
       AppView<AppInfoWithLiveness> appView,
       FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState predecessorState) {
+      ConcreteValueState predecessorState,
+      DexType outStaticType) {
     return predecessorState.asReferenceState().cast(appView, type);
   }
 
   @Override
-  public boolean verifyContainsBaseInFlow(BaseInFlow inFlow) {
-    assert inFlow.equals(this.inFlow);
-    return true;
-  }
-
-  @Override
-  public Iterable<BaseInFlow> getBaseInFlow() {
-    return Collections.singleton(inFlow);
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    return inFlow.traverseBaseInFlow(fn);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteValueState.java
index df2c3b8..065159e 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteValueState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteValueState.java
@@ -10,9 +10,11 @@
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.Action;
+import com.android.tools.r8.utils.TraversalContinuation;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 public abstract class ConcreteValueState extends NonEmptyValueState {
@@ -84,6 +86,18 @@
     return inFlow;
   }
 
+  public final <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    TraversalContinuation<TB, TC> traversalContinuation = TraversalContinuation.doContinue();
+    for (InFlow inFlow : getInFlow()) {
+      traversalContinuation = inFlow.traverseBaseInFlow(fn);
+      if (traversalContinuation.shouldBreak()) {
+        break;
+      }
+    }
+    return traversalContinuation;
+  }
+
   public abstract BottomValueState getCorrespondingBottom();
 
   public abstract ConcreteParameterStateKind getKind();
@@ -131,7 +145,7 @@
       DexType outStaticType,
       StateCloner cloner,
       Action onChangedAction) {
-    if (inState.isBottom()) {
+    if (inState.isBottom() || inState.isUnused()) {
       return this;
     }
     if (inState.isUnknown()) {
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/FieldValue.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/FieldValue.java
index 99fcf7b..c066906 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/FieldValue.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/FieldValue.java
@@ -3,11 +3,17 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.optimize.argumentpropagation.codescanner;
 
+import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.ir.analysis.value.AbstractValue;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.function.Function;
 
 // TODO(b/296030319): Change DexField to implement InFlow and use DexField in all places instead of
 //  FieldValue to avoid wrappers? This would also remove the need for the FieldValueFactory.
-public class FieldValue implements BaseInFlow {
+public class FieldValue implements BaseInFlow, ComputationTreeNode {
 
   private final DexField field;
 
@@ -15,6 +21,12 @@
     this.field = field;
   }
 
+  @Override
+  public AbstractValue evaluate(
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
+    return flowGraphStateProvider.getState(field).getAbstractValue(appView);
+  }
+
   public DexField getField() {
     return field;
   }
@@ -25,11 +37,21 @@
   }
 
   @Override
+  public BaseInFlow getSingleOpenVariable() {
+    return this;
+  }
+
+  @Override
   public int internalCompareToSameKind(InFlow other, InFlowComparator comparator) {
     return field.compareTo(other.asFieldValue().getField());
   }
 
   @Override
+  public boolean isComputationLeaf() {
+    return true;
+  }
+
+  @Override
   public boolean isFieldValue() {
     return true;
   }
@@ -45,6 +67,12 @@
   }
 
   @Override
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    return fn.apply(this);
+  }
+
+  @Override
   @SuppressWarnings("EqualsGetClass")
   public boolean equals(Object obj) {
     if (this == obj) {
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/FlowGraphStateProvider.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/FlowGraphStateProvider.java
index 385c957..639765b 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/FlowGraphStateProvider.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/FlowGraphStateProvider.java
@@ -5,6 +5,9 @@
 
 import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.ir.analysis.value.AbstractValue;
+import com.android.tools.r8.ir.optimize.info.MethodOptimizationInfo;
 import com.android.tools.r8.optimize.argumentpropagation.propagation.FlowGraph;
 import com.android.tools.r8.utils.InternalOptions;
 import java.util.function.Supplier;
@@ -18,7 +21,8 @@
     // If the abstract function needs to perform state lookups, we restrict state lookups to the
     // declared base in flow. This is important for arriving at the correct fix point.
     if (abstractFunction.usesFlowGraphStateProvider()) {
-      assert abstractFunction.isIfThenElseAbstractFunction()
+      assert abstractFunction.isAbstractComputation()
+          || abstractFunction.isIfThenElseAbstractFunction()
           || abstractFunction.isInstanceFieldReadAbstractFunction();
       return new FlowGraphStateProvider() {
 
@@ -38,10 +42,7 @@
     }
     // Otherwise, the abstract function is a canonical function, or the abstract function has a
     // single declared input, meaning we should never perform any state lookups.
-    assert abstractFunction.isIdentity()
-        || abstractFunction.isCastAbstractFunction()
-        || abstractFunction.isUnknownAbstractFunction()
-        || abstractFunction.isUpdateChangedFlagsAbstractFunction();
+    assert abstractFunction.isIdentity() || abstractFunction.isCastAbstractFunction();
     return new FlowGraphStateProvider() {
 
       @Override
@@ -57,6 +58,31 @@
     };
   }
 
+  static FlowGraphStateProvider createFromMethodOptimizationInfo(ProgramMethod method) {
+    return new FlowGraphStateProvider() {
+
+      @Override
+      public ValueState getState(DexField field) {
+        return ValueState.unknown();
+      }
+
+      @Override
+      public ValueState getState(
+          MethodParameter methodParameter, Supplier<ValueState> defaultStateProvider) {
+        if (methodParameter.getMethod().isNotIdenticalTo(method.getReference())) {
+          return ValueState.unknown();
+        }
+        MethodOptimizationInfo optimizationInfo = method.getOptimizationInfo();
+        AbstractValue abstractValue =
+            optimizationInfo.getArgumentInfos().getAbstractArgumentValue(methodParameter);
+        if (abstractValue.isUnknown()) {
+          return ValueState.unknown();
+        }
+        return ConcreteValueState.create(methodParameter.getType(), abstractValue);
+      }
+    };
+  }
+
   ValueState getState(DexField field);
 
   ValueState getState(MethodParameter methodParameter, Supplier<ValueState> defaultStateProvider);
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/IdentityAbstractFunction.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/IdentityAbstractFunction.java
index f31b821..8d37ae7 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/IdentityAbstractFunction.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/IdentityAbstractFunction.java
@@ -5,7 +5,10 @@
 
 import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.function.Function;
 
 public class IdentityAbstractFunction implements AbstractFunction {
 
@@ -21,17 +24,14 @@
   public ValueState apply(
       AppView<AppInfoWithLiveness> appView,
       FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState inState) {
+      ConcreteValueState inState,
+      DexType outStaticType) {
     return inState;
   }
 
   @Override
-  public boolean verifyContainsBaseInFlow(BaseInFlow inFlow) {
-    throw new Unreachable();
-  }
-
-  @Override
-  public Iterable<BaseInFlow> getBaseInFlow() {
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
     throw new Unreachable();
   }
 
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/IfThenElseAbstractFunction.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/IfThenElseAbstractFunction.java
index 4d21db9..a9b7dfa 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/IfThenElseAbstractFunction.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/IfThenElseAbstractFunction.java
@@ -9,9 +9,9 @@
 import com.android.tools.r8.ir.code.Position.SourcePosition;
 import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.android.tools.r8.utils.ListUtils;
-import java.util.List;
+import com.android.tools.r8.utils.TraversalContinuation;
 import java.util.Objects;
+import java.util.function.Function;
 
 /**
  * Represents a ternary expression (exp ? u : v). The {@link #condition} is an expression containing
@@ -19,6 +19,7 @@
  * true, then `u` is chosen. If the abstract value is false, then `v` is chosen. Otherwise, the
  * result is unknown.
  */
+// TODO(b/302281503): Replace this by a ComputationTreeNode.
 // TODO(b/302281503): Evaluate the impact of using the join of `u` and `v` instead of unknown when
 //  the condition does not evaluate to true or false.
 public class IfThenElseAbstractFunction implements AbstractFunction {
@@ -42,8 +43,9 @@
   public ValueState apply(
       AppView<AppInfoWithLiveness> appView,
       FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState inState) {
-    AbstractValue conditionValue = evaluateCondition(appView, flowGraphStateProvider);
+      ConcreteValueState inState,
+      DexType outStaticType) {
+    AbstractValue conditionValue = condition.evaluate(appView, flowGraphStateProvider);
     NonEmptyValueState resultState;
     if (conditionValue.isTrue()) {
       resultState = thenState;
@@ -60,30 +62,14 @@
     if (!concreteResultState.hasInFlow()) {
       return concreteResultState;
     }
-    return resolveInFlow(appView, flowGraphStateProvider, concreteResultState);
-  }
-
-  private AbstractValue evaluateCondition(
-      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
-    MethodParameter openVariable = condition.getSingleOpenVariable();
-    assert openVariable != null;
-    ValueState variableState = flowGraphStateProvider.getState(openVariable, () -> null);
-    if (variableState == null) {
-      // TODO(b/302281503): Conservatively return unknown for now. Investigate exactly when this
-      //  happens and whether we can return something more precise instead of unknown.
-      assert false;
-      return AbstractValue.unknown();
-    }
-    AbstractValue variableValue = variableState.getAbstractValue(appView);
-    // Since the condition is guaranteed to have a single open variable we simply return the
-    // `variableValue` for any given argument index.
-    return condition.evaluate(i -> variableValue, appView.abstractValueFactory());
+    return resolveInFlow(appView, flowGraphStateProvider, concreteResultState, outStaticType);
   }
 
   private ValueState resolveInFlow(
       AppView<AppInfoWithLiveness> appView,
       FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState resultStateWithInFlow) {
+      ConcreteValueState resultStateWithInFlow,
+      DexType outStaticType) {
     ValueState resultStateWithoutInFlow = resultStateWithInFlow.mutableCopyWithoutInFlow();
     for (InFlow inFlow : resultStateWithInFlow.getInFlow()) {
       // We currently only allow the primitive kinds of in flow (fields and method parameters) to
@@ -91,13 +77,11 @@
       assert inFlow.isBaseInFlow();
       ValueState inFlowState = flowGraphStateProvider.getState(inFlow.asBaseInFlow(), () -> null);
       if (inFlowState == null) {
-        assert false;
         return ValueState.unknown();
       }
       // TODO(b/302281503): The IfThenElseAbstractFunction is only used on input to base in flow.
       //  We should set  the `outStaticType` to the static type of the current field/parameter.
       DexType inStaticType = null;
-      DexType outStaticType = null;
       resultStateWithoutInFlow =
           resultStateWithoutInFlow.mutableJoin(
               appView, inFlowState, inStaticType, outStaticType, StateCloner.getCloner());
@@ -106,27 +90,25 @@
   }
 
   @Override
-  public boolean verifyContainsBaseInFlow(BaseInFlow inFlow) {
-    // TODO(b/302281503): Implement this.
-    return true;
-  }
-
-  @Override
-  public Iterable<BaseInFlow> getBaseInFlow() {
-    List<BaseInFlow> baseInFlow = ListUtils.newArrayList(condition.getSingleOpenVariable());
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    TraversalContinuation<TB, TC> traversalContinuation = condition.traverseBaseInFlow(fn);
+    if (traversalContinuation.shouldBreak()) {
+      return traversalContinuation;
+    }
     if (thenState.isConcrete()) {
-      for (InFlow inFlow : thenState.asConcrete().getInFlow()) {
-        assert inFlow.isBaseInFlow();
-        baseInFlow.add(inFlow.asBaseInFlow());
+      traversalContinuation = thenState.asConcrete().traverseBaseInFlow(fn);
+      if (traversalContinuation.shouldBreak()) {
+        return traversalContinuation;
       }
     }
     if (elseState.isConcrete()) {
-      for (InFlow inFlow : elseState.asConcrete().getInFlow()) {
-        assert inFlow.isBaseInFlow();
-        baseInFlow.add(inFlow.asBaseInFlow());
+      traversalContinuation = elseState.asConcrete().traverseBaseInFlow(fn);
+      if (traversalContinuation.shouldBreak()) {
+        return traversalContinuation;
       }
     }
-    return baseInFlow;
+    return traversalContinuation;
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlow.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlow.java
index 3c28c20..2eb6d54 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlow.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlow.java
@@ -5,7 +5,9 @@
 
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexMethod;
-import com.android.tools.r8.optimize.compose.UpdateChangedFlagsAbstractFunction;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.function.Function;
 
 public interface InFlow {
 
@@ -20,6 +22,14 @@
 
   InFlowKind getKind();
 
+  default boolean isAbstractComputation() {
+    return false;
+  }
+
+  default ComputationTreeNode asAbstractComputation() {
+    return null;
+  }
+
   default boolean isAbstractFunction() {
     return false;
   }
@@ -84,19 +94,10 @@
     return null;
   }
 
-  default OrAbstractFunction asOrAbstractFunction() {
-    return null;
-  }
-
-  default boolean isUnknownAbstractFunction() {
+  default boolean isUnknown() {
     return false;
   }
 
-  default boolean isUpdateChangedFlagsAbstractFunction() {
-    return false;
-  }
-
-  default UpdateChangedFlagsAbstractFunction asUpdateChangedFlagsAbstractFunction() {
-    return null;
-  }
+  <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn);
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlowComparator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlowComparator.java
index 82a1cc8..1327a44 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlowComparator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlowComparator.java
@@ -4,18 +4,29 @@
 package com.android.tools.r8.optimize.argumentpropagation.codescanner;
 
 import com.android.tools.r8.ir.code.Position.SourcePosition;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Map;
 
 public class InFlowComparator implements Comparator<InFlow> {
 
+  private final Map<ComputationTreeNode, SourcePosition> computationTreePositions;
   private final Map<IfThenElseAbstractFunction, SourcePosition> ifThenElsePositions;
 
-  private InFlowComparator(Map<IfThenElseAbstractFunction, SourcePosition> ifThenElsePositions) {
+  private InFlowComparator(
+      Map<ComputationTreeNode, SourcePosition> computationTreePositions,
+      Map<IfThenElseAbstractFunction, SourcePosition> ifThenElsePositions) {
+    this.computationTreePositions = computationTreePositions;
     this.ifThenElsePositions = ifThenElsePositions;
   }
 
+  public SourcePosition getComputationTreePosition(ComputationTreeNode computation) {
+    SourcePosition position = computationTreePositions.get(computation);
+    assert position != null;
+    return position;
+  }
+
   public SourcePosition getIfThenElsePosition(IfThenElseAbstractFunction fn) {
     SourcePosition position = ifThenElsePositions.get(fn);
     assert position != null;
@@ -37,9 +48,16 @@
 
   public static class Builder {
 
+    private final Map<ComputationTreeNode, SourcePosition> computationTreePositions =
+        new HashMap<>();
     private final Map<IfThenElseAbstractFunction, SourcePosition> ifThenElsePositions =
         new HashMap<>();
 
+    public void addComputationTreePosition(
+        ComputationTreeNode computation, SourcePosition position) {
+      computationTreePositions.put(computation, position);
+    }
+
     public void addIfThenElsePosition(IfThenElseAbstractFunction fn, SourcePosition position) {
       synchronized (ifThenElsePositions) {
         ifThenElsePositions.put(fn, position);
@@ -47,7 +65,7 @@
     }
 
     public InFlowComparator build() {
-      return new InFlowComparator(ifThenElsePositions);
+      return new InFlowComparator(computationTreePositions, ifThenElsePositions);
     }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlowKind.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlowKind.java
index 0714caf..6844aa2 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlowKind.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlowKind.java
@@ -4,13 +4,11 @@
 package com.android.tools.r8.optimize.argumentpropagation.codescanner;
 
 public enum InFlowKind {
+  ABSTRACT_COMPUTATION,
   ABSTRACT_FUNCTION_CAST,
   ABSTRACT_FUNCTION_IDENTITY,
   ABSTRACT_FUNCTION_IF_THEN_ELSE,
   ABSTRACT_FUNCTION_INSTANCE_FIELD_READ,
-  ABSTRACT_FUNCTION_OR,
-  ABSTRACT_FUNCTION_UNKNOWN,
-  ABSTRACT_FUNCTION_UPDATE_CHANGED_FLAGS,
   FIELD,
   METHOD_PARAMETER
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InstanceFieldReadAbstractFunction.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InstanceFieldReadAbstractFunction.java
index 382706b..022c2b4 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InstanceFieldReadAbstractFunction.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InstanceFieldReadAbstractFunction.java
@@ -5,9 +5,11 @@
 
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.google.common.collect.Lists;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.function.Function;
 
 public class InstanceFieldReadAbstractFunction implements AbstractFunction {
 
@@ -23,7 +25,8 @@
   public ValueState apply(
       AppView<AppInfoWithLiveness> appView,
       FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState predecessorState) {
+      ConcreteValueState predecessorState,
+      DexType outStaticType) {
     ValueState state = flowGraphStateProvider.getState(receiver, () -> ValueState.bottom(field));
     if (state.isBottom()) {
       return ValueState.bottom(field);
@@ -47,14 +50,13 @@
   }
 
   @Override
-  public boolean verifyContainsBaseInFlow(BaseInFlow inFlow) {
-    assert inFlow.equals(receiver) || inFlow.isFieldValue(field);
-    return true;
-  }
-
-  @Override
-  public Iterable<BaseInFlow> getBaseInFlow() {
-    return Lists.newArrayList(receiver, new FieldValue(field));
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    TraversalContinuation<TB, TC> traversalContinuation = fn.apply(receiver);
+    if (traversalContinuation.shouldContinue()) {
+      traversalContinuation = fn.apply(new FieldValue(field));
+    }
+    return traversalContinuation;
   }
 
   private ValueState getFallbackState(FlowGraphStateProvider flowGraphStateProvider) {
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodParameter.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodParameter.java
index ed49cea..4b6340b 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodParameter.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodParameter.java
@@ -4,15 +4,17 @@
 
 package com.android.tools.r8.optimize.argumentpropagation.codescanner;
 
+import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexClassAndMethod;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
-import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
 import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.BooleanUtils;
+import com.android.tools.r8.utils.TraversalContinuation;
 import java.util.Objects;
-import java.util.function.IntFunction;
+import java.util.function.Function;
 
 public class MethodParameter implements BaseInFlow, ComputationTreeNode {
 
@@ -30,8 +32,10 @@
     this.isMethodStatic = isMethodStatic;
   }
 
-  public static MethodParameter createStatic(DexMethod method, int index) {
-    return new MethodParameter(method, index, true);
+  @Override
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    return fn.apply(this);
   }
 
   @Override
@@ -48,7 +52,7 @@
   }
 
   @Override
-  public MethodParameter getSingleOpenVariable() {
+  public BaseInFlow getSingleOpenVariable() {
     return this;
   }
 
@@ -58,8 +62,9 @@
 
   @Override
   public AbstractValue evaluate(
-      IntFunction<AbstractValue> argumentAssignment, AbstractValueFactory abstractValueFactory) {
-    return argumentAssignment.apply(index);
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
+    ValueState state = flowGraphStateProvider.getState(this, () -> ValueState.bottom(getType()));
+    return state.getAbstractValue(appView);
   }
 
   @Override
@@ -78,6 +83,11 @@
   }
 
   @Override
+  public boolean isComputationLeaf() {
+    return true;
+  }
+
+  @Override
   public boolean isMethodParameter() {
     return true;
   }
@@ -93,13 +103,16 @@
   }
 
   @Override
-  @SuppressWarnings({"EqualsGetClass", "ReferenceEquality"})
+  @SuppressWarnings("EqualsGetClass")
   public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
     if (obj == null || getClass() != obj.getClass()) {
       return false;
     }
     MethodParameter methodParameter = (MethodParameter) obj;
-    return method == methodParameter.method && index == methodParameter.index;
+    return method.isIdenticalTo(methodParameter.method) && index == methodParameter.index;
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/OrAbstractFunction.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/OrAbstractFunction.java
deleted file mode 100644
index c8294a4..0000000
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/OrAbstractFunction.java
+++ /dev/null
@@ -1,96 +0,0 @@
-// Copyright (c) 2024, 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.codescanner;
-
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.ir.analysis.value.AbstractValue;
-import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
-import com.android.tools.r8.ir.analysis.value.arithmetic.AbstractCalculator;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.android.tools.r8.utils.IterableUtils;
-import java.util.Objects;
-
-/**
- * Encodes the `x | const` abstract function. This is currently used as part of the modeling of
- * updateChangedFlags, since the updateChangedFlags function is invoked with `changedFlags | 1` as
- * an argument.
- */
-public class OrAbstractFunction implements AbstractFunction {
-
-  public final BaseInFlow inFlow;
-  public final SingleNumberValue constant;
-
-  public OrAbstractFunction(BaseInFlow inFlow, SingleNumberValue constant) {
-    this.inFlow = inFlow;
-    this.constant = constant;
-  }
-
-  @Override
-  public OrAbstractFunction asOrAbstractFunction() {
-    return this;
-  }
-
-  @Override
-  public ValueState apply(
-      AppView<AppInfoWithLiveness> appView,
-      FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState inState) {
-    ConcretePrimitiveTypeValueState inPrimitiveState = inState.asPrimitiveState();
-    AbstractValue result =
-        AbstractCalculator.orIntegers(appView, inPrimitiveState.getAbstractValue(), constant);
-    return ConcretePrimitiveTypeValueState.create(result, inPrimitiveState.copyInFlow());
-  }
-
-  @Override
-  public boolean verifyContainsBaseInFlow(BaseInFlow otherInFlow) {
-    if (inFlow.isAbstractFunction()) {
-      assert inFlow.asAbstractFunction().verifyContainsBaseInFlow(otherInFlow);
-    } else {
-      assert inFlow.isBaseInFlow();
-      assert inFlow.equals(otherInFlow);
-    }
-    return true;
-  }
-
-  @Override
-  public Iterable<BaseInFlow> getBaseInFlow() {
-    if (inFlow.isAbstractFunction()) {
-      return inFlow.asAbstractFunction().getBaseInFlow();
-    }
-    return IterableUtils.singleton(inFlow);
-  }
-
-  @Override
-  public InFlowKind getKind() {
-    return InFlowKind.ABSTRACT_FUNCTION_OR;
-  }
-
-  @Override
-  public int internalCompareToSameKind(InFlow other, InFlowComparator comparator) {
-    OrAbstractFunction fn = other.asOrAbstractFunction();
-    int result = inFlow.compareTo(fn.inFlow, comparator);
-    if (result == 0) {
-      result = constant.getIntValue() - fn.constant.getIntValue();
-    }
-    return result;
-  }
-
-  @Override
-  @SuppressWarnings("EqualsGetClass")
-  public boolean equals(Object obj) {
-    if (this == obj) {
-      return true;
-    }
-    if (obj == null || getClass() != obj.getClass()) {
-      return false;
-    }
-    OrAbstractFunction fn = (OrAbstractFunction) obj;
-    return inFlow.equals(fn.inFlow) && constant.equals(fn.constant);
-  }
-
-  @Override
-  public int hashCode() {
-    return Objects.hash(getClass(), inFlow, constant);
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownAbstractFunction.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownAbstractFunction.java
deleted file mode 100644
index 726ade5..0000000
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownAbstractFunction.java
+++ /dev/null
@@ -1,53 +0,0 @@
-// Copyright (c) 2024, 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.codescanner;
-
-import com.android.tools.r8.errors.Unreachable;
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-
-public class UnknownAbstractFunction implements AbstractFunction {
-
-  private static final UnknownAbstractFunction INSTANCE = new UnknownAbstractFunction();
-
-  private UnknownAbstractFunction() {}
-
-  static UnknownAbstractFunction get() {
-    return INSTANCE;
-  }
-
-  @Override
-  public ValueState apply(
-      AppView<AppInfoWithLiveness> appView,
-      FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState inState) {
-    return ValueState.unknown();
-  }
-
-  @Override
-  public boolean verifyContainsBaseInFlow(BaseInFlow inFlow) {
-    throw new Unreachable();
-  }
-
-  @Override
-  public Iterable<BaseInFlow> getBaseInFlow() {
-    throw new Unreachable();
-  }
-
-  @Override
-  public InFlowKind getKind() {
-    return InFlowKind.ABSTRACT_FUNCTION_UNKNOWN;
-  }
-
-  @Override
-  public int internalCompareToSameKind(InFlow inFlow, InFlowComparator comparator) {
-    assert this == inFlow;
-    return 0;
-  }
-
-  @Override
-  public boolean isUnknownAbstractFunction() {
-    return true;
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedArrayTypeValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedArrayTypeValueState.java
new file mode 100644
index 0000000..e912bca
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedArrayTypeValueState.java
@@ -0,0 +1,45 @@
+// Copyright (c) 2024, 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.codescanner;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.Action;
+
+public class UnusedArrayTypeValueState extends UnusedValueState {
+
+  private static final UnusedArrayTypeValueState INSTANCE = new UnusedArrayTypeValueState();
+
+  private UnusedArrayTypeValueState() {}
+
+  public static UnusedArrayTypeValueState get() {
+    return INSTANCE;
+  }
+
+  @Override
+  public ValueState mutableJoin(
+      AppView<AppInfoWithLiveness> appView,
+      ValueState inState,
+      DexType inStaticType,
+      DexType outStaticType,
+      StateCloner cloner,
+      Action onChangedAction) {
+    if (inState.isBottom() || inState.isUnused()) {
+      return this;
+    }
+    if (inState.isUnknown()) {
+      return inState;
+    }
+    assert inState.isConcrete();
+    assert inState.asConcrete().isReferenceState();
+    return ValueState.bottomArrayTypeState()
+        .mutableJoin(appView, inState, inStaticType, outStaticType, cloner, onChangedAction);
+  }
+
+  @Override
+  public String toString() {
+    return "UNUSED(ARRAY)";
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedClassTypeValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedClassTypeValueState.java
new file mode 100644
index 0000000..9e55e43
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedClassTypeValueState.java
@@ -0,0 +1,45 @@
+// Copyright (c) 2024, 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.codescanner;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.Action;
+
+public class UnusedClassTypeValueState extends UnusedValueState {
+
+  private static final UnusedClassTypeValueState INSTANCE = new UnusedClassTypeValueState();
+
+  private UnusedClassTypeValueState() {}
+
+  public static UnusedClassTypeValueState get() {
+    return INSTANCE;
+  }
+
+  @Override
+  public ValueState mutableJoin(
+      AppView<AppInfoWithLiveness> appView,
+      ValueState inState,
+      DexType inStaticType,
+      DexType outStaticType,
+      StateCloner cloner,
+      Action onChangedAction) {
+    if (inState.isBottom() || inState.isUnused()) {
+      return this;
+    }
+    if (inState.isUnknown()) {
+      return inState;
+    }
+    assert inState.isConcrete();
+    assert inState.asConcrete().isReferenceState();
+    return ValueState.bottomClassTypeState()
+        .mutableJoin(appView, inState, inStaticType, outStaticType, cloner, onChangedAction);
+  }
+
+  @Override
+  public String toString() {
+    return "UNUSED(CLASS)";
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedPrimitiveTypeValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedPrimitiveTypeValueState.java
new file mode 100644
index 0000000..1ea268a
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedPrimitiveTypeValueState.java
@@ -0,0 +1,45 @@
+// Copyright (c) 2024, 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.codescanner;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.Action;
+
+public class UnusedPrimitiveTypeValueState extends UnusedValueState {
+
+  private static final UnusedPrimitiveTypeValueState INSTANCE = new UnusedPrimitiveTypeValueState();
+
+  private UnusedPrimitiveTypeValueState() {}
+
+  public static UnusedPrimitiveTypeValueState get() {
+    return INSTANCE;
+  }
+
+  @Override
+  public ValueState mutableJoin(
+      AppView<AppInfoWithLiveness> appView,
+      ValueState inState,
+      DexType inStaticType,
+      DexType outStaticType,
+      StateCloner cloner,
+      Action onChangedAction) {
+    if (inState.isBottom() || inState.isUnused()) {
+      return this;
+    }
+    if (inState.isUnknown()) {
+      return inState;
+    }
+    assert inState.isConcrete();
+    assert inState.asConcrete().isPrimitiveState();
+    return ValueState.bottomPrimitiveTypeState()
+        .mutableJoin(appView, inState, inStaticType, outStaticType, cloner, onChangedAction);
+  }
+
+  @Override
+  public String toString() {
+    return "UNUSED(PRIMITIVE)";
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedValueState.java
new file mode 100644
index 0000000..e4d7a81
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnusedValueState.java
@@ -0,0 +1,43 @@
+// Copyright (c) 2024, 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.codescanner;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.ir.analysis.value.AbstractValue;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+
+public abstract class UnusedValueState extends NonEmptyValueState {
+
+  UnusedValueState() {}
+
+  @Override
+  public final AbstractValue getAbstractValue(AppView<AppInfoWithLiveness> appView) {
+    return AbstractValue.bottom();
+  }
+
+  @Override
+  public final boolean isUnused() {
+    return true;
+  }
+
+  @Override
+  public final ValueState mutableCopy() {
+    return this;
+  }
+
+  @Override
+  public ValueState mutableCopyWithoutInFlow() {
+    return this;
+  }
+
+  @Override
+  public final boolean equals(Object obj) {
+    return this == obj;
+  }
+
+  @Override
+  public final int hashCode() {
+    return System.identityHashCode(this);
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ValueState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ValueState.java
index a0ed391..9876db3 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ValueState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ValueState.java
@@ -16,17 +16,20 @@
 public abstract class ValueState {
 
   public static BottomValueState bottom(ProgramField field) {
-    return bottom(field.getReference());
+    return bottom(field.getType());
   }
 
   public static BottomValueState bottom(DexField field) {
-    DexType fieldType = field.getType();
-    if (fieldType.isArrayType()) {
+    return bottom(field.getType());
+  }
+
+  public static BottomValueState bottom(DexType type) {
+    if (type.isArrayType()) {
       return bottomArrayTypeState();
-    } else if (fieldType.isClassType()) {
+    } else if (type.isClassType()) {
       return bottomClassTypeState();
     } else {
-      assert fieldType.isPrimitiveType();
+      assert type.isPrimitiveType();
       return bottomPrimitiveTypeState();
     }
   }
@@ -51,6 +54,29 @@
     return UnknownValueState.get();
   }
 
+  public static UnusedValueState unused(DexType type) {
+    if (type.isArrayType()) {
+      return unusedArrayTypeState();
+    } else if (type.isClassType()) {
+      return unusedClassTypeState();
+    } else {
+      assert type.isPrimitiveType();
+      return unusedPrimitiveTypeState();
+    }
+  }
+
+  public static UnusedArrayTypeValueState unusedArrayTypeState() {
+    return UnusedArrayTypeValueState.get();
+  }
+
+  public static UnusedClassTypeValueState unusedClassTypeState() {
+    return UnusedClassTypeValueState.get();
+  }
+
+  public static UnusedPrimitiveTypeValueState unusedPrimitiveTypeState() {
+    return UnusedPrimitiveTypeValueState.get();
+  }
+
   public abstract AbstractValue getAbstractValue(AppView<AppInfoWithLiveness> appView);
 
   public boolean isArrayState() {
@@ -117,6 +143,10 @@
     return false;
   }
 
+  public boolean isUnused() {
+    return false;
+  }
+
   public abstract ValueState mutableCopy();
 
   public abstract ValueState mutableCopyWithoutInFlow();
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComposableComputationTreeBuilder.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComposableComputationTreeBuilder.java
new file mode 100644
index 0000000..0e5f0f1
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComposableComputationTreeBuilder.java
@@ -0,0 +1,164 @@
+// Copyright (c) 2024, 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.computation;
+
+import static com.android.tools.r8.ir.code.Opcodes.AND;
+import static com.android.tools.r8.ir.code.Opcodes.ARGUMENT;
+import static com.android.tools.r8.ir.code.Opcodes.CONST_NUMBER;
+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.INVOKE_STATIC;
+import static com.android.tools.r8.ir.code.Opcodes.OR;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexClassAndMethod;
+import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.graph.DexMethod;
+import com.android.tools.r8.graph.MethodResolutionResult.SingleResolutionResult;
+import com.android.tools.r8.graph.ProgramField;
+import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.ir.analysis.path.PathConstraintSupplier;
+import com.android.tools.r8.ir.code.And;
+import com.android.tools.r8.ir.code.Argument;
+import com.android.tools.r8.ir.code.BasicBlock;
+import com.android.tools.r8.ir.code.ConstNumber;
+import com.android.tools.r8.ir.code.IRCode;
+import com.android.tools.r8.ir.code.If;
+import com.android.tools.r8.ir.code.InstanceGet;
+import com.android.tools.r8.ir.code.Instruction;
+import com.android.tools.r8.ir.code.InvokeStatic;
+import com.android.tools.r8.ir.code.Or;
+import com.android.tools.r8.ir.code.Phi;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldValueFactory;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameterFactory;
+import com.android.tools.r8.optimize.compose.ComputationTreeUnopUpdateChangedFlagsNode;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.google.common.collect.Sets;
+import java.util.Set;
+
+/**
+ * Similar to {@link DefaultComputationTreeBuilder} except that this also has support for
+ * int-valued, non-cyclic phis and logical OR instructions.
+ */
+public class ComposableComputationTreeBuilder extends ComputationTreeBuilder {
+
+  private final PathConstraintSupplier pathConstraintSupplier;
+
+  private final Set<Phi> seenPhis = Sets.newIdentityHashSet();
+
+  public ComposableComputationTreeBuilder(
+      AppView<AppInfoWithLiveness> appView,
+      IRCode code,
+      ProgramMethod method,
+      FieldValueFactory fieldValueFactory,
+      MethodParameterFactory methodParameterFactory,
+      PathConstraintSupplier pathConstraintSupplier) {
+    super(appView, code, method, fieldValueFactory, methodParameterFactory);
+    this.pathConstraintSupplier = pathConstraintSupplier;
+  }
+
+  @Override
+  ComputationTreeNode buildComputationTree(Instruction instruction) {
+    switch (instruction.opcode()) {
+      case AND:
+        {
+          And and = instruction.asAnd();
+          ComputationTreeNode left = getOrBuildComputationTree(and.leftValue());
+          ComputationTreeNode right = getOrBuildComputationTree(and.rightValue());
+          return ComputationTreeLogicalBinopAndNode.create(left, right);
+        }
+      case ARGUMENT:
+        {
+          Argument argument = instruction.asArgument();
+          if (argument.getOutType().isInt()) {
+            return methodParameterFactory.create(method, argument.getIndex());
+          }
+          break;
+        }
+      case CONST_NUMBER:
+        {
+          ConstNumber constNumber = instruction.asConstNumber();
+          if (constNumber.getOutType().isInt()) {
+            return constNumber.getAbstractValue(factory());
+          }
+          break;
+        }
+      case IF:
+        {
+          If theIf = instruction.asIf();
+          if (theIf.isZeroTest()) {
+            ComputationTreeNode operand = getOrBuildComputationTree(theIf.lhs());
+            return ComputationTreeUnopCompareNode.create(operand, theIf.getType());
+          }
+          break;
+        }
+      case INSTANCE_GET:
+        {
+          InstanceGet instanceGet = instruction.asInstanceGet();
+          DexField reference = instanceGet.getField();
+          if (instanceGet.object().isThis() && reference.getType().isIntType()) {
+            ProgramField field = instanceGet.resolveField(appView, method).getProgramField();
+            if (field != null) {
+              return fieldValueFactory.create(field);
+            }
+          }
+          break;
+        }
+      case INVOKE_STATIC:
+        {
+          InvokeStatic invoke = instruction.asInvokeStatic();
+          DexMethod reference = invoke.getInvokedMethod();
+          if (reference.getArity() != 1
+              || !reference.getParameter(0).isIntType()
+              || !reference.getReturnType().isIntType()) {
+            break;
+          }
+          SingleResolutionResult<?> resolutionResult =
+              invoke.resolveMethod(appView, method).asSingleResolution();
+          if (resolutionResult == null) {
+            break;
+          }
+          DexClassAndMethod singleTarget =
+              resolutionResult
+                  .lookupDispatchTarget(appView, invoke, method)
+                  .getSingleDispatchTarget();
+          if (singleTarget == null) {
+            break;
+          }
+          if (!singleTarget.getOptimizationInfo().getAbstractFunction().isUpdateChangedFlags()) {
+            break;
+          }
+          ComputationTreeNode operand = getOrBuildComputationTree(invoke.getFirstArgument());
+          return ComputationTreeUnopUpdateChangedFlagsNode.create(operand);
+        }
+      case OR:
+        {
+          Or or = instruction.asOr();
+          ComputationTreeNode left = getOrBuildComputationTree(or.leftValue());
+          ComputationTreeNode right = getOrBuildComputationTree(or.rightValue());
+          return ComputationTreeLogicalBinopOrNode.create(left, right);
+        }
+      default:
+        break;
+    }
+    return unknown();
+  }
+
+  @Override
+  ComputationTreeNode buildComputationTree(Phi phi) {
+    if (!seenPhis.add(phi) || phi.getOperands().size() != 2 || !phi.getType().isInt()) {
+      return unknown();
+    }
+    ComputationTreeNode left = getOrBuildComputationTree(phi.getOperand(0));
+    ComputationTreeNode right = getOrBuildComputationTree(phi.getOperand(1));
+    if (left.isUnknown() && right.isUnknown()) {
+      return unknown();
+    }
+    BasicBlock block = phi.getBlock();
+    ComputationTreeNode condition =
+        pathConstraintSupplier.getDifferentiatingPathConstraint(
+            block.getPredecessor(0), block.getPredecessor(1));
+    return ComputationTreeLogicalBinopIntPhiNode.create(condition, left, right);
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeBaseNode.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeBaseNode.java
index 2228c3e..f5c5870 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeBaseNode.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeBaseNode.java
@@ -3,14 +3,40 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.optimize.argumentpropagation.computation;
 
+import com.android.tools.r8.ir.code.Position.SourcePosition;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlow;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlowComparator;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlowKind;
+
 /**
  * Represents a computation tree with no open variables other than the arguments of a given method.
  */
 abstract class ComputationTreeBaseNode implements ComputationTreeNode {
 
   @Override
+  public InFlowKind getKind() {
+    return InFlowKind.ABSTRACT_COMPUTATION;
+  }
+
+  @Override
+  public int internalCompareToSameKind(InFlow inFlow, InFlowComparator comparator) {
+    SourcePosition position = comparator.getComputationTreePosition(this);
+    SourcePosition otherPosition =
+        comparator.getComputationTreePosition(inFlow.asAbstractComputation());
+    return position.compareTo(otherPosition);
+  }
+
+  @Override
+  public final boolean isComputationLeaf() {
+    return false;
+  }
+
+  @Override
   public abstract boolean equals(Object obj);
 
   @Override
   public abstract int hashCode();
+
+  @Override
+  public abstract String toString();
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeBuilder.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeBuilder.java
index 74804a2..b56c7c0 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeBuilder.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeBuilder.java
@@ -3,102 +3,80 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.optimize.argumentpropagation.computation;
 
-import static com.android.tools.r8.ir.code.Opcodes.AND;
-import static com.android.tools.r8.ir.code.Opcodes.ARGUMENT;
-import static com.android.tools.r8.ir.code.Opcodes.CONST_NUMBER;
-import static com.android.tools.r8.ir.code.Opcodes.IF;
-
+import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
 import com.android.tools.r8.ir.analysis.value.UnknownValue;
-import com.android.tools.r8.ir.code.And;
-import com.android.tools.r8.ir.code.Argument;
-import com.android.tools.r8.ir.code.ConstNumber;
-import com.android.tools.r8.ir.code.If;
+import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.Instruction;
+import com.android.tools.r8.ir.code.InstructionOrValue;
+import com.android.tools.r8.ir.code.Phi;
 import com.android.tools.r8.ir.code.Value;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldValueFactory;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameterFactory;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import java.util.IdentityHashMap;
 import java.util.Map;
 
-public class ComputationTreeBuilder {
+public abstract class ComputationTreeBuilder {
 
-  private final AbstractValueFactory abstractValueFactory;
-  private final ProgramMethod method;
-  private final MethodParameterFactory methodParameterFactory;
+  final AppView<AppInfoWithLiveness> appView;
+  final IRCode code;
+  final ProgramMethod method;
+  final FieldValueFactory fieldValueFactory;
+  final MethodParameterFactory methodParameterFactory;
 
-  private final Map<Instruction, ComputationTreeNode> cache = new IdentityHashMap<>();
+  private final Map<InstructionOrValue, ComputationTreeNode> cache = new IdentityHashMap<>();
 
   public ComputationTreeBuilder(
-      AbstractValueFactory abstractValueFactory,
+      AppView<AppInfoWithLiveness> appView,
+      IRCode code,
       ProgramMethod method,
+      FieldValueFactory fieldValueFactory,
       MethodParameterFactory methodParameterFactory) {
-    this.abstractValueFactory = abstractValueFactory;
+    this.appView = appView;
+    this.code = code;
     this.method = method;
+    this.fieldValueFactory = fieldValueFactory;
     this.methodParameterFactory = methodParameterFactory;
   }
 
+  AbstractValueFactory factory() {
+    return appView.abstractValueFactory();
+  }
+
   // TODO(b/302281503): "Long lived" computation trees (i.e., the ones that survive past the IR
   //  conversion of the current method) should be canonicalized.
-  public ComputationTreeNode getOrBuildComputationTree(Instruction instruction) {
-    ComputationTreeNode existing = cache.get(instruction);
+  public final ComputationTreeNode getOrBuildComputationTree(
+      InstructionOrValue instructionOrValue) {
+    ComputationTreeNode existing = cache.get(instructionOrValue);
     if (existing != null) {
       return existing;
     }
-    ComputationTreeNode result = buildComputationTree(instruction);
-    cache.put(instruction, result);
+    ComputationTreeNode result = buildComputationTree(instructionOrValue);
+    cache.put(instructionOrValue, result);
     return result;
   }
 
-  private ComputationTreeNode buildComputationTree(Instruction instruction) {
-    switch (instruction.opcode()) {
-      case AND:
-        {
-          And and = instruction.asAnd();
-          ComputationTreeNode left = buildComputationTreeFromValue(and.leftValue());
-          ComputationTreeNode right = buildComputationTreeFromValue(and.rightValue());
-          return ComputationTreeLogicalBinopAndNode.create(left, right);
-        }
-      case ARGUMENT:
-        {
-          Argument argument = instruction.asArgument();
-          if (argument.getOutType().isInt()) {
-            return methodParameterFactory.create(method, argument.getIndex());
-          }
-          break;
-        }
-      case CONST_NUMBER:
-        {
-          ConstNumber constNumber = instruction.asConstNumber();
-          if (constNumber.getOutType().isInt()) {
-            return constNumber.getAbstractValue(abstractValueFactory);
-          }
-          break;
-        }
-      case IF:
-        {
-          If theIf = instruction.asIf();
-          if (theIf.isZeroTest()) {
-            ComputationTreeNode operand = buildComputationTreeFromValue(theIf.lhs());
-            return ComputationTreeUnopCompareNode.create(operand, theIf.getType());
-          }
-          break;
-        }
-      default:
-        break;
+  private ComputationTreeNode buildComputationTree(InstructionOrValue instructionOrValue) {
+    if (instructionOrValue.isInstruction()) {
+      return buildComputationTree(instructionOrValue.asInstruction());
+    } else {
+      Value value = instructionOrValue.asValue();
+      if (value.isPhi()) {
+        return buildComputationTree(value.asPhi());
+      } else {
+        return buildComputationTree(value.getDefinition());
+      }
     }
-    return AbstractValue.unknown();
   }
 
-  private ComputationTreeNode buildComputationTreeFromValue(Value value) {
-    if (value.isPhi()) {
-      return unknown();
-    }
-    return getOrBuildComputationTree(value.getDefinition());
-  }
+  abstract ComputationTreeNode buildComputationTree(Instruction instruction);
 
-  private static UnknownValue unknown() {
+  abstract ComputationTreeNode buildComputationTree(Phi phi);
+
+  static UnknownValue unknown() {
     return AbstractValue.unknown();
   }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopAndNode.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopAndNode.java
index f18f081..9c60201 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopAndNode.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopAndNode.java
@@ -3,11 +3,12 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.optimize.argumentpropagation.computation;
 
+import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
-import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
 import com.android.tools.r8.ir.analysis.value.arithmetic.AbstractCalculator;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import java.util.Objects;
-import java.util.function.IntFunction;
 
 public class ComputationTreeLogicalBinopAndNode extends ComputationTreeLogicalBinopNode {
 
@@ -24,11 +25,17 @@
 
   @Override
   public AbstractValue evaluate(
-      IntFunction<AbstractValue> argumentAssignment, AbstractValueFactory abstractValueFactory) {
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
     assert getNumericType().isInt();
-    AbstractValue leftValue = left.evaluate(argumentAssignment, abstractValueFactory);
-    AbstractValue rightValue = right.evaluate(argumentAssignment, abstractValueFactory);
-    return AbstractCalculator.andIntegers(abstractValueFactory, leftValue, rightValue);
+    AbstractValue leftValue = left.evaluate(appView, flowGraphStateProvider);
+    if (leftValue.isBottom()) {
+      return leftValue;
+    }
+    AbstractValue rightValue = right.evaluate(appView, flowGraphStateProvider);
+    if (rightValue.isBottom()) {
+      return rightValue;
+    }
+    return AbstractCalculator.andIntegers(appView, leftValue, rightValue);
   }
 
   @Override
@@ -47,4 +54,9 @@
   public int hashCode() {
     return Objects.hash(getClass(), left, right);
   }
+
+  @Override
+  public String toString() {
+    return left.toStringWithParenthesis() + " & " + right.toStringWithParenthesis();
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopIntPhiNode.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopIntPhiNode.java
new file mode 100644
index 0000000..fb42934
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopIntPhiNode.java
@@ -0,0 +1,89 @@
+// Copyright (c) 2024, 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.computation;
+
+import com.android.tools.r8.graph.AppView;
+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.AbstractValueJoiner.AbstractValueConstantPropagationJoiner;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.Objects;
+import java.util.function.Function;
+
+public class ComputationTreeLogicalBinopIntPhiNode extends ComputationTreeLogicalBinopNode {
+
+  private final ComputationTreeNode condition;
+
+  private ComputationTreeLogicalBinopIntPhiNode(
+      ComputationTreeNode condition, ComputationTreeNode left, ComputationTreeNode right) {
+    super(left, right);
+    this.condition = condition;
+  }
+
+  public static ComputationTreeNode create(
+      ComputationTreeNode condition, ComputationTreeNode left, ComputationTreeNode right) {
+    if (left.isUnknown() && right.isUnknown()) {
+      return AbstractValue.unknown();
+    }
+    return new ComputationTreeLogicalBinopIntPhiNode(condition, left, right);
+  }
+
+  @Override
+  public AbstractValue evaluate(
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
+    AbstractValue result = condition.evaluate(appView, flowGraphStateProvider);
+    if (result.isBottom()) {
+      return AbstractValue.bottom();
+    } else if (result.isTrue()) {
+      return left.evaluate(appView, flowGraphStateProvider);
+    } else if (result.isFalse()) {
+      return right.evaluate(appView, flowGraphStateProvider);
+    } else {
+      AbstractValueConstantPropagationJoiner joiner =
+          appView.getAbstractValueConstantPropagationJoiner();
+      AbstractValue leftValue = left.evaluate(appView, flowGraphStateProvider);
+      AbstractValue rightValue = right.evaluate(appView, flowGraphStateProvider);
+      return joiner.join(leftValue, rightValue, TypeElement.getInt());
+    }
+  }
+
+  @Override
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    TraversalContinuation<TB, TC> traversalContinuation = condition.traverseBaseInFlow(fn);
+    if (traversalContinuation.shouldContinue()) {
+      traversalContinuation = super.traverseBaseInFlow(fn);
+    }
+    return traversalContinuation;
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
+    if (!(obj instanceof ComputationTreeLogicalBinopIntPhiNode)) {
+      return false;
+    }
+    ComputationTreeLogicalBinopIntPhiNode node = (ComputationTreeLogicalBinopIntPhiNode) obj;
+    return condition.equals(node.condition) && internalIsEqualTo(node);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(getClass(), condition, left, right);
+  }
+
+  @Override
+  public String toString() {
+    return condition.toStringWithParenthesis()
+        + " ? "
+        + left.toStringWithParenthesis()
+        + " : "
+        + right.toStringWithParenthesis();
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopNode.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopNode.java
index d7cf46b..a2504cb 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopNode.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopNode.java
@@ -4,7 +4,9 @@
 package com.android.tools.r8.optimize.argumentpropagation.computation;
 
 import com.android.tools.r8.ir.code.NumericType;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.function.Function;
 
 public abstract class ComputationTreeLogicalBinopNode extends ComputationTreeBaseNode {
 
@@ -17,13 +19,23 @@
     this.right = right;
   }
 
+  @Override
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    TraversalContinuation<TB, TC> traversalContinuation = left.traverseBaseInFlow(fn);
+    if (traversalContinuation.shouldContinue()) {
+      traversalContinuation = right.traverseBaseInFlow(fn);
+    }
+    return traversalContinuation;
+  }
+
   public NumericType getNumericType() {
     return NumericType.INT;
   }
 
   @Override
-  public final MethodParameter getSingleOpenVariable() {
-    MethodParameter openVariable = left.getSingleOpenVariable();
+  public final BaseInFlow getSingleOpenVariable() {
+    BaseInFlow openVariable = left.getSingleOpenVariable();
     if (openVariable != null) {
       return right.getSingleOpenVariable() == null ? openVariable : null;
     }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopOrNode.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopOrNode.java
new file mode 100644
index 0000000..3d7c3fe
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeLogicalBinopOrNode.java
@@ -0,0 +1,62 @@
+// Copyright (c) 2024, 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.computation;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.ir.analysis.value.AbstractValue;
+import com.android.tools.r8.ir.analysis.value.arithmetic.AbstractCalculator;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import java.util.Objects;
+
+public class ComputationTreeLogicalBinopOrNode extends ComputationTreeLogicalBinopNode {
+
+  private ComputationTreeLogicalBinopOrNode(ComputationTreeNode left, ComputationTreeNode right) {
+    super(left, right);
+  }
+
+  public static ComputationTreeNode create(ComputationTreeNode left, ComputationTreeNode right) {
+    if (left.isUnknown() && right.isUnknown()) {
+      return AbstractValue.unknown();
+    }
+    return new ComputationTreeLogicalBinopOrNode(left, right);
+  }
+
+  @Override
+  public AbstractValue evaluate(
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
+    assert getNumericType().isInt();
+    AbstractValue leftValue = left.evaluate(appView, flowGraphStateProvider);
+    if (leftValue.isBottom()) {
+      return leftValue;
+    }
+    AbstractValue rightValue = right.evaluate(appView, flowGraphStateProvider);
+    if (rightValue.isBottom()) {
+      return rightValue;
+    }
+    return AbstractCalculator.orIntegers(appView, leftValue, rightValue);
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
+    if (!(obj instanceof ComputationTreeLogicalBinopOrNode)) {
+      return false;
+    }
+    ComputationTreeLogicalBinopOrNode node = (ComputationTreeLogicalBinopOrNode) obj;
+    return internalIsEqualTo(node);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(getClass(), left, right);
+  }
+
+  @Override
+  public String toString() {
+    return left.toStringWithParenthesis() + " | " + right.toStringWithParenthesis();
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeNode.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeNode.java
index 1231d76..a22b633 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeNode.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeNode.java
@@ -3,23 +3,80 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.optimize.argumentpropagation.computation;
 
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.ir.analysis.type.DynamicType;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
-import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
-import java.util.function.IntFunction;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.AbstractFunction;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteClassTypeValueState;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcretePrimitiveTypeValueState;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteValueState;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.ValueState;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
 /**
  * Represents a computation tree with no open variables other than the arguments of a given method.
  */
-public interface ComputationTreeNode {
+public interface ComputationTreeNode extends AbstractFunction {
+
+  @Override
+  default ValueState apply(
+      AppView<AppInfoWithLiveness> appView,
+      FlowGraphStateProvider flowGraphStateProvider,
+      ConcreteValueState inState,
+      DexType outStaticType) {
+    AbstractValue abstractValue = evaluate(appView, flowGraphStateProvider);
+    if (abstractValue.isBottom()) {
+      return ValueState.bottom(outStaticType);
+    } else if (abstractValue.isUnknown()) {
+      return ValueState.unknown();
+    } else {
+      if (outStaticType.isArrayType()) {
+        // We currently do not track abstract values for array types.
+        return ValueState.unknown();
+      } else if (outStaticType.isClassType()) {
+        return ConcreteClassTypeValueState.create(abstractValue, DynamicType.unknown());
+      } else {
+        assert outStaticType.isPrimitiveType();
+        return ConcretePrimitiveTypeValueState.create(abstractValue);
+      }
+    }
+  }
 
   /** Evaluates the current computation tree on the given argument assignment. */
   AbstractValue evaluate(
-      IntFunction<AbstractValue> argumentAssignment, AbstractValueFactory abstractValueFactory);
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider);
 
-  MethodParameter getSingleOpenVariable();
+  BaseInFlow getSingleOpenVariable();
 
-  default boolean isUnknown() {
+  @Override
+  default boolean isAbstractComputation() {
+    return true;
+  }
+
+  @Override
+  default ComputationTreeNode asAbstractComputation() {
+    return this;
+  }
+
+  default boolean isArgumentBitSetCompareNode() {
     return false;
   }
+
+  boolean isComputationLeaf();
+
+  @Override
+  default boolean usesFlowGraphStateProvider() {
+    return true;
+  }
+
+  default String toStringWithParenthesis() {
+    if (isComputationLeaf()) {
+      return toString();
+    } else {
+      return "(" + this + ")";
+    }
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeUnopCompareNode.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeUnopCompareNode.java
index fb4c465..7b4705f 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeUnopCompareNode.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeUnopCompareNode.java
@@ -3,11 +3,14 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.optimize.argumentpropagation.computation;
 
+import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.ir.analysis.value.AbstractValue;
-import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
+import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
 import com.android.tools.r8.ir.code.IfType;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import java.util.Objects;
-import java.util.function.IntFunction;
 
 public class ComputationTreeUnopCompareNode extends ComputationTreeUnopNode {
 
@@ -27,9 +30,26 @@
 
   @Override
   public AbstractValue evaluate(
-      IntFunction<AbstractValue> argumentAssignment, AbstractValueFactory abstractValueFactory) {
-    AbstractValue operandValue = operand.evaluate(argumentAssignment, abstractValueFactory);
-    return type.evaluate(operandValue, abstractValueFactory);
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
+    AbstractValue operandValue = operand.evaluate(appView, flowGraphStateProvider);
+    if (operandValue.isBottom()) {
+      return operandValue;
+    }
+    return type.evaluate(operandValue, appView);
+  }
+
+  @Override
+  public boolean isArgumentBitSetCompareNode() {
+    if (!type.isEqualsOrNotEquals() || !(operand instanceof ComputationTreeLogicalBinopAndNode)) {
+      return false;
+    }
+    ComputationTreeLogicalBinopAndNode andOperand = (ComputationTreeLogicalBinopAndNode) operand;
+    return andOperand.left instanceof MethodParameter
+        && andOperand.right instanceof SingleNumberValue;
+  }
+
+  public ComputationTreeUnopCompareNode negate() {
+    return new ComputationTreeUnopCompareNode(operand, type.inverted());
   }
 
   @Override
@@ -48,4 +68,9 @@
   public int hashCode() {
     return Objects.hash(getClass(), operand, type);
   }
+
+  @Override
+  public String toString() {
+    return operand.toStringWithParenthesis() + " " + type.getSymbol() + " 0";
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeUnopNode.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeUnopNode.java
index 8ac1ca9..71bd020 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeUnopNode.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/ComputationTreeUnopNode.java
@@ -3,19 +3,27 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.optimize.argumentpropagation.computation;
 
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.function.Function;
 
 public abstract class ComputationTreeUnopNode extends ComputationTreeBaseNode {
 
-  final ComputationTreeNode operand;
+  protected final ComputationTreeNode operand;
 
-  ComputationTreeUnopNode(ComputationTreeNode operand) {
+  protected ComputationTreeUnopNode(ComputationTreeNode operand) {
     assert !operand.isUnknown();
     this.operand = operand;
   }
 
   @Override
-  public MethodParameter getSingleOpenVariable() {
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    return operand.traverseBaseInFlow(fn);
+  }
+
+  @Override
+  public BaseInFlow getSingleOpenVariable() {
     return operand.getSingleOpenVariable();
   }
 
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/DefaultComputationTreeBuilder.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/DefaultComputationTreeBuilder.java
new file mode 100644
index 0000000..e34cc2c
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/computation/DefaultComputationTreeBuilder.java
@@ -0,0 +1,80 @@
+// Copyright (c) 2024, 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.computation;
+
+import static com.android.tools.r8.ir.code.Opcodes.AND;
+import static com.android.tools.r8.ir.code.Opcodes.ARGUMENT;
+import static com.android.tools.r8.ir.code.Opcodes.CONST_NUMBER;
+import static com.android.tools.r8.ir.code.Opcodes.IF;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.ir.code.And;
+import com.android.tools.r8.ir.code.Argument;
+import com.android.tools.r8.ir.code.ConstNumber;
+import com.android.tools.r8.ir.code.IRCode;
+import com.android.tools.r8.ir.code.If;
+import com.android.tools.r8.ir.code.Instruction;
+import com.android.tools.r8.ir.code.Phi;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldValueFactory;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameterFactory;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+
+public class DefaultComputationTreeBuilder extends ComputationTreeBuilder {
+
+  public DefaultComputationTreeBuilder(
+      AppView<AppInfoWithLiveness> appView,
+      IRCode code,
+      ProgramMethod method,
+      FieldValueFactory fieldValueFactory,
+      MethodParameterFactory methodParameterFactory) {
+    super(appView, code, method, fieldValueFactory, methodParameterFactory);
+  }
+
+  @Override
+  ComputationTreeNode buildComputationTree(Instruction instruction) {
+    switch (instruction.opcode()) {
+      case AND:
+        {
+          And and = instruction.asAnd();
+          ComputationTreeNode left = getOrBuildComputationTree(and.leftValue());
+          ComputationTreeNode right = getOrBuildComputationTree(and.rightValue());
+          return ComputationTreeLogicalBinopAndNode.create(left, right);
+        }
+      case ARGUMENT:
+        {
+          Argument argument = instruction.asArgument();
+          if (argument.getOutType().isInt()) {
+            return methodParameterFactory.create(method, argument.getIndex());
+          }
+          break;
+        }
+      case CONST_NUMBER:
+        {
+          ConstNumber constNumber = instruction.asConstNumber();
+          if (constNumber.getOutType().isInt()) {
+            return constNumber.getAbstractValue(factory());
+          }
+          break;
+        }
+      case IF:
+        {
+          If theIf = instruction.asIf();
+          if (theIf.isZeroTest()) {
+            ComputationTreeNode operand = getOrBuildComputationTree(theIf.lhs());
+            return ComputationTreeUnopCompareNode.create(operand, theIf.getType());
+          }
+          break;
+        }
+      default:
+        break;
+    }
+    return unknown();
+  }
+
+  @Override
+  ComputationTreeNode buildComputationTree(Phi phi) {
+    return unknown();
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/FlowGraphBuilder.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/FlowGraphBuilder.java
index 5830d3f..c8389a1 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/FlowGraphBuilder.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/FlowGraphBuilder.java
@@ -16,7 +16,6 @@
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.conversion.IRConverter;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.AbstractFunction;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteMonomorphicMethodState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteValueState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldStateCollection;
@@ -159,38 +158,34 @@
 
   // Returns BREAK if the current node has been set to unknown.
   private TraversalContinuation<?, ?> addInFlow(InFlow inFlow, FlowGraphNode node) {
-    if (inFlow.isAbstractFunction()) {
-      return addInFlow(inFlow.asAbstractFunction(), node);
-    } else if (inFlow.isFieldValue()) {
+    if (inFlow.isFieldValue()) {
       return addInFlow(inFlow.asFieldValue(), node);
     } else if (inFlow.isMethodParameter()) {
       return addInFlow(inFlow.asMethodParameter(), node);
+    } else if (inFlow.isAbstractFunction()) {
+      return addInFlow(inFlow.asAbstractFunction(), node);
     } else {
       throw new Unreachable(inFlow.getClass().getTypeName());
     }
   }
 
   private TraversalContinuation<?, ?> addInFlow(AbstractFunction inFlow, FlowGraphNode node) {
-    for (BaseInFlow baseInFlow : inFlow.getBaseInFlow()) {
-      TraversalContinuation<?, ?> traversalContinuation;
-      if (baseInFlow.isFieldValue()) {
-        traversalContinuation = addInFlow(baseInFlow.asFieldValue(), node, inFlow);
-      } else {
-        assert baseInFlow.isMethodParameter();
-        traversalContinuation = addInFlow(baseInFlow.asMethodParameter(), node, inFlow);
-      }
-      if (traversalContinuation.shouldBreak()) {
-        return traversalContinuation;
-      }
-    }
-    return TraversalContinuation.doContinue();
+    return inFlow.traverseBaseInFlow(
+        baseInFlow -> {
+          if (baseInFlow.isFieldValue()) {
+            return addInFlow(baseInFlow.asFieldValue(), node, inFlow);
+          } else {
+            assert baseInFlow.isMethodParameter();
+            return addInFlow(baseInFlow.asMethodParameter(), node, inFlow);
+          }
+        });
   }
 
-  private TraversalContinuation<?, ?> addInFlow(FieldValue inFlow, FlowGraphNode node) {
+  private <TB, TC> TraversalContinuation<TB, TC> addInFlow(FieldValue inFlow, FlowGraphNode node) {
     return addInFlow(inFlow, node, AbstractFunction.identity());
   }
 
-  private TraversalContinuation<?, ?> addInFlow(
+  private <TB, TC> TraversalContinuation<TB, TC> addInFlow(
       FieldValue inFlow, FlowGraphNode node, AbstractFunction transferFunction) {
     assert !node.isUnknown();
 
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InFlowPropagator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InFlowPropagator.java
index 82f2350..fef91ce 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InFlowPropagator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InFlowPropagator.java
@@ -253,7 +253,8 @@
       FlowGraphStateProvider flowGraphStateProvider =
           FlowGraphStateProvider.create(flowGraph, transferFunction);
       ValueState transferState =
-          transferFunction.apply(appView, flowGraphStateProvider, stateToPropagate);
+          transferFunction.apply(
+              appView, flowGraphStateProvider, stateToPropagate, successorNode.getStaticType());
       ValueState oldSuccessorStateForDebugging =
           successorNode.getDebug() ? successorNode.getState().mutableCopy() : null;
       if (transferState.isBottom()) {
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InFlowPropagatorDebugUtils.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InFlowPropagatorDebugUtils.java
index a39b534..7687b99 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InFlowPropagatorDebugUtils.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InFlowPropagatorDebugUtils.java
@@ -6,11 +6,11 @@
 import com.android.tools.r8.graph.ProgramField;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.AbstractFunction;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteValueState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ValueState;
 import com.android.tools.r8.utils.StringUtils;
+import com.android.tools.r8.utils.TraversalContinuation;
 import com.android.tools.r8.utils.WorkList;
 import java.util.ArrayList;
 import java.util.List;
@@ -86,14 +86,16 @@
       List<String> transferFunctionDependencies = new ArrayList<>();
       transferFunctionDependencies.add("");
       transferFunctionDependencies.add("TRANSFER FN INPUTS:");
-      for (BaseInFlow transferFunctionDependency : transferFunction.getBaseInFlow()) {
-        if (!node.equalsBaseInFlow(transferFunctionDependency)) {
-          ValueState transferFunctionDependencyState =
-              flowGraphStateProvider.getState(transferFunctionDependency, null);
-          transferFunctionDependencies.add("  DEP: " + transferFunctionDependency);
-          transferFunctionDependencies.add("  DEP STATE: " + transferFunctionDependencyState);
-        }
-      }
+      transferFunction.traverseBaseInFlow(
+          transferFunctionDependency -> {
+            if (!node.equalsBaseInFlow(transferFunctionDependency)) {
+              ValueState transferFunctionDependencyState =
+                  flowGraphStateProvider.getState(transferFunctionDependency, null);
+              transferFunctionDependencies.add("  DEP: " + transferFunctionDependency);
+              transferFunctionDependencies.add("  DEP STATE: " + transferFunctionDependencyState);
+            }
+            return TraversalContinuation.doContinue();
+          });
       ValueState newSuccessorState = successorNode.getState();
       log(
           "PROPAGATE CONCRETE",
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/unusedarguments/EffectivelyUnusedArgumentsAnalysis.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/unusedarguments/EffectivelyUnusedArgumentsAnalysis.java
index 1eec104..564ef5f 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/unusedarguments/EffectivelyUnusedArgumentsAnalysis.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/unusedarguments/EffectivelyUnusedArgumentsAnalysis.java
@@ -11,12 +11,18 @@
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.graph.PrunedItems;
+import com.android.tools.r8.ir.analysis.path.PathConstraintSupplier;
+import com.android.tools.r8.ir.analysis.path.state.ConcretePathConstraintAnalysisState;
 import com.android.tools.r8.ir.code.Argument;
+import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.Instruction;
 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.optimize.argumentpropagation.codescanner.MethodParameter;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeUnopCompareNode;
 import com.android.tools.r8.optimize.argumentpropagation.utils.ParameterRemovalUtils;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.ListUtils;
@@ -28,6 +34,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Consumer;
 
 /**
  * Analysis to find arguments that are effectively unused. The analysis first computes the
@@ -68,6 +75,14 @@
 
   private final AppView<AppInfoWithLiveness> appView;
 
+  // If the condition for a given method parameter evaluates to true then the parameter is unused,
+  // regardless (!) of the constraints for the method parameter in `constraints`. If the condition
+  // evaluates to false, the method parameter can still be unused if all the constraints for the
+  // method parameter are unused.
+  //
+  // All constraints in this map are currently on the form `(arg & const) != 0`.
+  private final Map<MethodParameter, ComputationTreeNode> conditions = new ConcurrentHashMap<>();
+
   // Maps each method parameter p to the method parameters that must be effectively unused in order
   // for the method parameter p to be effectively unused.
   private final Map<MethodParameter, Set<MethodParameter>> constraints = new ConcurrentHashMap<>();
@@ -107,72 +122,16 @@
         });
   }
 
-  public void scan(ProgramMethod method, IRCode code) {
-    // If this method is not subject to optimization, then don't compute effectively unused
-    // constraints for the method parameters.
-    if (isUnoptimizable(method)) {
-      return;
-    }
-    Iterator<Argument> argumentIterator = code.argumentIterator();
-    while (argumentIterator.hasNext()) {
-      Argument argument = argumentIterator.next();
-      Value argumentValue = argument.outValue();
-      Set<MethodParameter> effectivelyUnusedConstraints =
-          computeEffectivelyUnusedConstraints(method, argument, argumentValue);
-      if (effectivelyUnusedConstraints != null && !effectivelyUnusedConstraints.isEmpty()) {
-        MethodParameter methodParameter = new MethodParameter(method, argument.getIndex());
-        assert !constraints.containsKey(methodParameter);
-        constraints.put(methodParameter, effectivelyUnusedConstraints);
-      }
-    }
-  }
-
-  private Set<MethodParameter> computeEffectivelyUnusedConstraints(
-      ProgramMethod method, Argument argument, Value argumentValue) {
-    if (method.getDefinition().isInstanceInitializer() && argumentValue.isThis()) {
-      return null;
-    }
-    if (!ParameterRemovalUtils.canRemoveUnusedParameter(appView, method, argument.getIndex())) {
-      return null;
-    }
-    if (!argumentValue.getType().isClassType()
-        || argumentValue.hasDebugUsers()
-        || argumentValue.hasPhiUsers()) {
-      return null;
-    }
-    Set<MethodParameter> effectivelyUnusedConstraints = new HashSet<>();
-    for (Instruction user : argumentValue.uniqueUsers()) {
-      if (user.isInvokeMethod()) {
-        InvokeMethod invoke = user.asInvokeMethod();
-        ProgramMethod resolvedMethod =
-            appView
-                .appInfo()
-                .unsafeResolveMethodDueToDexFormatLegacy(invoke.getInvokedMethod())
-                .getResolvedProgramMethod();
-        if (resolvedMethod == null || isUnoptimizable(resolvedMethod)) {
-          return null;
-        }
-        int dependentArgumentIndex =
-            ListUtils.uniqueIndexMatching(invoke.arguments(), value -> value == argumentValue);
-        if (dependentArgumentIndex < 0
-            || !ParameterRemovalUtils.canRemoveUnusedParameter(
-                appView, resolvedMethod, dependentArgumentIndex)) {
-          return null;
-        }
-        effectivelyUnusedConstraints.add(
-            new MethodParameter(resolvedMethod, dependentArgumentIndex));
-      } else {
-        return null;
-      }
-    }
-    return effectivelyUnusedConstraints;
+  public void scan(
+      ProgramMethod method, IRCode code, PathConstraintSupplier pathConstraintSupplier) {
+    new Analyzer(method, code, pathConstraintSupplier).analyze();
   }
 
   public void computeEffectivelyUnusedArguments(PrunedItems prunedItems) {
     // Build a graph where nodes are method parameters and there is an edge from method parameter p0
     // to method parameter p1 if the removal of p0 depends on the removal of p1.
     EffectivelyUnusedArgumentsGraph dependenceGraph =
-        EffectivelyUnusedArgumentsGraph.create(appView, constraints, prunedItems);
+        EffectivelyUnusedArgumentsGraph.create(appView, conditions, constraints, prunedItems);
 
     // Remove all unoptimizable method parameters from the graph, as well as all nodes that depend
     // on a node that is unoptimable.
@@ -202,17 +161,6 @@
     }
   }
 
-  private boolean isUnoptimizable(ProgramMethod method) {
-    if (method.getDefinition().belongsToDirectPool()) {
-      return !ParameterRemovalUtils.canRemoveUnusedParametersFrom(appView, method);
-    }
-    if (optimizableVirtualMethods.contains(method)) {
-      assert ParameterRemovalUtils.canRemoveUnusedParametersFrom(appView, method);
-      return false;
-    }
-    return true;
-  }
-
   public void onMethodPruned(ProgramMethod method) {
     onMethodCodePruned(method);
   }
@@ -221,7 +169,154 @@
     for (int argumentIndex = 0;
         argumentIndex < method.getDefinition().getNumberOfArguments();
         argumentIndex++) {
-      constraints.remove(new MethodParameter(method, argumentIndex));
+      MethodParameter methodParameter = new MethodParameter(method, argumentIndex);
+      conditions.remove(methodParameter);
+      constraints.remove(methodParameter);
+    }
+  }
+
+  private class Analyzer {
+
+    private final ProgramMethod method;
+    private final IRCode code;
+    private final PathConstraintSupplier pathConstraintSupplier;
+
+    private Analyzer(
+        ProgramMethod method, IRCode code, PathConstraintSupplier pathConstraintSupplier) {
+      this.method = method;
+      this.code = code;
+      this.pathConstraintSupplier = pathConstraintSupplier;
+    }
+
+    void analyze() {
+      // If this method is not subject to optimization, then don't compute effectively unused
+      // constraints for the method parameters.
+      if (isUnoptimizable(method)) {
+        return;
+      }
+      Iterator<Argument> argumentIterator = code.argumentIterator();
+      while (argumentIterator.hasNext()) {
+        Argument argument = argumentIterator.next();
+        Value argumentValue = argument.outValue();
+        computeEffectivelyUnusedConstraints(
+            argument,
+            argumentValue,
+            effectivelyUnusedCondition -> {
+              MethodParameter methodParameter = new MethodParameter(method, argument.getIndex());
+              assert !conditions.containsKey(methodParameter);
+              conditions.put(methodParameter, effectivelyUnusedCondition);
+            },
+            effectivelyUnusedConstraints -> {
+              MethodParameter methodParameter = new MethodParameter(method, argument.getIndex());
+              assert !constraints.containsKey(methodParameter);
+              constraints.put(methodParameter, effectivelyUnusedConstraints);
+            });
+      }
+    }
+
+    private void computeEffectivelyUnusedConstraints(
+        Argument argument,
+        Value argumentValue,
+        Consumer<ComputationTreeNode> effectivelyUnusedConditionsConsumer,
+        Consumer<Set<MethodParameter>> effectivelyUnusedConstraintsConsumer) {
+      if (!ParameterRemovalUtils.canRemoveUnusedParameter(appView, method, argument.getIndex())) {
+        return;
+      }
+      if (!argumentValue.getType().isClassType() || argumentValue.hasDebugUsers()) {
+        return;
+      }
+      Value usedValue;
+      if (argumentValue.hasPhiUsers()) {
+        // If the argument has one or more phi users, we check if there is a single phi due to a
+        // default value for the argument. If so, we record this condition and mark the users of the
+        // phi as effectively unused argument constraints for the current argument.
+        if (!computeEffectivelyUnusedCondition(
+            argumentValue, effectivelyUnusedConditionsConsumer)) {
+          return;
+        }
+        assert argumentValue.hasSingleUniquePhiUser();
+        Phi user = argumentValue.singleUniquePhiUser();
+        if (user.hasDebugUsers() || user.hasPhiUsers()) {
+          return;
+        }
+        usedValue = user;
+      } else {
+        usedValue = argumentValue;
+      }
+      Set<MethodParameter> effectivelyUnusedConstraints = new HashSet<>();
+      for (Instruction user : usedValue.uniqueUsers()) {
+        if (user.isInvokeMethod()) {
+          InvokeMethod invoke = user.asInvokeMethod();
+          ProgramMethod resolvedMethod =
+              appView
+                  .appInfo()
+                  .unsafeResolveMethodDueToDexFormatLegacy(invoke.getInvokedMethod())
+                  .getResolvedProgramMethod();
+          if (resolvedMethod == null || isUnoptimizable(resolvedMethod)) {
+            return;
+          }
+          int dependentArgumentIndex =
+              ListUtils.uniqueIndexMatching(invoke.arguments(), value -> value == argumentValue);
+          if (dependentArgumentIndex < 0
+              || !ParameterRemovalUtils.canRemoveUnusedParameter(
+                  appView, resolvedMethod, dependentArgumentIndex)) {
+            return;
+          }
+          effectivelyUnusedConstraints.add(
+              new MethodParameter(resolvedMethod, dependentArgumentIndex));
+        } else {
+          return;
+        }
+      }
+      if (!effectivelyUnusedConstraints.isEmpty()) {
+        effectivelyUnusedConstraintsConsumer.accept(effectivelyUnusedConstraints);
+      }
+    }
+
+    private boolean computeEffectivelyUnusedCondition(
+        Value argumentValue, Consumer<ComputationTreeNode> effectivelyUnusedConditionsConsumer) {
+      assert argumentValue.hasPhiUsers();
+      if (argumentValue.hasUsers() || !argumentValue.hasSingleUniquePhiUser()) {
+        return false;
+      }
+      Phi phi = argumentValue.singleUniquePhiUser();
+      if (phi.getOperands().size() != 2) {
+        return false;
+      }
+      BasicBlock block = phi.getBlock();
+      ConcretePathConstraintAnalysisState leftState =
+          pathConstraintSupplier.getPathConstraint(block.getPredecessor(0)).asConcreteState();
+      ConcretePathConstraintAnalysisState rightState =
+          pathConstraintSupplier.getPathConstraint(block.getPredecessor(1)).asConcreteState();
+      if (leftState == null || rightState == null) {
+        return false;
+      }
+      // Find a condition that can be used to distinguish program paths coming from the two
+      // predecessors.
+      ComputationTreeNode condition = leftState.getDifferentiatingPathConstraint(rightState);
+      if (!condition.isArgumentBitSetCompareNode()) {
+        return false;
+      }
+      // Extract the state corresponding to the program path where the argument is unused. If the
+      // condition evaluates to false on this program path then negate the condition.
+      ConcretePathConstraintAnalysisState unusedState =
+          phi.getOperand(0) == argumentValue ? rightState : leftState;
+      if (unusedState.isNegated(condition)) {
+        condition = ((ComputationTreeUnopCompareNode) condition).negate();
+      }
+      effectivelyUnusedConditionsConsumer.accept(condition);
+      return true;
+    }
+
+    private boolean isUnoptimizable(ProgramMethod method) {
+      if (method.getDefinition().belongsToDirectPool()) {
+        return !ParameterRemovalUtils.canRemoveUnusedParametersFrom(appView, method);
+      }
+      if (optimizableVirtualMethods.contains(method)) {
+        assert ParameterRemovalUtils.canRemoveUnusedParametersFrom(appView, method);
+        return false;
+      }
+      return true;
     }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/unusedarguments/EffectivelyUnusedArgumentsGraph.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/unusedarguments/EffectivelyUnusedArgumentsGraph.java
index 1f795ad..c2fa1f0 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/unusedarguments/EffectivelyUnusedArgumentsGraph.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/unusedarguments/EffectivelyUnusedArgumentsGraph.java
@@ -9,8 +9,11 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.graph.PrunedItems;
+import com.android.tools.r8.ir.analysis.value.AbstractValue;
 import com.android.tools.r8.ir.optimize.info.MethodOptimizationInfo;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.WorkList;
 import com.android.tools.r8.utils.dfs.DFSStack;
@@ -38,9 +41,34 @@
 
   public static EffectivelyUnusedArgumentsGraph create(
       AppView<AppInfoWithLiveness> appView,
+      Map<MethodParameter, ComputationTreeNode> conditions,
       Map<MethodParameter, Set<MethodParameter>> constraints,
       PrunedItems prunedItems) {
     EffectivelyUnusedArgumentsGraph graph = new EffectivelyUnusedArgumentsGraph(appView);
+    conditions.forEach(
+        (methodParameter, condition) -> {
+          ProgramMethod method =
+              asProgramMethodOrNull(appView.definitionFor(methodParameter.getMethod()));
+          if (method == null) {
+            assert false;
+            return;
+          }
+          // Evaluate the condition. If the condition evaluates to true, then create a graph node
+          // for the method parameter and delete all constraints. As a result, the node will not
+          // have any successors, meaning it is effectively unused.
+          if (method
+              .getOptimizationInfo()
+              .getArgumentInfos()
+              .isConcreteCallSiteOptimizationInfo()) {
+            FlowGraphStateProvider flowGraphStateProvider =
+                FlowGraphStateProvider.createFromMethodOptimizationInfo(method);
+            AbstractValue result = condition.evaluate(appView, flowGraphStateProvider);
+            if (result.isTrue()) {
+              graph.getOrCreateNode(methodParameter);
+              constraints.remove(methodParameter);
+            }
+          }
+        });
     constraints.forEach(
         (methodParameter, constraintsForMethodParameter) -> {
           EffectivelyUnusedArgumentsGraphNode node = graph.getOrCreateNode(methodParameter);
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorCodeScannerForComposableFunctions.java b/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorCodeScannerForComposableFunctions.java
deleted file mode 100644
index 371af3d..0000000
--- a/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorCodeScannerForComposableFunctions.java
+++ /dev/null
@@ -1,65 +0,0 @@
-// Copyright (c) 2023, 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.compose;
-
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.ir.analysis.path.PathConstraintSupplier;
-import com.android.tools.r8.ir.code.AbstractValueSupplier;
-import com.android.tools.r8.ir.code.IRCode;
-import com.android.tools.r8.ir.code.InvokeMethod;
-import com.android.tools.r8.optimize.argumentpropagation.ArgumentPropagatorCodeScanner;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.android.tools.r8.utils.Timing;
-
-public class ArgumentPropagatorCodeScannerForComposableFunctions
-    extends ArgumentPropagatorCodeScanner {
-
-  private final ComposableCallGraph callGraph;
-
-  public ArgumentPropagatorCodeScannerForComposableFunctions(
-      AppView<AppInfoWithLiveness> appView, ComposableCallGraph callGraph) {
-    super(appView);
-    this.callGraph = callGraph;
-  }
-
-  @Override
-  public void scan(
-      ProgramMethod method,
-      IRCode code,
-      AbstractValueSupplier abstractValueSupplier,
-      PathConstraintSupplier pathConstraintSupplier,
-      Timing timing) {
-    new CodeScanner(abstractValueSupplier, code, method, pathConstraintSupplier).scan(timing);
-  }
-
-  @Override
-  protected boolean isMethodParameterAlreadyUnknown(
-      DexType staticType, MethodParameter methodParameter, ProgramMethod method) {
-    // We haven't defined the virtual root mapping, so we can't tell.
-    return false;
-  }
-
-  private class CodeScanner extends ArgumentPropagatorCodeScanner.CodeScanner {
-
-    protected CodeScanner(
-        AbstractValueSupplier abstractValueSupplier,
-        IRCode code,
-        ProgramMethod method,
-        PathConstraintSupplier pathConstraintSupplier) {
-      super(abstractValueSupplier, code, method, pathConstraintSupplier);
-    }
-
-    @Override
-    protected void addTemporaryMethodState(
-        InvokeMethod invoke, ProgramMethod resolvedMethod, Timing timing) {
-      ComposableCallGraphNode node = callGraph.getNodes().get(resolvedMethod);
-      if (node != null && node.isComposable()) {
-        super.addTemporaryMethodState(invoke, resolvedMethod, timing);
-      }
-    }
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorComposeModeling.java b/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorComposeModeling.java
deleted file mode 100644
index 69a0331..0000000
--- a/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorComposeModeling.java
+++ /dev/null
@@ -1,230 +0,0 @@
-// Copyright (c) 2023, 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.compose;
-
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.DexClassAndMethod;
-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.graph.DexString;
-import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.MethodResolutionResult.SingleResolutionResult;
-import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.graph.lens.GraphLens;
-import com.android.tools.r8.ir.analysis.type.TypeElement;
-import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
-import com.android.tools.r8.ir.code.InstanceGet;
-import com.android.tools.r8.ir.code.Instruction;
-import com.android.tools.r8.ir.code.InvokeMethod;
-import com.android.tools.r8.ir.code.InvokeStatic;
-import com.android.tools.r8.ir.code.Or;
-import com.android.tools.r8.ir.code.Value;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcretePrimitiveTypeValueState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldValue;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.NonEmptyValueState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.OrAbstractFunction;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.android.tools.r8.utils.BooleanUtils;
-import com.google.common.collect.Iterables;
-
-public class ArgumentPropagatorComposeModeling {
-
-  private final AppView<AppInfoWithLiveness> appView;
-  private final ComposeReferences rewrittenComposeReferences;
-
-  private final DexType rewrittenFunction2Type;
-  private final DexString invokeName;
-
-  public ArgumentPropagatorComposeModeling(AppView<AppInfoWithLiveness> appView) {
-    assert appView
-        .options()
-        .getJetpackComposeOptions()
-        .isModelingChangedArgumentsToComposableFunctions();
-    this.appView = appView;
-    this.rewrittenComposeReferences =
-        appView
-            .getComposeReferences()
-            .rewrittenWithLens(appView.graphLens(), GraphLens.getIdentityLens());
-    DexItemFactory dexItemFactory = appView.dexItemFactory();
-    this.rewrittenFunction2Type =
-        appView
-            .graphLens()
-            .lookupType(
-                dexItemFactory.createType("Lkotlin/jvm/functions/Function2;"),
-                GraphLens.getIdentityLens());
-    this.invokeName = dexItemFactory.createString("invoke");
-  }
-
-  /**
-   * Models calls to @Composable functions from Compose restart lambdas.
-   *
-   * <p>The @Composable functions are static and should have one of the following two parameter
-   * lists:
-   *
-   * <ol>
-   *   <li>(..., Composable, int)
-   *   <li>(..., Composable, int, int)
-   * </ol>
-   *
-   * <p>The int argument after the Composable parameter is the $$changed parameter. The second int
-   * argument after the Composable parameter (if present) is the $$default parameter.
-   *
-   * <p>The call to a @Composable function from its restart lambda follows the following code
-   * pattern:
-   *
-   * <pre>
-   *   MyComposableFunction(
-   *       ..., this.composer, updateChangedFlags(this.$$changed) || 1, this.$$default)
-   * </pre>
-   *
-   * <p>The modeling performed by this method assumes that updateChangedFlags() does not have any
-   * impact on $$changed (see the current implementation below). The modeling also assumes that
-   * this.$$changed and this.$$default are captures of the $$changed and $$default parameters of
-   * the @Composable function.
-   *
-   * <pre>
-   *   internal fun updateChangedFlags(flags: Int): Int {
-   *     val lowBits = flags and changedLowBitMask
-   *     val highBits = flags and changedHighBitMask
-   *     return ((flags and changedMask) or
-   *         (lowBits or (highBits shr 1)) or ((lowBits shl 1) and highBits))
-   *   }
-   * </pre>
-   */
-  public NonEmptyValueState modelParameterStateForChangedOrDefaultArgumentToComposableFunction(
-      InvokeMethod invoke,
-      ProgramMethod singleTarget,
-      int argumentIndex,
-      Value argument,
-      ProgramMethod context) {
-    // TODO(b/302483644): Add some robust way of detecting restart lambda contexts.
-    if (!context.getHolder().getInterfaces().contains(rewrittenFunction2Type)
-        || !invoke.getPosition().getOutermostCaller().getMethod().getName().isEqualTo(invokeName)
-        || Iterables.isEmpty(
-            context
-                .getHolder()
-                .instanceFields(
-                    f -> f.getName().isIdenticalTo(rewrittenComposeReferences.changedFieldName)))) {
-      return null;
-    }
-
-    // First check if this is an invoke to a @Composable function.
-    if (singleTarget == null
-        || !singleTarget
-            .getDefinition()
-            .annotations()
-            .hasAnnotation(rewrittenComposeReferences.composableType)) {
-      return null;
-    }
-
-    // The @Composable function is expected to be static and have >= 2 parameters.
-    if (!invoke.isInvokeStatic()) {
-      return null;
-    }
-
-    DexMethod invokedMethod = invoke.getInvokedMethod();
-    if (invokedMethod.getArity() < 2) {
-      return null;
-    }
-
-    // Check if the parameters list is one of (..., Composer, int) or (..., Composer, int, int).
-    if (!invokedMethod.getParameter(invokedMethod.getArity() - 1).isIntType()) {
-      return null;
-    }
-
-    boolean hasDefaultParameter =
-        invokedMethod.getParameter(invokedMethod.getArity() - 2).isIntType();
-    if (hasDefaultParameter && invokedMethod.getArity() < 3) {
-      return null;
-    }
-
-    int composerParameterIndex =
-        invokedMethod.getArity() - 2 - BooleanUtils.intValue(hasDefaultParameter);
-    if (!invokedMethod
-        .getParameter(composerParameterIndex)
-        .isIdenticalTo(rewrittenComposeReferences.composerType)) {
-      return null;
-    }
-
-    // We only model the $$changed argument to the @Composable function.
-    if (argumentIndex != composerParameterIndex + 1) {
-      return null;
-    }
-
-    assert argument.getType().isInt();
-
-    DexField changedField =
-        appView
-            .dexItemFactory()
-            .createField(
-                context.getHolderType(),
-                appView.dexItemFactory().intType,
-                rewrittenComposeReferences.changedFieldName);
-
-    UpdateChangedFlagsAbstractFunction inFlow = null;
-    // We are looking at an argument to the $$changed parameter of the @Composable function.
-    // We generally expect this argument to be defined by a call to updateChangedFlags().
-    if (argument.isDefinedByInstructionSatisfying(Instruction::isInvokeStatic)) {
-      InvokeStatic invokeStatic = argument.getDefinition().asInvokeStatic();
-      SingleResolutionResult<?> resolutionResult =
-          invokeStatic.resolveMethod(appView, context).asSingleResolution();
-      if (resolutionResult == null) {
-        return null;
-      }
-      DexClassAndMethod invokeSingleTarget =
-          resolutionResult
-              .lookupDispatchTarget(appView, invokeStatic, context)
-              .getSingleDispatchTarget();
-      if (invokeSingleTarget == null) {
-        return null;
-      }
-      inFlow =
-          invokeSingleTarget
-              .getOptimizationInfo()
-              .getAbstractFunction()
-              .asUpdateChangedFlagsAbstractFunction();
-      if (inFlow == null) {
-        return null;
-      }
-      // By accounting for the abstract function we can safely strip the call.
-      argument = invokeStatic.getFirstArgument();
-    }
-    // Allow the argument to be defined by `this.$$changed | 1`.
-    if (argument.isDefinedByInstructionSatisfying(Instruction::isOr)) {
-      Or or = argument.getDefinition().asOr();
-      Value maybeNumberOperand = or.leftValue().isConstNumber() ? or.leftValue() : or.rightValue();
-      Value otherOperand = or.getOperand(1 - or.inValues().indexOf(maybeNumberOperand));
-      if (!maybeNumberOperand.isConstNumber(1)) {
-        return null;
-      }
-      // Strip the OR instruction.
-      argument = otherOperand;
-      // Update the model from bottom to a special value that effectively throws away any known
-      // information about the lowermost bit of $$changed.
-      SingleNumberValue one =
-          appView.abstractValueFactory().createSingleNumberValue(1, TypeElement.getInt());
-      inFlow =
-          new UpdateChangedFlagsAbstractFunction(
-              new OrAbstractFunction(new FieldValue(changedField), one));
-    } else {
-      inFlow = new UpdateChangedFlagsAbstractFunction(new FieldValue(changedField));
-    }
-
-    // At this point we expect that the restart lambda is reading this.$$changed using an
-    // instance-get.
-    if (!argument.isDefinedByInstructionSatisfying(Instruction::isInstanceGet)) {
-      return null;
-    }
-
-    // Check that the instance-get is reading the capture field that we expect it to.
-    InstanceGet instanceGet = argument.getDefinition().asInstanceGet();
-    if (!instanceGet.getField().isIdenticalTo(changedField)) {
-      return null;
-    }
-
-    // Return the argument model.
-    return new ConcretePrimitiveTypeValueState(inFlow);
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ComposableCallGraph.java b/src/main/java/com/android/tools/r8/optimize/compose/ComposableCallGraph.java
deleted file mode 100644
index 8364488..0000000
--- a/src/main/java/com/android/tools/r8/optimize/compose/ComposableCallGraph.java
+++ /dev/null
@@ -1,171 +0,0 @@
-// Copyright (c) 2023, 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.compose;
-
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.Code;
-import com.android.tools.r8.graph.DexEncodedMethod;
-import com.android.tools.r8.graph.DexField;
-import com.android.tools.r8.graph.DexMethod;
-import com.android.tools.r8.graph.DexProgramClass;
-import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.graph.UseRegistry;
-import com.android.tools.r8.graph.lens.GraphLens;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.android.tools.r8.utils.collections.ProgramMethodMap;
-import java.util.function.Consumer;
-
-/**
- * A partial call graph that stores call edges to @Composable functions. By processing all the call
- * sites of a given @Composable function we can reapply arguent propagation for the @Composable
- * function.
- */
-public class ComposableCallGraph {
-
-  private final ProgramMethodMap<ComposableCallGraphNode> nodes;
-
-  public ComposableCallGraph(ProgramMethodMap<ComposableCallGraphNode> nodes) {
-    this.nodes = nodes;
-  }
-
-  public static Builder builder(AppView<AppInfoWithLiveness> appView) {
-    return new Builder(appView);
-  }
-
-  public static ComposableCallGraph empty() {
-    return new ComposableCallGraph(ProgramMethodMap.empty());
-  }
-
-  public void forEachNode(Consumer<ComposableCallGraphNode> consumer) {
-    nodes.forEachValue(consumer);
-  }
-
-  public ProgramMethodMap<ComposableCallGraphNode> getNodes() {
-    return nodes;
-  }
-
-  public boolean isEmpty() {
-    return nodes.isEmpty();
-  }
-
-  public static class Builder {
-
-    private final AppView<AppInfoWithLiveness> appView;
-    private final ProgramMethodMap<ComposableCallGraphNode> nodes = ProgramMethodMap.create();
-
-    Builder(AppView<AppInfoWithLiveness> appView) {
-      this.appView = appView;
-    }
-
-    public ComposableCallGraph build() {
-      createCallGraphNodesForComposableFunctions();
-      if (!nodes.isEmpty()) {
-        addCallEdgesToComposableFunctions();
-      }
-      return new ComposableCallGraph(nodes);
-    }
-
-    private void createCallGraphNodesForComposableFunctions() {
-      ComposeReferences rewrittenComposeReferences =
-          appView
-              .getComposeReferences()
-              .rewrittenWithLens(appView.graphLens(), GraphLens.getIdentityLens());
-      for (DexProgramClass clazz : appView.appInfo().classes()) {
-        clazz.forEachProgramDirectMethodMatching(
-            method -> method.annotations().hasAnnotation(rewrittenComposeReferences.composableType),
-            method -> {
-              // TODO(b/302483644): Don't include kept @Composable functions, since we can't
-              //  optimize them anyway.
-              assert method.getAccessFlags().isStatic();
-              nodes.put(method, new ComposableCallGraphNode(method, true));
-            });
-      }
-    }
-
-    // TODO(b/302483644): Parallelize identification of @Composable call sites.
-    private void addCallEdgesToComposableFunctions() {
-      // Code is fully rewritten so no need to lens rewrite in registry.
-      assert appView.graphLens().isMemberRebindingIdentityLens();
-      assert appView.codeLens() == appView.graphLens().asNonIdentityLens().getPrevious();
-
-      for (DexProgramClass clazz : appView.appInfo().classes()) {
-        clazz.forEachProgramMethodMatching(
-            DexEncodedMethod::hasCode,
-            method -> {
-              Code code = method.getDefinition().getCode();
-
-              // TODO(b/302483644): Leverage LIR code constant pool for efficient checking.
-              // TODO(b/302483644): Maybe remove the possibility of CF/DEX at this point.
-              assert code.isLirCode()
-                  || code.isCfCode()
-                  || code.isDexCode()
-                  || code.isDefaultInstanceInitializerCode()
-                  || code.isThrowNullCode();
-
-              code.registerCodeReferences(
-                  method,
-                  new UseRegistry<>(appView, method) {
-
-                    private final AppView<AppInfoWithLiveness> appViewWithLiveness =
-                        appView.withLiveness();
-
-                    @Override
-                    public void registerInvokeStatic(DexMethod method) {
-                      ProgramMethod resolvedMethod =
-                          appViewWithLiveness
-                              .appInfo()
-                              .unsafeResolveMethodDueToDexFormat(method)
-                              .getResolvedProgramMethod();
-                      if (resolvedMethod == null) {
-                        return;
-                      }
-
-                      ComposableCallGraphNode callee = nodes.get(resolvedMethod);
-                      if (callee == null || !callee.isComposable()) {
-                        // Only record calls to Composable functions.
-                        return;
-                      }
-
-                      ComposableCallGraphNode caller =
-                          nodes.computeIfAbsent(
-                              getContext(), context -> new ComposableCallGraphNode(context, false));
-                      callee.addCaller(caller);
-                    }
-
-                    @Override
-                    public void registerInitClass(DexType type) {}
-
-                    @Override
-                    public void registerInvokeDirect(DexMethod method) {}
-
-                    @Override
-                    public void registerInvokeInterface(DexMethod method) {}
-
-                    @Override
-                    public void registerInvokeSuper(DexMethod method) {}
-
-                    @Override
-                    public void registerInvokeVirtual(DexMethod method) {}
-
-                    @Override
-                    public void registerInstanceFieldRead(DexField field) {}
-
-                    @Override
-                    public void registerInstanceFieldWrite(DexField field) {}
-
-                    @Override
-                    public void registerStaticFieldRead(DexField field) {}
-
-                    @Override
-                    public void registerStaticFieldWrite(DexField field) {}
-
-                    @Override
-                    public void registerTypeReference(DexType type) {}
-                  });
-            });
-      }
-    }
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ComposableCallGraphNode.java b/src/main/java/com/android/tools/r8/optimize/compose/ComposableCallGraphNode.java
deleted file mode 100644
index 13ec14db..0000000
--- a/src/main/java/com/android/tools/r8/optimize/compose/ComposableCallGraphNode.java
+++ /dev/null
@@ -1,53 +0,0 @@
-// Copyright (c) 2023, 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.compose;
-
-import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.utils.SetUtils;
-import java.util.Set;
-import java.util.function.Consumer;
-
-public class ComposableCallGraphNode {
-
-  private final ProgramMethod method;
-  private final boolean isComposable;
-
-  private final Set<ComposableCallGraphNode> callers = SetUtils.newIdentityHashSet();
-  private final Set<ComposableCallGraphNode> callees = SetUtils.newIdentityHashSet();
-
-  ComposableCallGraphNode(ProgramMethod method, boolean isComposable) {
-    this.method = method;
-    this.isComposable = isComposable;
-  }
-
-  public void addCaller(ComposableCallGraphNode caller) {
-    callers.add(caller);
-    caller.callees.add(this);
-  }
-
-  public void forEachComposableCallee(Consumer<ComposableCallGraphNode> consumer) {
-    for (ComposableCallGraphNode callee : callees) {
-      if (callee.isComposable()) {
-        consumer.accept(callee);
-      }
-    }
-  }
-
-  public Set<ComposableCallGraphNode> getCallers() {
-    return callers;
-  }
-
-  public ProgramMethod getMethod() {
-    return method;
-  }
-
-  public boolean isComposable() {
-    return isComposable;
-  }
-
-  @Override
-  public String toString() {
-    return method.toString();
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ComposableOptimizationPass.java b/src/main/java/com/android/tools/r8/optimize/compose/ComposableOptimizationPass.java
deleted file mode 100644
index c2890d8..0000000
--- a/src/main/java/com/android/tools/r8/optimize/compose/ComposableOptimizationPass.java
+++ /dev/null
@@ -1,114 +0,0 @@
-// Copyright (c) 2023, 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.compose;
-
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.ir.conversion.PrimaryR8IRConverter;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.android.tools.r8.utils.SetUtils;
-import com.android.tools.r8.utils.Timing;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Sets;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-
-public class ComposableOptimizationPass {
-
-  private final AppView<AppInfoWithLiveness> appView;
-  private final PrimaryR8IRConverter converter;
-
-  private ComposableOptimizationPass(
-      AppView<AppInfoWithLiveness> appView, PrimaryR8IRConverter converter) {
-    this.appView = appView;
-    this.converter = converter;
-  }
-
-  public static void run(
-      AppView<AppInfoWithLiveness> appView,
-      PrimaryR8IRConverter converter,
-      ExecutorService executorService,
-      Timing timing)
-      throws ExecutionException {
-    if (appView.options().getJetpackComposeOptions().isComposableOptimizationPassEnabled()) {
-      timing.time(
-          "ComposableOptimizationPass",
-          () -> new ComposableOptimizationPass(appView, converter).processWaves(executorService));
-    }
-  }
-
-  void processWaves(ExecutorService executorService) throws ExecutionException {
-    ComposableCallGraph callGraph = ComposableCallGraph.builder(appView).build();
-    ComposeMethodProcessor methodProcessor =
-        new ComposeMethodProcessor(appView, callGraph, converter);
-    Set<ComposableCallGraphNode> wave = createInitialWave(callGraph);
-    while (!wave.isEmpty()) {
-      Set<ComposableCallGraphNode> optimizedComposableFunctions =
-          methodProcessor.processWave(wave, executorService);
-      wave = createNextWave(methodProcessor, optimizedComposableFunctions);
-    }
-  }
-
-  // TODO(b/302483644): Should we skip root @Composable functions that don't have any nested
-  //  @Composable functions (?).
-  private Set<ComposableCallGraphNode> computeComposableRoots(ComposableCallGraph callGraph) {
-    Set<ComposableCallGraphNode> composableRoots = Sets.newIdentityHashSet();
-    callGraph.forEachNode(
-        node -> {
-          if (!node.isComposable()
-              || Iterables.any(node.getCallers(), ComposableCallGraphNode::isComposable)) {
-            // This is not a @Composable root.
-            return;
-          }
-          if (node.getCallers().isEmpty()) {
-            // Don't include root @Composable functions that are never called. These are either kept
-            // or will be removed in tree shaking.
-            return;
-          }
-          composableRoots.add(node);
-        });
-    return composableRoots;
-  }
-
-  private Set<ComposableCallGraphNode> createInitialWave(ComposableCallGraph callGraph) {
-    Set<ComposableCallGraphNode> wave = Sets.newIdentityHashSet();
-    Set<ComposableCallGraphNode> composableRoots = computeComposableRoots(callGraph);
-    composableRoots.forEach(composableRoot -> wave.addAll(composableRoot.getCallers()));
-    return wave;
-  }
-
-  // TODO(b/302483644): Consider repeatedly extracting the roots from the graph similar to the way
-  //  we extract leaves in the primary optimization pass.
-  private static Set<ComposableCallGraphNode> createNextWave(
-      ComposeMethodProcessor methodProcessor,
-      Set<ComposableCallGraphNode> optimizedComposableFunctions) {
-    Set<ComposableCallGraphNode> nextWave =
-        SetUtils.newIdentityHashSet(optimizedComposableFunctions);
-
-    // If the new wave contains two @Composable functions where one calls the other, then defer the
-    // processing of the callee to a later wave, to ensure that we have seen all of its callers
-    // before processing the callee.
-    List<ComposableCallGraphNode> deferredComposableFunctions = new ArrayList<>();
-    nextWave.forEach(
-        node -> {
-          if (SetUtils.containsAnyOf(nextWave, node.getCallers())) {
-            deferredComposableFunctions.add(node);
-          }
-        });
-    deferredComposableFunctions.forEach(nextWave::remove);
-
-    // To optimize the @Composable functions that are called from the @Composable functions of the
-    // next wave in the wave after that, we need to include their callers in the next wave as well.
-    Set<ComposableCallGraphNode> callersOfCalledComposableFunctions = Sets.newIdentityHashSet();
-    nextWave.forEach(
-        node ->
-            node.forEachComposableCallee(
-                callee -> callersOfCalledComposableFunctions.addAll(callee.getCallers())));
-    nextWave.addAll(callersOfCalledComposableFunctions);
-    nextWave.removeIf(methodProcessor::isProcessed);
-    return nextWave;
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ComposeMethodProcessor.java b/src/main/java/com/android/tools/r8/optimize/compose/ComposeMethodProcessor.java
deleted file mode 100644
index 3b51e7c..0000000
--- a/src/main/java/com/android/tools/r8/optimize/compose/ComposeMethodProcessor.java
+++ /dev/null
@@ -1,280 +0,0 @@
-// Copyright (c) 2023, 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.compose;
-
-import static com.android.tools.r8.graph.ProgramField.asProgramFieldOrNull;
-
-import com.android.tools.r8.contexts.CompilationContext.ProcessorContext;
-import com.android.tools.r8.errors.Unreachable;
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.graph.DexMethod;
-import com.android.tools.r8.graph.DexProgramClass;
-import com.android.tools.r8.graph.ProgramField;
-import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.ir.analysis.constant.SparseConditionalConstantPropagation;
-import com.android.tools.r8.ir.analysis.path.PathConstraintSupplier;
-import com.android.tools.r8.ir.analysis.value.AbstractValue;
-import com.android.tools.r8.ir.code.AbstractValueSupplier;
-import com.android.tools.r8.ir.code.IRCode;
-import com.android.tools.r8.ir.code.Value;
-import com.android.tools.r8.ir.conversion.MethodConversionOptions;
-import com.android.tools.r8.ir.conversion.MethodProcessor;
-import com.android.tools.r8.ir.conversion.MethodProcessorEventConsumer;
-import com.android.tools.r8.ir.conversion.PrimaryR8IRConverter;
-import com.android.tools.r8.ir.conversion.callgraph.CallSiteInformation;
-import com.android.tools.r8.ir.optimize.info.OptimizationFeedback;
-import com.android.tools.r8.optimize.argumentpropagation.ArgumentPropagatorCodeScanner;
-import com.android.tools.r8.optimize.argumentpropagation.ArgumentPropagatorOptimizationInfoPopulator;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteMonomorphicMethodState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcretePrimitiveTypeValueState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteValueState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldStateCollection;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlowComparator;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodStateCollectionByReference;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ValueState;
-import com.android.tools.r8.optimize.argumentpropagation.propagation.DefaultFieldValueJoiner;
-import com.android.tools.r8.optimize.argumentpropagation.propagation.FlowGraph;
-import com.android.tools.r8.optimize.argumentpropagation.propagation.InFlowPropagator;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.android.tools.r8.utils.IterableUtils;
-import com.android.tools.r8.utils.LazyBox;
-import com.android.tools.r8.utils.ThreadUtils;
-import com.android.tools.r8.utils.Timing;
-import com.android.tools.r8.utils.collections.ProgramMethodSet;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Sets;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-
-public class ComposeMethodProcessor extends MethodProcessor {
-
-  private final AppView<AppInfoWithLiveness> appView;
-  private final ArgumentPropagatorCodeScanner codeScanner;
-  private final PrimaryR8IRConverter converter;
-
-  private final Set<ComposableCallGraphNode> processed = Sets.newIdentityHashSet();
-
-  public ComposeMethodProcessor(
-      AppView<AppInfoWithLiveness> appView,
-      ComposableCallGraph callGraph,
-      PrimaryR8IRConverter converter) {
-    this.appView = appView;
-    this.codeScanner = new ArgumentPropagatorCodeScannerForComposableFunctions(appView, callGraph);
-    this.converter = converter;
-  }
-
-  public Set<ComposableCallGraphNode> processWave(
-      Set<ComposableCallGraphNode> wave, ExecutorService executorService)
-      throws ExecutionException {
-    ProcessorContext processorContext = appView.createProcessorContext();
-    ThreadUtils.processItems(
-        wave,
-        node -> {
-          assert !processed.contains(node);
-          converter.processDesugaredMethod(
-              node.getMethod(),
-              OptimizationFeedback.getIgnoreFeedback(),
-              this,
-              processorContext.createMethodProcessingContext(node.getMethod()),
-              MethodConversionOptions.forLirPhase(appView));
-        },
-        appView.options().getThreadingModule(),
-        executorService);
-    processed.addAll(wave);
-    return optimizeComposableFunctionsCalledFromWave(wave, executorService);
-  }
-
-  private Set<ComposableCallGraphNode> optimizeComposableFunctionsCalledFromWave(
-      Set<ComposableCallGraphNode> wave, ExecutorService executorService)
-      throws ExecutionException {
-    prepareForInFlowPropagator();
-
-    InFlowComparator emptyComparator = InFlowComparator.builder().build();
-    InFlowPropagator inFlowPropagator =
-        new InFlowPropagator(
-            appView,
-            null,
-            converter,
-            codeScanner.getFieldStates(),
-            codeScanner.getMethodStates(),
-            emptyComparator) {
-
-          @Override
-          protected DefaultFieldValueJoiner createDefaultFieldValueJoiner(
-              List<FlowGraph> flowGraphs) {
-            return new DefaultFieldValueJoiner(appView, null, fieldStates, flowGraphs) {
-
-              @Override
-              protected Map<DexProgramClass, List<ProgramField>> getFieldsOfInterest() {
-                // We do not rely on the optimization of any fields in the Composable optimization
-                // pass.
-                return Collections.emptyMap();
-              }
-            };
-          }
-        };
-    inFlowPropagator.run(executorService);
-
-    ArgumentPropagatorOptimizationInfoPopulator optimizationInfoPopulator =
-        new ArgumentPropagatorOptimizationInfoPopulator(appView, null, null, null, null);
-    Set<ComposableCallGraphNode> optimizedComposableFunctions = Sets.newIdentityHashSet();
-    wave.forEach(
-        node ->
-            node.forEachComposableCallee(
-                callee -> {
-                  if (Iterables.all(callee.getCallers(), this::isProcessed)) {
-                    optimizationInfoPopulator.setOptimizationInfo(
-                        callee.getMethod(), ProgramMethodSet.empty(), getMethodState(callee));
-                    // TODO(b/302483644): Only enqueue this callee if its optimization info changed.
-                    optimizedComposableFunctions.add(callee);
-                  }
-                }));
-    return optimizedComposableFunctions;
-  }
-
-  private void prepareForInFlowPropagator() {
-    FieldStateCollection fieldStates = codeScanner.getFieldStates();
-
-    // Set all field states to unknown since we are not guaranteed to have processes all field
-    // writes.
-    fieldStates.forEach(
-        (field, fieldState) ->
-            fieldStates.addTemporaryFieldState(
-                appView, field, ValueState::unknown, Timing.empty()));
-
-    // Widen all parameter states that have in-flow to unknown, except when the in-flow is an
-    // update-changed-flags abstract function.
-    MethodStateCollectionByReference methodStates = codeScanner.getMethodStates();
-    methodStates.forEach(
-        (method, methodState) -> {
-          if (!methodState.isMonomorphic()) {
-            assert methodState.isUnknown();
-            return;
-          }
-          ConcreteMonomorphicMethodState monomorphicMethodState = methodState.asMonomorphic();
-          for (int parameterIndex = 0;
-              parameterIndex < monomorphicMethodState.size();
-              parameterIndex++) {
-            ValueState parameterState = monomorphicMethodState.getParameterState(parameterIndex);
-            if (parameterState.isConcrete()) {
-              ConcreteValueState concreteParameterState = parameterState.asConcrete();
-              prepareParameterStateForInFlowPropagator(
-                  method, monomorphicMethodState, parameterIndex, concreteParameterState);
-            }
-          }
-        });
-  }
-
-  private void prepareParameterStateForInFlowPropagator(
-      DexMethod method,
-      ConcreteMonomorphicMethodState methodState,
-      int parameterIndex,
-      ConcreteValueState parameterState) {
-    if (!parameterState.hasInFlow()) {
-      return;
-    }
-
-    UpdateChangedFlagsAbstractFunction transferFunction = null;
-    if (parameterState.getInFlow().size() == 1) {
-      transferFunction =
-          Iterables.getOnlyElement(parameterState.getInFlow())
-              .asUpdateChangedFlagsAbstractFunction();
-    }
-    if (transferFunction == null) {
-      methodState.setParameterState(parameterIndex, ValueState.unknown());
-      return;
-    }
-
-    // This is a call to a composable function from a restart function.
-    Iterable<BaseInFlow> baseInFlow = transferFunction.getBaseInFlow();
-    assert Iterables.size(baseInFlow) == 1;
-    BaseInFlow singleBaseInFlow = IterableUtils.first(baseInFlow);
-    assert singleBaseInFlow.isFieldValue();
-
-    ProgramField field =
-        asProgramFieldOrNull(appView.definitionFor(singleBaseInFlow.asFieldValue().getField()));
-    assert field != null;
-
-    // If the only input to the $$changed parameter of the Composable function is in-flow then skip.
-    if (methodState.getParameterState(parameterIndex).getAbstractValue(appView).isBottom()) {
-      methodState.setParameterState(parameterIndex, ValueState.unknown());
-      return;
-    }
-
-    codeScanner
-        .getFieldStates()
-        .addTemporaryFieldState(
-            appView,
-            field,
-            () ->
-                new ConcretePrimitiveTypeValueState(
-                    MethodParameter.createStatic(method, parameterIndex)),
-            Timing.empty());
-  }
-
-  private MethodState getMethodState(ComposableCallGraphNode node) {
-    assert processed.containsAll(node.getCallers());
-    return codeScanner.getMethodStates().get(node.getMethod());
-  }
-
-  public void scan(ProgramMethod method, IRCode code, Timing timing) {
-    LazyBox<Map<Value, AbstractValue>> abstractValues =
-        new LazyBox<>(() -> new SparseConditionalConstantPropagation(appView).analyze(code));
-    AbstractValueSupplier abstractValueSupplier =
-        value -> {
-          AbstractValue abstractValue = abstractValues.computeIfAbsent().get(value);
-          assert abstractValue != null;
-          return abstractValue;
-        };
-    PathConstraintSupplier pathConstraintSupplier =
-        new PathConstraintSupplier(appView, code, codeScanner.getMethodParameterFactory());
-    codeScanner.scan(method, code, abstractValueSupplier, pathConstraintSupplier, timing);
-  }
-
-  public boolean isProcessed(ComposableCallGraphNode node) {
-    return processed.contains(node);
-  }
-
-  @Override
-  public CallSiteInformation getCallSiteInformation() {
-    return CallSiteInformation.empty();
-  }
-
-  @Override
-  public MethodProcessorEventConsumer getEventConsumer() {
-    throw new Unreachable();
-  }
-
-  @Override
-  public boolean isComposeMethodProcessor() {
-    return true;
-  }
-
-  @Override
-  public ComposeMethodProcessor asComposeMethodProcessor() {
-    return this;
-  }
-
-  @Override
-  public boolean isProcessedConcurrently(ProgramMethod method) {
-    return false;
-  }
-
-  @Override
-  public void scheduleDesugaredMethodForProcessing(ProgramMethod method) {
-    throw new Unreachable();
-  }
-
-  @Override
-  public boolean shouldApplyCodeRewritings(ProgramMethod method) {
-    return false;
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ComposeReferences.java b/src/main/java/com/android/tools/r8/optimize/compose/ComposeReferences.java
index e38f9a4..633008d 100644
--- a/src/main/java/com/android/tools/r8/optimize/compose/ComposeReferences.java
+++ b/src/main/java/com/android/tools/r8/optimize/compose/ComposeReferences.java
@@ -7,6 +7,7 @@
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexString;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.ProgramDefinition;
 import com.android.tools.r8.graph.lens.GraphLens;
 
 public class ComposeReferences {
@@ -42,6 +43,11 @@
     this.updatedChangedFlagsMethod = updatedChangedFlagsMethod;
   }
 
+  public boolean isComposable(ProgramDefinition definition) {
+    return definition.isProgramMethod()
+        && definition.asProgramMethod().getAnnotations().hasAnnotation(composableType);
+  }
+
   public ComposeReferences rewrittenWithLens(GraphLens graphLens, GraphLens codeLens) {
     return new ComposeReferences(
         changedFieldName,
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ComputationTreeUnopUpdateChangedFlagsNode.java b/src/main/java/com/android/tools/r8/optimize/compose/ComputationTreeUnopUpdateChangedFlagsNode.java
new file mode 100644
index 0000000..15ca10f
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/compose/ComputationTreeUnopUpdateChangedFlagsNode.java
@@ -0,0 +1,136 @@
+// Copyright (c) 2024, 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.compose;
+
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.ir.analysis.value.AbstractValue;
+import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
+import com.android.tools.r8.ir.analysis.value.DefiniteBitsNumberValue;
+import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
+import com.android.tools.r8.ir.analysis.value.arithmetic.AbstractCalculator;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
+import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeUnopNode;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.TraversalContinuation;
+import java.util.Objects;
+import java.util.function.Function;
+
+public class ComputationTreeUnopUpdateChangedFlagsNode extends ComputationTreeUnopNode {
+
+  private static final int changedLowBitMask = 0b001_001_001_001_001_001_001_001_001_001_0;
+  private static final int changedHighBitMask = changedLowBitMask << 1;
+  private static final int changedMask = ~(changedLowBitMask | changedHighBitMask);
+
+  public ComputationTreeUnopUpdateChangedFlagsNode(ComputationTreeNode operand) {
+    super(operand);
+  }
+
+  public static ComputationTreeNode create(ComputationTreeNode operand) {
+    if (operand.isUnknown()) {
+      return AbstractValue.unknown();
+    }
+    return new ComputationTreeUnopUpdateChangedFlagsNode(operand);
+  }
+
+  @Override
+  public AbstractValue evaluate(
+      AppView<AppInfoWithLiveness> appView, FlowGraphStateProvider flowGraphStateProvider) {
+    AbstractValue operandValue = operand.evaluate(appView, flowGraphStateProvider);
+    if (operandValue.isBottom()) {
+      return operandValue;
+    } else if (operandValue.isSingleNumberValue()) {
+      return evaluateConcrete(appView, operandValue.asSingleNumberValue().getIntValue());
+    } else if (operandValue.isDefiniteBitsNumberValue()) {
+      return evaluateAbstract(appView, operandValue.asDefiniteBitsNumberValue());
+    } else {
+      assert !operandValue.hasDefinitelySetAndUnsetBitsInformation();
+      return AbstractValue.unknown();
+    }
+  }
+
+  /**
+   * Applies the following function to the given {@param abstractValue}.
+   *
+   * <pre>
+   * private const val changedLowBitMask = 0b001_001_001_001_001_001_001_001_001_001_0
+   * private const val changedHighBitMask = changedLowBitMask shl 1
+   * private const val changedMask = (changedLowBitMask or changedHighBitMask).inv()
+   *
+   * internal fun updateChangedFlags(flags: Int): Int {
+   *     val lowBits = flags and changedLowBitMask
+   *     val highBits = flags and changedHighBitMask
+   *     return ((flags and changedMask) or
+   *         (lowBits or (highBits shr 1)) or ((lowBits shl 1) and highBits))
+   * }
+   * </pre>
+   */
+  private AbstractValue evaluateAbstract(
+      AppView<AppInfoWithLiveness> appView, DefiniteBitsNumberValue flags) {
+    AbstractValueFactory factory = appView.abstractValueFactory();
+    // Load constants.
+    AbstractValue changedLowBitMaskValue =
+        factory.createUncheckedSingleNumberValue(changedLowBitMask);
+    AbstractValue changedHighBitMaskValue =
+        factory.createUncheckedSingleNumberValue(changedHighBitMask);
+    AbstractValue changedMaskValue = factory.createUncheckedSingleNumberValue(changedMask);
+    // Evaluate expression.
+    AbstractValue lowBitsValue =
+        AbstractCalculator.andIntegers(appView, flags, changedLowBitMaskValue);
+    AbstractValue highBitsValue =
+        AbstractCalculator.andIntegers(appView, flags, changedHighBitMaskValue);
+    AbstractValue changedBitsValue =
+        AbstractCalculator.andIntegers(appView, flags, changedMaskValue);
+    return AbstractCalculator.orIntegers(
+        appView,
+        changedBitsValue,
+        lowBitsValue,
+        AbstractCalculator.shrIntegers(appView, highBitsValue, 1),
+        AbstractCalculator.andIntegers(
+            appView, AbstractCalculator.shlIntegers(appView, lowBitsValue, 1), highBitsValue));
+  }
+
+  private SingleNumberValue evaluateConcrete(AppView<AppInfoWithLiveness> appView, int flags) {
+    int lowBits = flags & changedLowBitMask;
+    int highBits = flags & changedHighBitMask;
+    int changedBits = flags & changedMask;
+    int result = changedBits | lowBits | (highBits >> 1) | ((lowBits << 1) & highBits);
+    return appView.abstractValueFactory().createUncheckedSingleNumberValue(result);
+  }
+
+  @Override
+  public boolean isUpdateChangedFlags() {
+    return true;
+  }
+
+  @Override
+  public <TB, TC> TraversalContinuation<TB, TC> traverseBaseInFlow(
+      Function<? super BaseInFlow, TraversalContinuation<TB, TC>> fn) {
+    return operand.traverseBaseInFlow(fn);
+  }
+
+  @Override
+  @SuppressWarnings("EqualsGetClass")
+  public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
+    if (obj == null || getClass() != obj.getClass()) {
+      return false;
+    }
+    ComputationTreeUnopUpdateChangedFlagsNode fn = (ComputationTreeUnopUpdateChangedFlagsNode) obj;
+    return operand.equals(fn.operand);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(getClass(), operand);
+  }
+
+  @Override
+  public String toString() {
+    return "UpdateChangedFlags(" + operand + ")";
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/JetpackComposeOptions.java b/src/main/java/com/android/tools/r8/optimize/compose/JetpackComposeOptions.java
index 69ec5e6..65a8199 100644
--- a/src/main/java/com/android/tools/r8/optimize/compose/JetpackComposeOptions.java
+++ b/src/main/java/com/android/tools/r8/optimize/compose/JetpackComposeOptions.java
@@ -10,10 +10,6 @@
 
   private final InternalOptions options;
 
-  public boolean enableComposableOptimizationPass =
-      SystemPropertyUtils.parseSystemPropertyOrDefault(
-          "com.android.tools.r8.jetpackcompose.enableComposableOptimizationPass", false);
-
   public boolean enableModelingOfChangedArguments =
       SystemPropertyUtils.parseSystemPropertyOrDefault(
           "com.android.tools.r8.jetpackcompose.enableModelingOfChangedArguments", false);
@@ -22,20 +18,6 @@
     this.options = options;
   }
 
-  public void enableAllOptimizations(boolean enable) {
-    enableComposableOptimizationPass = enable;
-    enableModelingOfChangedArguments = enable;
-  }
-
-  public boolean isAnyOptimizationsEnabled() {
-    return isComposableOptimizationPassEnabled()
-        || isModelingChangedArgumentsToComposableFunctions();
-  }
-
-  public boolean isComposableOptimizationPassEnabled() {
-    return isModelingChangedArgumentsToComposableFunctions() && enableComposableOptimizationPass;
-  }
-
   public boolean isModelingChangedArgumentsToComposableFunctions() {
     return options.isOptimizing() && options.isShrinking() && enableModelingOfChangedArguments;
   }
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/UpdateChangedFlagsAbstractFunction.java b/src/main/java/com/android/tools/r8/optimize/compose/UpdateChangedFlagsAbstractFunction.java
deleted file mode 100644
index 91476d3..0000000
--- a/src/main/java/com/android/tools/r8/optimize/compose/UpdateChangedFlagsAbstractFunction.java
+++ /dev/null
@@ -1,166 +0,0 @@
-// Copyright (c) 2024, 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.compose;
-
-import com.android.tools.r8.graph.AppView;
-import com.android.tools.r8.ir.analysis.value.AbstractValue;
-import com.android.tools.r8.ir.analysis.value.AbstractValueFactory;
-import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
-import com.android.tools.r8.ir.analysis.value.arithmetic.AbstractCalculator;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.AbstractFunction;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.BaseInFlow;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcretePrimitiveTypeValueState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteValueState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.FlowGraphStateProvider;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlow;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlowComparator;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlowKind;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.OrAbstractFunction;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ValueState;
-import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import com.android.tools.r8.utils.IterableUtils;
-import java.util.Objects;
-
-public class UpdateChangedFlagsAbstractFunction implements AbstractFunction {
-
-  private static final int changedLowBitMask = 0b001_001_001_001_001_001_001_001_001_001_0;
-  private static final int changedHighBitMask = changedLowBitMask << 1;
-  private static final int changedMask = ~(changedLowBitMask | changedHighBitMask);
-
-  private final InFlow inFlow;
-
-  public UpdateChangedFlagsAbstractFunction(InFlow inFlow) {
-    this.inFlow = inFlow;
-  }
-
-  @Override
-  public ValueState apply(
-      AppView<AppInfoWithLiveness> appView,
-      FlowGraphStateProvider flowGraphStateProvider,
-      ConcreteValueState baseInState) {
-    ValueState inState;
-    if (inFlow.isAbstractFunction()) {
-      AbstractFunction orFunction = inFlow.asAbstractFunction();
-      assert orFunction instanceof OrAbstractFunction;
-      inState = orFunction.apply(appView, flowGraphStateProvider, baseInState);
-    } else {
-      inState = baseInState;
-    }
-    if (!inState.isPrimitiveState()) {
-      assert inState.isBottom() || inState.isUnknown();
-      return inState;
-    }
-    AbstractValue result = apply(appView, inState.asPrimitiveState().getAbstractValue());
-    return ConcretePrimitiveTypeValueState.create(result);
-  }
-
-  /**
-   * Applies the following function to the given {@param abstractValue}.
-   *
-   * <pre>
-   * private const val changedLowBitMask = 0b001_001_001_001_001_001_001_001_001_001_0
-   * private const val changedHighBitMask = changedLowBitMask shl 1
-   * private const val changedMask = (changedLowBitMask or changedHighBitMask).inv()
-   *
-   * internal fun updateChangedFlags(flags: Int): Int {
-   *     val lowBits = flags and changedLowBitMask
-   *     val highBits = flags and changedHighBitMask
-   *     return ((flags and changedMask) or
-   *         (lowBits or (highBits shr 1)) or ((lowBits shl 1) and highBits))
-   * }
-   * </pre>
-   */
-  private AbstractValue apply(AppView<AppInfoWithLiveness> appView, AbstractValue flagsValue) {
-    if (flagsValue.isSingleNumberValue()) {
-      return apply(appView, flagsValue.asSingleNumberValue().getIntValue());
-    }
-    AbstractValueFactory factory = appView.abstractValueFactory();
-    // Load constants.
-    AbstractValue changedLowBitMaskValue =
-        factory.createUncheckedSingleNumberValue(changedLowBitMask);
-    AbstractValue changedHighBitMaskValue =
-        factory.createUncheckedSingleNumberValue(changedHighBitMask);
-    AbstractValue changedMaskValue = factory.createUncheckedSingleNumberValue(changedMask);
-    // Evaluate expression.
-    AbstractValue lowBitsValue =
-        AbstractCalculator.andIntegers(appView, flagsValue, changedLowBitMaskValue);
-    AbstractValue highBitsValue =
-        AbstractCalculator.andIntegers(appView, flagsValue, changedHighBitMaskValue);
-    AbstractValue changedBitsValue =
-        AbstractCalculator.andIntegers(appView, flagsValue, changedMaskValue);
-    return AbstractCalculator.orIntegers(
-        appView,
-        changedBitsValue,
-        lowBitsValue,
-        AbstractCalculator.shrIntegers(appView, highBitsValue, 1),
-        AbstractCalculator.andIntegers(
-            appView, AbstractCalculator.shlIntegers(appView, lowBitsValue, 1), highBitsValue));
-  }
-
-  private SingleNumberValue apply(AppView<AppInfoWithLiveness> appView, int flags) {
-    int lowBits = flags & changedLowBitMask;
-    int highBits = flags & changedHighBitMask;
-    int changedBits = flags & changedMask;
-    int result = changedBits | lowBits | (highBits >> 1) | ((lowBits << 1) & highBits);
-    return appView.abstractValueFactory().createUncheckedSingleNumberValue(result);
-  }
-
-  @Override
-  public boolean verifyContainsBaseInFlow(BaseInFlow otherInFlow) {
-    if (inFlow.isAbstractFunction()) {
-      assert inFlow.asAbstractFunction().verifyContainsBaseInFlow(otherInFlow);
-    } else {
-      assert inFlow.isBaseInFlow();
-      assert inFlow.equals(otherInFlow);
-    }
-    return true;
-  }
-
-  @Override
-  public Iterable<BaseInFlow> getBaseInFlow() {
-    if (inFlow.isAbstractFunction()) {
-      return inFlow.asAbstractFunction().getBaseInFlow();
-    }
-    assert inFlow.isBaseInFlow();
-    return IterableUtils.singleton(inFlow.asBaseInFlow());
-  }
-
-  @Override
-  public InFlowKind getKind() {
-    return InFlowKind.ABSTRACT_FUNCTION_UPDATE_CHANGED_FLAGS;
-  }
-
-  @Override
-  public int internalCompareToSameKind(InFlow other, InFlowComparator comparator) {
-    return inFlow.compareTo(other.asUpdateChangedFlagsAbstractFunction().inFlow, comparator);
-  }
-
-  @Override
-  public boolean isUpdateChangedFlagsAbstractFunction() {
-    return true;
-  }
-
-  @Override
-  public UpdateChangedFlagsAbstractFunction asUpdateChangedFlagsAbstractFunction() {
-    return this;
-  }
-
-  @Override
-  @SuppressWarnings("EqualsGetClass")
-  public boolean equals(Object obj) {
-    if (this == obj) {
-      return true;
-    }
-    if (obj == null || getClass() != obj.getClass()) {
-      return false;
-    }
-    UpdateChangedFlagsAbstractFunction fn = (UpdateChangedFlagsAbstractFunction) obj;
-    return inFlow.equals(fn.inFlow);
-  }
-
-  @Override
-  public int hashCode() {
-    return Objects.hash(getClass(), inFlow);
-  }
-}
diff --git a/src/main/java/com/android/tools/r8/profile/rewriting/ProfileRewritingCovariantReturnTypeAnnotationTransformerEventConsumer.java b/src/main/java/com/android/tools/r8/profile/rewriting/ProfileRewritingCovariantReturnTypeAnnotationTransformerEventConsumer.java
index ef07715..9186c46 100644
--- a/src/main/java/com/android/tools/r8/profile/rewriting/ProfileRewritingCovariantReturnTypeAnnotationTransformerEventConsumer.java
+++ b/src/main/java/com/android/tools/r8/profile/rewriting/ProfileRewritingCovariantReturnTypeAnnotationTransformerEventConsumer.java
@@ -4,8 +4,8 @@
 
 package com.android.tools.r8.profile.rewriting;
 
+import com.android.tools.r8.desugar.covariantreturntype.CovariantReturnTypeAnnotationTransformerEventConsumer;
 import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.ir.desugar.CovariantReturnTypeAnnotationTransformerEventConsumer;
 
 public class ProfileRewritingCovariantReturnTypeAnnotationTransformerEventConsumer
     implements CovariantReturnTypeAnnotationTransformerEventConsumer {
diff --git a/src/main/java/com/android/tools/r8/shaking/AnnotationRemover.java b/src/main/java/com/android/tools/r8/shaking/AnnotationRemover.java
index bcfa888..0e2508f 100644
--- a/src/main/java/com/android/tools/r8/shaking/AnnotationRemover.java
+++ b/src/main/java/com/android/tools/r8/shaking/AnnotationRemover.java
@@ -142,7 +142,7 @@
             .startsWith(options.itemFactory.dalvikAnnotationOptimizationPrefix)) {
           return true;
         }
-        if (isComposableAnnotationToRetain(appView, annotation, kind, mode, options)) {
+        if (isComposableAnnotationToRetain(appView, annotation, kind, mode)) {
           return true;
         }
         return shouldKeepNormalAnnotation(
@@ -309,13 +309,8 @@
   }
 
   private static boolean isComposableAnnotationToRetain(
-      AppView<?> appView,
-      DexAnnotation annotation,
-      AnnotatedKind kind,
-      Mode mode,
-      InternalOptions options) {
-    return options.getJetpackComposeOptions().isAnyOptimizationsEnabled()
-        && mode.isInitialTreeShaking()
+      AppView<?> appView, DexAnnotation annotation, AnnotatedKind kind, Mode mode) {
+    return mode.isInitialTreeShaking()
         && kind.isMethod()
         && annotation
             .getAnnotationType()
diff --git a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
index 5ca4dbe..1781300 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -22,6 +22,7 @@
 import com.android.tools.r8.cf.code.CfInvoke;
 import com.android.tools.r8.contexts.CompilationContext.MethodProcessingContext;
 import com.android.tools.r8.contexts.CompilationContext.ProcessorContext;
+import com.android.tools.r8.desugar.covariantreturntype.CovariantReturnTypeEnqueuerExtension;
 import com.android.tools.r8.dex.IndexedItemCollection;
 import com.android.tools.r8.dex.code.CfOrDexInstruction;
 import com.android.tools.r8.errors.InterfaceDesugarMissingTypeDiagnostic;
@@ -537,6 +538,7 @@
           shrinker -> registerAnalysis(shrinker.createEnqueuerAnalysis()));
       IsolatedFeatureSplitsChecker.register(appView, this);
       ResourceAccessAnalysis.register(appView, this);
+      CovariantReturnTypeEnqueuerExtension.register(appView, this);
     }
 
     targetedMethods = new LiveMethodsSet(graphReporter::registerMethod);
@@ -4807,7 +4809,9 @@
 
         // Notify each analysis that a fixpoint has been reached, and give each analysis an
         // opportunity to add items to the worklist.
-        analyses.forEach(analysis -> analysis.notifyFixpoint(this, worklist, timing));
+        for (EnqueuerAnalysis analysis : analyses) {
+          analysis.notifyFixpoint(this, worklist, executorService, timing);
+        }
         if (!worklist.isEmpty()) {
           continue;
         }
diff --git a/src/main/java/com/android/tools/r8/shaking/KeepInfo.java b/src/main/java/com/android/tools/r8/shaking/KeepInfo.java
index 23be490..d37185a 100644
--- a/src/main/java/com/android/tools/r8/shaking/KeepInfo.java
+++ b/src/main/java/com/android/tools/r8/shaking/KeepInfo.java
@@ -544,6 +544,10 @@
       return builder.isCheckDiscardedEnabled();
     }
 
+    public boolean isMinificationAllowed() {
+      return builder.isMinificationAllowed();
+    }
+
     public boolean isOptimizationAllowed() {
       return builder.isOptimizationAllowed();
     }
diff --git a/src/main/java/com/android/tools/r8/utils/FunctionUtils.java b/src/main/java/com/android/tools/r8/utils/FunctionUtils.java
index 2f19c5c..f7da18f 100644
--- a/src/main/java/com/android/tools/r8/utils/FunctionUtils.java
+++ b/src/main/java/com/android/tools/r8/utils/FunctionUtils.java
@@ -41,4 +41,8 @@
   public static <T, R> Function<T, R> ignoreArgument(Supplier<R> supplier) {
     return ignore -> supplier.get();
   }
+
+  public static <T, R> Function<T, R> supplyValue(R value) {
+    return ignore -> value;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/utils/InternalOptions.java b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
index dac665d..44c289f 100644
--- a/src/main/java/com/android/tools/r8/utils/InternalOptions.java
+++ b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
@@ -372,7 +372,8 @@
     if (isAndroidPlatformBuild) {
       apiModelingOptions().disableApiModeling();
       disableBackportsAndReportIfTriggered = true;
-      addAndroidPlatformBuildToMarker = isAndroidPlatformBuild;
+      addAndroidPlatformBuildToMarker = true;
+      processCovariantReturnTypeAnnotations = true;
     }
   }
 
@@ -436,6 +437,8 @@
   public boolean enableSideEffectAnalysis = true;
   public boolean enableDeterminismAnalysis = true;
   public boolean enableServiceLoaderRewriting = true;
+  public boolean allowServiceLoaderRewritingPinnedTypes =
+      System.getProperty("com.android.tools.r8.allowServiceLoaderRewritingPinnedTypes") != null;
   public boolean enableNameReflectionOptimization = true;
   public boolean enableStringConcatenationOptimization = true;
   // Enabled only for R8 (not D8).
@@ -778,7 +781,7 @@
   public OffOrAuto tryWithResourcesDesugaring = OffOrAuto.Auto;
   // Flag to turn on/off processing of @dalvik.annotation.codegen.CovariantReturnType and
   // @dalvik.annotation.codegen.CovariantReturnType$CovariantReturnTypes.
-  public boolean processCovariantReturnTypeAnnotations = true;
+  public boolean processCovariantReturnTypeAnnotations = false;
 
   public boolean loadAllClassDefinitions = false;
 
@@ -2460,7 +2463,6 @@
     public boolean enableDeadSwitchCaseElimination = true;
     public boolean enableEnqueuerDeferredTracingForReferenceFields =
         System.getProperty("com.android.tools.r8.disableEnqueuerDeferredTracing") == null;
-    public boolean enableIfThenElseAbstractFunction = false;
     public boolean enableInvokeSuperToInvokeVirtualRewriting = true;
     public boolean enableLegacyClassDefOrdering =
         System.getProperty("com.android.tools.r8.enableLegacyClassDefOrdering") != null;
diff --git a/src/main/java/com/android/tools/r8/utils/SetUtils.java b/src/main/java/com/android/tools/r8/utils/SetUtils.java
index a7c3366..ea46d64 100644
--- a/src/main/java/com/android/tools/r8/utils/SetUtils.java
+++ b/src/main/java/com/android/tools/r8/utils/SetUtils.java
@@ -5,6 +5,7 @@
 package com.android.tools.r8.utils;
 
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 import java.util.Collection;
 import java.util.Collections;
@@ -149,6 +150,12 @@
     return set;
   }
 
+  public static <T> Set<T> unionHashSet(Iterable<? extends T> one, Iterable<? extends T> other) {
+    Set<T> union = Sets.newHashSet(one);
+    Iterables.addAll(union, other);
+    return union;
+  }
+
   public static <T> Set<T> unionIdentityHashSet(Set<T> one, Set<T> other) {
     Set<T> union = Sets.newIdentityHashSet();
     union.addAll(one);
diff --git a/src/test/java/com/android/tools/r8/androidapi/GenerateCovariantReturnTypeMethodsTest.java b/src/test/java/com/android/tools/r8/androidapi/GenerateCovariantReturnTypeMethodsTest.java
index a06547e..76c0cde 100644
--- a/src/test/java/com/android/tools/r8/androidapi/GenerateCovariantReturnTypeMethodsTest.java
+++ b/src/test/java/com/android/tools/r8/androidapi/GenerateCovariantReturnTypeMethodsTest.java
@@ -5,7 +5,6 @@
 package com.android.tools.r8.androidapi;
 
 import static com.android.tools.r8.apimodel.JavaSourceCodePrinter.Type.fromType;
-import static com.android.tools.r8.ir.desugar.CovariantReturnTypeAnnotationTransformer.isCovariantReturnTypeAnnotation;
 import static com.android.tools.r8.utils.MapUtils.ignoreKey;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
@@ -23,7 +22,6 @@
 import com.android.tools.r8.graph.DexAnnotation;
 import com.android.tools.r8.graph.DexAnnotationElement;
 import com.android.tools.r8.graph.DexEncodedAnnotation;
-import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexValue;
 import com.android.tools.r8.graph.DexValue.DexValueAnnotation;
 import com.android.tools.r8.graph.DexValue.DexValueArray;
@@ -63,6 +61,11 @@
 @RunWith(Parameterized.class)
 public class GenerateCovariantReturnTypeMethodsTest extends TestBase {
 
+  private static final String COVARIANT_RETURN_TYPE_ANNOTATION_NAME =
+      "dalvik.annotation.codegen.CovariantReturnType";
+  private static final String COVARIANT_RETURN_TYPES_ANNOTATION_NAME =
+      "dalvik.annotation.codegen.CovariantReturnType$CovariantReturnTypes";
+
   private static final String CLASS_NAME = "CovariantReturnTypeMethods";
   private static final String PACKAGE_NAME = "com.android.tools.r8.androidapi";
   // When updating to support a new api level build libcore in aosp and update the cloud dependency.
@@ -200,20 +203,22 @@
     public static CovariantMethodsInJarResult create() throws Exception {
       Map<ClassReference, List<MethodReferenceWithApiLevel>> methodReferenceMap = new HashMap<>();
       CodeInspector inspector = new CodeInspector(PATH_TO_CORE_JAR);
-      DexItemFactory factory = inspector.getFactory();
       for (FoundClassSubject clazz : inspector.allClasses()) {
         clazz.forAllMethods(
             method -> {
               List<DexAnnotation> covariantAnnotations =
                   inspector.findAnnotations(
                       method.getMethod().annotations(),
-                      annotation ->
-                          isCovariantReturnTypeAnnotation(annotation.annotation, factory));
+                      annotation -> {
+                        String typeName = annotation.getAnnotationType().getTypeName();
+                        return typeName.equals(COVARIANT_RETURN_TYPE_ANNOTATION_NAME)
+                            || typeName.equals(COVARIANT_RETURN_TYPES_ANNOTATION_NAME);
+                      });
               if (!covariantAnnotations.isEmpty()) {
                 MethodReference methodReference = method.asMethodReference();
                 for (DexAnnotation covariantAnnotation : covariantAnnotations) {
                   createCovariantMethodReference(
-                      factory, methodReference, covariantAnnotation.annotation, methodReferenceMap);
+                      methodReference, covariantAnnotation.annotation, methodReferenceMap);
                 }
               }
             });
@@ -222,19 +227,21 @@
     }
 
     private static void createCovariantMethodReference(
-        DexItemFactory factory,
         MethodReference methodReference,
         DexEncodedAnnotation covariantAnnotation,
         Map<ClassReference, List<MethodReferenceWithApiLevel>> methodReferenceMap) {
-      if (covariantAnnotation.type == factory.annotationCovariantReturnType) {
+      if (covariantAnnotation
+          .getType()
+          .getTypeName()
+          .equals(COVARIANT_RETURN_TYPE_ANNOTATION_NAME)) {
         DexAnnotationElement returnTypeElement = covariantAnnotation.getElement(0);
-        assert returnTypeElement.name.toString().equals("returnType");
+        assert returnTypeElement.getName().toString().equals("returnType");
         DexValueType newReturnType = returnTypeElement.getValue().asDexValueType();
         DexAnnotationElement presentAfterElement = covariantAnnotation.getElement(1);
-        assert presentAfterElement.name.toString().equals("presentAfter");
+        assert presentAfterElement.getName().toString().equals("presentAfter");
         AndroidApiLevel apiLevel =
             AndroidApiLevel.getAndroidApiLevel(
-                presentAfterElement.getValue().asDexValueInt().value);
+                presentAfterElement.getValue().asDexValueInt().getValue());
         methodReferenceMap
             .computeIfAbsent(methodReference.getHolderClass(), ignoreKey(ArrayList::new))
             .add(
@@ -243,20 +250,23 @@
                         methodReference.getHolderClass(),
                         methodReference.getMethodName(),
                         methodReference.getFormalTypes(),
-                        newReturnType.value.asClassReference()),
+                        newReturnType.getValue().asClassReference()),
                     apiLevel));
       } else {
-        assert covariantAnnotation.type == factory.annotationCovariantReturnTypes;
+        assert covariantAnnotation
+            .getType()
+            .getTypeName()
+            .equals(COVARIANT_RETURN_TYPES_ANNOTATION_NAME);
         DexAnnotationElement valuesElement = covariantAnnotation.getElement(0);
-        assert valuesElement.name.toString().equals("value");
-        DexValueArray array = valuesElement.value.asDexValueArray();
+        assertEquals("value", valuesElement.getName().toString());
+        DexValueArray array = valuesElement.getValue().asDexValueArray();
         if (array == null) {
           fail(
               String.format(
                   "Expected element \"value\" of CovariantReturnTypes annotation to "
                       + "be an array (method: \"%s\", was: %s)",
                   methodReference.toSourceString(),
-                  valuesElement.value.getClass().getCanonicalName()));
+                  valuesElement.getValue().getClass().getCanonicalName()));
         }
 
         // Handle the inner dalvik.annotation.codegen.CovariantReturnType annotations recursively.
@@ -264,7 +274,7 @@
           assert value.isDexValueAnnotation();
           DexValueAnnotation innerAnnotation = value.asDexValueAnnotation();
           createCovariantMethodReference(
-              factory, methodReference, innerAnnotation.value, methodReferenceMap);
+              methodReference, innerAnnotation.value, methodReferenceMap);
         }
       }
     }
diff --git a/src/test/java/com/android/tools/r8/annotations/DalvikAnnotationOptimizationTest.java b/src/test/java/com/android/tools/r8/annotations/DalvikAnnotationOptimizationTest.java
index f14c979..c6c128f 100644
--- a/src/test/java/com/android/tools/r8/annotations/DalvikAnnotationOptimizationTest.java
+++ b/src/test/java/com/android/tools/r8/annotations/DalvikAnnotationOptimizationTest.java
@@ -12,7 +12,6 @@
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.ToolHelper;
-import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.origin.Origin;
 import com.android.tools.r8.transformers.MethodTransformer;
 import com.android.tools.r8.utils.AndroidApiLevel;
@@ -56,8 +55,7 @@
         BooleanUtils.values());
   }
 
-  private static final String dalvikOptimizationPrefix =
-      DexItemFactory.dalvikAnnotationOptimizationPrefixString;
+  private static final String dalvikOptimizationPrefix = "Ldalvik/annotation/optimization/";
   private static final String dalvikCodegenPrefix = "Ldalvik/annotation/codegen/";
   private static final String ourClassName =
       DescriptorUtils.javaTypeToDescriptor(DalvikAnnotationOptimizationTest.class.getTypeName());
diff --git a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkConfig.java b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkConfig.java
index 0eab29b..bff9585 100644
--- a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkConfig.java
+++ b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkConfig.java
@@ -168,8 +168,23 @@
       return this;
     }
 
-    public Builder measureComposableCodeSize() {
-      metrics.add(BenchmarkMetric.ComposableCodeSize);
+    public Builder measureInstructionCodeSize() {
+      metrics.add(BenchmarkMetric.InstructionCodeSize);
+      return this;
+    }
+
+    public Builder measureComposableInstructionCodeSize() {
+      metrics.add(BenchmarkMetric.ComposableInstructionCodeSize);
+      return this;
+    }
+
+    public Builder measureDexSegmentsCodeSize() {
+      metrics.add(BenchmarkMetric.DexSegmentsCodeSize);
+      return this;
+    }
+
+    public Builder measureDex2OatCodeSize() {
+      metrics.add(BenchmarkMetric.Dex2OatCodeSize);
       return this;
     }
 
@@ -248,7 +263,7 @@
   }
 
   public boolean containsComposableCodeSizeMetric() {
-    return containsMetric(BenchmarkMetric.ComposableCodeSize);
+    return containsMetric(BenchmarkMetric.ComposableInstructionCodeSize);
   }
 
   public BenchmarkIdentifier getIdentifier() {
diff --git a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsCollection.java b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsCollection.java
index 8fc8508..dc841ba 100644
--- a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsCollection.java
+++ b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsCollection.java
@@ -3,7 +3,9 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.benchmarks;
 
+import com.android.tools.r8.DexSegments.SegmentInfo;
 import com.android.tools.r8.errors.Unimplemented;
+import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import java.io.PrintStream;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -23,24 +25,40 @@
 
   @Override
   public void addRuntimeResult(long result) {
-    throw new BenchmarkConfigError(
-        "Unexpected attempt to add a result to a the root of a benchmark with sub-benchmarks");
+    throw error();
   }
 
   @Override
   public void addCodeSizeResult(long result) {
-    throw new BenchmarkConfigError(
-        "Unexpected attempt to add a result to a the root of a benchmark with sub-benchmarks");
+    throw error();
   }
 
   @Override
-  public void addComposableCodeSizeResult(long result) {
-    throw new BenchmarkConfigError(
-        "Unexpected attempt to add a result to a the root of a benchmark with sub-benchmarks");
+  public void addInstructionCodeSizeResult(long result) {
+    throw error();
+  }
+
+  @Override
+  public void addComposableInstructionCodeSizeResult(long result) {
+    throw error();
+  }
+
+  @Override
+  public void addDexSegmentsSizeResult(Int2ReferenceMap<SegmentInfo> result) {
+    throw error();
+  }
+
+  @Override
+  public void addDex2OatSizeResult(long result) {
+    throw error();
   }
 
   @Override
   public void addResourceSizeResult(long result) {
+    throw error();
+  }
+
+  private BenchmarkConfigError error() {
     throw new BenchmarkConfigError(
         "Unexpected attempt to add a result to a the root of a benchmark with sub-benchmarks");
   }
diff --git a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsSingle.java b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsSingle.java
index ad6b1b3..7553ea8 100644
--- a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsSingle.java
+++ b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsSingle.java
@@ -3,12 +3,19 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.benchmarks;
 
+import com.android.tools.r8.DexSegments.SegmentInfo;
+import com.android.tools.r8.dex.DexSection;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import it.unimi.dsi.fastutil.longs.LongArrayList;
 import it.unimi.dsi.fastutil.longs.LongList;
 import java.io.PrintStream;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
 import java.util.Set;
+import java.util.function.IntToLongFunction;
 import java.util.function.LongConsumer;
 
 public class BenchmarkResultsSingle implements BenchmarkResults {
@@ -16,8 +23,13 @@
   private final String name;
   private final Set<BenchmarkMetric> metrics;
   private final LongList runtimeResults = new LongArrayList();
+
+  // Consider using LongSet to eliminate duplicate results for size.
   private final LongList codeSizeResults = new LongArrayList();
-  private final LongList composableCodeSizeResults = new LongArrayList();
+  private final LongList instructionCodeSizeResults = new LongArrayList();
+  private final LongList composableInstructionCodeSizeResults = new LongArrayList();
+  private final LongList dex2OatSizeResult = new LongArrayList();
+  private final List<Int2ReferenceMap<SegmentInfo>> dexSegmentsSizeResults = new ArrayList<>();
 
   public BenchmarkResultsSingle(String name, Set<BenchmarkMetric> metrics) {
     this.name = name;
@@ -32,8 +44,20 @@
     return codeSizeResults;
   }
 
-  public LongList getComposableCodeSizeResults() {
-    return composableCodeSizeResults;
+  public LongList getInstructionCodeSizeResults() {
+    return instructionCodeSizeResults;
+  }
+
+  public LongList getComposableInstructionCodeSizeResults() {
+    return composableInstructionCodeSizeResults;
+  }
+
+  public List<Int2ReferenceMap<SegmentInfo>> getDexSegmentsSizeResults() {
+    return dexSegmentsSizeResults;
+  }
+
+  public LongList getDex2OatSizeResult() {
+    return dex2OatSizeResult;
   }
 
   public LongList getRuntimeResults() {
@@ -53,12 +77,37 @@
   }
 
   @Override
-  public void addComposableCodeSizeResult(long result) {
+  public void addInstructionCodeSizeResult(long result) {
     verifyMetric(
-        BenchmarkMetric.ComposableCodeSize,
-        metrics.contains(BenchmarkMetric.ComposableCodeSize),
+        BenchmarkMetric.InstructionCodeSize,
+        metrics.contains(BenchmarkMetric.InstructionCodeSize),
         true);
-    composableCodeSizeResults.add(result);
+    instructionCodeSizeResults.add(result);
+  }
+
+  @Override
+  public void addComposableInstructionCodeSizeResult(long result) {
+    verifyMetric(
+        BenchmarkMetric.ComposableInstructionCodeSize,
+        metrics.contains(BenchmarkMetric.ComposableInstructionCodeSize),
+        true);
+    composableInstructionCodeSizeResults.add(result);
+  }
+
+  @Override
+  public void addDexSegmentsSizeResult(Int2ReferenceMap<SegmentInfo> result) {
+    verifyMetric(
+        BenchmarkMetric.DexSegmentsCodeSize,
+        metrics.contains(BenchmarkMetric.DexSegmentsCodeSize),
+        true);
+    dexSegmentsSizeResults.add(result);
+  }
+
+  @Override
+  public void addDex2OatSizeResult(long result) {
+    verifyMetric(
+        BenchmarkMetric.Dex2OatCodeSize, metrics.contains(BenchmarkMetric.Dex2OatCodeSize), true);
+    dex2OatSizeResult.add(result);
   }
 
   @Override
@@ -91,12 +140,24 @@
         !runtimeResults.isEmpty());
     verifyMetric(
         BenchmarkMetric.CodeSize,
-        metrics.contains(BenchmarkMetric.CodeSize),
+        isBenchmarkingCodeSize() && metrics.contains(BenchmarkMetric.CodeSize),
         !codeSizeResults.isEmpty());
     verifyMetric(
-        BenchmarkMetric.ComposableCodeSize,
-        metrics.contains(BenchmarkMetric.ComposableCodeSize),
-        !composableCodeSizeResults.isEmpty());
+        BenchmarkMetric.InstructionCodeSize,
+        isBenchmarkingCodeSize() && metrics.contains(BenchmarkMetric.InstructionCodeSize),
+        !instructionCodeSizeResults.isEmpty());
+    verifyMetric(
+        BenchmarkMetric.ComposableInstructionCodeSize,
+        isBenchmarkingCodeSize() && metrics.contains(BenchmarkMetric.ComposableInstructionCodeSize),
+        !composableInstructionCodeSizeResults.isEmpty());
+    verifyMetric(
+        BenchmarkMetric.DexSegmentsCodeSize,
+        isBenchmarkingCodeSize() && metrics.contains(BenchmarkMetric.DexSegmentsCodeSize),
+        !dexSegmentsSizeResults.isEmpty());
+    verifyMetric(
+        BenchmarkMetric.Dex2OatCodeSize,
+        isBenchmarkingCodeSize() && metrics.contains(BenchmarkMetric.Dex2OatCodeSize),
+        !dex2OatSizeResult.isEmpty());
   }
 
   private void printRunTime(long duration) {
@@ -105,12 +166,29 @@
   }
 
   private void printCodeSize(long bytes) {
-    System.out.println(BenchmarkResults.prettyMetric(name, BenchmarkMetric.CodeSize, "" + bytes));
+    System.out.println(BenchmarkResults.prettyMetric(name, BenchmarkMetric.CodeSize, bytes));
   }
 
-  private void printComposableCodeSize(long bytes) {
+  private void printInstructionCodeSize(long bytes) {
     System.out.println(
-        BenchmarkResults.prettyMetric(name, BenchmarkMetric.ComposableCodeSize, "" + bytes));
+        BenchmarkResults.prettyMetric(name, BenchmarkMetric.InstructionCodeSize, bytes));
+  }
+
+  private void printComposableInstructionCodeSize(long bytes) {
+    System.out.println(
+        BenchmarkResults.prettyMetric(name, BenchmarkMetric.ComposableInstructionCodeSize, bytes));
+  }
+
+  private void printDexSegmentSize(int section, long bytes) {
+    System.out.println(
+        BenchmarkResults.prettyMetric(
+            name,
+            BenchmarkMetric.DexSegmentsCodeSize + ", " + DexSection.typeName(section),
+            bytes));
+  }
+
+  private void printDex2OatSize(long bytes) {
+    System.out.println(BenchmarkResults.prettyMetric(name, BenchmarkMetric.Dex2OatCodeSize, bytes));
   }
 
   @Override
@@ -123,23 +201,45 @@
     }
     printCodeSizeResults(codeSizeResults, failOnCodeSizeDifferences, this::printCodeSize);
     printCodeSizeResults(
-        composableCodeSizeResults, failOnCodeSizeDifferences, this::printComposableCodeSize);
+        instructionCodeSizeResults, failOnCodeSizeDifferences, this::printInstructionCodeSize);
+    printCodeSizeResults(
+        composableInstructionCodeSizeResults,
+        failOnCodeSizeDifferences,
+        this::printComposableInstructionCodeSize);
+    for (int section : DexSection.getConstants()) {
+      printCodeSizeResults(
+          dexSegmentsSizeResults,
+          i -> dexSegmentsSizeResults.get(i).get(section).getSegmentSize(),
+          failOnCodeSizeDifferences,
+          result -> printDexSegmentSize(section, result));
+    }
+    printCodeSizeResults(dex2OatSizeResult, failOnCodeSizeDifferences, this::printDex2OatSize);
   }
 
   private static void printCodeSizeResults(
       LongList codeSizeResults, boolean failOnCodeSizeDifferences, LongConsumer printer) {
-    if (!codeSizeResults.isEmpty()) {
-      long size = codeSizeResults.getLong(0);
-      if (failOnCodeSizeDifferences) {
-        for (int i = 1; i < codeSizeResults.size(); i++) {
-          if (size != codeSizeResults.getLong(i)) {
-            throw new RuntimeException(
-                "Unexpected code size difference: " + size + " and " + codeSizeResults.getLong(i));
-          }
+    printCodeSizeResults(
+        codeSizeResults, codeSizeResults::getLong, failOnCodeSizeDifferences, printer);
+  }
+
+  private static void printCodeSizeResults(
+      Collection<?> codeSizeResults,
+      IntToLongFunction getter,
+      boolean failOnCodeSizeDifferences,
+      LongConsumer printer) {
+    if (codeSizeResults.isEmpty()) {
+      return;
+    }
+    long size = getter.applyAsLong(0);
+    if (failOnCodeSizeDifferences) {
+      for (int i = 1; i < codeSizeResults.size(); i++) {
+        if (size != getter.applyAsLong(i)) {
+          throw new RuntimeException(
+              "Unexpected code size difference: " + size + " and " + getter.applyAsLong(i));
         }
       }
-      printer.accept(size);
     }
+    printer.accept(size);
   }
 
   @Override
diff --git a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsSingleAdapter.java b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsSingleAdapter.java
index 68d6fb8..86cfe62 100644
--- a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsSingleAdapter.java
+++ b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsSingleAdapter.java
@@ -3,13 +3,16 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.benchmarks;
 
-import com.android.tools.r8.utils.ListUtils;
+import com.android.tools.r8.dex.DexSection;
+import com.google.common.base.CaseFormat;
 import com.google.gson.JsonArray;
 import com.google.gson.JsonElement;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonSerializationContext;
 import com.google.gson.JsonSerializer;
 import java.lang.reflect.Type;
+import java.util.Collection;
+import java.util.function.IntToLongFunction;
 
 public class BenchmarkResultsSingleAdapter implements JsonSerializer<BenchmarkResultsSingle> {
 
@@ -17,22 +20,74 @@
   public JsonElement serialize(
       BenchmarkResultsSingle result, Type type, JsonSerializationContext jsonSerializationContext) {
     JsonArray resultsArray = new JsonArray();
-    ListUtils.forEachWithIndex(
-        result.getCodeSizeResults(),
-        (codeSizeResult, iteration) -> {
-          JsonObject resultObject = new JsonObject();
-          resultObject.addProperty("code_size", codeSizeResult);
-          if (!result.getComposableCodeSizeResults().isEmpty()) {
-            resultObject.addProperty(
-                "composable_code_size", result.getComposableCodeSizeResults().getLong(iteration));
-          }
-          resultObject.addProperty("runtime", result.getRuntimeResults().getLong(iteration));
+    for (int iteration = 0; iteration < result.getCodeSizeResults().size(); iteration++) {
+      JsonObject resultObject = new JsonObject();
+      addPropertyIfValueDifferentFromRepresentative(
+          resultObject,
+          "code_size",
+          iteration,
+          result.getCodeSizeResults(),
+          i -> result.getCodeSizeResults().getLong(i));
+      addPropertyIfValueDifferentFromRepresentative(
+          resultObject,
+          "ins_code_size",
+          iteration,
+          result.getInstructionCodeSizeResults(),
+          i -> result.getInstructionCodeSizeResults().getLong(i));
+      addPropertyIfValueDifferentFromRepresentative(
+          resultObject,
+          "composable_code_size",
+          iteration,
+          result.getComposableInstructionCodeSizeResults(),
+          i -> result.getComposableInstructionCodeSizeResults().getLong(i));
+      for (int section : DexSection.getConstants()) {
+        String sectionName = DexSection.typeName(section);
+        String sectionNameUnderscore =
+            CaseFormat.UPPER_CAMEL.to(CaseFormat.UPPER_UNDERSCORE, sectionName);
+        addPropertyIfValueDifferentFromRepresentative(
+            resultObject,
+            "dex_" + sectionNameUnderscore + "_size",
+            iteration,
+            result.getDexSegmentsSizeResults(),
+            i -> result.getDexSegmentsSizeResults().get(i).get(section).getSegmentSize());
+      }
+      addPropertyIfValueDifferentFromRepresentative(
+          resultObject,
+          "oat_code_size",
+          iteration,
+          result.getDex2OatSizeResult(),
+          i -> result.getDex2OatSizeResult().getLong(i));
+      addPropertyIfValueDifferentFromRepresentative(
+          resultObject,
+          "runtime",
+          iteration,
+          result.getRuntimeResults(),
+          i -> result.getRuntimeResults().getLong(i));
           resultsArray.add(resultObject);
-        });
+    }
 
     JsonObject benchmarkObject = new JsonObject();
     benchmarkObject.addProperty("benchmark_name", result.getName());
     benchmarkObject.add("results", resultsArray);
     return benchmarkObject;
   }
+
+  private void addPropertyIfValueDifferentFromRepresentative(
+      JsonObject resultObject,
+      String propertyName,
+      int iteration,
+      Collection<?> results,
+      IntToLongFunction getter) {
+    if (results.isEmpty()) {
+      return;
+    }
+    long result = getter.applyAsLong(iteration);
+    if (iteration != 0) {
+      long representativeResult = getter.applyAsLong(0);
+      if (result == representativeResult) {
+        return;
+      }
+    }
+    resultObject.addProperty(propertyName, result);
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsWarmup.java b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsWarmup.java
index 19c1de8..4f16cc4 100644
--- a/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsWarmup.java
+++ b/src/test/java/com/android/tools/r8/benchmarks/BenchmarkResultsWarmup.java
@@ -3,7 +3,10 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.benchmarks;
 
+import com.android.tools.r8.DexSegments.SegmentInfo;
 import com.android.tools.r8.errors.Unimplemented;
+import com.android.tools.r8.errors.Unreachable;
+import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import it.unimi.dsi.fastutil.longs.LongArrayList;
 import it.unimi.dsi.fastutil.longs.LongList;
 import java.io.PrintStream;
@@ -13,7 +16,6 @@
   private final String name;
   private final LongList runtimeResults = new LongArrayList();
   private long codeSizeResult = -1;
-  private long composableCodeSizeResult = -1;
   private long resourceSizeResult = -1;
 
   public BenchmarkResultsWarmup(String name) {
@@ -37,17 +39,27 @@
   }
 
   @Override
-  public void addComposableCodeSizeResult(long result) {
-    if (composableCodeSizeResult == -1) {
-      composableCodeSizeResult = result;
-    }
-    if (composableCodeSizeResult != result) {
-      throw new RuntimeException(
-          "Unexpected Composable code size difference: "
-              + result
-              + " and "
-              + composableCodeSizeResult);
-    }
+  public void addInstructionCodeSizeResult(long result) {
+    throw addSizeResultError();
+  }
+
+  @Override
+  public void addComposableInstructionCodeSizeResult(long result) {
+    throw addSizeResultError();
+  }
+
+  @Override
+  public void addDexSegmentsSizeResult(Int2ReferenceMap<SegmentInfo> result) {
+    throw addSizeResultError();
+  }
+
+  @Override
+  public void addDex2OatSizeResult(long result) {
+    throw addSizeResultError();
+  }
+
+  private Unreachable addSizeResultError() {
+    throw new Unreachable("Unexpected attempt to add size result for warmup run");
   }
 
   @Override
@@ -68,6 +80,11 @@
   }
 
   @Override
+  public boolean isBenchmarkingCodeSize() {
+    return false;
+  }
+
+  @Override
   public void printResults(ResultMode mode, boolean failOnCodeSizeDifferences) {
     if (runtimeResults.isEmpty()) {
       throw new BenchmarkConfigError("Expected runtime results for warmup run");
diff --git a/src/test/java/com/android/tools/r8/benchmarks/appdumps/AppDumpBenchmarkBuilder.java b/src/test/java/com/android/tools/r8/benchmarks/appdumps/AppDumpBenchmarkBuilder.java
index 09934bc..1ba1fd0 100644
--- a/src/test/java/com/android/tools/r8/benchmarks/appdumps/AppDumpBenchmarkBuilder.java
+++ b/src/test/java/com/android/tools/r8/benchmarks/appdumps/AppDumpBenchmarkBuilder.java
@@ -118,7 +118,10 @@
         .addDependency(dumpDependency)
         .measureRunTime()
         .measureCodeSize()
-        .measureComposableCodeSize()
+        .measureInstructionCodeSize()
+        .measureComposableInstructionCodeSize()
+        .measureDexSegmentsCodeSize()
+        .measureDex2OatCodeSize()
         .setTimeout(10, TimeUnit.MINUTES)
         .build();
   }
@@ -306,7 +309,9 @@
                           r ->
                               r.benchmarkCompile(results)
                                   .benchmarkCodeSize(results)
-                                  .benchmarkComposableCodeSize(results));
+                                  .benchmarkInstructionCodeSize(results)
+                                  .benchmarkDexSegmentsCodeSize(results)
+                                  .benchmarkDex2OatCodeSize(results));
                 });
   }
 
diff --git a/src/test/java/com/android/tools/r8/compose/NestedComposableArgumentPropagationTest.java b/src/test/java/com/android/tools/r8/compose/NestedComposableArgumentPropagationTest.java
index 7c40462..b8f63c1 100644
--- a/src/test/java/com/android/tools/r8/compose/NestedComposableArgumentPropagationTest.java
+++ b/src/test/java/com/android/tools/r8/compose/NestedComposableArgumentPropagationTest.java
@@ -3,8 +3,8 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.compose;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertTrue;
 
 import com.android.tools.r8.R8TestCompileResult;
 import com.android.tools.r8.TestBase;
@@ -12,6 +12,7 @@
 import com.android.tools.r8.TestParametersCollection;
 import com.android.tools.r8.ThrowableConsumer;
 import com.android.tools.r8.ToolHelper;
+import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.references.Reference;
 import com.android.tools.r8.utils.AndroidApiLevel;
@@ -40,7 +41,19 @@
   enum ComposableFunction {
     A,
     B,
-    C
+    C;
+
+    public int getExpectedNumberOfIfInstructions() {
+      switch (this) {
+        case A:
+        case B:
+          return 5;
+        case C:
+          return 3;
+        default:
+          throw new Unreachable();
+      }
+    }
   }
 
   static class CodeStats {
@@ -76,24 +89,15 @@
 
   @Test
   public void test() throws Exception {
-    EnumMap<ComposableFunction, CodeStats> defaultCodeStats = build(false);
-    EnumMap<ComposableFunction, CodeStats> optimizedCodeStats = build(true);
+    EnumMap<ComposableFunction, CodeStats> result = build();
     for (ComposableFunction composableFunction : ComposableFunction.values()) {
-      CodeStats defaultCodeStatsForFunction = defaultCodeStats.get(composableFunction);
-      CodeStats optimizedCodeStatsForFunction = optimizedCodeStats.get(composableFunction);
-      assertTrue(
-          composableFunction
-              + ": "
-              + defaultCodeStatsForFunction.numberOfIfInstructions
-              + " vs "
-              + optimizedCodeStatsForFunction.numberOfIfInstructions,
-          defaultCodeStatsForFunction.numberOfIfInstructions
-              > optimizedCodeStatsForFunction.numberOfIfInstructions);
+      CodeStats codeStats = result.get(composableFunction);
+      assertEquals(
+          composableFunction.getExpectedNumberOfIfInstructions(), codeStats.numberOfIfInstructions);
     }
   }
 
-  private EnumMap<ComposableFunction, CodeStats> build(boolean enableComposeOptimizations)
-      throws Exception {
+  private EnumMap<ComposableFunction, CodeStats> build() throws Exception {
     Box<ClassReference> mainActivityKtClassReference =
         new Box<>(Reference.classFromTypeName("com.example.MainActivityKt"));
     R8TestCompileResult compileResult =
@@ -114,12 +118,7 @@
                 updateMainActivityKt(
                     MinificationInspector::getTarget, mainActivityKtClassReference, true))
             .addOptionsModification(
-                options -> {
-                  options.getOpenClosedInterfacesOptions().suppressAllOpenInterfaces();
-                  options
-                      .getJetpackComposeOptions()
-                      .enableAllOptimizations(enableComposeOptimizations);
-                })
+                options -> options.getOpenClosedInterfacesOptions().suppressAllOpenInterfaces())
             .setMinApi(AndroidApiLevel.N)
             .allowDiagnosticMessages()
             .allowUnnecessaryDontWarnWildcards()
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/DateTimeStandaloneDayTest.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/DateTimeStandaloneDayTest.java
new file mode 100644
index 0000000..1cab68d
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/DateTimeStandaloneDayTest.java
@@ -0,0 +1,83 @@
+// Copyright (c) 2024, 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.desugar.desugaredlibrary;
+
+import static com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification.DEFAULT_SPECIFICATIONS;
+import static com.android.tools.r8.desugar.desugaredlibrary.test.LibraryDesugaringSpecification.getJdk8Jdk11;
+
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.desugar.desugaredlibrary.test.CompilationSpecification;
+import com.android.tools.r8.desugar.desugaredlibrary.test.LibraryDesugaringSpecification;
+import com.android.tools.r8.utils.AndroidApiLevel;
+import com.android.tools.r8.utils.StringUtils;
+import java.time.DayOfWeek;
+import java.time.format.TextStyle;
+import java.util.List;
+import java.util.Locale;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class DateTimeStandaloneDayTest extends DesugaredLibraryTestBase {
+
+  // TODO(b/362277530): Replace expected output when desugared library is updated.
+  private static final String EXPECTED_OUTPUT_TO_FIX =
+      StringUtils.lines("1", "2", "3", "4", "5", "6", "7");
+  private static final String EXPECTED_OUTPUT =
+      StringUtils.lines(
+          "Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday");
+
+  private final TestParameters parameters;
+  private final CompilationSpecification compilationSpecification;
+  private final LibraryDesugaringSpecification libraryDesugaringSpecification;
+
+  @Parameters(name = "{0}, spec: {1}, {2}")
+  public static List<Object[]> data() {
+    return buildParameters(
+        getTestParameters().withDexRuntimes().withAllApiLevels().build(),
+        getJdk8Jdk11(),
+        DEFAULT_SPECIFICATIONS);
+  }
+
+  public DateTimeStandaloneDayTest(
+      TestParameters parameters,
+      LibraryDesugaringSpecification libraryDesugaringSpecification,
+      CompilationSpecification compilationSpecification) {
+    this.parameters = parameters;
+    this.compilationSpecification = compilationSpecification;
+    this.libraryDesugaringSpecification = libraryDesugaringSpecification;
+  }
+
+  @Test
+  public void testStandaloneDay() throws Throwable {
+    testForDesugaredLibrary(parameters, libraryDesugaringSpecification, compilationSpecification)
+        .addInnerClasses(getClass())
+        .addKeepMainRule(TestClass.class)
+        .run(parameters.getRuntime(), TestClass.class)
+        .assertSuccessWithOutputIf(
+            parameters.getApiLevel().isGreaterThanOrEqualTo(AndroidApiLevel.O), EXPECTED_OUTPUT)
+        .assertSuccessWithOutputIf(
+            parameters.getApiLevel().isLessThan(AndroidApiLevel.O), EXPECTED_OUTPUT_TO_FIX);
+  }
+
+  public static void main(String[] args) {
+    TestClass.main(args);
+  }
+
+  static class TestClass {
+
+    public static void main(String[] args) {
+      System.out.println(DayOfWeek.MONDAY.getDisplayName(TextStyle.FULL_STANDALONE, Locale.US));
+      System.out.println(DayOfWeek.TUESDAY.getDisplayName(TextStyle.FULL_STANDALONE, Locale.US));
+      System.out.println(DayOfWeek.WEDNESDAY.getDisplayName(TextStyle.FULL_STANDALONE, Locale.US));
+      System.out.println(DayOfWeek.THURSDAY.getDisplayName(TextStyle.FULL_STANDALONE, Locale.US));
+      System.out.println(DayOfWeek.FRIDAY.getDisplayName(TextStyle.FULL_STANDALONE, Locale.US));
+      System.out.println(DayOfWeek.SATURDAY.getDisplayName(TextStyle.FULL_STANDALONE, Locale.US));
+      System.out.println(DayOfWeek.SUNDAY.getDisplayName(TextStyle.FULL_STANDALONE, Locale.US));
+    }
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/dex/DexCodeDeduppingTest.java b/src/test/java/com/android/tools/r8/dex/DexCodeDeduppingTest.java
index d5b8f32..55f23f9 100644
--- a/src/test/java/com/android/tools/r8/dex/DexCodeDeduppingTest.java
+++ b/src/test/java/com/android/tools/r8/dex/DexCodeDeduppingTest.java
@@ -239,7 +239,7 @@
   public SegmentInfo getCodeSegmentInfo(Path path)
       throws CompilationFailedException, ResourceException, IOException {
     Command.Builder builder = Command.builder().addProgramFiles(path);
-    Map<Integer, SegmentInfo> segmentInfoMap = DexSegments.run(builder.build());
+    Map<Integer, SegmentInfo> segmentInfoMap = DexSegments.runForTesting(builder.build());
     return segmentInfoMap.get(Constants.TYPE_CODE_ITEM);
   }
 
diff --git a/src/test/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysisUnitTest.java b/src/test/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysisUnitTest.java
index 8107aff..a174673 100644
--- a/src/test/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysisUnitTest.java
+++ b/src/test/java/com/android/tools/r8/ir/analysis/path/PathConstraintAnalysisUnitTest.java
@@ -18,6 +18,7 @@
 import com.android.tools.r8.ir.analysis.path.state.PathConstraintKind;
 import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.IRCode;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.FieldValueFactory;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameterFactory;
 import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeNode;
 import com.android.tools.r8.optimize.argumentpropagation.computation.ComputationTreeUnopCompareNode;
@@ -56,7 +57,8 @@
     IRCode code =
         inspector.clazz(Main.class).uniqueMethodWithOriginalName("greet").buildIR(appView);
     PathConstraintAnalysis analysis =
-        new PathConstraintAnalysis(appView, code, new MethodParameterFactory());
+        new PathConstraintAnalysis(
+            appView, code, new FieldValueFactory(), new MethodParameterFactory());
     DataflowAnalysisResult result = analysis.run(code.entryBlock());
     assertTrue(result.isSuccessfulAnalysisResult());
     SuccessfulDataflowAnalysisResult<BasicBlock, PathConstraintAnalysisState> successfulResult =
diff --git a/src/test/java/com/android/tools/r8/ir/desugar/annotations/CovariantReturnTypeAnnotationTransformerR8Test.java b/src/test/java/com/android/tools/r8/ir/desugar/annotations/CovariantReturnTypeAnnotationTransformerR8Test.java
new file mode 100644
index 0000000..37f133f
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/ir/desugar/annotations/CovariantReturnTypeAnnotationTransformerR8Test.java
@@ -0,0 +1,201 @@
+// Copyright (c) 2024, 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.ir.desugar.annotations;
+
+import static com.android.tools.r8.DiagnosticsMatcher.diagnosticMessage;
+import static com.android.tools.r8.DiagnosticsMatcher.diagnosticType;
+import static com.android.tools.r8.utils.codeinspector.AssertUtils.assertFailsCompilation;
+import static org.hamcrest.CoreMatchers.allOf;
+import static org.hamcrest.CoreMatchers.equalTo;
+
+import com.android.tools.r8.Diagnostic;
+import com.android.tools.r8.R8FullTestBuilder;
+import com.android.tools.r8.R8TestCompileResult;
+import com.android.tools.r8.TestBase;
+import com.android.tools.r8.TestCompilerBuilder.DiagnosticsConsumer;
+import com.android.tools.r8.TestDiagnosticMessages;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.TestParametersCollection;
+import com.android.tools.r8.ThrowableConsumer;
+import com.android.tools.r8.ToolHelper;
+import com.android.tools.r8.errors.NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic;
+import com.android.tools.r8.ir.desugar.annotations.CovariantReturnType.CovariantReturnTypes;
+import com.android.tools.r8.references.MethodReference;
+import com.android.tools.r8.references.Reference;
+import com.android.tools.r8.utils.ListUtils;
+import com.android.tools.r8.utils.MethodReferenceUtils;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import java.util.List;
+import java.util.Map;
+import org.hamcrest.Matcher;
+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 CovariantReturnTypeAnnotationTransformerR8Test extends TestBase {
+
+  private static final String covariantReturnTypeDescriptor =
+      "Ldalvik/annotation/codegen/CovariantReturnType;";
+  private static final String covariantReturnTypesDescriptor =
+      "Ldalvik/annotation/codegen/CovariantReturnType$CovariantReturnTypes;";
+
+  private static final Map<String, String> descriptorTransformation =
+      ImmutableMap.of(
+          descriptor(com.android.tools.r8.ir.desugar.annotations.version2.B.class),
+          descriptor(B.class),
+          descriptor(com.android.tools.r8.ir.desugar.annotations.version2.C.class),
+          descriptor(C.class),
+          descriptor(com.android.tools.r8.ir.desugar.annotations.version2.E.class),
+          descriptor(E.class),
+          descriptor(com.android.tools.r8.ir.desugar.annotations.version2.F.class),
+          descriptor(F.class),
+          descriptor(CovariantReturnType.class),
+          covariantReturnTypeDescriptor,
+          descriptor(CovariantReturnTypes.class),
+          covariantReturnTypesDescriptor);
+
+  @Parameter(0)
+  public TestParameters parameters;
+
+  @Parameters(name = "{0}")
+  public static TestParametersCollection data() {
+    return getTestParameters().withDexRuntimes().withMaximumApiLevel().build();
+  }
+
+  @Test
+  public void testDontObfuscateDontOptimizeDontShrink() throws Exception {
+    R8TestCompileResult r8CompileResult =
+        compileWithR8(
+            testBuilder -> testBuilder.addDontObfuscate().addDontOptimize().addDontShrink());
+    testOnRuntime(r8CompileResult);
+  }
+
+  @Test
+  public void testUnconditionalKeepAllPublicMethods() throws Exception {
+    R8TestCompileResult r8CompileResult =
+        compileWithR8(
+            testBuilder -> testBuilder.addKeepRules("-keep public class * { public <methods>; }"));
+    testOnRuntime(r8CompileResult);
+  }
+
+  @Test
+  public void testUnconditionalKeepAllPublicMethodsAllowObfuscation() throws Exception {
+    assertFailsCompilation(
+        () ->
+            compileWithR8(
+                testBuilder ->
+                    testBuilder.addKeepRules(
+                        "-keep,allowobfuscation public class * { public <methods>; }"),
+                this::inspectDiagnostics));
+  }
+
+  @Test
+  public void testUnconditionalKeepAllPublicMethodsAllowOptimization() throws Exception {
+    assertFailsCompilation(
+        () ->
+            compileWithR8(
+                testBuilder ->
+                    testBuilder.addKeepRules(
+                        "-keep,allowoptimization public class * { public <methods>; }"),
+                this::inspectDiagnostics));
+  }
+
+  @Test
+  public void testConditionalKeepAllPublicMethods() throws Exception {
+    assertFailsCompilation(
+        () ->
+            compileWithR8(
+                testBuilder ->
+                    testBuilder.addKeepRules(
+                        "-if public class * -keep class <1> { public <methods>; }",
+                        "-keep public class *"),
+                this::inspectDiagnostics));
+  }
+
+  private R8TestCompileResult compileWithR8(
+      ThrowableConsumer<? super R8FullTestBuilder> configuration) throws Exception {
+    return compileWithR8(configuration, TestDiagnosticMessages::assertNoMessages);
+  }
+
+  private R8TestCompileResult compileWithR8(
+      ThrowableConsumer<? super R8FullTestBuilder> configuration,
+      DiagnosticsConsumer<?> diagnosticsConsumer)
+      throws Exception {
+    return testForR8(parameters.getBackend())
+        .addProgramClasses(A.class, D.class)
+        .addProgramClassFileData(
+            transformer(com.android.tools.r8.ir.desugar.annotations.version2.B.class)
+                .replaceClassDescriptorInAnnotations(descriptorTransformation)
+                .replaceClassDescriptorInMethodInstructions(descriptorTransformation)
+                .setClassDescriptor(descriptor(B.class))
+                .transform(),
+            transformer(com.android.tools.r8.ir.desugar.annotations.version2.C.class)
+                .replaceClassDescriptorInAnnotations(descriptorTransformation)
+                .replaceClassDescriptorInMethodInstructions(descriptorTransformation)
+                .setClassDescriptor(descriptor(C.class))
+                .transform(),
+            transformer(com.android.tools.r8.ir.desugar.annotations.version2.E.class)
+                .replaceClassDescriptorInAnnotations(descriptorTransformation)
+                .replaceClassDescriptorInMethodInstructions(descriptorTransformation)
+                .setClassDescriptor(descriptor(E.class))
+                .transform(),
+            transformer(com.android.tools.r8.ir.desugar.annotations.version2.F.class)
+                .replaceClassDescriptorInAnnotations(descriptorTransformation)
+                .replaceClassDescriptorInMethodInstructions(descriptorTransformation)
+                .setClassDescriptor(descriptor(F.class))
+                .transform(),
+            transformer(CovariantReturnType.class)
+                .replaceClassDescriptorInAnnotations(descriptorTransformation)
+                .setClassDescriptor(covariantReturnTypeDescriptor)
+                .transform(),
+            transformer(CovariantReturnTypes.class)
+                .replaceClassDescriptorInAnnotations(descriptorTransformation)
+                .replaceClassDescriptorInMembers(
+                    descriptor(CovariantReturnType.class), covariantReturnTypeDescriptor)
+                .setClassDescriptor(covariantReturnTypesDescriptor)
+                .transform())
+        .addLibraryFiles(ToolHelper.getMostRecentAndroidJar())
+        .addOptionsModification(options -> options.processCovariantReturnTypeAnnotations = true)
+        .apply(configuration)
+        .setMinApi(parameters)
+        .compileWithExpectedDiagnostics(diagnosticsConsumer);
+  }
+
+  private void inspectDiagnostics(TestDiagnosticMessages diagnostics) throws Exception {
+    List<MethodReference> methods =
+        ImmutableList.of(
+            Reference.methodFromMethod(B.class.getDeclaredMethod("method")),
+            Reference.methodFromMethod(C.class.getDeclaredMethod("method")),
+            Reference.methodFromMethod(F.class.getDeclaredMethod("method")));
+    List<String> messages =
+        ListUtils.map(
+            methods,
+            method ->
+                "Methods with @CovariantReturnType annotations should be kept, but was not: "
+                    + MethodReferenceUtils.toSourceString(method));
+    List<Matcher<Diagnostic>> matchers =
+        ListUtils.map(
+            messages,
+            message ->
+                allOf(
+                    diagnosticType(NonKeptMethodWithCovariantReturnTypeAnnotationDiagnostic.class),
+                    diagnosticMessage(equalTo(message))));
+    diagnostics.assertErrorsMatch(matchers);
+  }
+
+  private void testOnRuntime(R8TestCompileResult r8CompileResult) throws Exception {
+    testForD8()
+        .addProgramClasses(Client.class)
+        .addClasspathClasses(A.class, B.class, C.class, D.class, E.class, F.class)
+        .setMinApi(parameters)
+        .compile()
+        .addRunClasspathFiles(r8CompileResult.writeToZip())
+        .run(parameters.getRuntime(), Client.class)
+        .assertSuccessWithOutputLines("a=A", "b=B", "c=C", "d=F", "e=F", "f=F");
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/ir/desugar/annotations/version2/C.java b/src/test/java/com/android/tools/r8/ir/desugar/annotations/version2/C.java
index 52fe17f..058e13e 100644
--- a/src/test/java/com/android/tools/r8/ir/desugar/annotations/version2/C.java
+++ b/src/test/java/com/android/tools/r8/ir/desugar/annotations/version2/C.java
@@ -5,6 +5,7 @@
 package com.android.tools.r8.ir.desugar.annotations.version2;
 
 import com.android.tools.r8.ir.desugar.annotations.A;
+import com.android.tools.r8.ir.desugar.annotations.B;
 import com.android.tools.r8.ir.desugar.annotations.CovariantReturnType;
 
 public class C extends B {
diff --git a/src/test/java/com/android/tools/r8/ir/desugar/annotations/version2/F.java b/src/test/java/com/android/tools/r8/ir/desugar/annotations/version2/F.java
index a133826..5423b42 100644
--- a/src/test/java/com/android/tools/r8/ir/desugar/annotations/version2/F.java
+++ b/src/test/java/com/android/tools/r8/ir/desugar/annotations/version2/F.java
@@ -6,6 +6,7 @@
 
 import com.android.tools.r8.ir.desugar.annotations.CovariantReturnType;
 import com.android.tools.r8.ir.desugar.annotations.D;
+import com.android.tools.r8.ir.desugar.annotations.E;
 
 public class F extends E {
   @CovariantReturnType(returnType = E.class, presentAfter = 25)
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/callsites/RestartLambdaPropagationWithDefaultArgumentTest.java b/src/test/java/com/android/tools/r8/ir/optimize/callsites/RestartLambdaPropagationWithDefaultArgumentTest.java
index 93c4be1..63f24f0 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/callsites/RestartLambdaPropagationWithDefaultArgumentTest.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/callsites/RestartLambdaPropagationWithDefaultArgumentTest.java
@@ -41,8 +41,6 @@
     testForR8(parameters.getBackend())
         .addInnerClasses(getClass())
         .addKeepMainRule(Main.class)
-        .addOptionsModification(
-            options -> options.getTestingOptions().enableIfThenElseAbstractFunction = true)
         .enableInliningAnnotations()
         .setMinApi(parameters)
         .compile()
@@ -71,11 +69,10 @@
                       .streamInstructions()
                       .noneMatch(
                           instruction -> instruction.isConstString("DefaultValueNeverUsed")));
-              // TODO(b/302281503): This argument is never used and should be removed.
               assertTrue(
                   mainMethodSubject
                       .streamInstructions()
-                      .anyMatch(
+                      .noneMatch(
                           instruction ->
                               instruction.isConstString("Unused[DefaultValueAlwaysUsed]")));
 
diff --git a/src/test/java/com/android/tools/r8/kotlin/KotlinClassInlinerTest.java b/src/test/java/com/android/tools/r8/kotlin/KotlinClassInlinerTest.java
index e3944a9..4014e78 100644
--- a/src/test/java/com/android/tools/r8/kotlin/KotlinClassInlinerTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/KotlinClassInlinerTest.java
@@ -8,6 +8,8 @@
 import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLINC_1_5_0;
 import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLINC_1_6_0;
 import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLINC_1_9_21;
+import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLINC_2_0_20;
+import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLIN_DEV;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isAbsent;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isAbsentIf;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
@@ -212,12 +214,26 @@
                                     kotlinParameters
                                         .getCompiler()
                                         .getCompilerVersion()
-                                        .isGreaterThanOrEqualTo(KOTLINC_1_9_21),
+                                        .isBetweenBothIncluded(KOTLINC_1_9_21, KOTLINC_2_0_20),
                                     i ->
                                         i.assertIsCompleteMergeGroup(
                                             SyntheticItemsTestUtils.syntheticLambdaClass(mainKt, 0),
                                             SyntheticItemsTestUtils.syntheticLambdaClass(
-                                                mainKt, 1)))
+                                                mainKt, 1)),
+                                    kotlinParameters
+                                        .getCompiler()
+                                        .getCompilerVersion()
+                                        .isEqualTo(KOTLIN_DEV),
+                                    i -> {
+                                      ClassReference sequencesKt =
+                                          Reference.classFromTypeName(
+                                              "kotlin.sequences.SequencesKt__SequencesKt");
+                                      i.assertIsCompleteMergeGroup(
+                                          SyntheticItemsTestUtils.syntheticLambdaClass(mainKt, 0),
+                                          SyntheticItemsTestUtils.syntheticLambdaClass(mainKt, 1),
+                                          SyntheticItemsTestUtils.syntheticLambdaClass(
+                                              sequencesKt, 0));
+                                    })
                                 .assertNoOtherClassesMerged();
                           } else {
                             assert kotlinParameters.getLambdaGeneration().isInvokeDynamic()
diff --git a/src/test/java/com/android/tools/r8/kotlin/lambda/KotlinLambdaMergingSingletonTest.java b/src/test/java/com/android/tools/r8/kotlin/lambda/KotlinLambdaMergingSingletonTest.java
index 988651b..ce259f2 100644
--- a/src/test/java/com/android/tools/r8/kotlin/lambda/KotlinLambdaMergingSingletonTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/lambda/KotlinLambdaMergingSingletonTest.java
@@ -4,12 +4,13 @@
 
 package com.android.tools.r8.kotlin.lambda;
 
+import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLINC_1_9_21;
+import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLINC_2_0_20;
 import static junit.framework.TestCase.assertEquals;
 import static org.hamcrest.CoreMatchers.containsString;
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
-import com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion;
 import com.android.tools.r8.KotlinTestBase;
 import com.android.tools.r8.KotlinTestParameters;
 import com.android.tools.r8.TestParameters;
@@ -70,13 +71,13 @@
         KotlinLambdasInInput.create(getProgramFiles(), getTestName());
     assertEquals(
         kotlinParameters.getLambdaGeneration().isInvokeDynamic()
-                && kotlinParameters.getCompilerVersion() == KotlinCompilerVersion.KOTLIN_DEV
+                && kotlinParameters.getCompilerVersion().isLessThanOrEqualTo(KOTLINC_2_0_20)
             ? 8
             : 2,
         lambdasInInput.getNumberOfJStyleLambdas());
     assertEquals(
         kotlinParameters.getLambdaGeneration().isInvokeDynamic()
-            ? (kotlinParameters.getCompilerVersion() == KotlinCompilerVersion.KOTLIN_DEV ? 0 : 6)
+            ? (kotlinParameters.getCompilerVersion().isLessThanOrEqualTo(KOTLINC_2_0_20) ? 0 : 6)
             : 7,
         lambdasInInput.getNumberOfKStyleLambdas());
 
@@ -99,9 +100,7 @@
   private void inspect(
       HorizontallyMergedClassesInspector inspector, KotlinLambdasInInput lambdasInInput) {
     // All J-style Kotlin lambdas should be merged into one class.
-    if (kotlinParameters
-        .getCompilerVersion()
-        .isLessThanOrEqualTo(KotlinCompilerVersion.KOTLINC_1_9_21)) {
+    if (kotlinParameters.getCompilerVersion().isLessThanOrEqualTo(KOTLINC_1_9_21)) {
       inspector.assertIsCompleteMergeGroup(lambdasInInput.getJStyleLambdas());
     } else {
       assertEquals(
diff --git a/src/test/java/com/android/tools/r8/kotlin/metadata/KotlinMetadataTestBase.java b/src/test/java/com/android/tools/r8/kotlin/metadata/KotlinMetadataTestBase.java
index e2339f6..055f31c 100644
--- a/src/test/java/com/android/tools/r8/kotlin/metadata/KotlinMetadataTestBase.java
+++ b/src/test/java/com/android/tools/r8/kotlin/metadata/KotlinMetadataTestBase.java
@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.kotlin.metadata;
 
+import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLINC_2_0_20;
 import static com.android.tools.r8.utils.FunctionUtils.ignoreArgument;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
 import static junit.framework.TestCase.assertNotNull;
@@ -181,13 +182,16 @@
 
   protected String unresolvedReferenceMessage(KotlinTestParameters param, String ref) {
     if (param.isKotlinDev()) {
+      return "unresolved reference 'class " + ref + "'";
+    } else if (param.getCompilerVersion().isGreaterThanOrEqualTo(KOTLINC_2_0_20)) {
       return "unresolved reference '" + ref + "'";
+    } else {
+      return "unresolved reference: " + ref;
     }
-    return "unresolved reference: " + ref;
   }
 
   protected String cannotAccessMessage(KotlinTestParameters param, String ref) {
-    if (param.isKotlinDev()) {
+    if (param.isNewerThanOrEqualTo(KOTLINC_2_0_20)) {
       return "cannot access 'class " + ref + " : Any'";
     }
     return "cannot access '" + ref + "'";
diff --git a/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInSealedClassTest.java b/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInSealedClassTest.java
index 24f257b..d6b8c34 100644
--- a/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInSealedClassTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInSealedClassTest.java
@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.kotlin.metadata;
 
+import static com.android.tools.r8.KotlinCompilerTool.KotlinCompilerVersion.KOTLINC_2_0_20;
 import static com.android.tools.r8.utils.DescriptorUtils.descriptorToJavaType;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isExtensionFunction;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
@@ -178,8 +179,11 @@
   private String sealedClassErrorMessage() {
     if (kotlinParameters.isKotlinDev()) {
       return "a class can only extend a sealed class or interface declared in the same package";
+    } else if (kotlinParameters.is(KOTLINC_2_0_20)) {
+      return "extending sealed classes or interfaces from a different module is prohibited";
+    } else {
+      return "inheritance of sealed classes or interfaces from different module is prohibited";
     }
-    return "inheritance of sealed classes or interfaces from different module is prohibited";
   }
 
   private void inspectInvalid(CodeInspector inspector) {
diff --git a/src/test/java/com/android/tools/r8/kotlin/optimize/switches/KotlinEnumSwitchTest.java b/src/test/java/com/android/tools/r8/kotlin/optimize/switches/KotlinEnumSwitchTest.java
index bcb01da..4824bf1 100644
--- a/src/test/java/com/android/tools/r8/kotlin/optimize/switches/KotlinEnumSwitchTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/optimize/switches/KotlinEnumSwitchTest.java
@@ -55,6 +55,8 @@
     testForR8(parameters.getBackend())
         // Use android.jar with java.lang.ClassValue.
         .addLibraryFiles(ToolHelper.getAndroidJar(AndroidApiLevel.U))
+        // Add java.lang.invoke.LambdaMetafactory for class file generation.
+        .applyIf(parameters.isCfRuntime(), b -> b.addLibraryFiles(ToolHelper.getCoreLambdaStubs()))
         .addProgramFiles(
             kotlinJars.getForConfiguration(kotlinParameters), kotlinc.getKotlinAnnotationJar())
         .addKeepMainRule("enumswitch.EnumSwitchKt")
diff --git a/src/test/java/com/android/tools/r8/kotlin/reflection/ReflectiveConstructionWithInlineClassTest.java b/src/test/java/com/android/tools/r8/kotlin/reflection/ReflectiveConstructionWithInlineClassTest.java
index a5c6dd3..4d28c48 100644
--- a/src/test/java/com/android/tools/r8/kotlin/reflection/ReflectiveConstructionWithInlineClassTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/reflection/ReflectiveConstructionWithInlineClassTest.java
@@ -119,6 +119,8 @@
   @Test
   public void testR8KeepDataClass() throws Exception {
     configureR8(testForR8(parameters.getBackend()).addDontObfuscate())
+        // Add java.lang.invoke.LambdaMetafactory for class file generation.
+        .applyIf(parameters.isCfRuntime(), b -> b.addLibraryFiles(ToolHelper.getCoreLambdaStubs()))
         .compile()
         .assertNoErrorMessages()
         .apply(KotlinMetadataTestBase::verifyExpectedWarningsFromKotlinReflectAndStdLib)
@@ -130,6 +132,8 @@
   public void testR8KeepDataClassAndInlineClass() throws Exception {
     configureR8(testForR8(parameters.getBackend()))
         .addKeepRules("-keep class " + PKG + ".Value { *; }")
+        // Add java.lang.invoke.LambdaMetafactory for class file generation.
+        .applyIf(parameters.isCfRuntime(), b -> b.addLibraryFiles(ToolHelper.getCoreLambdaStubs()))
         .compile()
         .assertNoErrorMessages()
         .apply(KotlinMetadataTestBase::verifyExpectedWarningsFromKotlinReflectAndStdLib)
diff --git a/src/test/java/com/android/tools/r8/naming/sourcefile/StringPoolSizeWithLazyDexStringsTest.java b/src/test/java/com/android/tools/r8/naming/sourcefile/StringPoolSizeWithLazyDexStringsTest.java
index 4598b8a..bfc8f65 100644
--- a/src/test/java/com/android/tools/r8/naming/sourcefile/StringPoolSizeWithLazyDexStringsTest.java
+++ b/src/test/java/com/android/tools/r8/naming/sourcefile/StringPoolSizeWithLazyDexStringsTest.java
@@ -66,7 +66,7 @@
 
   private void checkStringSegmentSize(R8TestCompileResult result) throws Exception {
     Map<Integer, SegmentInfo> segments =
-        DexSegments.run(Command.builder().addProgramFiles(result.writeToZip()).build());
+        DexSegments.runForTesting(Command.builder().addProgramFiles(result.writeToZip()).build());
     SegmentInfo stringInfo = segments.get(Constants.TYPE_STRING_ID_ITEM);
     assertEquals(8, stringInfo.getItemCount());
   }
diff --git a/src/test/java/com/android/tools/r8/repackage/RepackageMissingTypeCollisionTest.java b/src/test/java/com/android/tools/r8/repackage/RepackageMissingTypeCollisionTest.java
index e1f0cb0..fff87b6 100644
--- a/src/test/java/com/android/tools/r8/repackage/RepackageMissingTypeCollisionTest.java
+++ b/src/test/java/com/android/tools/r8/repackage/RepackageMissingTypeCollisionTest.java
@@ -43,7 +43,7 @@
             transformer(A.class).setClassDescriptor(newADescriptor).transform(),
             transformer(Anno.class)
                 .replaceClassDescriptorInMembers(descriptor(Missing.class), newMissingDescriptor)
-                .replaceClassDescriptorInAnnotationDefault(
+                .replaceClassDescriptorInAnnotations(
                     descriptor(Missing.class), newMissingDescriptor)
                 .transform(),
             transformer(Main.class)
@@ -96,7 +96,7 @@
             transformer(A.class).setClassDescriptor(newADescriptor).transform(),
             transformer(Anno.class)
                 .replaceClassDescriptorInMembers(descriptor(Missing.class), newMissingDescriptor)
-                .replaceClassDescriptorInAnnotationDefault(
+                .replaceClassDescriptorInAnnotations(
                     descriptor(Missing.class), newMissingDescriptor)
                 .transform(),
             transformer(Main.class)
diff --git a/src/test/testbase/java/com/android/tools/r8/Dex2OatTestRunResult.java b/src/test/testbase/java/com/android/tools/r8/Dex2OatTestRunResult.java
index 13ae71f..6b89f23 100644
--- a/src/test/testbase/java/com/android/tools/r8/Dex2OatTestRunResult.java
+++ b/src/test/testbase/java/com/android/tools/r8/Dex2OatTestRunResult.java
@@ -8,14 +8,20 @@
 
 import com.android.tools.r8.ToolHelper.ProcessResult;
 import com.android.tools.r8.utils.AndroidApp;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import org.hamcrest.CoreMatchers;
 import org.hamcrest.Matcher;
 
 public class Dex2OatTestRunResult extends SingleTestRunResult<Dex2OatTestRunResult> {
 
+  private final Path oat;
+
   public Dex2OatTestRunResult(
-      AndroidApp app, TestRuntime runtime, ProcessResult result, TestState state) {
+      AndroidApp app, Path oat, TestRuntime runtime, ProcessResult result, TestState state) {
     super(app, runtime, result, state);
+    this.oat = oat;
   }
 
   @Override
@@ -44,13 +50,7 @@
     return self();
   }
 
-  public Dex2OatTestRunResult assertVerificationErrors() {
-    assertSuccess();
-    Matcher<? super String> matcher = CoreMatchers.containsString("Verification error");
-    assertThat(
-        errorMessage("Run dex2oat did not produce verification errors.", matcher.toString()),
-        getStdErr(),
-        matcher);
-    return self();
+  public long getOatSizeOrDefault(long defaultValue) throws IOException {
+    return Files.exists(oat) ? Files.size(oat) : defaultValue;
   }
 }
diff --git a/src/test/testbase/java/com/android/tools/r8/KotlinCompilerTool.java b/src/test/testbase/java/com/android/tools/r8/KotlinCompilerTool.java
index 8b7b4dd..eb7eb2d 100644
--- a/src/test/testbase/java/com/android/tools/r8/KotlinCompilerTool.java
+++ b/src/test/testbase/java/com/android/tools/r8/KotlinCompilerTool.java
@@ -77,10 +77,11 @@
     KOTLINC_1_7_0("kotlin-compiler-1.7.0", KotlinLambdaGeneration.CLASS),
     KOTLINC_1_8_0("kotlin-compiler-1.8.0", KotlinLambdaGeneration.CLASS),
     KOTLINC_1_9_21("kotlin-compiler-1.9.21", KotlinLambdaGeneration.CLASS),
+    KOTLINC_2_0_20("kotlin-compiler-2.0.20", KotlinLambdaGeneration.INVOKE_DYNAMIC),
     KOTLIN_DEV("kotlin-compiler-dev", KotlinLambdaGeneration.INVOKE_DYNAMIC);
 
-    public static final KotlinCompilerVersion MIN_SUPPORTED_VERSION = KOTLINC_1_7_0;
-    public static final KotlinCompilerVersion MAX_SUPPORTED_VERSION = KOTLINC_1_9_21;
+    public static final KotlinCompilerVersion MIN_SUPPORTED_VERSION = KOTLINC_2_0_20;
+    public static final KotlinCompilerVersion MAX_SUPPORTED_VERSION = KOTLINC_2_0_20;
 
     private final String folder;
     private final KotlinLambdaGeneration defaultLambdaGeneration;
diff --git a/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java b/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
index 3ff7487..e9956ad 100644
--- a/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
+++ b/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
@@ -8,13 +8,18 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 
+import com.android.tools.r8.DexSegments.SegmentInfo;
+import com.android.tools.r8.TestRuntime.DexRuntime;
+import com.android.tools.r8.ToolHelper.DexVm;
 import com.android.tools.r8.ToolHelper.ProcessResult;
 import com.android.tools.r8.androidresources.AndroidResourceTestingUtils;
 import com.android.tools.r8.androidresources.AndroidResourceTestingUtils.ResourceTableInspector;
 import com.android.tools.r8.benchmarks.BenchmarkResults;
+import com.android.tools.r8.benchmarks.InstructionCodeSizeResult;
 import com.android.tools.r8.dexsplitter.SplitterTestBase.SplitRunner;
+import com.android.tools.r8.graph.DexEncodedMethod;
+import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.metadata.R8BuildMetadata;
 import com.android.tools.r8.profile.art.model.ExternalArtProfile;
 import com.android.tools.r8.profile.art.utils.ArtProfileInspector;
@@ -30,6 +35,8 @@
 import com.android.tools.r8.utils.codeinspector.CodeInspector;
 import com.android.tools.r8.utils.codeinspector.FoundClassSubject;
 import com.android.tools.r8.utils.graphinspector.GraphInspector;
+import it.unimi.dsi.fastutil.ints.Int2ReferenceLinkedOpenHashMap;
+import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
@@ -314,34 +321,71 @@
   @Override
   public R8TestCompileResult benchmarkCodeSize(BenchmarkResults results)
       throws IOException, ResourceException {
-    int applicationSizeWithFeatures =
-        AndroidApp.builder(app).addProgramFiles(features).build().applicationSize();
-    results.addCodeSizeResult(applicationSizeWithFeatures);
+    if (results.isBenchmarkingCodeSize()) {
+      int applicationSizeWithFeatures =
+          AndroidApp.builder(app).addProgramFiles(features).build().applicationSize();
+      results.addCodeSizeResult(applicationSizeWithFeatures);
+    }
     return self();
   }
 
   @Override
-  public R8TestCompileResult benchmarkComposableCodeSize(BenchmarkResults results)
+  public R8TestCompileResult benchmarkInstructionCodeSize(BenchmarkResults results)
       throws IOException {
-    int composableCodeSize = getComposableCodeSize(inspector());
-    for (Path feature : features) {
-      composableCodeSize += getComposableCodeSize(featureInspector(feature));
+    if (results.isBenchmarkingCodeSize()) {
+      InstructionCodeSizeResult result = getComposableCodeSize(inspector());
+      for (Path feature : features) {
+        result.add(getComposableCodeSize(featureInspector(feature)));
+      }
+      results.addInstructionCodeSizeResult(result.instructionCodeSize);
+      results.addComposableInstructionCodeSizeResult(result.composableInstructionCodeSize);
     }
-    results.addComposableCodeSizeResult(composableCodeSize);
     return self();
   }
 
-  private int getComposableCodeSize(CodeInspector inspector) {
+  private InstructionCodeSizeResult getComposableCodeSize(CodeInspector inspector) {
     DexType composableType =
         inspector.getFactory().createType("Landroidx/compose/runtime/Composable;");
-    int composableCodeSize = 0;
+    InstructionCodeSizeResult result = new InstructionCodeSizeResult();
     for (FoundClassSubject classSubject : inspector.allClasses()) {
-      for (ProgramMethod method : classSubject.getDexProgramClass().directProgramMethods()) {
-        if (method.getAnnotations().hasAnnotation(composableType)) {
-          composableCodeSize += method.getDefinition().getCode().asDexCode().codeSizeInBytes();
+      DexProgramClass clazz = classSubject.getDexProgramClass();
+      for (DexEncodedMethod method : clazz.methods(DexEncodedMethod::hasCode)) {
+        int instructionCodeSize = method.getCode().asDexCode().codeSizeInBytes();
+        result.instructionCodeSize += instructionCodeSize;
+        if (method.annotations().hasAnnotation(composableType)) {
+          result.composableInstructionCodeSize += instructionCodeSize;
         }
       }
     }
-    return composableCodeSize;
+    return result;
+  }
+
+  @Override
+  public R8TestCompileResult benchmarkDexSegmentsCodeSize(BenchmarkResults results)
+      throws IOException, ResourceException {
+    if (results.isBenchmarkingCodeSize()) {
+      AndroidApp appWithFeatures =
+          features.isEmpty() ? app : AndroidApp.builder(app).addProgramFiles(features).build();
+      results.addDexSegmentsSizeResult(runDexSegments(appWithFeatures));
+    }
+    return self();
+  }
+
+  private Int2ReferenceMap<SegmentInfo> runDexSegments(AndroidApp app)
+      throws IOException, ResourceException {
+    Map<Integer, SegmentInfo> result = DexSegments.runForTesting(app);
+    Int2ReferenceMap<SegmentInfo> rewrittenResult = new Int2ReferenceLinkedOpenHashMap<>();
+    rewrittenResult.putAll(result);
+    return rewrittenResult;
+  }
+
+  @Override
+  public R8TestCompileResult benchmarkDex2OatCodeSize(BenchmarkResults results) throws IOException {
+    if (results.isBenchmarkingCodeSize()) {
+      Dex2OatTestRunResult dex2OatTestRunResult =
+          runDex2Oat(new DexRuntime(DexVm.Version.LATEST_DEX2OAT));
+      results.addDex2OatSizeResult(dex2OatTestRunResult.getOatSizeOrDefault(0));
+    }
+    return self();
   }
 }
diff --git a/src/test/testbase/java/com/android/tools/r8/TestCompileResult.java b/src/test/testbase/java/com/android/tools/r8/TestCompileResult.java
index 1f50960..80cf267 100644
--- a/src/test/testbase/java/com/android/tools/r8/TestCompileResult.java
+++ b/src/test/testbase/java/com/android/tools/r8/TestCompileResult.java
@@ -715,7 +715,7 @@
     Path oatFile = tmp.resolve("out.oat");
     app.writeToZipForTesting(jarFile, OutputMode.DexIndexed);
     return new Dex2OatTestRunResult(
-        app, runtime, ToolHelper.runDex2OatRaw(jarFile, oatFile, vm), state);
+        app, oatFile, runtime, ToolHelper.runDex2OatRaw(jarFile, oatFile, vm), state);
   }
 
   public CR benchmarkCodeSize(BenchmarkResults results) throws IOException, ResourceException {
@@ -723,7 +723,17 @@
     return self();
   }
 
-  public CR benchmarkComposableCodeSize(BenchmarkResults results)
+  public CR benchmarkInstructionCodeSize(BenchmarkResults results)
+      throws IOException, ResourceException {
+    throw new Unimplemented();
+  }
+
+  public CR benchmarkDexSegmentsCodeSize(BenchmarkResults results)
+      throws IOException, ResourceException {
+    throw new Unimplemented();
+  }
+
+  public CR benchmarkDex2OatCodeSize(BenchmarkResults results)
       throws IOException, ResourceException {
     throw new Unimplemented();
   }
diff --git a/src/test/testbase/java/com/android/tools/r8/ToolHelper.java b/src/test/testbase/java/com/android/tools/r8/ToolHelper.java
index 2061f8f..d588f80 100644
--- a/src/test/testbase/java/com/android/tools/r8/ToolHelper.java
+++ b/src/test/testbase/java/com/android/tools/r8/ToolHelper.java
@@ -102,6 +102,13 @@
 public class ToolHelper {
 
   public static String getProjectRoot() {
+    String oracle = System.getProperty("REPO_ROOT");
+    if (oracle != null) {
+      if (!oracle.endsWith("/")) {
+        oracle = oracle + "/";
+      }
+      return oracle;
+    }
     String current = System.getProperty("user.dir");
     if (!current.contains("d8_r8")) {
       return "";
@@ -388,6 +395,8 @@
       // TODO(b/204855476): Rename to DEFAULT alias once the checked in VM is removed.
       public static final Version NEW_DEFAULT = DEFAULT;
 
+      public static final Version LATEST_DEX2OAT = V12_0_0;
+
       Version(String shortName) {
         this.shortName = shortName;
       }
diff --git a/src/test/testbase/java/com/android/tools/r8/benchmarks/BenchmarkMetric.java b/src/test/testbase/java/com/android/tools/r8/benchmarks/BenchmarkMetric.java
index abf93ce..11ffbcd 100644
--- a/src/test/testbase/java/com/android/tools/r8/benchmarks/BenchmarkMetric.java
+++ b/src/test/testbase/java/com/android/tools/r8/benchmarks/BenchmarkMetric.java
@@ -6,7 +6,10 @@
 public enum BenchmarkMetric {
   RunTimeRaw,
   CodeSize,
-  ComposableCodeSize,
+  InstructionCodeSize,
+  ComposableInstructionCodeSize,
+  DexSegmentsCodeSize,
+  Dex2OatCodeSize,
   StartupTime;
 
   public String getDartType() {
diff --git a/src/test/testbase/java/com/android/tools/r8/benchmarks/BenchmarkResults.java b/src/test/testbase/java/com/android/tools/r8/benchmarks/BenchmarkResults.java
index a925368..6f8d5b6 100644
--- a/src/test/testbase/java/com/android/tools/r8/benchmarks/BenchmarkResults.java
+++ b/src/test/testbase/java/com/android/tools/r8/benchmarks/BenchmarkResults.java
@@ -3,17 +3,26 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.benchmarks;
 
+import com.android.tools.r8.DexSegments.SegmentInfo;
 import com.android.tools.r8.utils.StringUtils;
+import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import java.io.PrintStream;
 
 public interface BenchmarkResults {
+
   // Append a runtime result. This may be summed or averaged depending on the benchmark set up.
   void addRuntimeResult(long result);
 
   // Append a code size result. This is always assumed to be identical if called multiple times.
   void addCodeSizeResult(long result);
 
-  void addComposableCodeSizeResult(long result);
+  void addInstructionCodeSizeResult(long result);
+
+  void addComposableInstructionCodeSizeResult(long result);
+
+  void addDexSegmentsSizeResult(Int2ReferenceMap<SegmentInfo> result);
+
+  void addDex2OatSizeResult(long result);
 
   // Append a resource size result. This is always assumed to be identical if called multiple times.
   void addResourceSizeResult(long result);
@@ -22,6 +31,10 @@
   // This will throw if called on a benchmark without sub-benchmarks.
   BenchmarkResults getSubResults(String name);
 
+  default boolean isBenchmarkingCodeSize() {
+    return true;
+  }
+
   void printResults(ResultMode resultMode, boolean failOnCodeSizeDifferences);
 
   void writeResults(PrintStream out);
@@ -30,8 +43,20 @@
     return "" + (nanoTime / 1000000) + " ms";
   }
 
+  static String prettyMetric(String name, BenchmarkMetric metric, long value) {
+    return prettyMetric(name, metric.name(), Long.toString(value));
+  }
+
   static String prettyMetric(String name, BenchmarkMetric metric, String value) {
-    return name + "(" + metric.name() + "): " + value;
+    return prettyMetric(name, metric.name(), value);
+  }
+
+  static String prettyMetric(String name, String metricName, long value) {
+    return prettyMetric(name, metricName, Long.toString(value));
+  }
+
+  static String prettyMetric(String name, String metricName, String value) {
+    return name + "(" + metricName + "): " + value;
   }
 
   enum ResultMode {
diff --git a/src/test/testbase/java/com/android/tools/r8/benchmarks/InstructionCodeSizeResult.java b/src/test/testbase/java/com/android/tools/r8/benchmarks/InstructionCodeSizeResult.java
new file mode 100644
index 0000000..20af572
--- /dev/null
+++ b/src/test/testbase/java/com/android/tools/r8/benchmarks/InstructionCodeSizeResult.java
@@ -0,0 +1,15 @@
+// Copyright (c) 2024, 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.benchmarks;
+
+public class InstructionCodeSizeResult {
+
+  public long instructionCodeSize;
+  public long composableInstructionCodeSize;
+
+  public void add(InstructionCodeSizeResult result) {
+    instructionCodeSize += result.instructionCodeSize;
+    composableInstructionCodeSize += result.composableInstructionCodeSize;
+  }
+}
diff --git a/src/test/testbase/java/com/android/tools/r8/transformers/ClassFileTransformer.java b/src/test/testbase/java/com/android/tools/r8/transformers/ClassFileTransformer.java
index 72bc22a..69bd9b2 100644
--- a/src/test/testbase/java/com/android/tools/r8/transformers/ClassFileTransformer.java
+++ b/src/test/testbase/java/com/android/tools/r8/transformers/ClassFileTransformer.java
@@ -1089,27 +1089,67 @@
         });
   }
 
-  public ClassFileTransformer replaceClassDescriptorInAnnotationDefault(
+  public ClassFileTransformer replaceClassDescriptorInAnnotations(
       String oldDescriptor, String newDescriptor) {
-    return addMethodTransformer(
-        new MethodTransformer() {
+    return replaceClassDescriptorInAnnotations(ImmutableMap.of(oldDescriptor, newDescriptor));
+  }
 
-          @Override
-          public AnnotationVisitor visitAnnotationDefault() {
-            return new AnnotationVisitor(ASM_VERSION, super.visitAnnotationDefault()) {
-              @Override
-              public void visit(String name, Object value) {
-                super.visit(name, value);
-              }
+  public ClassFileTransformer replaceClassDescriptorInAnnotations(Map<String, String> map) {
+    class AnnotationTransformer extends AnnotationVisitor {
 
-              @Override
-              public void visitEnum(String name, String descriptor, String value) {
-                super.visitEnum(
-                    name, descriptor.equals(oldDescriptor) ? newDescriptor : descriptor, value);
-              }
-            };
+      protected AnnotationTransformer(AnnotationVisitor annotationVisitor) {
+        super(ASM_VERSION, annotationVisitor);
+      }
+
+      @Override
+      public AnnotationVisitor visitAnnotation(String name, String descriptor) {
+        return new AnnotationTransformer(
+            super.visitAnnotation(name, map.getOrDefault(descriptor, descriptor)));
+      }
+
+      @Override
+      public AnnotationVisitor visitArray(String name) {
+        return new AnnotationTransformer(super.visitArray(name));
+      }
+
+      @Override
+      public void visitEnum(String name, String descriptor, String value) {
+        super.visitEnum(name, map.getOrDefault(descriptor, descriptor), value);
+      }
+
+      @Override
+      public void visit(String name, Object value) {
+        if (value instanceof Type) {
+          Type type = (Type) value;
+          if (map.containsKey(type.getDescriptor())) {
+            value = Type.getType(map.get(type.getDescriptor()));
           }
-        });
+        }
+        super.visit(name, value);
+      }
+    }
+    return addClassTransformer(
+            new ClassTransformer() {
+              @Override
+              public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
+                return new AnnotationTransformer(
+                    super.visitAnnotation(map.getOrDefault(descriptor, descriptor), visible));
+              }
+            })
+        .addMethodTransformer(
+            new MethodTransformer() {
+
+              @Override
+              public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
+                return new AnnotationTransformer(
+                    super.visitAnnotation(map.getOrDefault(descriptor, descriptor), visible));
+              }
+
+              @Override
+              public AnnotationVisitor visitAnnotationDefault() {
+                return new AnnotationTransformer(super.visitAnnotationDefault());
+              }
+            });
   }
 
   public ClassFileTransformer replaceClassDescriptorInMethodInstructions(
diff --git a/src/test/testbase/java/com/android/tools/r8/utils/codeinspector/HorizontallyMergedClassesInspector.java b/src/test/testbase/java/com/android/tools/r8/utils/codeinspector/HorizontallyMergedClassesInspector.java
index b3ca5e2..04663af 100644
--- a/src/test/testbase/java/com/android/tools/r8/utils/codeinspector/HorizontallyMergedClassesInspector.java
+++ b/src/test/testbase/java/com/android/tools/r8/utils/codeinspector/HorizontallyMergedClassesInspector.java
@@ -136,6 +136,19 @@
     return this;
   }
 
+  public HorizontallyMergedClassesInspector applyIf(
+      boolean condition,
+      ThrowableConsumer<HorizontallyMergedClassesInspector> thenConsumer,
+      boolean elseIfCondition,
+      ThrowableConsumer<HorizontallyMergedClassesInspector> elseIfThenConsumer) {
+    if (condition) {
+      thenConsumer.acceptWithRuntimeException(this);
+    } else if (elseIfCondition) {
+      elseIfThenConsumer.acceptWithRuntimeException(this);
+    }
+    return this;
+  }
+
   public HorizontallyMergedClassesInspector assertMergedInto(Class<?> from, Class<?> target) {
     return assertMergedInto(toDexType(from), toDexType(target));
   }
diff --git a/third_party/kotlin/kotlin-compiler-2.0.20.tar.gz.sha1 b/third_party/kotlin/kotlin-compiler-2.0.20.tar.gz.sha1
new file mode 100644
index 0000000..93c9bdc
--- /dev/null
+++ b/third_party/kotlin/kotlin-compiler-2.0.20.tar.gz.sha1
@@ -0,0 +1 @@
+4016121311b1e63f1fda5651fb92926028b63aaa
\ No newline at end of file
diff --git a/tools/perf/benchmark_data.json b/tools/perf/benchmark_data.json
index 9d18466..e54df9e 100644
--- a/tools/perf/benchmark_data.json
+++ b/tools/perf/benchmark_data.json
@@ -7,17 +7,48 @@
     "NowInAndroidApp": {
       "benchmark_name": "NowInAndroidApp",
       "results": [
-        { "code_size": 42, "runtime": 1 },
-        { "code_size": 42, "runtime": 2 },
-        { "code_size": 42, "runtime": 3 }
+        {
+          "code_size": 1,
+          "instruction_code_size": 2,
+          "composable_code_size": 3,
+          "oat_code_size": 4,
+          "runtime": 1
+        },
+        {
+          "code_size": 1,
+          "instruction_code_size": 2,
+          "composable_code_size": 3,
+          "oat_code_size": 4,
+          "runtime": 1
+        },
+        {
+          "code_size": 1,
+          "instruction_code_size": 2,
+          "composable_code_size": 3,
+          "oat_code_size": 4,
+          "runtime": 1
+        }
       ]
     },
     "TiviApp": {
       "benchmark_name": "TiviApp",
       "results": [
-        { "code_size": 84, "runtime": 4 },
-        { "code_size": 84, "runtime": 5 },
-        { "code_size": 84, "runtime": 6 }
+        {
+          "code_size": 1,
+          "instruction_code_size": 2,
+          "composable_code_size": 3,
+          "oat_code_size": 4,
+          "runtime": 1
+        },
+        {
+        },
+        {
+          "code_size": 1,
+          "instruction_code_size": 2,
+          "composable_code_size": 3,
+          "oat_code_size": 4,
+          "runtime": 1
+        }
       ]
     }
   }
diff --git a/tools/perf/index.html b/tools/perf/index.html
index 73c2931..8d4a6ef 100644
--- a/tools/perf/index.html
+++ b/tools/perf/index.html
@@ -121,6 +121,29 @@
       benchmarkSelectors.appendChild(label);
     }
 
+    function getSingleResult(benchmark, commit, resultName, resultIteration = 0) {
+      if (!(benchmark in commit.benchmarks)) {
+        return NaN;
+      }
+      const allResults = commit.benchmarks[benchmark].results;
+      const resultsForIteration = allResults[resultIteration];
+      // If a given iteration does not declare a result, then the result
+      // was the same as the first run.
+      if (resultIteration > 0 && !(resultName in resultsForIteration)) {
+        return allResults.first()[resultName];
+      }
+      return resultsForIteration[resultName];
+    }
+
+    function getAllResults(benchmark, commit, resultName) {
+      const result = [];
+      const allResults = commit.benchmarks[benchmark].results;
+      for (var iteration = 0; iteration < allResults.length; iteration++) {
+        result.push(getSingleResult(benchmark, commit, resultName, iteration));
+      }
+      return result;
+    }
+
     // Chart data provider.
     function getData() {
       const filteredCommits = commits.slice(zoom.left, zoom.right);
@@ -129,40 +152,28 @@
       for (const selectedBenchmark of selectedBenchmarks.values()) {
         const codeSizeData =
             filteredCommits.map(
-                (c, i) =>
-                    selectedBenchmark in filteredCommits[i].benchmarks
-                        ? filteredCommits[i]
-                            .benchmarks[selectedBenchmark]
-                            .results
-                            .first()
-                            .code_size
-                        : NaN);
-        const composableCodeSizeData =
+                (c, i) => getSingleResult(selectedBenchmark, filteredCommits[i], "code_size"));
+        const instructionCodeSizeData =
             filteredCommits.map(
-                (c, i) =>
-                    selectedBenchmark in filteredCommits[i].benchmarks
-                        ? filteredCommits[i]
-                            .benchmarks[selectedBenchmark]
-                            .results
-                            .first()
-                            .composable_code_size
-                        : NaN);
+                (c, i) => getSingleResult(selectedBenchmark, filteredCommits[i], "ins_code_size"));
+        const composableInstructionCodeSizeData =
+            filteredCommits.map(
+                (c, i) => getSingleResult(selectedBenchmark, filteredCommits[i], "composable_code_size"));
+        const oatCodeSizeData =
+            filteredCommits.map(
+                (c, i) => getSingleResult(selectedBenchmark, filteredCommits[i], "oat_code_size"));
         const codeSizeScatterData = [];
         for (const commit of filteredCommits.values()) {
           if (!(selectedBenchmark in commit.benchmarks)) {
             continue;
           }
-          const codeSizes =
-              commit.benchmarks[selectedBenchmark].results.map(result => result.code_size)
-          const expectedCodeSize = codeSizes.first();
-          if (codeSizes.any(codeSize => codeSize != expectedCodeSize)) {
-            const seen = new Set();
-            seen.add(expectedCodeSize);
-            for (const codeSize of codeSizes.values()) {
-              if (!seen.has(codeSize)) {
-                codeSizeScatterData.push({ x: commit.index, y: codeSize });
-                seen.add(codeSize);
-              }
+          const seen = new Set();
+          seen.add(getSingleResult(selectedBenchmark, commit, "code_size"));
+          const codeSizes = getAllResults(selectedBenchmark, commit, "code_size")
+          for (const codeSize of codeSizes.values()) {
+            if (!seen.has(codeSize)) {
+              codeSizeScatterData.push({ x: commit.index, y: codeSize });
+              seen.add(codeSize);
             }
           }
         }
@@ -170,10 +181,7 @@
             filteredCommits.map(
                 (c, i) =>
                     selectedBenchmark in filteredCommits[i].benchmarks
-                        ? filteredCommits[i]
-                            .benchmarks[selectedBenchmark]
-                            .results
-                            .map(result => result.runtime)
+                        ? getAllResults(selectedBenchmark, filteredCommits[i], "runtime")
                             .min()
                             .ns_to_s()
                         : NaN);
@@ -182,8 +190,7 @@
           if (!(selectedBenchmark in commit.benchmarks)) {
             continue;
           }
-          const runtimes =
-              commit.benchmarks[selectedBenchmark].results.map(result => result.runtime)
+          const runtimes = getAllResults(selectedBenchmark, commit, "runtime")
           for (const runtime of runtimes.values()) {
             runtimeScatterData.push({ x: commit.index, y: runtime.ns_to_s() });
           }
@@ -194,7 +201,7 @@
           {
             benchmark: selectedBenchmark,
             type: 'line',
-            label: 'Code size',
+            label: 'Dex size',
             data: codeSizeData,
             datalabels: {
               align: 'end',
@@ -216,14 +223,58 @@
           {
             benchmark: selectedBenchmark,
             type: 'line',
+            label: 'Instruction size',
+            data: instructionCodeSizeData,
+            datalabels: {
+              align: 'end',
+              anchor: 'end'
+            },
+            tension: 0.1,
+            yAxisID: 'y_ins_code_size',
+            segment: {
+              borderColor: ctx =>
+                  skipped(
+                      ctx,
+                      myChart
+                          ? myChart.data.datasets[ctx.datasetIndex].backgroundColor
+                          : undefined),
+              borderDash: ctx => skipped(ctx, [6, 6]),
+            },
+            spanGaps: true
+          },
+          {
+            benchmark: selectedBenchmark,
+            type: 'line',
             label: 'Composable size',
-            data: composableCodeSizeData,
+            data: composableInstructionCodeSizeData,
             datalabels: {
               align: 'start',
               anchor: 'start'
             },
             tension: 0.1,
-            yAxisID: 'y_composable_code_size',
+            yAxisID: 'y_ins_code_size',
+            segment: {
+              borderColor: ctx =>
+                  skipped(
+                      ctx,
+                      myChart
+                          ? myChart.data.datasets[ctx.datasetIndex].backgroundColor
+                          : undefined),
+              borderDash: ctx => skipped(ctx, [6, 6]),
+            },
+            spanGaps: true
+          },
+          {
+            benchmark: selectedBenchmark,
+            type: 'line',
+            label: 'Oat size',
+            data: oatCodeSizeData,
+            datalabels: {
+              align: 'start',
+              anchor: 'start'
+            },
+            tension: 0.1,
+            yAxisID: 'y_oat_code_size',
             segment: {
               borderColor: ctx =>
                   skipped(
@@ -294,7 +345,7 @@
 
     // Legend tracking.
     const legends =
-        new Set(['Code size', 'Composable size', 'Nondeterminism', 'Runtime', 'Runtime variance']);
+        new Set(['Dex size', 'Instruction size', 'Composable size', 'Oat size', 'Nondeterminism', 'Runtime', 'Runtime variance']);
     const selectedLegends =
         new Set(
               unescape(window.location.hash.substring(1))
@@ -408,7 +459,7 @@
           position: 'left',
           title: {
             display: true,
-            text: 'Code size (bytes)'
+            text: 'Dex size (bytes)'
           }
         },
         y_runtime: {
@@ -418,13 +469,20 @@
             text: 'Runtime (seconds)'
           }
         },
-        y_composable_code_size: {
+        y_ins_code_size: {
           position: 'left',
           title: {
             display: true,
-            text: 'Composable code size (bytes)'
+            text: 'Instruction size (bytes)'
           }
         },
+        y_oat_code_size: {
+          position: 'left',
+          title: {
+            display: true,
+            text: 'Oat size (bytes)'
+          }
+        }
       }
     };
 
@@ -522,8 +580,10 @@
         }
 
         // Update scales.
-        options.scales.y.display = selectedLegends.has('Code size');
-        options.scales.y_composable_code_size.display = selectedLegends.has('Composable size');
+        options.scales.y.display = selectedLegends.has('Dex size');
+        options.scales.y_ins_code_size.display =
+            selectedLegends.has('Instruction size') || selectedLegends.has('Composable size');
+        options.scales.y_oat_code_size.display = selectedLegends.has('Oat size');
         options.scales.y_runtime.display =
             selectedLegends.has('Runtime') || selectedLegends.has('Runtime variance');
 
diff --git a/tools/r8_release.py b/tools/r8_release.py
index e524f5c..ad22cb0 100755
--- a/tools/r8_release.py
+++ b/tools/r8_release.py
@@ -17,7 +17,7 @@
 
 import utils
 
-R8_DEV_BRANCH = '8.7'
+R8_DEV_BRANCH = '8.8'
 R8_VERSION_FILE = os.path.join('src', 'main', 'java', 'com', 'android', 'tools',
                                'r8', 'Version.java')
 THIS_FILE_RELATIVE = os.path.join('tools', 'r8_release.py')
diff --git a/tools/run_benchmark.py b/tools/run_benchmark.py
index f01535f..783e695 100755
--- a/tools/run_benchmark.py
+++ b/tools/run_benchmark.py
@@ -168,7 +168,9 @@
         jdk.GetJavaExecutable(jdkhome), '-Xms8g', '-Xmx8g',
         '-XX:+TieredCompilation', '-XX:TieredStopAtLevel=4',
         '-DBENCHMARK_IGNORE_CODE_SIZE_DIFFERENCES',
-        f'-DBUILD_PROP_KEEPANNO_RUNTIME_PATH={utils.REPO_ROOT}/d8_r8/keepanno/build/classes/java/main'
+        f'-DBUILD_PROP_KEEPANNO_RUNTIME_PATH={utils.REPO_ROOT}/d8_r8/keepanno/build/classes/java/main',
+        # Since we change the working directory to a temp folder.
+        f'-DREPO_ROOT={utils.REPO_ROOT}'
     ]
     if options.enable_assertions:
         cmd.append('-ea')
diff --git a/tools/test.py b/tools/test.py
index 973f8a4..aa63a0f 100755
--- a/tools/test.py
+++ b/tools/test.py
@@ -319,9 +319,19 @@
     if not os.path.exists("tools/linux/art-7.0.0/lib/libncurses.so.5"):
         os.symlink("/usr/lib/i386-linux-gnu/libncurses.so.6",
                    art7 + "/lib/libncurses.so.5")
+    if not os.path.exists("tools/linux/art-7.0.0/lib/libtinfo.so.5"):
+        # We don't have libtinfo, but this is never used, so we just need
+        # a valid 32bit library file.
+        os.symlink("/usr/lib/i386-linux-gnu/libncurses.so.6",
+                   art7 + "/lib/libtinfo.so.5")
     if not os.path.exists("tools/linux/art-7.0.0/lib64/libncurses.so.5"):
         os.symlink("/usr/lib/x86_64-linux-gnu/libncurses.so.6",
                    art7 + "/lib64/libncurses.so.5")
+    if not os.path.exists("tools/linux/art-7.0.0/lib64/libtinfo.so.5"):
+        # We don't have libtinfo, but this is never used, so we just need
+        # a valid 64bit library file.
+        os.symlink("/usr/lib/x86_64-linux-gnu/libncurses.so.6",
+                   art7 + "/lib64/libtinfo.so.5")
 
 def Main():
     (options, args) = ParseOptions()