content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.tooling.muzzle.generation;
  
  import static java.util.logging.Level.INFO;
  
  import com.google.errorprone.annotations.CanIgnoreReturnValue;
  import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
  import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
  import io.opentelemetry.javaagent.extension.instrumentation.internal.AsmApi;
  import io.opentelemetry.javaagent.tooling.muzzle.HelperResource;
  import io.opentelemetry.javaagent.tooling.muzzle.HelperResourceBuilderImpl;
  import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationModuleMuzzle;
  import io.opentelemetry.javaagent.tooling.muzzle.ReferenceCollector;
  import io.opentelemetry.javaagent.tooling.muzzle.VirtualFieldMappings;
  import io.opentelemetry.javaagent.tooling.muzzle.references.ClassRef;
  import io.opentelemetry.javaagent.tooling.muzzle.references.ClassRefBuilder;
  import io.opentelemetry.javaagent.tooling.muzzle.references.FieldRef;
  import io.opentelemetry.javaagent.tooling.muzzle.references.Flag;
  import io.opentelemetry.javaagent.tooling.muzzle.references.MethodRef;
  import io.opentelemetry.javaagent.tooling.muzzle.references.Source;
  import java.net.URLClassLoader;
  import java.util.Collection;
  import java.util.List;
  import java.util.Set;
  import java.util.logging.Logger;
  import java.util.regex.Matcher;
  import java.util.regex.Pattern;
  import net.bytebuddy.asm.AsmVisitorWrapper;
  import net.bytebuddy.description.field.FieldDescription;
  import net.bytebuddy.description.field.FieldList;
  import net.bytebuddy.description.method.MethodList;
  import net.bytebuddy.description.type.TypeDescription;
  import net.bytebuddy.implementation.Implementation;
  import net.bytebuddy.pool.TypePool;
  import org.objectweb.asm.ClassVisitor;
  import org.objectweb.asm.ClassWriter;
  import org.objectweb.asm.MethodVisitor;
  import org.objectweb.asm.Opcodes;
  import org.objectweb.asm.Type;
  
  /**
   * This class generates the actual implementation of the {@code
   * InstrumentationModule#getMuzzleReferences()} method. It collects references from all advice
   * classes defined in an instrumentation and writes them as Java bytecode in the generated {@code
   * InstrumentationModule#getMuzzleReferences()} method.
   *
   * <p>This class is run at compile time by the {@link MuzzleCodeGenerationPlugin} ByteBuddy plugin.
   */
  final class MuzzleCodeGenerator implements AsmVisitorWrapper {
    private static final Logger logger = Logger.getLogger(MuzzleCodeGenerator.class.getName());
  
    private static final String MUZZLE_REFERENCES_METHOD_NAME = "getMuzzleReferences";
    private static final String MUZZLE_HELPER_CLASSES_METHOD_NAME = "getMuzzleHelperClassNames";
    private static final String MUZZLE_VIRTUAL_FIELDS_METHOD_NAME = "registerMuzzleVirtualFields";
    private final URLClassLoader classLoader;
  
    public MuzzleCodeGenerator(URLClassLoader classLoader) {
      this.classLoader = classLoader;
    }
  
    @Override
    public int mergeWriter(int flags) {
      return flags | ClassWriter.COMPUTE_MAXS;
    }
  
    @Override
    @CanIgnoreReturnValue
    public int mergeReader(int flags) {
      return flags;
    }
  
    @Override
    public ClassVisitor wrap(
        TypeDescription instrumentedType,
        ClassVisitor classVisitor,
        Implementation.Context implementationContext,
        TypePool typePool,
        FieldList<FieldDescription.InDefinedShape> fields,
        MethodList<?> methods,
        int writerFlags,
        int readerFlags) {
      return new GenerateMuzzleMethodsAndFields(classVisitor, classLoader);
    }
  
    private static class GenerateMuzzleMethodsAndFields extends ClassVisitor {
      private final String[] defaultInterfaces =
          new String[] {Utils.getInternalName(InstrumentationModuleMuzzle.class)};
  
      private final URLClassLoader classLoader;
      private String instrumentationClassName;
      private InstrumentationModule instrumentationModule;
  
      private boolean generateReferencesMethod = true;
      private boolean generateHelperClassNamesMethod = true;
      private boolean generateVirtualFieldsMethod = true;
  
      public GenerateMuzzleMethodsAndFields(ClassVisitor classVisitor, URLClassLoader classLoader) {
        super(AsmApi.VERSION, classVisitor);
        this.classLoader = classLoader;
      }
  
      @Override
      public void visit(
          int version,
          int access,
          String name,
          String signature,
          String superName,
          String[] interfaces) {
        this.instrumentationClassName = name;
        try {
          instrumentationModule =
              (InstrumentationModule)
                  classLoader
                      .loadClass(Utils.getClassName(instrumentationClassName))
                      .getDeclaredConstructor()
                      .newInstance();
        } catch (Exception e) {
          throw new IllegalStateException(e);
        }
  
        super.visit(version, access, name, signature, superName, addMuzzleInterface(interfaces));
      }
  
      private String[] addMuzzleInterface(String[] interfaces) {
        if (interfaces == null) {
          return defaultInterfaces;
        }
  
        String[] allInterfaces = new String[interfaces.length + 1];
        allInterfaces[0] = Utils.getInternalName(InstrumentationModuleMuzzle.class);
        System.arraycopy(interfaces, 0, allInterfaces, 1, interfaces.length);
        return allInterfaces;
      }
  
      @Override
      public MethodVisitor visitMethod(
          int access, String name, String descriptor, String signature, String[] exceptions) {
        if (MUZZLE_REFERENCES_METHOD_NAME.equals(name)) {
          generateReferencesMethod = false;
          logMethodAlreadyExistsMessage(MUZZLE_REFERENCES_METHOD_NAME);
        }
        if (MUZZLE_HELPER_CLASSES_METHOD_NAME.equals(name)) {
          generateHelperClassNamesMethod = false;
          logMethodAlreadyExistsMessage(MUZZLE_HELPER_CLASSES_METHOD_NAME);
        }
        if (MUZZLE_VIRTUAL_FIELDS_METHOD_NAME.equals(name)) {
          generateVirtualFieldsMethod = false;
          logMethodAlreadyExistsMessage(MUZZLE_VIRTUAL_FIELDS_METHOD_NAME);
        }
        return super.visitMethod(access, name, descriptor, signature, exceptions);
      }
  
      private void logMethodAlreadyExistsMessage(String muzzleVirtualFieldsMethodName) {
        if (logger.isLoggable(INFO)) {
          logger.log(
              INFO,
              "The \"{0}\" method was already found in class \"{1}\". Muzzle will not generate it again",
              new Object[] {muzzleVirtualFieldsMethodName, instrumentationClassName});
        }
      }
  
      @Override
      public void visitEnd() {
        ReferenceCollector collector = collectReferences();
        if (generateReferencesMethod) {
          generateMuzzleReferencesMethod(collector);
        }
        if (generateHelperClassNamesMethod) {
          generateMuzzleHelperClassNamesMethod(collector);
        }
        if (generateVirtualFieldsMethod) {
          generateMuzzleVirtualFieldsMethod(collector);
        }
        super.visitEnd();
      }
  
      private ReferenceCollector collectReferences() {
        AdviceClassNameCollector adviceClassNameCollector = new AdviceClassNameCollector();
        for (TypeInstrumentation typeInstrumentation : instrumentationModule.typeInstrumentations()) {
          typeInstrumentation.transform(adviceClassNameCollector);
        }
  
        // the class loader has a parent including the Gradle classpath, such as buildSrc
        // dependencies.
        // These may have resources take precedence over ones we define, so we need to make sure to
        // not include them when loading resources.
        // TODO analyze anew if this is needed
        ClassLoader resourceLoader = new URLClassLoader(classLoader.getURLs(), null);
        ReferenceCollector collector =
            new ReferenceCollector(instrumentationModule::isHelperClass, resourceLoader);
        for (String adviceClass : adviceClassNameCollector.getAdviceClassNames()) {
          collector.collectReferencesFromAdvice(adviceClass);
        }
        HelperResourceBuilderImpl helperResourceBuilder = new HelperResourceBuilderImpl();
        instrumentationModule.registerHelperResources(helperResourceBuilder);
        for (HelperResource resource : helperResourceBuilder.getResources()) {
          collector.collectReferencesFromResource(resource);
        }
        collector.prune();
        return collector;
      }
  
      private void generateMuzzleReferencesMethod(ReferenceCollector collector) {
        Type referenceType = Type.getType(ClassRef.class);
        Type referenceBuilderType = Type.getType(ClassRefBuilder.class);
        Type referenceFlagType = Type.getType(Flag.class);
        Type referenceFlagArrayType = Type.getType(Flag[].class);
        Type referenceSourceArrayType = Type.getType(Source[].class);
        Type stringType = Type.getType(String.class);
        Type typeType = Type.getType(Type.class);
        Type typeArrayType = Type.getType(Type[].class);
  
        /*
         * public Map<String, ClassRef> getMuzzleReferences() {
         *   Map<String, ClassRef> references = new HashMap<>(...);
         *   references.put("reference class name", ClassRef.builder(...)
         *       ...
         *       .build());
         *   return references;
         * }
         */
        MethodVisitor mv =
            super.visitMethod(
                Opcodes.ACC_PUBLIC, MUZZLE_REFERENCES_METHOD_NAME, "()Ljava/util/Map;", null, null);
        mv.visitCode();
  
        Collection<ClassRef> references = collector.getReferences().values();
  
        writeNewMap(mv, references.size());
        // stack: map
        mv.visitVarInsn(Opcodes.ASTORE, 1);
        // stack: <empty>
  
        references.forEach(
            reference -> {
              mv.visitVarInsn(Opcodes.ALOAD, 1);
              // stack: map
              mv.visitLdcInsn(reference.getClassName());
              // stack: map, className
  
              mv.visitLdcInsn(reference.getClassName());
              mv.visitMethodInsn(
                  Opcodes.INVOKESTATIC,
                  referenceType.getInternalName(),
                  "builder",
                  Type.getMethodDescriptor(referenceBuilderType, stringType),
                  /* isInterface= */ false);
              // stack: map, className, builder
  
              for (Source source : reference.getSources()) {
                mv.visitLdcInsn(source.getName());
                mv.visitLdcInsn(source.getLine());
                mv.visitMethodInsn(
                    Opcodes.INVOKEVIRTUAL,
                    referenceBuilderType.getInternalName(),
                    "addSource",
                    Type.getMethodDescriptor(referenceBuilderType, stringType, Type.INT_TYPE),
                    /* isInterface= */ false);
              }
              // stack: map, className, builder
              for (Flag flag : reference.getFlags()) {
                String enumClassName = getEnumClassInternalName(flag);
                mv.visitFieldInsn(
                    Opcodes.GETSTATIC, enumClassName, flag.name(), "L" + enumClassName + ";");
                mv.visitMethodInsn(
                    Opcodes.INVOKEVIRTUAL,
                    referenceBuilderType.getInternalName(),
                    "addFlag",
                    Type.getMethodDescriptor(referenceBuilderType, referenceFlagType),
                    /* isInterface= */ false);
              }
              // stack: map, className, builder
              if (null != reference.getSuperClassName()) {
                mv.visitLdcInsn(reference.getSuperClassName());
                mv.visitMethodInsn(
                    Opcodes.INVOKEVIRTUAL,
                    referenceBuilderType.getInternalName(),
                    "setSuperClassName",
                    Type.getMethodDescriptor(referenceBuilderType, stringType),
                    /* isInterface= */ false);
              }
              // stack: map, className, builder
              for (String interfaceName : reference.getInterfaceNames()) {
                mv.visitLdcInsn(interfaceName);
                mv.visitMethodInsn(
                    Opcodes.INVOKEVIRTUAL,
                    referenceBuilderType.getInternalName(),
                    "addInterfaceName",
                    Type.getMethodDescriptor(referenceBuilderType, stringType),
                    /* isInterface= */ false);
              }
              // stack: map, className, builder
              for (FieldRef field : reference.getFields()) {
                writeSourcesArray(mv, field.getSources());
                writeFlagsArray(mv, field.getFlags());
                // field name
                mv.visitLdcInsn(field.getName());
                writeType(mv, field.getDescriptor());
                // declared flag
                mv.visitLdcInsn(field.isDeclared());
  
                mv.visitMethodInsn(
                    Opcodes.INVOKEVIRTUAL,
                    referenceBuilderType.getInternalName(),
                    "addField",
                    Type.getMethodDescriptor(
                        referenceBuilderType,
                        referenceSourceArrayType,
                        referenceFlagArrayType,
                        stringType,
                        typeType,
                        Type.BOOLEAN_TYPE),
                    /* isInterface= */ false);
              }
              // stack: map, className, builder
              for (MethodRef method : reference.getMethods()) {
                writeSourcesArray(mv, method.getSources());
                writeFlagsArray(mv, method.getFlags());
                // method name
                mv.visitLdcInsn(method.getName());
                // method return and argument types
                {
                  // we cannot pass the whole method descriptor string as it won't be shaded, so
                  // we
                  // have to pass the return and parameter types separately - strings in
                  // Type.getType()
                  // calls will be shaded correctly
                  Type methodType = Type.getMethodType(method.getDescriptor());
  
                  writeType(mv, methodType.getReturnType().getDescriptor());
  
                  mv.visitLdcInsn(methodType.getArgumentTypes().length);
                  mv.visitTypeInsn(Opcodes.ANEWARRAY, typeType.getInternalName());
                  int i = 0;
                  for (Type parameterType : methodType.getArgumentTypes()) {
                    mv.visitInsn(Opcodes.DUP);
                    mv.visitLdcInsn(i);
                    writeType(mv, parameterType.getDescriptor());
                    mv.visitInsn(Opcodes.AASTORE);
                    i++;
                  }
                }
  
                mv.visitMethodInsn(
                    Opcodes.INVOKEVIRTUAL,
                    referenceBuilderType.getInternalName(),
                    "addMethod",
                    Type.getMethodDescriptor(
                        referenceBuilderType,
                        referenceSourceArrayType,
                        referenceFlagArrayType,
                        stringType,
                        typeType,
                        typeArrayType),
                    /* isInterface= */ false);
              }
              // stack: map, className, builder
              mv.visitMethodInsn(
                  Opcodes.INVOKEVIRTUAL,
                  referenceBuilderType.getInternalName(),
                  "build",
                  Type.getMethodDescriptor(referenceType),
                  /* isInterface= */ false);
              // stack: map, className, classRef
  
              mv.visitMethodInsn(
                  Opcodes.INVOKEINTERFACE,
                  "java/util/Map",
                  "put",
                  "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;",
                  /* isInterface= */ true);
              // stack: previousValue
              mv.visitInsn(Opcodes.POP);
              // stack: <empty>
            });
  
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        // stack: map
        mv.visitInsn(Opcodes.ARETURN);
  
        mv.visitMaxs(0, 0); // recomputed
        mv.visitEnd();
      }
  
      private static void writeNewMap(MethodVisitor mv, int size) {
        mv.visitTypeInsn(Opcodes.NEW, "java/util/HashMap");
        // stack: map
        mv.visitInsn(Opcodes.DUP);
        // stack: map, map
        // pass bigger size to avoid resizes; same formula as in e.g. HashSet(Collection)
        // 0.75 is the default load factor
        mv.visitLdcInsn((int) (size / 0.75f) + 1);
        // stack: map, map, size
        mv.visitLdcInsn(0.75f);
        // stack: map, map, size, loadFactor
        mv.visitMethodInsn(
            Opcodes.INVOKESPECIAL, "java/util/HashMap", "<init>", "(IF)V", /* isInterface= */ false);
      }
  
      private static void writeSourcesArray(MethodVisitor mv, Set<Source> sources) {
        Type referenceSourceType = Type.getType(Source.class);
  
        mv.visitLdcInsn(sources.size());
        mv.visitTypeInsn(Opcodes.ANEWARRAY, referenceSourceType.getInternalName());
  
        int i = 0;
        for (Source source : sources) {
          mv.visitInsn(Opcodes.DUP);
          mv.visitLdcInsn(i);
  
          mv.visitTypeInsn(Opcodes.NEW, referenceSourceType.getInternalName());
          mv.visitInsn(Opcodes.DUP);
          mv.visitLdcInsn(source.getName());
          mv.visitLdcInsn(source.getLine());
          mv.visitMethodInsn(
              Opcodes.INVOKESPECIAL,
              referenceSourceType.getInternalName(),
              "<init>",
              "(Ljava/lang/String;I)V",
              /* isInterface= */ false);
  
          mv.visitInsn(Opcodes.AASTORE);
          ++i;
        }
      }
  
      private static void writeFlagsArray(MethodVisitor mv, Set<Flag> flags) {
        Type referenceFlagType = Type.getType(Flag.class);
  
        mv.visitLdcInsn(flags.size());
        mv.visitTypeInsn(Opcodes.ANEWARRAY, referenceFlagType.getInternalName());
  
        int i = 0;
        for (Flag flag : flags) {
          mv.visitInsn(Opcodes.DUP);
          mv.visitLdcInsn(i);
          String enumClassName = getEnumClassInternalName(flag);
          mv.visitFieldInsn(Opcodes.GETSTATIC, enumClassName, flag.name(), "L" + enumClassName + ";");
          mv.visitInsn(Opcodes.AASTORE);
          ++i;
        }
      }
  
      private static final Pattern ANONYMOUS_ENUM_CONSTANT_CLASS =
          Pattern.compile("(?<enumClass>.*)\\$[0-9]+$");
  
      // drops "$1" suffix for enum constants that override/implement super class methods
      private static String getEnumClassInternalName(Flag flag) {
        String fullInternalName = Utils.getInternalName(flag.getClass());
        Matcher m = ANONYMOUS_ENUM_CONSTANT_CLASS.matcher(fullInternalName);
        return m.matches() ? m.group("enumClass") : fullInternalName;
      }
  
      private static void writeType(MethodVisitor mv, String descriptor) {
        Type typeType = Type.getType(Type.class);
  
        mv.visitLdcInsn(descriptor);
        mv.visitMethodInsn(
            Opcodes.INVOKESTATIC,
            typeType.getInternalName(),
            "getType",
            Type.getMethodDescriptor(typeType, Type.getType(String.class)),
            /* isInterface= */ false);
      }
  
      private void generateMuzzleHelperClassNamesMethod(ReferenceCollector collector) {
        /*
         * public List<String> getMuzzleHelperClassNames() {
         *   List<String> helperClassNames = new ArrayList<>(...);
         *   helperClassNames.add(...);
         *   return helperClassNames;
         * }
         */
        MethodVisitor mv =
            super.visitMethod(
                Opcodes.ACC_PUBLIC,
                MUZZLE_HELPER_CLASSES_METHOD_NAME,
                "()Ljava/util/List;",
                null,
                null);
        mv.visitCode();
  
        List<String> helperClassNames = collector.getSortedHelperClasses();
  
        mv.visitTypeInsn(Opcodes.NEW, "java/util/ArrayList");
        // stack: list
        mv.visitInsn(Opcodes.DUP);
        // stack: list, list
        mv.visitLdcInsn(helperClassNames.size());
        // stack: list, list, size
        mv.visitMethodInsn(
            Opcodes.INVOKESPECIAL, "java/util/ArrayList", "<init>", "(I)V", /* isInterface= */ false);
        // stack: list
        mv.visitVarInsn(Opcodes.ASTORE, 1);
        // stack: <empty>
  
        helperClassNames.forEach(
            helperClassName -> {
              mv.visitVarInsn(Opcodes.ALOAD, 1);
              // stack: list
              mv.visitLdcInsn(helperClassName);
              // stack: list, helperClassName
              mv.visitMethodInsn(
                  Opcodes.INVOKEINTERFACE,
                  "java/util/List",
                  "add",
                  "(Ljava/lang/Object;)Z",
                  /* isInterface= */ true);
              // stack: added
              mv.visitInsn(Opcodes.POP);
              // stack: <empty>
            });
  
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        // stack: list
        mv.visitInsn(Opcodes.ARETURN);
  
        mv.visitMaxs(0, 0);
        mv.visitEnd();
      }
  
      private void generateMuzzleVirtualFieldsMethod(ReferenceCollector collector) {
        /*
         * public void registerMuzzleVirtualFields(VirtualFieldMappingsBuilder builder) {
         *   builder.register(..., ...);
         * }
         */
        MethodVisitor mv =
            super.visitMethod(
                Opcodes.ACC_PUBLIC,
                MUZZLE_VIRTUAL_FIELDS_METHOD_NAME,
                "(Lio/opentelemetry/javaagent/tooling/muzzle/VirtualFieldMappingsBuilder;)V",
                null,
                null);
        mv.visitCode();
  
        VirtualFieldMappings virtualFieldMappings = collector.getVirtualFieldMappings();
  
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        // stack: builder
        virtualFieldMappings.forEach(
            (typeName, fieldTypeName) -> {
              mv.visitLdcInsn(typeName);
              // stack: builder, typeName
              mv.visitLdcInsn(fieldTypeName);
              // stack: builder, typeName, fieldTypeName
              mv.visitMethodInsn(
                  Opcodes.INVOKEINTERFACE,
                  "io/opentelemetry/javaagent/tooling/muzzle/VirtualFieldMappingsBuilder",
                  "register",
                  "(Ljava/lang/String;Ljava/lang/String;)Lio/opentelemetry/javaagent/tooling/muzzle/VirtualFieldMappingsBuilder;",
                  /* isInterface= */ true);
              // stack: builder
            });
        mv.visitInsn(Opcodes.POP);
        // stack: <empty>
        mv.visitInsn(Opcodes.RETURN);
  
        mv.visitMaxs(0, 0);
        mv.visitEnd();
      }
    }
  }
  