Add additional instructions for class injector.

raphw · Oct 31, 2024 · 5a45f35 · 5a45f35
1 parent 1ba91d0
commit 5a45f35
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 1 deletion.
diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/DynamicType.java b/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/DynamicType.java
@@ -86,6 +86,13 @@ public interface DynamicType extends ClassFileLocator {
      */
     byte[] getBytes();
 
+    /**
+     * Returns a set of all types that are represented by this dynamic type.
+     *
+     * @return A set of all represented types.
+     */
+    Set<TypeDescription> getTypeDescriptions();
+
     /**
      * <p>
      * Returns a map of all auxiliary types that are required for making use of the main type.
@@ -6123,6 +6130,18 @@ public void close() {
             /* do nothing */
         }
 
+        /**
+         * {@inheritDoc}
+         */
+        public Set<TypeDescription> getTypeDescriptions() {
+            Set<TypeDescription> types = new LinkedHashSet<TypeDescription>();
+            types.add(typeDescription);
+            for (DynamicType auxiliaryType : auxiliaryTypes) {
+                types.addAll(auxiliaryType.getTypeDescriptions());
+            }
+            return types;
+        }
+
         /**
          * {@inheritDoc}
          */

diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/loading/ClassInjector.java b/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/loading/ClassInjector.java
@@ -97,6 +97,8 @@ public interface ClassInjector {
      */
     Map<String, Class<?>> inject(Set<String> names, ClassFileLocator classFileLocator);
 
+    Map<TypeDescription, Class<?>> inject(List<? extends DynamicType> dynamicTypes);
+
     /**
      * Injects the given types into the represented class loader.
      *
@@ -118,6 +120,23 @@ public interface ClassInjector {
      */
     abstract class AbstractBase implements ClassInjector {
 
+        /**
+         * {@inheritDoc}
+         */
+        public Map<TypeDescription, Class<?>> inject(List<? extends DynamicType> dynamicTypes) {
+            Map<String, TypeDescription> types = new LinkedHashMap<String, TypeDescription>();
+            for (DynamicType dynamicType : dynamicTypes) {
+                for (TypeDescription typeDescription : dynamicType.getTypeDescriptions()) {
+                    types.put(typeDescription.getName(), typeDescription);
+                }
+            }
+            Map<TypeDescription, Class<?>> result = new HashMap<TypeDescription, Class<?>>();
+            for (Map.Entry<String, Class<?>> entry : inject(types.keySet(), new ClassFileLocator.Compound(dynamicTypes)).entrySet()) {
+                result.put(types.get(entry.getKey()), entry.getValue());
+            }
+            return result;
+        }
+
         /**
          * {@inheritDoc}
          */
@@ -127,7 +146,7 @@ public Map<TypeDescription, Class<?>> inject(Map<? extends TypeDescription, byte
                 binaryRepresentations.put(entry.getKey().getName(), entry.getValue());
             }
             Map<String, Class<?>> loadedTypes = injectRaw(binaryRepresentations);
-            Map<TypeDescription, Class<?>> result = new LinkedHashMap<TypeDescription, Class<?>>();
+            Map<TypeDescription, Class<?>> result = new HashMap<TypeDescription, Class<?>>();
             for (TypeDescription typeDescription : types.keySet()) {
                 result.put(typeDescription, loadedTypes.get(typeDescription.getName()));
             }