Fixed subclasses in decompilers

src/me/ByteEdit/boxes/CompilationBox.java

@@ -15,6 +15,8 @@ import java.io.StringWriter;
 import java.net.URI;
 import java.nio.file.Files;
 import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
 
 import javax.swing.JComponent;
 import javax.swing.JFrame;
@@ -122,7 +124,13 @@ public class CompilationBox extends JFrame {
 								ClassNode node = new ClassNode();
 								read.accept(node, 0);
 								EnumDecompiler decompiler = EnumDecompiler.BYTEEDIT;
-								String dis = decompiler.getDecompiler().decompile(node);
+								ClassNode mainNode = node;
+
+								Map<String, ClassNode> classNodes;
+								synchronized (Main.classNodes) {
+									classNodes = new HashMap<>(Main.classNodes);
+								}
+								classNodes.put("Compiled.class", node);
 								File[] files = tmpFolder.listFiles();
 								if (files != null)
 									for (File f : files) {
@@ -137,21 +145,26 @@ public class CompilationBox extends JFrame {
 										read = new ClassReader(baos.toByteArray());
 										node = new ClassNode();
 										read.accept(node, 0);
-										dis += "\n" + decompiler.getDecompiler().decompile(node);
+										classNodes.put(f.getName(), node);
 									}
-								compSuccess.textArea.setText(dis);
+								compSuccess.textArea
+										.setText(decompiler.getDecompiler().decompile(mainNode, classNodes));
 								compSuccess.setVisible(true);
 							} catch (IOException e2) {
 								e2.printStackTrace();
 							}
+							try {
+								for (File f : tmpFolder.listFiles()) {
+									Files.delete(f.toPath());
+								}
+								Files.delete(tmpFolder.toPath());
+							} catch (IOException e2) {
+								e2.printStackTrace();
+							}
 						});
 					} else {
 						JOptionPane.showMessageDialog(CompilationBox.this, res, "Error!", JOptionPane.ERROR_MESSAGE);
 					}
-					for (File f : tmpFolder.listFiles()) {
-						Files.delete(f.toPath());
-					}
-					Files.delete(tmpFolder.toPath());
 				} catch (IOException e1) {
 					e1.printStackTrace();
 				}

src/me/ByteEdit/decompiler/CFRDecompiler.java

@@ -3,12 +3,18 @@ package me.ByteEdit.decompiler;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.lang.reflect.Field;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.Map;
 
 import org.benf.cfr.reader.PluginRunner;
 import org.benf.cfr.reader.api.ClassFileSource;
 import org.benf.cfr.reader.bytecode.analysis.parse.utils.Pair;
