7 files changed, 111 insertions, 31 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/EntryReference.java b/src/main/java/cuchaz/enigma/analysis/EntryReference.java
index 101729d..df36c23 100644
--- a/src/main/java/cuchaz/enigma/analysis/EntryReference.java
+++ b/src/main/java/cuchaz/enigma/analysis/EntryReference.java
@@ -70,7 +70,7 @@ public class EntryReference<E extends Entry, C extends Entry> {
                 return entry;
         }
 
-        public String getNamableName() {
+        public String getNameableName() {
                 if (getNameableEntry() instanceof ClassEntry) {
                         ClassEntry classEntry = (ClassEntry) getNameableEntry();
                         if (classEntry.isInnerClass()) {
diff --git a/src/main/java/cuchaz/enigma/analysis/IndexClassVisitor.java b/src/main/java/cuchaz/enigma/analysis/IndexClassVisitor.java
index 69fe54f..4d5e803 100644
--- a/src/main/java/cuchaz/enigma/analysis/IndexClassVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/IndexClassVisitor.java
@@ -14,6 +14,11 @@ public class IndexClassVisitor extends ClassVisitor {
                 this.index = index;
         }
 
+        public IndexClassVisitor(JarIndex index, int api, ClassVisitor cv) {
+                super(api, cv);
+                this.index = index;
+        }
+
         @Override
         public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
                 this.classEntry = this.index.indexClass(access, name, signature, superName, interfaces);
diff --git a/src/main/java/cuchaz/enigma/analysis/IndexInnerClassVisitor.java b/src/main/java/cuchaz/enigma/analysis/IndexInnerClassVisitor.java
index 0474227..b6ab2d5 100644
--- a/src/main/java/cuchaz/enigma/analysis/IndexInnerClassVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/IndexInnerClassVisitor.java
@@ -11,6 +11,11 @@ public class IndexInnerClassVisitor extends ClassVisitor {
                 this.index = index;
         }
 
+        public IndexInnerClassVisitor(JarIndex index, int api, ClassVisitor cv) {
+                super(api, cv);
+                this.index = index;
+        }
+
         @Override
         public void visitInnerClass(String name, String outerName, String innerName, int access) {
                 ClassEntry entry = new ClassEntry(name);
@@ -19,5 +24,6 @@ public class IndexInnerClassVisitor extends ClassVisitor {
                         ClassEntry outerEntry = new ClassEntry(outerName);
                         index.indexInnerClass(entry, outerEntry);
                 }
+                super.visitInnerClass(name, outerName, innerName, access);
         }
 }
diff --git a/src/main/java/cuchaz/enigma/analysis/JarIndex.java b/src/main/java/cuchaz/enigma/analysis/JarIndex.java
index 158df4b..e8bda8e 100644
--- a/src/main/java/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/JarIndex.java
@@ -15,6 +15,8 @@ import com.google.common.collect.*;
 import cuchaz.enigma.bytecode.AccessFlags;
 import cuchaz.enigma.mapping.*;
 import cuchaz.enigma.mapping.entry.*;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.Opcodes;
 
 import java.util.*;
@@ -62,10 +64,15 @@ public class JarIndex {
                 obfClassEntries.addAll(jar.getClassEntries());
 
                 // step 2: index classes, fields, methods, interfaces
-                jar.visit(node -> node.accept(new IndexClassVisitor(this, Opcodes.ASM5)));
+                if (buildInnerClasses) {
+                        // + step 5: index inner classes
+                        jar.visitReader(name -> new IndexClassVisitor(this, Opcodes.ASM5, new IndexInnerClassVisitor(this, Opcodes.ASM5)), ClassReader.SKIP_CODE);
+                } else {
+                        jar.visitReader(name -> new IndexClassVisitor(this, Opcodes.ASM5), ClassReader.SKIP_CODE);
+                }
 
                 // step 3: index field, method, constructor references
-                jar.visit(node -> node.accept(new IndexReferenceVisitor(this, Opcodes.ASM5)));
+                jar.visitReader(name -> new IndexReferenceVisitor(this, Opcodes.ASM5), ClassReader.SKIP_FRAMES);
 
                 // step 4: index access and bridged methods
                 for (MethodDefEntry methodEntry : methods.values()) {
@@ -79,9 +86,6 @@ public class JarIndex {
                 }
 
                 if (buildInnerClasses) {
-                        // step 5: index inner classes and anonymous classes
-                        jar.visit(node -> node.accept(new IndexInnerClassVisitor(this, Opcodes.ASM5)));
-
                         // step 6: update other indices with inner class info
                         Map<String, String> renames = Maps.newHashMap();
                         for (ClassEntry innerClassEntry : this.innerClassesByOuter.values()) {
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;
+    }
 }
diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
index 78195cb..4c84e69 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
@@ -22,10 +22,7 @@ import com.strobel.decompiler.languages.java.ast.Identifier;
 import com.strobel.decompiler.languages.java.ast.TypeDeclaration;
 import cuchaz.enigma.mapping.entry.Entry;
 
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.TreeMap;
+import java.util.*;
 import java.util.regex.Pattern;
 
 public class SourceIndex {
@@ -48,9 +45,12 @@ public class SourceIndex {
                 this.tokenToReference = Maps.newTreeMap();
                 this.referenceToTokens = HashMultimap.create();
                 this.declarationToToken = Maps.newHashMap();
-                this.lineOffsets = Lists.newArrayList();
+                calculateLineOffsets();
+        }
 
+        private void calculateLineOffsets() {
                 // count the lines
+                this.lineOffsets = Lists.newArrayList();
                 this.lineOffsets.add(0);
                 for (int i = 0; i < source.length(); i++) {
                         if (source.charAt(i) == '\n') {
@@ -59,6 +59,32 @@ public class SourceIndex {
                 }
         }
 
+        public void remap(String source, Map<Token, Token> tokenMap) {
+                this.source = source;
+                calculateLineOffsets();
+
+                for (Entry entry : Lists.newArrayList(declarationToToken.keySet())) {
+                        Token token = declarationToToken.get(entry);
+                        declarationToToken.put(entry, tokenMap.getOrDefault(token, token));
+                }
+
+                for (Token token : Lists.newArrayList(tokenToReference.keySet())) {
+                        EntryReference<Entry, Entry> e = tokenToReference.remove(token);
+                        tokenToReference.put(tokenMap.getOrDefault(token, token), e);
+                }
+
+                for (EntryReference<Entry, Entry> ref : Lists.newArrayList(referenceToTokens.keySet())) {
+                        List<Token> newTokens = new ArrayList<>();
+
+                        for (Token token : referenceToTokens.get(ref)) {
+                                newTokens.add(tokenMap.getOrDefault(token, token));
+                        }
+
+                        referenceToTokens.removeAll(ref);
+                        referenceToTokens.putAll(ref, newTokens);
+                }
+        }
+
         public String getSource() {
                 return this.source;
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/Token.java b/src/main/java/cuchaz/enigma/analysis/Token.java
index 266d202..14fa7ca 100644
--- a/src/main/java/cuchaz/enigma/analysis/Token.java
+++ b/src/main/java/cuchaz/enigma/analysis/Token.java
@@ -25,6 +25,25 @@ public class Token implements Comparable<Token> {
                 }
         }
 
+        public int getRenameOffset(String to) {
+                int length = this.end - this.start;
+                return to.length() - length;
+        }
+
+        public String rename(String source, String to) {
+                int oldEnd = this.end;
+                this.text = to;
+                this.end = this.start + to.length();
+
+                return source.substring(0, this.start) + to + source.substring(oldEnd);
+        }
+
+        public Token move(int offset) {
+                Token token = new Token(this.start + offset, this.end + offset, null);
+                token.text = text;
+                return token;
+        }
+
         public boolean contains(int pos) {
                 return pos >= start && pos <= end;
         }
@@ -41,7 +60,7 @@ public class Token implements Comparable<Token> {
 
         @Override
         public int hashCode() {
-                return Integer.hashCode(start) + Integer.hashCode(end) + (text != null ? text.hashCode() : 0);
+                return start * 37 + end;
         }
 
         public boolean equals(Token other) {