8 files changed, 158 insertions, 35 deletions

diff --git a/.gitignore b/.gitignore
index 1729e4d6..662e408b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -150,6 +150,7 @@ flycheck_*.el
 
 # IntelliJ
 /out/
+.idea
 
 # mpeltonen/sbt-idea plugin
 .idea_modules/
diff --git a/src/main/java/cuchaz/enigma/CachingClasspathTypeLoader.java b/src/main/java/cuchaz/enigma/CachingClasspathTypeLoader.java
new file mode 100644
index 00000000..fca2a47e
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/CachingClasspathTypeLoader.java
@@ -0,0 +1,20 @@
+package cuchaz.enigma;
+
+import com.strobel.assembler.metadata.Buffer;
+import com.strobel.assembler.metadata.ClasspathTypeLoader;
+import com.strobel.assembler.metadata.ITypeLoader;
+
+/**
+ * Caching version of {@link ClasspathTypeLoader}
+ */
+class CachingClasspathTypeLoader extends CachingTypeLoader {
+        private final ITypeLoader classpathLoader = new ClasspathTypeLoader();
+
+        protected byte[] doLoad(String className) {
+                Buffer parentBuf = new Buffer();
+                if (classpathLoader.tryLoadType(className, parentBuf)) {
+                        return parentBuf.array();
+                }
+                return EMPTY_ARRAY;//need to return *something* as null means no store
+        }
+}
diff --git a/src/main/java/cuchaz/enigma/CachingTypeLoader.java b/src/main/java/cuchaz/enigma/CachingTypeLoader.java
new file mode 100644
index 00000000..22c31c63
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/CachingTypeLoader.java
@@ -0,0 +1,38 @@
+package cuchaz.enigma;
+
+import com.google.common.collect.Maps;
+import com.strobel.assembler.metadata.Buffer;
+import com.strobel.assembler.metadata.ITypeLoader;
+
+import java.util.Map;
+
+/**
+ * Common cache functions
+ */
+public abstract class CachingTypeLoader implements ITypeLoader {
+        protected static final byte[] EMPTY_ARRAY = {};
+
+        private final Map<String, byte[]> cache = Maps.newHashMap();
+
+        protected abstract byte[] doLoad(String className);
+
+        @Override
+        public boolean tryLoadType(String className, Buffer out) {
+
+                // check the cache
+                byte[] data = this.cache.computeIfAbsent(className, this::doLoad);
+
+                if (data == EMPTY_ARRAY) {
+                        return false;
+                }
+
+                out.reset(data.length);
+                System.arraycopy(data, 0, out.array(), out.position(), data.length);
+                out.position(0);
+                return true;
+        }
+
+        public void clearCache() {
+                this.cache.clear();
+        }
+}
diff --git a/src/main/java/cuchaz/enigma/Deobfuscator.java b/src/main/java/cuchaz/enigma/Deobfuscator.java
index 0e037536..b2cecfe2 100644
--- a/src/main/java/cuchaz/enigma/Deobfuscator.java
+++ b/src/main/java/cuchaz/enigma/Deobfuscator.java
@@ -163,6 +163,10 @@ public class Deobfuscator {
         }
 
         public CompilationUnit getSourceTree(String className) {
+                return getSourceTree(className, createTypeLoader());
+        }
+
+        public CompilationUnit getSourceTree(String className, ITranslatingTypeLoader loader) {
 
                 // we don't know if this class name is obfuscated or deobfuscated
                 // we need to tell the decompiler the deobfuscated name so it doesn't get freaked out
@@ -178,7 +182,6 @@ public class Deobfuscator {
                 }
 
                 // set the desc loader
-                TranslatingTypeLoader loader = createTypeLoader();
                 this.settings.setTypeLoader(loader);
 
                 // see if procyon can find the desc
@@ -267,6 +270,10 @@ public class Deobfuscator {
                         progress.init(classEntries.size(), "Decompiling classes...");
                 }
 
+                //create a common instance outside the loop as mappings shouldn't be changing while this is happening
+                //synchronized to make sure the parallelStream doesn't CME with the cache
+                ITranslatingTypeLoader typeLoader = new SynchronizedTypeLoader(createTypeLoader());
+
                 // DEOBFUSCATE ALL THE THINGS!! @_@
                 Stopwatch stopwatch = Stopwatch.createStarted();
                 AtomicInteger count = new AtomicInteger();
@@ -278,7 +285,7 @@ public class Deobfuscator {
 
                         try {
                                 // get the source
-                                CompilationUnit sourceTree = getSourceTree(obfClassEntry.getName());
+                                CompilationUnit sourceTree = getSourceTree(obfClassEntry.getName(), typeLoader);
 
                                 // write the file
                                 File file = new File(dirOut, deobfClassEntry.getName().replace('.', '/') + ".java");
@@ -295,7 +302,7 @@ public class Deobfuscator {
                         }
                 });
                 stopwatch.stop();
-                System.out.println("Done in : " + stopwatch.toString());
+                System.out.println("writeSources Done in : " + stopwatch.toString());
                 if (progress != null) {
                         progress.onProgress(count.get(), "Done:");
                 }
@@ -658,4 +665,5 @@ public class Deobfuscator {
         public interface ClassTransformer {
                 String transform(ClassNode node, ClassWriter writer);
         }
+
 }
diff --git a/src/main/java/cuchaz/enigma/ITranslatingTypeLoader.java b/src/main/java/cuchaz/enigma/ITranslatingTypeLoader.java
new file mode 100644
index 00000000..547ed0b2
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/ITranslatingTypeLoader.java
@@ -0,0 +1,19 @@
+package cuchaz.enigma;
+
+import com.strobel.assembler.metadata.ITypeLoader;
+import cuchaz.enigma.mapping.entry.ClassEntry;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.tree.ClassNode;
+
+import java.util.List;
+
+/**
+ * For delegation of TranslatingTypeLoader without needing the subclass the whole thing
+ */
+public interface ITranslatingTypeLoader extends ITypeLoader {
+        List<String> getClassNamesToTry(String className);
+
+        List<String> getClassNamesToTry(ClassEntry obfClassEntry);
+
+        String transformInto(ClassNode node, ClassWriter writer);
+}
diff --git a/src/main/java/cuchaz/enigma/SynchronizedTypeLoader.java b/src/main/java/cuchaz/enigma/SynchronizedTypeLoader.java
new file mode 100644
index 00000000..9aa3c5e4
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/SynchronizedTypeLoader.java
@@ -0,0 +1,39 @@
+package cuchaz.enigma;
+
+import com.strobel.assembler.metadata.Buffer;
+import cuchaz.enigma.mapping.entry.ClassEntry;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.tree.ClassNode;
+
+import java.util.List;
+
+/**
+ * Typeloader with synchronized tryLoadType method
+ */
+public class SynchronizedTypeLoader implements ITranslatingTypeLoader {
+        private final TranslatingTypeLoader delegate;
+
+        SynchronizedTypeLoader(TranslatingTypeLoader delegate) {
+                this.delegate = delegate;
+        }
+
+        @Override
+        public List<String> getClassNamesToTry(String className) {
+                return delegate.getClassNamesToTry(className);
+        }
+
+        @Override
+        public List<String> getClassNamesToTry(ClassEntry obfClassEntry) {
+                return delegate.getClassNamesToTry(obfClassEntry);
+        }
+
+        @Override
+        public String transformInto(ClassNode node, ClassWriter writer) {
+                return delegate.transformInto(node, writer);
+        }
+
+        @Override
+        public synchronized boolean tryLoadType(String internalName, Buffer buffer) {
+                return delegate.tryLoadType(internalName, buffer);
+        }
+}
diff --git a/src/main/java/cuchaz/enigma/TranslatingTypeLoader.java b/src/main/java/cuchaz/enigma/TranslatingTypeLoader.java
index 34901d52..eb780ee9 100644
--- a/src/main/java/cuchaz/enigma/TranslatingTypeLoader.java
+++ b/src/main/java/cuchaz/enigma/TranslatingTypeLoader.java
@@ -12,9 +12,7 @@
 package cuchaz.enigma;
 
 import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
 import com.strobel.assembler.metadata.Buffer;
-import com.strobel.assembler.metadata.ClasspathTypeLoader;
 import com.strobel.assembler.metadata.ITypeLoader;
 import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.analysis.ParsedJar;
@@ -27,17 +25,16 @@ import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.ClassNode;
 
 import java.util.List;
-import java.util.Map;
 
-public class TranslatingTypeLoader implements ITypeLoader {
+public class TranslatingTypeLoader extends CachingTypeLoader implements ITranslatingTypeLoader {
+        //Store one instance as the classpath shouldnt change during load
+        private static final ITypeLoader defaultTypeLoader = new CachingClasspathTypeLoader();
 
         private final ParsedJar jar;
         private final JarIndex jarIndex;
         private final ReferencedEntryPool entryPool;
         private final Translator obfuscatingTranslator;
         private final Translator deobfuscatingTranslator;
-        private final Map<String, byte[]> cache;
-        private final ClasspathTypeLoader defaultTypeLoader;
 
         public TranslatingTypeLoader(ParsedJar jar, JarIndex jarIndex, ReferencedEntryPool entryPool, Translator obfuscatingTranslator, Translator deobfuscatingTranslator) {
                 this.jar = jar;
@@ -45,37 +42,19 @@ public class TranslatingTypeLoader implements ITypeLoader {
                 this.entryPool = entryPool;
                 this.obfuscatingTranslator = obfuscatingTranslator;
                 this.deobfuscatingTranslator = deobfuscatingTranslator;
-                this.cache = Maps.newHashMap();
-                this.defaultTypeLoader = new ClasspathTypeLoader();
-
-        }
-
-        public void clearCache() {
-                this.cache.clear();
         }
 
-        @Override
-        public boolean tryLoadType(String className, Buffer out) {
-
-                // check the cache
-                byte[] data;
-                if (this.cache.containsKey(className)) {
-                        data = this.cache.get(className);
-                } else {
-                        data = loadType(className);
-                        this.cache.put(className, data);
-                }
-
+        protected byte[] doLoad(String className){
+                byte[] data = loadType(className);
                 if (data == null) {
                         // chain to default desc loader
-                        return this.defaultTypeLoader.tryLoadType(className, out);
+                        Buffer parentBuf = new Buffer();
+                        if (defaultTypeLoader.tryLoadType(className, parentBuf)){
+                                return parentBuf.array();
+                        }
+                        return EMPTY_ARRAY;//need to return *something* as null means no store
                 }
-
-                // send the class to the decompiler
-                out.reset(data.length);
-                System.arraycopy(data, 0, out.array(), out.position(), data.length);
-                out.position(0);
-                return true;
+                return data;
         }
 
         private byte[] loadType(String className) {
@@ -131,10 +110,12 @@ public class TranslatingTypeLoader implements ITypeLoader {
                 return null;
         }
 
+        @Override
         public List<String> getClassNamesToTry(String className) {
                 return getClassNamesToTry(this.obfuscatingTranslator.getTranslatedClass(new ClassEntry(className)));
         }
 
+        @Override
         public List<String> getClassNamesToTry(ClassEntry obfClassEntry) {
                 List<String> classNamesToTry = Lists.newArrayList();
                 classNamesToTry.add(obfClassEntry.getName());
@@ -145,8 +126,10 @@ public class TranslatingTypeLoader implements ITypeLoader {
                 return classNamesToTry;
         }
 
+        @Override
         public String transformInto(ClassNode node, ClassWriter writer) {
                 node.accept(new TranslationClassVisitor(deobfuscatingTranslator, jarIndex, entryPool, Opcodes.ASM5, writer));
                 return deobfuscatingTranslator.getTranslatedClass(new ClassEntry(node.name)).getName();
         }
+
 }
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationSignatureVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationSignatureVisitor.java
index 928c6d19..e66b085f 100644
--- a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationSignatureVisitor.java
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationSignatureVisitor.java
@@ -3,12 +3,14 @@ package cuchaz.enigma.bytecode.translators;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.signature.SignatureVisitor;
 
+import java.util.Stack;
 import java.util.function.Function;
 
 public class TranslationSignatureVisitor extends SignatureVisitor {
         private final Function<String, String> remapper;
 
         private final SignatureVisitor sv;
+        private final Stack<String> classStack = new Stack<>();
 
         public TranslationSignatureVisitor(Function<String, String> remapper, SignatureVisitor sv) {
                 super(Opcodes.ASM5);
@@ -18,13 +20,24 @@ public class TranslationSignatureVisitor extends SignatureVisitor {
 
         @Override
         public void visitClassType(String name) {
+                classStack.push(name);
                 String translatedEntry = this.remapper.apply(name);
                 this.sv.visitClassType(translatedEntry);
         }
 
         @Override
         public void visitInnerClassType(String name) {
+                String lastClass = classStack.pop();
+                if (!name.startsWith(lastClass+"$")){//todo see if there's a way to base this on whether there were type params or not
+                        name = lastClass+"$"+name;
+                }
                 String translatedEntry = this.remapper.apply(name);
+                if (translatedEntry.contains("/")){
+                        translatedEntry = translatedEntry.substring(translatedEntry.lastIndexOf("/")+1);
+                }
+                if (translatedEntry.contains("$")){
+                        translatedEntry = translatedEntry.substring(translatedEntry.lastIndexOf("$")+1);
+                }
                 this.sv.visitInnerClassType(translatedEntry);
         }
 
@@ -105,6 +118,8 @@ public class TranslationSignatureVisitor extends SignatureVisitor {
         @Override
         public void visitEnd() {
                 this.sv.visitEnd();
+                if (!classStack.empty())
+                        classStack.pop();
         }
 
         @Override