Adjust memory implementation of in memory target to support multi-rel…

…ease jars.
raphw · Sep 24, 2024 · e06d8ea · e06d8ea
1 parent 258ea66
commit e06d8ea
Showing 1 changed file with 43 additions and 8 deletions.
diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/build/Plugin.java b/byte-buddy-dep/src/main/java/net/bytebuddy/build/Plugin.java
@@ -3442,31 +3442,59 @@ public void close() {
             }
 
             /**
-             * A sink that stores all elements in a memory map. In case of multi-release jars, this memory
-             * storage aims to retain the non-versioned class file.
+             * A sink that stores all elements in a memory map.
              */
-            @HashCodeAndEqualsPlugin.Enhance
             class InMemory implements Target, Sink {
 
+                @MaybeNull
+                private final ClassFileVersion classFileVersion;
+
                 /**
                  * The map for storing all elements being received.
                  */
                 private final Map<String, byte[]> storage;
 
+                private final Map<String, Integer> versions;
+
                 /**
                  * Creates a new in-memory storage.
                  */
                 public InMemory() {
                     this(new HashMap<String, byte[]>());
                 }
 
+                /**
+                 * Creates a new in-memory storage.
+                 *
+                 * @param classFileVersion The class file version to consider as the maximum version when accepting
+                 *                         multi-release classes or {@code null} if multi-release versions should
+                 *                         be discarded.
+                 */
+                public InMemory(@MaybeNull ClassFileVersion classFileVersion) {
+                    this(classFileVersion, new HashMap<String, byte[]>());
+                }
+
                 /**
                  * Creates a new in-memory storage.
                  *
                  * @param storage The map for storing all elements being received.
                  */
                 public InMemory(Map<String, byte[]> storage) {
+                    this(null, storage);
+                }
+
+                /**
+                 * Creates a new in-memory storage.
+                 *
+                 * @param classFileVersion The class file version to consider as the maximum version when accepting
+                 *                         multi-release classes or {@code null} if multi-release versions should
+                 *                         be discarded.
+                 * @param storage          The map for storing all elements being received.
+                 */
+                public InMemory(@MaybeNull ClassFileVersion classFileVersion, Map<String, byte[]> storage) {
+                    this.classFileVersion = classFileVersion;
                     this.storage = storage;
+                    versions = new HashMap<String, Integer>();
                 }
 
                 /**
@@ -3490,18 +3518,25 @@ public Sink write(@MaybeNull Manifest manifest) throws IOException {
                  */
                 public void store(Map<TypeDescription, byte[]> binaryRepresentations) {
                     for (Map.Entry<TypeDescription, byte[]> entry : binaryRepresentations.entrySet()) {
-                        storage.put(entry.getKey().getInternalName() + CLASS_FILE_EXTENSION, entry.getValue());
+                        String name = entry.getKey().getInternalName() + CLASS_FILE_EXTENSION;
+                        if (!storage.containsKey(name)) {
+                            storage.put(name, entry.getValue());
+                        }
                     }
                 }
 
                 /**
                  * {@inheritDoc}
                  */
                 public void store(int version, Map<TypeDescription, byte[]> binaryRepresentations) throws IOException {
-                    for (Map.Entry<TypeDescription, byte[]> entry : binaryRepresentations.entrySet()) {
-                        String name = entry.getKey().getInternalName() + CLASS_FILE_EXTENSION;
-                        if (!storage.containsKey(name)) {
-                            storage.put(name, entry.getValue());
+                    if (classFileVersion != null && version <= classFileVersion.getJavaVersion()) {
+                        for (Map.Entry<TypeDescription, byte[]> entry : binaryRepresentations.entrySet()) {
+                            String name = entry.getKey().getInternalName() + CLASS_FILE_EXTENSION;
+                            Integer current = versions.get(name);
+                            if (current == null || current < version) {
+                                versions.put(name, version);
+                                storage.put(name, entry.getValue());
+                            }
                         }
                     }
                 }