src/main/java/com/android/tools/r8/keepradius/RootSetKeepRadius.java

// Copyright (c) 2026, 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.keepradius;

import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
import com.android.tools.r8.graph.AppView;
import com.android.tools.r8.graph.DexField;
import com.android.tools.r8.graph.DexMethod;
import com.android.tools.r8.graph.DexReference;
import com.android.tools.r8.graph.DexType;
import com.android.tools.r8.shaking.Enqueuer;
import com.android.tools.r8.shaking.KeepClassInfo;
import com.android.tools.r8.shaking.KeepClassMembersNoShrinkingOfInitializerOnSubclassesFakeProguardRule;
import com.android.tools.r8.shaking.KeepFieldInfo;
import com.android.tools.r8.shaking.KeepInfo;
import com.android.tools.r8.shaking.KeepInfoCollectionEventConsumer;
import com.android.tools.r8.shaking.KeepMethodInfo;
import com.android.tools.r8.shaking.ProguardKeepRuleBase;
import com.android.tools.r8.shaking.rules.KeepAnnotationFakeProguardRule;
import com.android.tools.r8.utils.internal.ListUtils;
import java.util.Collection;
import java.util.Comparator;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

public class RootSetKeepRadius {

  private final Map<ProguardKeepRuleBase, RootSetKeepRadiusForRule> keepRadius;

  private RootSetKeepRadius(Map<ProguardKeepRuleBase, RootSetKeepRadiusForRule> keepRadius) {
    this.keepRadius = keepRadius;
  }

  public static Builder builder(
      AppView<? extends AppInfoWithClassHierarchy> appView, Enqueuer.Mode mode) {
    return mode.isInitialTreeShaking() && appView.options().getKeepRadiusOptions().isEnabled()
        ? new Builder()
        : null;
  }

  public Collection<RootSetKeepRadiusForRule> getKeepRadius() {
    return keepRadius.values();
  }

  public Collection<RootSetKeepRadiusForRule> getKeepRadiusWithDeterministicOrder() {
    // TODO(b/441055269): Sorting by source is not guaranteed to be deterministic.
    return ListUtils.sort(getKeepRadius(), Comparator.comparing(x -> x.getRule().getSource()));
  }

  public Map<RootSetKeepRadiusForRule, Collection<RootSetKeepRadiusForRule>> getSubsumedByInfo(
      KeepRadiusOptions options) {
    return new KeepRuleSubsumptionAnalysis(this).run(options);
  }

  public static class Builder implements KeepInfoCollectionEventConsumer {

    private final Map<ProguardKeepRuleBase, RootSetKeepRadiusForRule> keepRadius =
        new IdentityHashMap<>();

    @Override
    public void acceptKeepClassInfo(
        DexType type, Consumer<? super KeepClassInfo.Joiner> keepInfoEffect) {
      acceptKeepInfo(
          type,
          keepInfoEffect,
          KeepClassInfo.newEmptyJoiner(),
          RootSetKeepRadiusForRule::addMatchedClass);
    }

    @Override
    public void acceptKeepFieldInfo(
        DexField field, Consumer<? super KeepFieldInfo.Joiner> keepInfoEffect) {
      acceptKeepInfo(
          field,
          keepInfoEffect,
          KeepFieldInfo.newEmptyJoiner(),
          RootSetKeepRadiusForRule::addMatchedField);
    }

    @Override
    public void acceptKeepMethodInfo(
        DexMethod method, Consumer<? super KeepMethodInfo.Joiner> keepInfoEffect) {
      acceptKeepInfo(
          method,
          keepInfoEffect,
          KeepMethodInfo.newEmptyJoiner(),
          RootSetKeepRadiusForRule::addMatchedMethod);
    }

    private <R extends DexReference, J extends KeepInfo.Joiner<?, ?, ?>> void acceptKeepInfo(
        R reference,
        Consumer<? super J> keepInfoEffect,
        J emptyJoiner,
        BiConsumer<RootSetKeepRadiusForRule, R> addReferenceToRuleKeepRadius) {
      keepInfoEffect.accept(emptyJoiner);
      for (ProguardKeepRuleBase rule : emptyJoiner.getRules()) {
        if (rule.isProguardIfRule()) {
          // Perform attribution to the root -if rule.
          rule = rule.asProguardIfRule().getParentOrThis();
        }
        RootSetKeepRadiusForRule ruleKeepRadius =
            keepRadius.computeIfAbsent(rule, RootSetKeepRadiusForRule::new);
        addReferenceToRuleKeepRadius.accept(ruleKeepRadius, reference);
      }
    }

    public RootSetKeepRadius build(AppView<? extends AppInfoWithClassHierarchy> appView) {
      // Add all rules so that the keep radius result also contains empty rules.
      for (var rule : appView.options().getProguardConfiguration().getRules()) {
        if (rule instanceof ProguardKeepRuleBase) {
          keepRadius.computeIfAbsent((ProguardKeepRuleBase) rule, RootSetKeepRadiusForRule::new);
        }
      }
      // Remove fake rules from output.
      keepRadius.keySet().removeIf(Builder::isFakeKeepRule);
      return new RootSetKeepRadius(keepRadius);
    }

    private static boolean isFakeKeepRule(ProguardKeepRuleBase rule) {
      return rule instanceof KeepAnnotationFakeProguardRule
          || rule instanceof KeepClassMembersNoShrinkingOfInitializerOnSubclassesFakeProguardRule;
    }
  }
}