Allow bytecode transformers to run in parallel when using shamrock:run

core/deployment/src/main/java/org/jboss/shamrock/deployment/BuildTimeGenerator.java

@@ -34,6 +34,7 @@
 import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -86,7 +87,7 @@ public class BuildTimeGenerator {
     private final DeploymentProcessorInjection injection;
     private final ClassLoader classLoader;
     private final boolean useStaticInit;
-    private final List<Function<String, Function<ClassVisitor, ClassVisitor>>> bytecodeTransformers = new ArrayList<>();
+    private final Map<String, List<BiFunction<String, ClassVisitor, ClassVisitor>>> byteCodeTransformers = new HashMap<>();
     private final Set<String> applicationArchiveMarkers;
     private final ArchiveContextBuilder archiveContextBuilder;
     private final Set<String> capabilities;
@@ -110,8 +111,8 @@ public BuildTimeGenerator(ClassOutput classOutput, ClassLoader cl, boolean useSt
         this.capabilities = new HashSet<>(setupContext.capabilities);
     }
 
-    public List<Function<String, Function<ClassVisitor, ClassVisitor>>> getBytecodeTransformers() {
-        return bytecodeTransformers;
+    public Map<String, List<BiFunction<String, ClassVisitor, ClassVisitor>>> getByteCodeTransformers() {
+        return byteCodeTransformers;
     }
 
     public void run(Path root) throws IOException {
@@ -342,8 +343,8 @@ public void createResource(String name, byte[] data) throws IOException {
         }
 
         @Override
-        public void addByteCodeTransformer(Function<String, Function<ClassVisitor, ClassVisitor>> visitorFunction) {
-            bytecodeTransformers.add(visitorFunction);
+        public void addByteCodeTransformer(String className, BiFunction<String, ClassVisitor, ClassVisitor> visitorFunction) {
+            byteCodeTransformers.computeIfAbsent(className, (e) -> new ArrayList<>()).add(visitorFunction);
         }
 
         @Override
@@ -438,7 +439,6 @@ void writeMainClass() throws IOException {
                 ResultHandle dup = mv.newInstance(ofConstructor(holder.className));
                 mv.invokeInterfaceMethod(ofMethod(StartupTask.class, "deploy", void.class, StartupContext.class), dup, startupContext);
             }
-
             mv.invokeStaticMethod(ofMethod(Timing.class, "printStartupTime", void.class));
             mv.returnValue(null);
 

core/deployment/src/main/java/org/jboss/shamrock/deployment/ProcessorContext.java

@@ -3,6 +3,7 @@
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 import org.jboss.jandex.FieldInfo;
@@ -41,9 +42,9 @@ public interface ProcessorContext {
      * This method is used to indicate that a given class requires reflection.
      * <p>
      * It is used in the graal output to allow the class to be reflected when running on substrate VM
-     *
+     * <p>
      * This will add all constructors, as well as all methods and fields if the appropriate fields are set.
-     *
+     * <p>
      * Where possible consider using the more fine grained addReflective* variants
      *
      * @param className The class name
@@ -60,56 +61,56 @@ public interface ProcessorContext {
 
     /**
      * Attempts to register a complete type hierarchy for reflection.
-     *
+     * <p>
      * This is intended to be used to register types that are going to be serialized,
      * e.g. by Jackson or some other JSON mapper.
-     *
+     * <p>
      * This will do 'smart discovery' and in addition to registering the type itself it will also attempt to
      * register the following:
-     *
+     * <p>
      * - Superclasses
      * - Component types of collections
      * - Types used in bean properties if (if method reflection is enabled)
      * - Field types (if field reflection is enabled)
-     *
+     * <p>
      * This discovery is applied recursively, so any additional types that are registered will also have their dependencies
      * discovered
-     *
      */
     void addReflectiveHierarchy(Type type);
 
 
     /**
-     *
      * @param applicationClass If this class should be loaded by the application class loader when in runtime mode
-     * @param name The class name
-     * @param classData The class bytes
+     * @param name             The class name
+     * @param classData        The class bytes
      * @throws IOException
      */
     void addGeneratedClass(boolean applicationClass, String name, byte[] classData) throws IOException;
 
     /**
      * Creates a resources with the provided contents
+     *
      * @param name
      * @param data
      * @throws IOException
      */
     void createResource(String name, byte[] data) throws IOException;
 
     /**
-     * Adds a bytecode transformer that can transform application classes.
+     * Adds a bytecode transformer that can transform application classes
+     * <p>
+     * This is added on a per-class basis, by specifying the class name. The transformer is a function that
+     * can be used to wrap an ASM {@link ClassVisitor}.
      * <p>
-     * This takes the form of a function that takes the class name as a String, and returns a Function that wraps an
-     * ASM {@link ClassVisitor}.
-     *
-     * If this function returns null then no transform is applied. If it returns a function then it will be transformed.
-     *
      * The transformation is applied by calling each function that has been registered it turn to create a chain
      * of visitors. These visitors are then applied and the result is saved to the output.
-     *
+     * <p>
      * At present these transformations are only applied to application classes, not classes provided by dependencies
+     * <p>
+     * These transformations may be run concurrently in multiple threads, so if a function is registered for multiple
+     * classes it must be thread safe
      */
-    void addByteCodeTransformer(Function<String, Function<ClassVisitor, ClassVisitor>> visitorFunction);
+    void addByteCodeTransformer(String classToTransform, BiFunction<String, ClassVisitor, ClassVisitor> visitorFunction);
 
     /**
      * Adds a resource to the image that will be accessible when running under substrate.
@@ -126,25 +127,24 @@ public interface ProcessorContext {
      *
      * @param classes The classes to lazily init
      */
-    void addRuntimeInitializedClasses(String ... classes);
+    void addRuntimeInitializedClasses(String... classes);
 
     /**
      * Adds a proxy definition to allow proxies to be created using {@link java.lang.reflect.Proxy}
      *
      * @param proxyClasses The interface names that this proxy will implement
      */
-    void addProxyDefinition(String ... proxyClasses);
+    void addProxyDefinition(String... proxyClasses);
 
     /**
      * Set a system property to be passed in to the native image tool.
      *
-     * @param name the property name (must not be {@code null})
+     * @param name  the property name (must not be {@code null})
      * @param value the property value
      */
     void addNativeImageSystemProperty(String name, String value);
 
     /**
-     *
      * @param capability
      * @return if the given capability is present
      * @see Capabilities