Skip to content

Commit

Permalink
use cached class nodes to not have to read twice
Browse files Browse the repository at this point in the history
  • Loading branch information
@wagyourtail
wagyourtail committed May 13, 2024
1 parent a293ba7 commit b6ec1b1
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 27 deletions.
37 changes: 20 additions & 17 deletions src/main/java/xyz/wagyourtail/jvmdg/compile/ApiShader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,16 +67,16 @@ public static void shadeApis(int target, String prefix, File input, File output,
shadeApis(prefix, input.toPath(), output.toPath(), downgradedApiPath);
}

public static void shadeApis(final String prefix, Path input, Path output, Path downgradedApi) throws IOException {
public static void shadeApis(final String prefix, Path input, final Path output, Path downgradedApi) throws IOException {
if (!prefix.endsWith("/")) throw new IllegalArgumentException("prefix must end with /");
// step 2: collect classes in the api and their references
ReferenceGraph refs = new ReferenceGraph();
final ReferenceGraph apiRefs = new ReferenceGraph();
final Set<Type> apiClasses;
try (FileSystem apiFs = Utils.openZipFileSystem(downgradedApi, new HashMap<String, Object>())) {
final Path apiRoot = apiFs.getPath("/");
Map<Path, Type> preScan = refs.preScan(apiFs.getPath("/"));
Map<Path, Type> preScan = apiRefs.preScan(apiFs.getPath("/"));
apiClasses = new HashSet<>(preScan.values());
refs.scan(preScan, new ReferenceGraph.Filter() {
apiRefs.scan(preScan, new ReferenceGraph.Filter() {
@Override
public boolean shouldInclude(FullyQualifiedMemberNameAndDesc member) {
return apiClasses.contains(member.getOwner());
Expand All @@ -85,7 +85,7 @@ public boolean shouldInclude(FullyQualifiedMemberNameAndDesc member) {
try (final FileSystem inputFs = Utils.openZipFileSystem(input, new HashMap<String, Object>())) {
final Path inputRoot = inputFs.getPath("/");
// step 3: traverse the input classes for references to the api
ReferenceGraph inputRefs = new ReferenceGraph();
final ReferenceGraph inputRefs = new ReferenceGraph();
inputRefs.scan(inputFs.getPath("/"), new ReferenceGraph.Filter() {
@Override
public boolean shouldInclude(FullyQualifiedMemberNameAndDesc member) {
Expand All @@ -97,7 +97,7 @@ public boolean shouldInclude(FullyQualifiedMemberNameAndDesc member) {
try (final FileSystem outputFs = Utils.openZipFileSystem(output, env)) {
final Path outputRoot = outputFs.getPath("/");
// step 4: create remapper for api classes to prefixed api classes
Pair<Set<FullyQualifiedMemberNameAndDesc>, Set<String>> required = refs.recursiveResolveFrom(inputRefs.getAllRefs());
Pair<Set<FullyQualifiedMemberNameAndDesc>, Set<String>> required = apiRefs.recursiveResolveFrom(inputRefs.getAllRefs());
final Map<Type, Set<MemberNameAndDesc>> byType = byType(required.getFirst());
final Map<String, String> remap = new HashMap<>();
for (Type type : byType.keySet()) {
Expand All @@ -109,17 +109,13 @@ public boolean shouldInclude(FullyQualifiedMemberNameAndDesc member) {
Future<Void> apiWrite = AsyncUtils.forEachAsync(byType.entrySet(), new IOConsumer<Map.Entry<Type, Set<MemberNameAndDesc>>>() {
@Override
public void accept(Map.Entry<Type, Set<MemberNameAndDesc>> type) throws IOException {
Path inPath = apiRoot.resolve(type.getKey().getInternalName() + ".class");
Path outPath = outputRoot.resolve(prefix + type.getKey().getInternalName() + ".class");
Path parent = outPath.getParent();
if (parent != null) {
Files.createDirectories(outputRoot.resolve(parent));
}
// load api class as a ClassNode
ClassNode node = new ClassNode();
try (InputStream stream = Files.newInputStream(inPath)) {
new ClassReader(stream).accept(node, 0);
}
ClassNode node = apiRefs.getClassFor(type.getKey());
if ((node.access & Opcodes.ACC_ENUM) == 0) {
// remove unused members
Set<MemberNameAndDesc> members = type.getValue();
Expand All @@ -141,7 +137,7 @@ public void accept(Map.Entry<Type, Set<MemberNameAndDesc>> type) throws IOExcept
// write the class to the output
ClassWriter writer = new ClassWriter(0);
node.accept(new ClassRemapper(writer, remapper));
Files.write(outPath, writer.toByteArray());
Files.write(outPath, writer.toByteArray(), StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
}

});
Expand Down Expand Up @@ -170,12 +166,19 @@ public Boolean apply(Path path) throws IOException {
}, new IOConsumer<Path>() {
@Override
public void accept(Path path) throws IOException {
Path output = outputRoot.resolve(inputRoot.relativize(path));
if (path.toString().endsWith(".class")) {
byte[] data = Files.readAllBytes(path);
Path relPath = inputRoot.relativize(path);
Path output = outputRoot.resolve(relPath);
String pathStr = relPath.toString();
if (!pathStr.startsWith("META-INF/versions") && pathStr.endsWith(".class") && !pathStr.equals("module-info.class")) {
// byte[] data = Files.readAllBytes(path);
ClassWriter writer = new ClassWriter(0);
new ClassReader(data).accept(new ClassRemapper(writer, remapper), 0);
Files.write(output, writer.toByteArray());
pathStr = pathStr.substring(0, pathStr.length() - 6);
ClassNode cn = inputRefs.getClassFor(Type.getObjectType(pathStr));
if (cn == null) {
throw new IllegalStateException("Class not found from ref cache? " + pathStr);
}
cn.accept(new ClassRemapper(writer, remapper));
Files.write(output, writer.toByteArray(), StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
} else {
Files.copy(path, output, StandardCopyOption.REPLACE_EXISTING);
}
Expand Down
41 changes: 32 additions & 9 deletions src/main/java/xyz/wagyourtail/jvmdg/compile/shade/ReferenceGraph.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,23 @@

public class ReferenceGraph {
private final Map<Type, References> references = new ConcurrentHashMap<>();
private final Map<Type, ClassNode> classNodes = new ConcurrentHashMap<>();
private final boolean keepClassNodes;

public ReferenceGraph() {
keepClassNodes = true;
}

public ReferenceGraph(boolean keepClassNodes) {
this.keepClassNodes = keepClassNodes;
}

public ClassNode getClassFor(Type type) {
if (!keepClassNodes) {
throw new IllegalStateException();
}
return classNodes.get(type);
}

public void scan(Path root, Filter filter) throws IOException, ExecutionException, InterruptedException {
scan(preScan(root), filter);
Expand Down Expand Up @@ -58,17 +75,23 @@ public void accept(Path path) throws IOException {
return newScanTargets;
}

public void scan(final Map<Path, Type> newScanTargets, Filter filter) throws IOException {
for (Path path : newScanTargets.keySet()) {
try (InputStream stream = Files.newInputStream(path)) {
ClassNode node = ASMUtils.bytesToClassNode(Utils.readAllBytes(stream));
Type type = Type.getObjectType(node.name);
if (!type.equals(newScanTargets.get(path))) {
throw new IllegalStateException("Expected path to match class name: " + path + " != " + type.getInternalName());
public void scan(final Map<Path, Type> newScanTargets, final Filter filter) throws IOException, ExecutionException, InterruptedException {
AsyncUtils.forEachAsync(newScanTargets.keySet(), new IOConsumer<Path>() {
@Override
public void accept(Path path) throws IOException {
try (InputStream stream = Files.newInputStream(path)) {
ClassNode node = ASMUtils.bytesToClassNode(Utils.readAllBytes(stream));
Type type = Type.getObjectType(node.name);
if (!type.equals(newScanTargets.get(path))) {
throw new IllegalStateException("Expected path to match class name: " + path + " != " + type.getInternalName());
}
if (keepClassNodes) {
classNodes.put(type, node);
}
references.get(type).scan(node, filter);
}
references.get(type).scan(node, filter);
}
}
}).get();
}

public void debugPrint() {
Expand Down
2 changes: 1 addition & 1 deletion src/shared/java/xyz/wagyourtail/jvmdg/util/AsyncUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
import java.util.concurrent.atomic.AtomicReference;

public class AsyncUtils {
private static final ForkJoinPool pool = new ForkJoinPool(Runtime.getRuntime().availableProcessors(), ForkJoinPool.defaultForkJoinWorkerThreadFactory, null, false);
private static final ForkJoinPool pool = new ForkJoinPool(Math.min(Math.max(1, Runtime.getRuntime().availableProcessors() - 1), Short.MAX_VALUE), ForkJoinPool.defaultForkJoinWorkerThreadFactory, null, true);

public static <T> Future<List<T>> waitForFutures(Future<T>... futures) {
return waitForFutures(new ArrayDeque<>(Arrays.asList(futures)));
Expand Down

0 comments on commit b6ec1b1

Please sign in to comment.