// 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.OrAbstractFunction;
import com.android.tools.r8.optimize.argumentpropagation.codescanner.ValueState;
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 ValueState 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);
  }
}