// Copyright (c) 2017, 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 stringconcat;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.invoke.CallSite;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Handle;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

public class TestGenerator {
  private static final String RECIPE_PREFIX = "RECIPE:";

  private static final String STRING_CONCAT_FACTORY = "java/lang/invoke/StringConcatFactory";

  private static final Handle MAKE_CONCAT_WITH_CONSTANTS = new Handle(
      Opcodes.H_INVOKESTATIC, STRING_CONCAT_FACTORY, "makeConcatWithConstants",
      MethodType.methodType(CallSite.class, MethodHandles.Lookup.class, String.class,
          MethodType.class, String.class, Object[].class).toMethodDescriptorString(),
      false);

  private static final Handle MAKE_CONCAT = new Handle(
      Opcodes.H_INVOKESTATIC, STRING_CONCAT_FACTORY, "makeConcat",
      MethodType.methodType(CallSite.class, MethodHandles.Lookup.class,
          String.class, MethodType.class).toMethodDescriptorString(),
      false);

  public static void main(String[] args) throws IOException {
    assert args.length == 1;
    generateTests(Paths.get(args[0],
        TestGenerator.class.getPackage().getName(),
        StringConcat.class.getSimpleName() + ".class"));
  }

  private static void generateTests(Path classNamePath) throws IOException {
    try (InputStream input = Files.newInputStream(classNamePath)) {
      ClassReader cr = new ClassReader(input);
      ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
      cr.accept(
          new ClassVisitor(Opcodes.ASM7, cw) {
            @Override
            public MethodVisitor visitMethod(int access,
                final String methodName, String desc, String signature, String[] exceptions) {
              MethodVisitor mv = super.visitMethod(access, methodName, desc, signature, exceptions);
              return new MethodVisitor(Opcodes.ASM7, mv) {
                private List<Object> recentConstants = new ArrayList<>();

                @Override
                public void visitLdcInsn(Object cst) {
                  if (!recentConstants.isEmpty() ||
                      (cst instanceof String && ((String) cst).startsWith(RECIPE_PREFIX))) {
                    // Add the constant, don't push anything on stack.
                    recentConstants.add(cst);
                    return;
                  }
                  super.visitLdcInsn(cst);
                }

                @Override
                public void visitMethodInsn(
                    int opcode, String owner, String name, String desc, boolean itf) {
                  // Replace calls to 'makeConcat(...)' with appropriate `invokedynamic`.
                  if (opcode == Opcodes.INVOKESTATIC && name.equals("makeConcat")) {
                    mv.visitInvokeDynamicInsn(MAKE_CONCAT.getName(), desc, MAKE_CONCAT);
                    recentConstants.clear();
                    return;
                  }

                  // Replace calls to 'makeConcat(...)' with appropriate `invokedynamic`.
                  if (opcode == Opcodes.INVOKESTATIC && name.equals("makeConcatWithConstants")) {
                    if (recentConstants.isEmpty()) {
                      throw new AssertionError("No constants detected in `" +
                          methodName + "`: call to " + name + desc);
                    }
                    recentConstants.set(0,
                        ((String) recentConstants.get(0)).substring(RECIPE_PREFIX.length()));

                    mv.visitInvokeDynamicInsn(
                        MAKE_CONCAT_WITH_CONSTANTS.getName(),
                        removeLastParams(desc, recentConstants.size()),
                        MAKE_CONCAT_WITH_CONSTANTS,
                        recentConstants.toArray());
                    recentConstants.clear();
                    return;
                  }

                  // Otherwise fall back to default implementation.
                  super.visitMethodInsn(opcode, owner, name, desc, itf);
                }

                private String removeLastParams(String descr, int paramsToRemove) {
                  MethodType methodType =
                      MethodType.fromMethodDescriptorString(
                          descr, this.getClass().getClassLoader());
                  return methodType
                      .dropParameterTypes(
                          methodType.parameterCount() - paramsToRemove,
                          methodType.parameterCount())
                      .toMethodDescriptorString();
                }

                @Override
                public void visitInsn(int opcode) {
                  switch (opcode) {
                    case Opcodes.ICONST_0:
                      if (!recentConstants.isEmpty()) {
                        recentConstants.add(0);
                        return;
                      }
                      break;
                    case Opcodes.ICONST_1:
                      if (!recentConstants.isEmpty()) {
                        recentConstants.add(1);
                        return;
                      }
                      break;
                    case Opcodes.ICONST_2:
                      if (!recentConstants.isEmpty()) {
                        recentConstants.add(2);
                        return;
                      }
                      break;
                    case Opcodes.ICONST_3:
                      if (!recentConstants.isEmpty()) {
                        recentConstants.add(3);
                        return;
                      }
                      break;
                    case Opcodes.ICONST_4:
                      if (!recentConstants.isEmpty()) {
                        recentConstants.add(4);
                        return;
                      }
                      break;
                    case Opcodes.ICONST_5:
                      if (!recentConstants.isEmpty()) {
                        recentConstants.add(5);
                        return;
                      }
                      break;
                    case Opcodes.ICONST_M1:
                      if (!recentConstants.isEmpty()) {
                        recentConstants.add(-1);
                        return;
                      }
                      break;
                    default:
                      recentConstants.clear();
                      break;
                  }
                  super.visitInsn(opcode);
                }

                @Override
                public void visitIntInsn(int opcode, int operand) {
                  recentConstants.clear();
                  super.visitIntInsn(opcode, operand);
                }

                @Override
                public void visitVarInsn(int opcode, int var) {
                  recentConstants.clear();
                  super.visitVarInsn(opcode, var);
                }

                @Override
                public void visitTypeInsn(int opcode, String type) {
                  recentConstants.clear();
                  super.visitTypeInsn(opcode, type);
                }

                @Override
                public void visitFieldInsn(int opcode, String owner, String name, String desc) {
                  recentConstants.clear();
                  super.visitFieldInsn(opcode, owner, name, desc);
                }

                @Override
                public void visitJumpInsn(int opcode, Label label) {
                  recentConstants.clear();
                  super.visitJumpInsn(opcode, label);
                }

                @Override
                public void visitIincInsn(int var, int increment) {
                  recentConstants.clear();
                  super.visitIincInsn(var, increment);
                }

                @Override
                public void visitTableSwitchInsn(int min, int max, Label dflt, Label... labels) {
                  recentConstants.clear();
                  super.visitTableSwitchInsn(min, max, dflt, labels);
                }

                @Override
                public void visitLookupSwitchInsn(Label dflt, int[] keys, Label[] labels) {
                  recentConstants.clear();
                  super.visitLookupSwitchInsn(dflt, keys, labels);
                }

                @Override
                public void visitMultiANewArrayInsn(String desc, int dims) {
                  recentConstants.clear();
                  super.visitMultiANewArrayInsn(desc, dims);
                }
              };
            }
          }, 0);
      try (OutputStream output = Files.newOutputStream(classNamePath)) {
        output.write(cw.toByteArray());
      }
    }
  }
}