Add class file locator abstraction to dynamic types. Allow explicit t…

…ype injection from build plugins.

byte-buddy-dep/src/main/java/net/bytebuddy/build/Plugin.java

@@ -97,6 +97,20 @@ interface WithPreprocessor extends Plugin {
         void onPreprocess(TypeDescription typeDescription, ClassFileLocator classFileLocator);
     }
 
+    /**
+     * Allows for the generation of types before a plugin is applied.
+     */
+    interface WithInitialization extends Plugin {
+
+        /**
+         * Returns a mapping of classes that should be created before discovering any types.
+         *
+         * @param classFileLocator A class file locator that can locate other types in the scope of the project.
+         * @return A mapping of types to their binary representation.
+         */
+        Map<TypeDescription, byte[]> initialize(ClassFileLocator classFileLocator);
+    }
+
     /**
      * A factory for providing a build plugin.
      */
@@ -110,7 +124,7 @@ interface Factory {
         Plugin make();
 
         /**
-         * A simple factory that returns a preconstructed plugin instance..
+         * A simple factory that returns a preconstructed plugin instance.
          */
         @HashCodeAndEqualsPlugin.Enhance
         class Simple implements Factory {
@@ -4476,6 +4490,7 @@ public Summary apply(Source source, Target target, List<? extends Plugin.Factory
                 List<String> unresolved = new ArrayList<String>();
                 Throwable rethrown = null;
                 List<Plugin> plugins = new ArrayList<Plugin>(factories.size());
+                List<WithInitialization> initializers = new ArrayList<WithInitialization>();
                 List<WithPreprocessor> preprocessors = new ArrayList<WithPreprocessor>();
                 try {
                     for (Plugin.Factory factory : factories) {
@@ -4484,6 +4499,9 @@ public Summary apply(Source source, Target target, List<? extends Plugin.Factory
                         if (plugin instanceof WithPreprocessor) {
                             preprocessors.add((WithPreprocessor) plugin);
                         }
+                        if (plugin instanceof WithInitialization) {
+                            initializers.add((WithInitialization) plugin);
+                        }
                     }
                     Source.Origin origin = source.read();
                     try {
@@ -4493,6 +4511,9 @@ public Summary apply(Source source, Target target, List<? extends Plugin.Factory
                         listener.onManifest(manifest);
                         Target.Sink sink = target.write(manifest);
                         try {
+                            for (WithInitialization initializer : initializers) {
+                                sink.store(initializer.initialize(classFileLocator));
+                            }
                             Dispatcher dispatcher = dispatcherFactory.make(sink, transformed, failed, unresolved);
                             try {
                                 for (Source.Element element : origin) {

byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/DynamicType.java

@@ -59,7 +59,7 @@
  * Note that the {@link TypeDescription}s will represent their
  * unloaded forms and therefore differ from the loaded types, especially with regards to annotations.
  */
-public interface DynamicType {
+public interface DynamicType extends ClassFileLocator {
 
     /**
      * <p>
@@ -5922,6 +5922,26 @@ public Default(TypeDescription typeDescription,
             this.auxiliaryTypes = auxiliaryTypes;
         }
 
+        /**
+         * {@inheritDoc}
+         */
+        public Resolution locate(String name) throws IOException {
+            if (typeDescription.getName().equals(name)) {
+                return new Resolution.Explicit(binaryRepresentation);
+            }
+            for (DynamicType auxiliaryType : auxiliaryTypes) {
+                Resolution resolution = auxiliaryType.locate(name);
+                if (resolution.isResolved()) {
+                    return resolution;
+                }
+            }
+            return new Resolution.Illegal(name);
+        }
+
+        public void close() {
+            /* do nothing */
+        }
+
         /**
          * {@inheritDoc}
          */

byte-buddy-dep/src/test/java/net/bytebuddy/build/PluginEngineDefaultTest.java

@@ -24,6 +24,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Map;
 import java.util.jar.*;
 
 import static net.bytebuddy.test.utility.FieldByFieldComparison.hasPrototype;
@@ -83,6 +84,35 @@ public void testSimpleTransformation() throws Exception {
         verifyNoMoreInteractions(listener);
     }
 
+    @Test
+    public void testSimpleTransformationWithInitialization() throws Exception {
+        Plugin.Engine.Listener listener = mock(Plugin.Engine.Listener.class);
+        Plugin plugin = eager
+                ? new InitializingPlugin(new SimplePlugin())
+                : new PreprocessingPlugin(new InitializingPlugin(new SimplePlugin()));
+        Plugin.Engine.Source source = Plugin.Engine.Source.InMemory.ofTypes(Sample.class);
+        Plugin.Engine.Target.InMemory target = new Plugin.Engine.Target.InMemory();
+        Plugin.Engine.Summary summary = new Plugin.Engine.Default()
+                .with(listener)
+                .with(ClassFileLocator.ForClassLoader.of(SimplePlugin.class.getClassLoader()))
+                .with(dispatcherFactory)
+                .apply(source, target, new Plugin.Factory.Simple(plugin));
+        ClassLoader classLoader = new ByteArrayClassLoader(ClassLoadingStrategy.BOOTSTRAP_LOADER, target.toTypeMap());
+        Class<?> type = classLoader.loadClass(Sample.class.getName());
+        assertThat(type.getDeclaredField(FOO).getType(), is((Object) Void.class));
+        Class<?> otherType = classLoader.loadClass(OtherSample.class.getName());
+        assertThat(otherType.getDeclaredFields().length, is(0));
+        assertThat(summary.getTransformed(), hasItems(TypeDescription.ForLoadedType.of(Sample.class)));
+        assertThat(summary.getFailed().size(), is(0));
+        assertThat(summary.getUnresolved().size(), is(0));
+        verify(listener).onManifest(Plugin.Engine.Source.Origin.NO_MANIFEST);
+        verify(listener).onDiscovery(Sample.class.getName());
+        verify(listener).onTransformation(TypeDescription.ForLoadedType.of(Sample.class), plugin);
+        verify(listener).onTransformation(TypeDescription.ForLoadedType.of(Sample.class), Collections.singletonList(plugin));
+        verify(listener).onComplete(TypeDescription.ForLoadedType.of(Sample.class));
+        verifyNoMoreInteractions(listener);
+    }
+
     @Test
     public void testSimpleTransformationIgnoredByPlugin() throws Exception {
         Plugin.Engine.Listener listener = mock(Plugin.Engine.Listener.class);
@@ -378,6 +408,10 @@ private static class Sample {
         /* empty */
     }
 
+    private static class OtherSample {
+        /* empty */
+    }
+
     private static class SimplePlugin implements Plugin {
 
         public DynamicType.Builder<?> apply(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassFileLocator classFileLocator) {
@@ -452,7 +486,7 @@ public void close() {
         }
     }
 
-    private static class PreprocessingPlugin implements Plugin.WithPreprocessor {
+    private static class PreprocessingPlugin implements Plugin.WithPreprocessor, Plugin.WithInitialization {
 
         private final Plugin plugin;
 
@@ -464,6 +498,39 @@ public void onPreprocess(TypeDescription typeDescription, ClassFileLocator class
             /* empty */
         }
 
+        public Map<TypeDescription, byte[]> initialize(ClassFileLocator classFileLocator) {
+            if (plugin instanceof WithInitialization) {
+                return ((WithInitialization) plugin).initialize(classFileLocator);
+            } else {
+                return Collections.emptyMap();
+            }
+        }
+
+        public DynamicType.Builder<?> apply(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassFileLocator classFileLocator) {
+            return plugin.apply(builder, typeDescription, classFileLocator);
+        }
+
+        public boolean matches(TypeDescription target) {
+            return plugin.matches(target);
+        }
+
+        public void close() throws IOException {
+            plugin.close();
+        }
+    }
+
+    private static class InitializingPlugin implements Plugin.WithInitialization {
+
+        private final Plugin plugin;
+
+        private InitializingPlugin(Plugin plugin) {
+            this.plugin = plugin;
+        }
+
+        public Map<TypeDescription, byte[]> initialize(ClassFileLocator classFileLocator) {
+            return Collections.singletonMap(TypeDescription.ForLoadedType.of(OtherSample.class), ClassFileLocator.ForClassLoader.read(OtherSample.class));
+        }
+
         public DynamicType.Builder<?> apply(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassFileLocator classFileLocator) {
             return plugin.apply(builder, typeDescription, classFileLocator);
         }

byte-buddy-dep/src/test/java/net/bytebuddy/dynamic/DynamicTypeDefaultTest.java

@@ -324,4 +324,11 @@ public void testIterationOrder() throws Exception {
         assertThat(types.next(), is(auxiliaryTypeDescription));
         assertThat(types.hasNext(), is(false));
     }
+
+    @Test
+    public void testClassFileLocator() throws Exception {
+        assertThat(dynamicType.locate(FOOBAR.replace('/', '.')).isResolved(), is(true));
+        assertThat(dynamicType.locate(QUXBAZ.replace('/', '.')).isResolved(), is(true));
+        assertThat(dynamicType.locate(BARBAZ.replace('/', '.')).isResolved(), is(false));
+    }
 }