+import org.benf.cfr.reader.state.ClassFileSourceImpl;
+import org.benf.cfr.reader.util.MiscConstants.Version;
+import org.benf.cfr.reader.util.getopt.OptionsImpl;
 import org.objectweb.asm.tree.ClassNode;
 
 public class CFRDecompiler implements IDecompiler {
@@ -68,17 +74,24 @@ public class CFRDecompiler implements IDecompiler {
 		options.put("sugarenums", "true");
 		options.put("tidymonitors", "true");
 		options.put("usenametable", "true");
+
+		try {
+			Field f = Version.class.getDeclaredField("version");
+			f.setAccessible(true);
+			f.set(null, "ByteEdit");
+		} catch (Exception ex) {
+		}
 	}
 
 	@Override
-	public String decompile(ClassNode cn) {
-		return doDecompilation(cn, getBytes(cn));
+	public String decompile(ClassNode cn, Map<String, ClassNode> classNodes) {
+		return doDecompilation(cn, IDecompiler.getBytes(cn), classNodes);
 	}
 
-	private static String doDecompilation(ClassNode cn, byte[] b) {
+	private static String doDecompilation(ClassNode cn, byte[] b, Map<String, ClassNode> classNodes) {
 		try {
 			HashMap<String, String> ops = options;
-			ClassFileSource cfs = new ClassFileSource() {
+			ClassFileSource cfs = new ClassFileSourceImpl(new OptionsImpl(Collections.EMPTY_MAP)) {
 
 				@Override
 				public void informAnalysisRelativePathDetail(String a, String b) {
@@ -92,10 +105,16 @@ public class CFRDecompiler implements IDecompiler {
 				@Override
 				public Pair<byte[], String> getClassFileContent(String path) throws IOException {
 					String name = path.substring(0, path.length() - 6);
-					if (name.equals(cn.name)) {
+					if (name.equals(cn.name))
 						return Pair.make(b, name);
+					// try loading other classes from jar
+					ClassNode node;
+					synchronized (classNodes) {
+						node = classNodes.get(name + ".class");
 					}
-					return null; // cfr loads unnecessary classes
+					if (node != null)
+						return Pair.make(IDecompiler.getBytes(node), name);
+					return super.getClassFileContent(path);
 				}
 
 				@Override

src/me/ByteEdit/decompiler/FernflowerDecompiler.java

@@ -5,6 +5,7 @@ import java.io.IOException;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.HashMap;
+import java.util.Map;
 import java.util.Map.Entry;
 import java.util.jar.Manifest;
 
@@ -51,11 +52,11 @@ public class FernflowerDecompiler implements IBytecodeProvider, IResultSaver, ID
 	}
 
 	@Override
-	public String decompile(ClassNode cn) {
-		return doDecompilation(cn, getBytes(cn));
+	public String decompile(ClassNode cn, Map<String, ClassNode> classNodes) {
+		return doDecompilation(cn, IDecompiler.getBytes(cn), classNodes);
 	}
 
-	public String doDecompilation(ClassNode cn, byte[] b) {
+	public String doDecompilation(ClassNode cn, byte[] b, Map<String, ClassNode> classNodes) {
 		try {
 			// TODO decompile method only
 			this.bytes = b;
@@ -87,7 +88,7 @@ public class FernflowerDecompiler implements IBytecodeProvider, IResultSaver, ID
 					sc.getDecompiledData());
 			sc.getUnits().put(fakePath, unit);
 			unit.addClass(cl, "none.class");
-			sc.getLoader().addClassLink(cn.name, new LazyLoader.Link(LazyLoader.Link.CLASS, fakePath, null));
+			sc.getLoader().addClassLink(cn.name, new LazyLoader.Link(LazyLoader.Link.CLASS, fakePath, cn.name));
 
 			f.decompileContext();
 			return returned;

src/me/ByteEdit/decompiler/IDecompiler.java

@@ -1,13 +1,15 @@
 package me.ByteEdit.decompiler;
 
+import java.util.Map;
+
 import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.tree.ClassNode;
 
 public interface IDecompiler {
 
-	String decompile(ClassNode cn);
+	String decompile(ClassNode cn, Map<String, ClassNode> classNodes);
 
-	default byte[] getBytes(ClassNode cn) {
+	static byte[] getBytes(ClassNode cn) {
 		ClassWriter classWriter = new ClassWriter(0);
 		cn.accept(classWriter);
 		return classWriter.toByteArray();

src/me/ByteEdit/decompiler/JDDecompiler.java

@@ -1,5 +1,7 @@
 package me.ByteEdit.decompiler;
 
+import java.util.Map;
+
 import org.jd.core.v1.ClassFileToJavaSourceDecompiler;
 import org.jd.core.v1.api.loader.Loader;
 import org.jd.core.v1.api.loader.LoaderException;
@@ -7,12 +9,19 @@ import org.jd.core.v1.api.printer.Printer;
 import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.tree.ClassNode;
 
-import me.ByteEdit.main.Main;
-
 public class JDDecompiler implements Loader, Printer, IDecompiler {
 
+	protected static final String TAB = "    ";
+	protected static final String NEWLINE = "\n";
+
+	protected int indentationCount = 0;
+	protected StringBuilder sb = new StringBuilder();
+
+	private Map<String, ClassNode> classNodes;
+
 	@Override
-	public String decompile(ClassNode cn) {
+	public String decompile(ClassNode cn, Map<String, ClassNode> classNodes) {
+		this.classNodes = classNodes;
 		ClassFileToJavaSourceDecompiler decompiler = new ClassFileToJavaSourceDecompiler();
 		try {
 			decompiler.decompile(this, this, cn.name);
@@ -32,22 +41,21 @@ public class JDDecompiler implements Loader, Printer, IDecompiler {
 
 	@Override
 	public boolean canLoad(String internalName) {
-		return Main.classNodes.containsKey(internalName + ".class");
+		synchronized (classNodes) {
+			return classNodes.containsKey(internalName + ".class");
+		}
 	}
 
 	@Override
 	public byte[] load(String internalName) throws LoaderException {
-		ClassNode cl = Main.classNodes.get(internalName + ".class");
-		ClassWriter cw = new ClassWriter(0);
-		cl.accept(cw);
-		return cw.toByteArray();
+		ClassNode node;
+		synchronized (classNodes) {
+			node = classNodes.get(internalName + ".class");
+		}
+		if (node != null)
+			return IDecompiler.getBytes(node);
+		throw new LoaderException("Failed to load '" + internalName + "'");
 	}
-
-	protected static final String TAB = "    ";
-	protected static final String NEWLINE = "\n";
-
-	protected int indentationCount = 0;
-	protected StringBuilder sb = new StringBuilder();
 
 	@Override
 	public String toString() {

src/me/ByteEdit/decompiler/ProcyonDecompiler.java

@@ -3,6 +3,7 @@ package me.ByteEdit.decompiler;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.lang.reflect.Field;
+import java.util.Map;
 
 import org.objectweb.asm.tree.ClassNode;
 
@@ -19,11 +20,11 @@ import com.strobel.decompiler.PlainTextOutput;
 public class ProcyonDecompiler implements IDecompiler {
 
 	@Override
-	public String decompile(ClassNode cn) {
-		return doDecompilation(cn, getBytes(cn));
+	public String decompile(ClassNode cn, Map<String, ClassNode> classNodes) {
+		return doDecompilation(cn, IDecompiler.getBytes(cn), classNodes);
 	}
 
-	private static String doDecompilation(ClassNode cn, byte[] b) {
+	private static String doDecompilation(ClassNode cn, byte[] b, Map<String, ClassNode> classNodes) {
 		try {
 			// TODO decompile method only
 			DecompilerSettings settings = new DecompilerSettings();
@@ -48,6 +49,16 @@ public class ProcyonDecompiler implements IDecompiler {
 						buffer.position(0);
 						return true;
 					}
+					ClassNode node;
+					synchronized (classNodes) {
+						node = classNodes.get(s + ".class");
+					}
+					if (node != null) {
+						byte[] data = IDecompiler.getBytes(node);
+						buffer.putByteArray(data, 0, data.length);
+						buffer.position(0);
+						return true;
+					}
 					return backLoader.tryLoadType(s, buffer);
 				}
 			});

src/me/ByteEdit/edit/Disassembler.java

@@ -21,6 +21,8 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -69,8 +71,21 @@ public class Disassembler implements IDecompiler {
 			.newFixedThreadPool((Runtime.getRuntime().availableProcessors() + 1) / 2);
 
 	@Override
-	public String decompile(ClassNode cn) {
-		return disassemble(cn);
+	public String decompile(ClassNode cn, Map<String, ClassNode> classNodes) {
+		String name = cn.name;
+		StringBuilder dis = new StringBuilder(disassemble(cn));
+		synchronized (classNodes) {
+			for (Entry<String, ClassNode> entry : classNodes.entrySet()) {
+				if (entry.getKey().contains("$")) {
+					String[] split = entry.getKey().split("\\$");
+					if (split[0].equals(name)) {
+						dis.append('\n');
+						dis.append(disassemble(entry.getValue()));
+					}
+				}
+			}
+		}
+		return dis.toString();
 	}
 
 	private static String disassemble(ClassNode classNode) {

src/me/ByteEdit/main/Main.java

@@ -786,23 +786,11 @@ public class Main extends JFrame {
 			String current = currentNodeName;
 			ClassNode classNode = classNodes.get(current);
 			int prev = txtByteEditView.getCaretPosition();
-			String dis = decompiler.getDecompiler().decompile(classNode);
-			String substr = current.substring(0, current.length() - 6);
-			synchronized (Main.classNodes) {
-				for (Entry<String, ClassNode> entry : Main.classNodes.entrySet()) {
-					if (entry.getKey().contains("$")) {
-						String[] split = entry.getKey().split("\\$");
-						if (split[0].equals(substr)) {
-							dis += "\n" + decompiler.getDecompiler().decompile(entry.getValue());
-						}
-					}
-				}
-			}
-			String fdis = dis;
+			String dis = decompiler.getDecompiler().decompile(classNode, Main.classNodes);
 			EventQueue.invokeAndWait(() -> {
 				synchronized (treeLock) {
-					txtByteEditView.setText(fdis);
-					if (fdis.length() > prev)
+					txtByteEditView.setText(dis);
+					if (dis.length() > prev)
 						txtByteEditView.setCaretPosition(prev);
 				}
 			});