content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.tooling.field;
  
  import static io.opentelemetry.javaagent.tooling.field.GeneratedVirtualFieldNames.getRealGetterName;
  import static io.opentelemetry.javaagent.tooling.field.GeneratedVirtualFieldNames.getRealSetterName;
  import static io.opentelemetry.javaagent.tooling.field.GeneratedVirtualFieldNames.getVirtualFieldImplementationClassName;
  
  import com.google.errorprone.annotations.CanIgnoreReturnValue;
  import io.opentelemetry.instrumentation.api.internal.cache.Cache;
  import io.opentelemetry.instrumentation.api.util.VirtualField;
  import io.opentelemetry.javaagent.extension.instrumentation.internal.AsmApi;
  import io.opentelemetry.javaagent.tooling.Utils;
  import io.opentelemetry.javaagent.tooling.muzzle.VirtualFieldMappings;
  import java.util.HashMap;
  import java.util.Map;
  import net.bytebuddy.ByteBuddy;
  import net.bytebuddy.ClassFileVersion;
  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.modifier.SyntheticState;
  import net.bytebuddy.description.modifier.TypeManifestation;
  import net.bytebuddy.description.modifier.Visibility;
  import net.bytebuddy.description.type.TypeDescription;
  import net.bytebuddy.dynamic.DynamicType;
  import net.bytebuddy.implementation.Implementation;
  import net.bytebuddy.pool.TypePool;
  import org.objectweb.asm.ClassVisitor;
  import org.objectweb.asm.ClassWriter;
  import org.objectweb.asm.Label;
  import org.objectweb.asm.MethodVisitor;
  import org.objectweb.asm.Opcodes;
  
  final class VirtualFieldImplementationsGenerator {
  
    private final ByteBuddy byteBuddy;
  
    VirtualFieldImplementationsGenerator(ByteBuddy byteBuddy) {
      this.byteBuddy = byteBuddy;
    }
  
    VirtualFieldImplementations generateClasses(
        VirtualFieldMappings virtualFieldMappings, FieldAccessorInterfaces fieldAccessorInterfaces) {
      Map<String, DynamicType.Unloaded<?>> virtualFieldImplementations =
          new HashMap<>(virtualFieldMappings.size());
      for (Map.Entry<String, String> entry : virtualFieldMappings.entrySet()) {
        DynamicType.Unloaded<?> type =
            makeVirtualFieldImplementationClass(
                entry.getKey(), entry.getValue(), fieldAccessorInterfaces);
        virtualFieldImplementations.put(type.getTypeDescription().getName(), type);
      }
      return new VirtualFieldImplementations(virtualFieldImplementations);
    }
  
    /**
     * Generate an 'implementation' of a context store class for given key class name and context
     * class name.
     *
     * @param typeName key class name
     * @param fieldTypeName context class name
     * @return unloaded dynamic type containing generated class
     */
    private DynamicType.Unloaded<?> makeVirtualFieldImplementationClass(
        String typeName, String fieldTypeName, FieldAccessorInterfaces fieldAccessorInterfaces) {
      return byteBuddy
          .rebase(VirtualFieldImplementationTemplate.class)
          .modifiers(Visibility.PUBLIC, TypeManifestation.FINAL, SyntheticState.SYNTHETIC)
          .name(getVirtualFieldImplementationClassName(typeName, fieldTypeName))
          .visit(
              getVirtualFieldImplementationVisitor(typeName, fieldTypeName, fieldAccessorInterfaces))
          .make();
    }
  
    /**
     * Returns a visitor that 'fills in' missing methods into concrete implementation of {@link
     * VirtualFieldImplementationsGenerator.VirtualFieldImplementationTemplate} for given key class
     * name and context class name.
     *
     * @param typeName key class name
     * @param fieldTypeName context class name
     * @return visitor that adds implementation for methods that need to be generated
     */
    private AsmVisitorWrapper getVirtualFieldImplementationVisitor(
        String typeName, String fieldTypeName, FieldAccessorInterfaces fieldAccessorInterfaces) {
      return new AsmVisitorWrapper() {
  
        @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 ClassVisitor(AsmApi.VERSION, classVisitor) {
  
            private final TypeDescription accessorInterface =
                fieldAccessorInterfaces.find(typeName, fieldTypeName);
            private final String accessorInterfaceInternalName = accessorInterface.getInternalName();
            private final String instrumentedTypeInternalName = instrumentedType.getInternalName();
            private final boolean frames =
                implementationContext.getClassFileVersion().isAtLeast(ClassFileVersion.JAVA_V6);
  
            @Override
            public MethodVisitor visitMethod(
                int access, String name, String descriptor, String signature, String[] exceptions) {
              if ("realGet".equals(name)) {
                generateRealGetMethod(name);
                return null;
              } else if ("realPut".equals(name)) {
                generateRealPutMethod(name);
                return null;
              } else {
                return super.visitMethod(access, name, descriptor, signature, exceptions);
              }
            }
  
            /**
             * Provides implementation for {@code realGet} method that looks like below.
             *
             * <blockquote>
             *
             * <pre>
             * private Object realGet(final Object key) {
             *   if (key instanceof $accessorInterfaceInternalName) {
             *     return (($accessorInterfaceInternalName) key).$getterName();
             *   } else {
             *     return mapGet(key);
             *   }
             * }
             * </pre>
             *
             * </blockquote>
             *
             * @param name name of the method being visited
             */
            private void generateRealGetMethod(String name) {
              String getterName = getRealGetterName(typeName, fieldTypeName);
              Label elseLabel = new Label();
              MethodVisitor mv = getMethodVisitor(name);
              mv.visitCode();
              mv.visitVarInsn(Opcodes.ALOAD, 1);
              mv.visitTypeInsn(Opcodes.INSTANCEOF, accessorInterfaceInternalName);
              mv.visitJumpInsn(Opcodes.IFEQ, elseLabel);
              mv.visitVarInsn(Opcodes.ALOAD, 1);
              mv.visitTypeInsn(Opcodes.CHECKCAST, accessorInterfaceInternalName);
              mv.visitMethodInsn(
                  Opcodes.INVOKEINTERFACE,
                  accessorInterfaceInternalName,
                  getterName,
                  Utils.getMethodDefinition(accessorInterface, getterName).getDescriptor(),
                  /* isInterface= */ true);
              mv.visitInsn(Opcodes.ARETURN);
              mv.visitLabel(elseLabel);
              if (frames) {
                mv.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
              }
              mv.visitVarInsn(Opcodes.ALOAD, 0);
              mv.visitVarInsn(Opcodes.ALOAD, 1);
              mv.visitMethodInsn(
                  Opcodes.INVOKESPECIAL,
                  instrumentedTypeInternalName,
                  "mapGet",
                  Utils.getMethodDefinition(instrumentedType, "mapGet").getDescriptor(),
                  /* isInterface= */ false);
              mv.visitInsn(Opcodes.ARETURN);
              mv.visitMaxs(0, 0);
              mv.visitEnd();
            }
  
            /**
             * Provides implementation for {@code realPut} method that looks like below.
             *
             * <blockquote>
             *
             * <pre>
             * private void realPut(final Object key, final Object value) {
             *   if (key instanceof $accessorInterfaceInternalName) {
             *     (($accessorInterfaceInternalName) key).$setterName(value);
             *   } else {
             *     mapPut(key, value);
             *   }
             * }
             * </pre>
             *
             * </blockquote>
             *
             * @param name name of the method being visited
             */
            private void generateRealPutMethod(String name) {
              String setterName = getRealSetterName(typeName, fieldTypeName);
              Label elseLabel = new Label();
              Label endLabel = new Label();
              MethodVisitor mv = getMethodVisitor(name);
              mv.visitCode();
              mv.visitVarInsn(Opcodes.ALOAD, 1);
              mv.visitTypeInsn(Opcodes.INSTANCEOF, accessorInterfaceInternalName);
              mv.visitJumpInsn(Opcodes.IFEQ, elseLabel);
              mv.visitVarInsn(Opcodes.ALOAD, 1);
              mv.visitTypeInsn(Opcodes.CHECKCAST, accessorInterfaceInternalName);
              mv.visitVarInsn(Opcodes.ALOAD, 2);
              mv.visitMethodInsn(
                  Opcodes.INVOKEINTERFACE,
                  accessorInterfaceInternalName,
                  setterName,
                  Utils.getMethodDefinition(accessorInterface, setterName).getDescriptor(),
                  /* isInterface= */ true);
              mv.visitJumpInsn(Opcodes.GOTO, endLabel);
              mv.visitLabel(elseLabel);
              if (frames) {
                mv.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
              }
              mv.visitVarInsn(Opcodes.ALOAD, 0);
              mv.visitVarInsn(Opcodes.ALOAD, 1);
              mv.visitVarInsn(Opcodes.ALOAD, 2);
              mv.visitMethodInsn(
                  Opcodes.INVOKESPECIAL,
                  instrumentedTypeInternalName,
                  "mapPut",
                  Utils.getMethodDefinition(instrumentedType, "mapPut").getDescriptor(),
                  /* isInterface= */ false);
              mv.visitLabel(endLabel);
              if (frames) {
                mv.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
              }
              mv.visitInsn(Opcodes.RETURN);
              mv.visitMaxs(0, 0);
              mv.visitEnd();
            }
  
            private MethodVisitor getMethodVisitor(String methodName) {
              return cv.visitMethod(
                  Opcodes.ACC_PRIVATE,
                  methodName,
                  Utils.getMethodDefinition(instrumentedType, methodName).getDescriptor(),
                  null,
                  null);
            }
          };
        }
      };
    }
  
    /**
     * Template class used to generate the class that accesses stored context using either key
     * instance's own injected field or global hash map if field is not available.
     */
    // Called from generated code
    @SuppressWarnings({"UnusedMethod", "UnusedVariable", "MethodCanBeStatic"})
    static final class VirtualFieldImplementationTemplate extends VirtualField<Object, Object> {
      private static final VirtualFieldImplementationTemplate INSTANCE =
          new VirtualFieldImplementationTemplate(Cache.weak());
  
      private final Cache<Object, Object> map;
  
      private VirtualFieldImplementationTemplate(Cache<Object, Object> map) {
        this.map = map;
      }
  
      @Override
      public Object get(Object object) {
        return realGet(object);
      }
  
      @Override
      public void set(Object object, Object fieldValue) {
        realPut(object, fieldValue);
      }
  
      private Object realGet(Object key) {
        // to be generated
        return null;
      }
  
      private void realPut(Object key, Object value) {
        // to be generated
      }
  
      private Object mapGet(Object key) {
        return map.get(key);
      }
  
      private void mapPut(Object key, Object value) {
        if (value == null) {
          map.remove(key);
        } else {
          map.put(key, value);
        }
      }
  
      public static VirtualField<?, ?> getVirtualField(Class<?> keyClass, Class<?> contextClass) {
        // We do not actually check the keyClass here - but that should be fine since compiler would
        // check things for us.
        return INSTANCE;
      }
    }
  }
  