From 4bc3afe4ff08b9f0c08952ec7f6e0ac930280cc5 Mon Sep 17 00:00:00 2001
From: asie
Date: Sat, 8 Dec 2018 11:21:18 +0100
Subject: add barebones plugin framework, cleanup

---
 .../java/cuchaz/enigma/analysis/ParsedJar.java     | 58 +++++++++++++++-------
 1 file changed, 39 insertions(+), 19 deletions(-)

(limited to 'src/main/java/cuchaz/enigma/analysis/ParsedJar.java')

diff --git a/src/main/java/cuchaz/enigma/analysis/ParsedJar.java b/src/main/java/cuchaz/enigma/analysis/ParsedJar.java
index 55f2141..86655d0 100644
--- a/src/main/java/cuchaz/enigma/analysis/ParsedJar.java
+++ b/src/main/java/cuchaz/enigma/analysis/ParsedJar.java
@@ -11,8 +11,10 @@
 
 package cuchaz.enigma.analysis;
 
+import com.google.common.io.ByteStreams;
 import cuchaz.enigma.mapping.entry.ClassEntry;
 import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.tree.ClassNode;
 
 import java.io.BufferedInputStream;
@@ -20,14 +22,17 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.util.*;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 import java.util.jar.JarInputStream;
 
 public class ParsedJar {
-	private final Map<String, ClassNode> nodes = new LinkedHashMap<>();
+	private final Map<String, byte[]> classBytes;
+	private final Map<String, ClassNode> nodeCache = new HashMap<>();
 
 	public ParsedJar(JarFile jar) throws IOException {
+		Map<String, byte[]> uClassBytes = new LinkedHashMap<>();;
 		try {
 			// get the jar entries that correspond to classes
 			Enumeration<JarEntry> entries = jar.entries();
@@ -36,60 +41,75 @@ public class ParsedJar {
 				// is this a class file?
 				if (entry.getName().endsWith(".class")) {
 					try (InputStream input = new BufferedInputStream(jar.getInputStream(entry))) {
-						// read the ClassNode from the jar
-						ClassReader reader = new ClassReader(input);
-						ClassNode node = new ClassNode();
-						reader.accept(node, 0);
 						String path = entry.getName().substring(0, entry.getName().length() - ".class".length());
-						nodes.put(path, node);
+						uClassBytes.put(path, ByteStreams.toByteArray(input));
 					}
 				}
 			}
 		} finally {
 			jar.close();
+			classBytes = Collections.unmodifiableMap(uClassBytes);
 		}
 	}
 
 	public ParsedJar(JarInputStream jar) throws IOException {
+		Map<String, byte[]> uClassBytes = new LinkedHashMap<>();
 		try {
 			// get the jar entries that correspond to classes
 			JarEntry entry;
 			while ((entry = jar.getNextJarEntry()) != null) {
 				// is this a class file?
 				if (entry.getName().endsWith(".class")) {
-					// read the ClassNode from the jar
-					ClassReader reader = new ClassReader(jar);
-					ClassNode node = new ClassNode();
-					reader.accept(node, 0);
 					String path = entry.getName().substring(0, entry.getName().length() - ".class".length());
-					nodes.put(path, node);
+					uClassBytes.put(path, ByteStreams.toByteArray(jar));
 					jar.closeEntry();
 				}
 			}
 		} finally {
 			jar.close();
+			classBytes = Collections.unmodifiableMap(uClassBytes);
 		}
 	}
 
-	public void visit(Consumer<ClassNode> visitor) {
-		for (ClassNode node : nodes.values()) {
-			visitor.accept(node);
+	public void visitReader(Function<String, ClassVisitor> visitorFunction, int options) {
+		for (String s : classBytes.keySet()) {
+			ClassNode nodeCached = nodeCache.get(s);
+			if (nodeCached != null) {
+				nodeCached.accept(visitorFunction.apply(s));
+			} else {
+				new ClassReader(classBytes.get(s)).accept(visitorFunction.apply(s), options);
+			}
+		}
+	}
+
+	public void visitNode(Consumer<ClassNode> consumer) {
+		for (String s : classBytes.keySet()) {
+			consumer.accept(getClassNode(s));
 		}
 	}
 
 	public int getClassCount() {
-		return nodes.size();
+		return classBytes.size();
 	}
 
 	public List<ClassEntry> getClassEntries() {
-		List<ClassEntry> entries = new ArrayList<>(nodes.size());
-		for (ClassNode node : nodes.values()) {
-			entries.add(new ClassEntry(node.name));
+		List<ClassEntry> entries = new ArrayList<>(classBytes.size());
+		for (String s : classBytes.keySet()) {
+			entries.add(new ClassEntry(s));
 		}
 		return entries;
 	}
 
 	public ClassNode getClassNode(String name) {
-		return nodes.get(name);
+		return nodeCache.computeIfAbsent(name, (n) -> {
+			ClassReader reader = new ClassReader(classBytes.get(name));
+			ClassNode node = new ClassNode();
+			reader.accept(node, 0);
+			return node;
+		});
 	}
+
+    public Map<String, byte[]> getClassDataMap() {
+		return classBytes;
+    }
 }
-- 
cgit v1.2.3