speed up Deobfuscator's getSources by using a single TranslatingTypeloader and caching the ClassLoaderTypeloader

5 files changed, 131 insertions, 35 deletions

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..cb02ffab 100644
--- a/src/main/java/cuchaz/enigma/Deobfuscator.java
+++ b/src/main/java/cuchaz/enigma/Deobfuscator.java
@@ -15,6 +15,7 @@ import com.google.common.base.Stopwatch;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
+import com.strobel.assembler.metadata.Buffer;
 import com.strobel.assembler.metadata.MetadataSystem;
 import com.strobel.assembler.metadata.TypeDefinition;
 import com.strobel.assembler.metadata.TypeReference;
@@ -163,6 +164,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 +183,6 @@ public class Deobfuscator {
                 }
 
                 // set the desc loader
-                TranslatingTypeLoader loader = createTypeLoader();
                 this.settings.setTypeLoader(loader);
 
                 // see if procyon can find the desc
@@ -267,6 +271,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 +286,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 +303,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 +666,32 @@ public class Deobfuscator {
         public interface ClassTransformer {
                 String transform(ClassNode node, ClassWriter writer);
         }
+
+        private class SynchronizedTypeLoader implements ITranslatingTypeLoader {
+                private final TranslatingTypeLoader delegate;
+
+                private 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/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/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();
         }
+
 }