// Copyright (c) 2022, 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.profile.art;

import static com.android.tools.r8.synthesis.SyntheticNaming.COMPANION_CLASS_SUFFIX;
import static com.android.tools.r8.synthesis.SyntheticNaming.EXTERNAL_SYNTHETIC_CLASS_SEPARATOR;

import com.android.tools.r8.TextInputStream;
import com.android.tools.r8.errors.Unreachable;
import com.android.tools.r8.references.ClassReference;
import com.android.tools.r8.references.MethodReference;
import com.android.tools.r8.references.Reference;
import com.android.tools.r8.startup.StartupProfileBuilder;
import com.android.tools.r8.utils.MethodReferenceUtils;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.util.function.Consumer;

public class ArtProfileBuilderUtils {

  public interface SyntheticToSyntheticContextGeneralization {

    ClassReference getSyntheticContextReference(ClassReference classReference);

    /**
     * When a startup profile is given to D8, the program input should be dex and the startup
     * profile should have been generated by launching the dex that is given on input. Therefore,
     * synthetic items in the ART profile should be present in the program input to D8 with the
     * exact same synthetic names as in the ART profile. This means that there is no need to
     * generalize synthetic items to their synthetic context.
     */
    static SyntheticToSyntheticContextGeneralization createForD8() {
      return classReference -> null;
    }

    /**
     * When a startup profile is given to R8, the program input is class files and the startup
     * profile should have been generated by dexing the program input (in release and intermediate
     * mode) and then launching the resulting app. The synthetic items in the resulting ART profile
     * do not exist in the program input to R8 (and the same synthetics may receive different names
     * in the R8 compilation). Therefore, synthetic items in the ART profile are generalized into
     * matching all synthetics from their synthetic context.
     */
    static SyntheticToSyntheticContextGeneralization createForR8() {
      return classReference -> {
        // TODO(b/243777722): Move this logic into synthetic items and extend the mapping from
        //  synthetic classes to their synthetic context to all synthetic kinds.
        String classDescriptor = classReference.getDescriptor();
        for (int i = 1; i < classDescriptor.length() - 1; i++) {
          if (classDescriptor.charAt(i) != '$') {
            continue;
          }
          if (classDescriptor.regionMatches(
                  i, COMPANION_CLASS_SUFFIX, 0, COMPANION_CLASS_SUFFIX.length())
              || classDescriptor.regionMatches(
                  i,
                  EXTERNAL_SYNTHETIC_CLASS_SEPARATOR,
                  0,
                  EXTERNAL_SYNTHETIC_CLASS_SEPARATOR.length())) {
            return Reference.classFromDescriptor(classDescriptor.substring(0, i) + ";");
          }
        }
        return null;
      };
    }
  }

  /**
   * Helper for creating an {@link ArtProfileBuilder} that performs callbacks on the given {@param
   * startupProfileBuilder}.
   */
  public static ArtProfileBuilder createBuilderForArtProfileToStartupProfileConversion(
      StartupProfileBuilder startupProfileBuilder,
      SyntheticToSyntheticContextGeneralization syntheticToSyntheticContextGeneralization) {
    return new ArtProfileBuilder() {

      @Override
      public ArtProfileBuilder addClassRule(
          Consumer<ArtProfileClassRuleBuilder> classRuleBuilderConsumer) {
        MutableArtProfileClassRule classRule = new MutableArtProfileClassRule();
        classRuleBuilderConsumer.accept(classRule);
        ClassReference syntheticContextReference =
            syntheticToSyntheticContextGeneralization.getSyntheticContextReference(
                classRule.getClassReference());
        if (syntheticContextReference == null) {
          startupProfileBuilder.addStartupClass(
              startupClassBuilder ->
                  startupClassBuilder.setClassReference(classRule.getClassReference()));
        } else {
          startupProfileBuilder.addSyntheticStartupMethod(
              syntheticStartupMethodBuilder ->
                  syntheticStartupMethodBuilder.setSyntheticContextReference(
                      syntheticContextReference));
        }
        return this;
      }

      @Override
      public ArtProfileBuilder addMethodRule(
          Consumer<ArtProfileMethodRuleBuilder> methodRuleBuilderConsumer) {
        MutableArtProfileMethodRule methodRule = new MutableArtProfileMethodRule();
        methodRuleBuilderConsumer.accept(methodRule);
        ClassReference syntheticContextReference =
            syntheticToSyntheticContextGeneralization.getSyntheticContextReference(
                methodRule.getMethodReference().getHolderClass());
        if (syntheticContextReference == null) {
          startupProfileBuilder.addStartupMethod(
              startupMethodBuilder ->
                  startupMethodBuilder.setMethodReference(methodRule.getMethodReference()));
        } else {
          startupProfileBuilder.addSyntheticStartupMethod(
              syntheticStartupMethodBuilder ->
                  syntheticStartupMethodBuilder.setSyntheticContextReference(
                      syntheticContextReference));
        }
        return this;
      }

      @Override
      public ArtProfileBuilder addHumanReadableArtProfile(
          TextInputStream textInputStream,
          Consumer<HumanReadableArtProfileParserBuilder> parserBuilderConsumer) {
        // The ART profile parser never calls addHumanReadableArtProfile().
        throw new Unreachable();
      }
    };
  }

  static class MutableArtProfileClassRule implements ArtProfileClassRuleBuilder {

    private ClassReference classReference;

    MutableArtProfileClassRule() {}

    public ClassReference getClassReference() {
      return classReference;
    }

    @Override
    public ArtProfileClassRuleBuilder setClassReference(ClassReference classReference) {
      this.classReference = classReference;
      return this;
    }

    public ArtProfileClassRuleInfo getClassRuleInfo() {
      return ArtProfileClassRuleInfoImpl.empty();
    }

    public void writeHumanReadableRuleString(OutputStreamWriter writer) throws IOException {
      writer.write(classReference.getDescriptor());
    }
  }

  static class MutableArtProfileMethodRule implements ArtProfileMethodRuleBuilder {

    private MethodReference methodReference;
    private ArtProfileMethodRuleInfoImpl methodRuleInfo = ArtProfileMethodRuleInfoImpl.empty();

    MutableArtProfileMethodRule() {}

    public MethodReference getMethodReference() {
      return methodReference;
    }

    public ArtProfileMethodRuleInfo getMethodRuleInfo() {
      return methodRuleInfo;
    }

    @Override
    public ArtProfileMethodRuleBuilder setMethodReference(MethodReference methodReference) {
      this.methodReference = methodReference;
      return this;
    }

    @Override
    public ArtProfileMethodRuleBuilder setMethodRuleInfo(
        Consumer<ArtProfileMethodRuleInfoBuilder> methodRuleInfoBuilderConsumer) {
      ArtProfileMethodRuleInfoImpl.Builder methodRuleInfoBuilder =
          ArtProfileMethodRuleInfoImpl.builder();
      methodRuleInfoBuilderConsumer.accept(methodRuleInfoBuilder);
      methodRuleInfo = methodRuleInfoBuilder.build();
      return this;
    }

    public void writeHumanReadableRuleString(OutputStreamWriter writer) throws IOException {
      methodRuleInfo.writeHumanReadableFlags(writer);
      writer.write(MethodReferenceUtils.toSmaliString(methodReference));
    }
  }
}