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

import com.android.tools.r8.graph.AppView;
import com.android.tools.r8.graph.DexEncodedField;
import com.android.tools.r8.graph.DexEncodedMethod;
import com.android.tools.r8.graph.DexType;
import com.android.tools.r8.ir.analysis.type.TypeLatticeElement;
import com.android.tools.r8.ir.code.Assume;
import com.android.tools.r8.ir.code.Assume.DynamicTypeAssumption;
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.InstructionListIterator;
import com.android.tools.r8.ir.code.InvokeMethod;
import com.android.tools.r8.ir.code.JumpInstruction;
import com.android.tools.r8.ir.code.Position;
import com.android.tools.r8.ir.code.StaticGet;
import com.android.tools.r8.ir.code.Value;
import com.android.tools.r8.ir.optimize.info.MethodOptimizationInfo;
import com.android.tools.r8.shaking.AppInfoWithLiveness;
import com.google.common.base.Predicates;
import java.util.ArrayList;
import java.util.List;
import java.util.ListIterator;
import java.util.function.Predicate;

public class DynamicTypeOptimization {

  private final AppView<AppInfoWithLiveness> appView;

  public DynamicTypeOptimization(AppView<AppInfoWithLiveness> appView) {
    this.appView = appView;
  }

  public void insertAssumeDynamicTypeInstructions(IRCode code) {
    insertAssumeDynamicTypeInstructionsInBlocks(code, code.listIterator(), Predicates.alwaysTrue());
  }

  public void insertAssumeDynamicTypeInstructionsInBlocks(
      IRCode code, ListIterator<BasicBlock> blockIterator, Predicate<BasicBlock> blockTester) {
    while (blockIterator.hasNext()) {
      BasicBlock block = blockIterator.next();
      if (blockTester.test(block)) {
        insertAssumeDynamicTypeInstructionsInBlock(code, blockIterator, block);
      }
    }
  }

  // TODO(b/127461806): Should also insert AssumeDynamicType instructions after instanceof
  //  instructions.
  private void insertAssumeDynamicTypeInstructionsInBlock(
      IRCode code, ListIterator<BasicBlock> blockIterator, BasicBlock block) {
    InstructionListIterator instructionIterator = block.listIterator(code);
    while (instructionIterator.hasNext()) {
      Instruction current = instructionIterator.next();
      Value outValue = current.outValue();
      if (outValue == null) {
        continue;
      }

      TypeLatticeElement dynamicType;
      if (current.isInvokeMethod()) {
        InvokeMethod invoke = current.asInvokeMethod();

        DexType staticReturnTypeRaw = invoke.getInvokedMethod().proto.returnType;
        if (!staticReturnTypeRaw.isReferenceType()) {
          continue;
        }

        DexEncodedMethod singleTarget =
            invoke.lookupSingleTarget(appView, code.method.method.holder);
        if (singleTarget == null) {
          continue;
        }

        MethodOptimizationInfo optimizationInfo = singleTarget.getOptimizationInfo();
        if (optimizationInfo.returnsArgument()) {
          continue;
        }

        dynamicType = optimizationInfo.getDynamicReturnType();
      } else if (current.isStaticGet()) {
        StaticGet staticGet = current.asStaticGet();
        DexEncodedField encodedField = appView.appInfo().resolveField(staticGet.getField());
        if (encodedField == null) {
          continue;
        }

        dynamicType = encodedField.getOptimizationInfo().getDynamicType();
      } else {
        continue;
      }

      if (dynamicType == null
          || !dynamicType.strictlyLessThan(outValue.getTypeLattice(), appView)) {
        continue;
      }

      // Split block if needed (only debug instructions are allowed after the throwing
      // instruction, if any).
      BasicBlock insertionBlock =
          block.hasCatchHandlers() ? instructionIterator.split(code, blockIterator) : block;

      // Replace usages of out-value by the out-value of the AssumeDynamicType instruction.
      Value specializedOutValue =
          code.createValue(outValue.getTypeLattice(), outValue.getLocalInfo());
      outValue.replaceUsers(specializedOutValue);

      // Insert AssumeDynamicType instruction.
      Assume<DynamicTypeAssumption> assumeInstruction =
          Assume.createAssumeDynamicTypeInstruction(
              dynamicType, specializedOutValue, outValue, current, appView);
      assumeInstruction.setPosition(
          appView.options().debug ? current.getPosition() : Position.none());
      if (insertionBlock == block) {
        instructionIterator.add(assumeInstruction);
      } else {
        insertionBlock.listIterator(code).add(assumeInstruction);
      }
    }
  }

  /**
   * Computes the dynamic return type of the given method.
   *
   * <p>If the method has no normal exits, then null is returned.
   */
  public TypeLatticeElement computeDynamicReturnType(DexEncodedMethod method, IRCode code) {
    assert method.method.proto.returnType.isReferenceType();
    List<TypeLatticeElement> returnedTypes = new ArrayList<>();
    for (BasicBlock block : code.blocks) {
      JumpInstruction exitInstruction = block.exit();
      if (exitInstruction.isReturn()) {
        returnedTypes.add(exitInstruction.asReturn().getReturnType());
      }
    }
    return returnedTypes.isEmpty() ? null : TypeLatticeElement.join(returnedTypes, appView);
  }
}