Starting implementing local variables (#33)

TODO:
- Store format (need to be defined)
- Implement some translate operations

This commit also fix some cases where argument tokens are not selected

7 files changed, 219 insertions, 5 deletions

diff --git a/src/main/java/cuchaz/enigma/Deobfuscator.java b/src/main/java/cuchaz/enigma/Deobfuscator.java
index e9a998d..22194ad 100644
--- a/src/main/java/cuchaz/enigma/Deobfuscator.java
+++ b/src/main/java/cuchaz/enigma/Deobfuscator.java
@@ -533,6 +533,10 @@ public class Deobfuscator {
             return false;
         } else if (obfEntry instanceof ArgumentEntry) {
             return translator.translate((ArgumentEntry) obfEntry) != null;
+        } else if (obfEntry instanceof LocalVariableEntry) {
+            // TODO: Implement it
+            //return translator.translate((LocalVariableEntry)obfEntry) != null;
+            return false;
         } else {
             throw new Error("Unknown entry type: " + obfEntry.getClass().getName());
         }
@@ -549,6 +553,8 @@ public class Deobfuscator {
             throw new IllegalArgumentException("Cannot rename constructors");
         } else if (obfEntry instanceof ArgumentEntry) {
             this.renamer.setArgumentTreeName((ArgumentEntry) obfEntry, newName);
+        } else if (obfEntry instanceof LocalVariableEntry) {
+            // TODO: Implement it
         } else {
             throw new Error("Unknown entry type: " + obfEntry.getClass().getName());
         }
@@ -587,6 +593,8 @@ public class Deobfuscator {
             throw new IllegalArgumentException("Cannot rename constructors");
         } else if (obfEntry instanceof ArgumentEntry) {
             this.renamer.markArgumentAsDeobfuscated((ArgumentEntry) obfEntry);
+        } else if (obfEntry instanceof LocalVariableEntry) {
+            // TODO: Implement it
         } else {
             throw new Error("Unknown entry type: " + obfEntry);
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/JarIndex.java b/src/main/java/cuchaz/enigma/analysis/JarIndex.java
index 619d862..f955a40 100644
--- a/src/main/java/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/JarIndex.java
@@ -791,6 +791,9 @@ public class JarIndex {
             return containsObfBehavior((BehaviorEntry) obfEntry);
         } else if (obfEntry instanceof ArgumentEntry) {
             return containsObfArgument((ArgumentEntry) obfEntry);
+        } else if (obfEntry instanceof LocalVariableEntry) {
+            // TODO: Implement it
+            return false;
         } else {
             throw new Error("Entry type not supported: " + obfEntry.getClass().getName());
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java b/src/main/java/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
index c328f5f..e690abd 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
@@ -10,6 +10,9 @@
  ******************************************************************************/
 package cuchaz.enigma.analysis;
 
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
 import com.strobel.assembler.metadata.MemberReference;
 import com.strobel.assembler.metadata.MethodReference;
 import com.strobel.assembler.metadata.ParameterDefinition;
@@ -17,6 +20,7 @@ import com.strobel.assembler.metadata.TypeReference;
 import com.strobel.decompiler.languages.TextLocation;
 import com.strobel.decompiler.languages.java.ast.*;
 import cuchaz.enigma.mapping.*;
+import javassist.bytecode.Descriptor;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -27,11 +31,14 @@ public class SourceIndexBehaviorVisitor extends SourceIndexVisitor {
 
     // TODO: Really fix Procyon index problem with inner classes
     private int argumentPosition;
-    private Map<String, Entry> argumentCache = new HashMap<>();
+    private int localsPosition;
+    private Multimap<String, Identifier> unmatchedIdentifier = HashMultimap.create();
+    private Map<String, Entry> identifierEntryCache = new HashMap<>();
 
     public SourceIndexBehaviorVisitor(BehaviorEntry behaviorEntry) {
         this.behaviorEntry = behaviorEntry;
         this.argumentPosition = 0;
+        this.localsPosition = 0;
     }
 
     @Override
@@ -66,6 +73,9 @@ public class SourceIndexBehaviorVisitor extends SourceIndexVisitor {
             }
         }
 
+        // Check for identifier
+        node.getArguments().stream().filter(expression -> expression instanceof IdentifierExpression)
+                .forEach(expression -> this.checkIdentifier((IdentifierExpression) expression, index));
         return recurse(node, index);
     }
 
@@ -106,7 +116,7 @@ public class SourceIndexBehaviorVisitor extends SourceIndexVisitor {
                     argumentPosition++, node.getName());
             Identifier identifier = node.getNameToken();
             // cache the argument entry and the identifier
-            argumentCache.put(identifier.getName(), argumentEntry);
+            identifierEntryCache.put(identifier.getName(), argumentEntry);
             index.addDeclaration(identifier, argumentEntry);
         }
 
@@ -121,12 +131,31 @@ public class SourceIndexBehaviorVisitor extends SourceIndexVisitor {
             FieldEntry fieldEntry = new FieldEntry(classEntry, ref.getName(), new Type(ref.getErasedSignature()));
             index.addReference(node.getIdentifierToken(), fieldEntry, this.behaviorEntry);
         }
-        else if (argumentCache.containsKey(node.getIdentifier())) // If it's in the argument cache, create a token!
-            index.addDeclaration(node.getIdentifierToken(), argumentCache.get(node.getIdentifier()));
-
+        else
+            this.checkIdentifier(node, index);
         return recurse(node, index);
     }
 
+    private void checkIdentifier(IdentifierExpression node, SourceIndex index)
+    {
+        if (identifierEntryCache.containsKey(node.getIdentifier())) // If it's in the argument cache, create a token!
+            index.addDeclaration(node.getIdentifierToken(), identifierEntryCache.get(node.getIdentifier()));
+        else
+            unmatchedIdentifier.put(node.getIdentifier(), node.getIdentifierToken()); // Not matched actually, put it!
+    }
+
+    private void addDeclarationToUnmatched(String key, SourceIndex index)
+    {
+        Entry entry = identifierEntryCache.get(key);
+
+        // This cannot happened in theory
+        if (entry == null)
+            return;
+        for (Identifier identifier : unmatchedIdentifier.get(key))
+            index.addDeclaration(identifier, entry);
+        unmatchedIdentifier.removeAll(key);
+    }
+
     @Override
     public Void visitObjectCreationExpression(ObjectCreationExpression node, SourceIndex index) {
         MemberReference ref = node.getUserData(Keys.MEMBER_REFERENCE);
@@ -141,4 +170,43 @@ public class SourceIndexBehaviorVisitor extends SourceIndexVisitor {
 
         return recurse(node, index);
     }
+
+    @Override
+    public Void visitForEachStatement(ForEachStatement node, SourceIndex index) {
+        if (node.getVariableType() instanceof SimpleType)
+        {
+            SimpleType type = (SimpleType) node.getVariableType();
+            TypeReference typeReference = type.getUserData(Keys.TYPE_REFERENCE);
+            Identifier identifier = node.getVariableNameToken();
+            String signature = Descriptor.of(typeReference.getErasedDescription());
+            LocalVariableEntry localVariableEntry = new LocalVariableEntry(behaviorEntry, localsPosition++, identifier.getName(), new Type(signature));
+            identifierEntryCache.put(identifier.getName(), localVariableEntry);
+            addDeclarationToUnmatched(identifier.getName(), index);
+            index.addDeclaration(identifier, localVariableEntry);
+        }
+        return recurse(node, index);
+    }
+
+    @Override
+    public Void visitVariableDeclaration(VariableDeclarationStatement node, SourceIndex index) {
+        AstNodeCollection<VariableInitializer> variables = node.getVariables();
+
+        // Single assignation
+        if (variables.size() == 1)
+        {
+            VariableInitializer initializer = variables.firstOrNullObject();
+            if (initializer != null && node.getType() instanceof SimpleType)
+            {
+                SimpleType type = (SimpleType) node.getType();
+                TypeReference typeReference = type.getUserData(Keys.TYPE_REFERENCE);
+                String signature = Descriptor.of(typeReference.getErasedDescription());
+                Identifier identifier = initializer.getNameToken();
+                LocalVariableEntry localVariableEntry = new LocalVariableEntry(behaviorEntry, localsPosition++, initializer.getName(), new Type(signature));
+                identifierEntryCache.put(identifier.getName(), localVariableEntry);
+                addDeclarationToUnmatched(identifier.getName(), index);
+                index.addDeclaration(identifier, localVariableEntry);
+            }
+        }
+        return recurse(node, index);
+    }
 }
diff --git a/src/main/java/cuchaz/enigma/analysis/TranslationIndex.java b/src/main/java/cuchaz/enigma/analysis/TranslationIndex.java
index 25edaef..d51131f 100644
--- a/src/main/java/cuchaz/enigma/analysis/TranslationIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/TranslationIndex.java
@@ -181,6 +181,8 @@ public class TranslationIndex {
             return behaviorExists((BehaviorEntry) entry);
         } else if (entry instanceof ArgumentEntry) {
             return behaviorExists(((ArgumentEntry) entry).getBehaviorEntry());
+        } else if (entry instanceof LocalVariableEntry) {
+            return behaviorExists(((LocalVariableEntry) entry).getBehaviorEntry());
         }
         throw new IllegalArgumentException("Cannot check existence for " + entry.getClass());
     }
diff --git a/src/main/java/cuchaz/enigma/gui/Gui.java b/src/main/java/cuchaz/enigma/gui/Gui.java
index d3f7406..2a35507 100644
--- a/src/main/java/cuchaz/enigma/gui/Gui.java
+++ b/src/main/java/cuchaz/enigma/gui/Gui.java
@@ -442,6 +442,8 @@ public class Gui {
             showConstructorEntry((ConstructorEntry) m_reference.entry);
         } else if (m_reference.entry instanceof ArgumentEntry) {
             showArgumentEntry((ArgumentEntry) m_reference.entry);
+        } else if (m_reference.entry instanceof LocalVariableEntry) {
+            showLocalVariableEntry((LocalVariableEntry) m_reference.entry);
         } else {
             throw new Error("Unknown entry type: " + m_reference.entry.getClass().getName());
         }
@@ -449,6 +451,14 @@ public class Gui {
         redraw();
     }
 
+    private void showLocalVariableEntry(LocalVariableEntry entry) {
+        addNameValue(m_infoPanel, "Variable", entry.getName());
+        addNameValue(m_infoPanel, "Class", entry.getClassEntry().getName());
+        addNameValue(m_infoPanel, "Method", entry.getBehaviorEntry().getName());
+        addNameValue(m_infoPanel, "Index", Integer.toString(entry.getIndex()));
+        addNameValue(m_infoPanel, "Type", entry.getType().toString());
+    }
+
     private void showClassEntry(ClassEntry entry) {
         addNameValue(m_infoPanel, "Class", entry.getName());
     }
diff --git a/src/main/java/cuchaz/enigma/mapping/LocalVariableEntry.java b/src/main/java/cuchaz/enigma/mapping/LocalVariableEntry.java
new file mode 100644
index 0000000..8bbaaaf
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/mapping/LocalVariableEntry.java
@@ -0,0 +1,104 @@
+package cuchaz.enigma.mapping;
+
+import cuchaz.enigma.utils.Utils;
+
+/**
+ * Desc...
+ * Created by Thog
+ * 19/10/2016
+ */
+public class LocalVariableEntry implements Entry
+{
+
+    protected final BehaviorEntry behaviorEntry;
+    protected final String name;
+    protected final Type type;
+    protected final int index;
+
+    public LocalVariableEntry(BehaviorEntry behaviorEntry, int index, String name, Type type) {
+        if (behaviorEntry == null) {
+            throw new IllegalArgumentException("Behavior cannot be null!");
+        }
+        if (index < 0) {
+            throw new IllegalArgumentException("Index must be non-negative!");
+        }
+        if (name == null) {
+            throw new IllegalArgumentException("Variable name cannot be null!");
+        }
+        if (type == null) {
+            throw new IllegalArgumentException("Variable type cannot be null!");
+        }
+
+        this.behaviorEntry = behaviorEntry;
+        this.name = name;
+        this.type = type;
+        this.index = index;
+    }
+
+
+    public LocalVariableEntry(LocalVariableEntry other, ClassEntry newClassEntry) {
+        this.behaviorEntry = (BehaviorEntry) other.behaviorEntry.cloneToNewClass(newClassEntry);
+        this.name = other.name;
+        this.type = other.type;
+        this.index = other.index;
+    }
+
+    public BehaviorEntry getBehaviorEntry() {
+        return this.behaviorEntry;
+    }
+
+    public Type getType() {
+        return type;
+    }
+
+    public int getIndex() {
+        return index;
+    }
+
+    @Override
+    public String getName() {
+        return this.name;
+    }
+
+    @Override
+    public ClassEntry getClassEntry() {
+        return this.behaviorEntry.getClassEntry();
+    }
+
+    @Override
+    public String getClassName() {
+        return this.behaviorEntry.getClassName();
+    }
+
+    @Override
+    public LocalVariableEntry cloneToNewClass(ClassEntry classEntry) {
+        return new LocalVariableEntry(this, classEntry);
+    }
+
+    public String getMethodName() {
+        return this.behaviorEntry.getName();
+    }
+
+    public Signature getMethodSignature() {
+        return this.behaviorEntry.getSignature();
+    }
+
+    @Override
+    public int hashCode() {
+        return Utils.combineHashesOrdered(this.behaviorEntry, this.type.hashCode(), this.name.hashCode(), Integer.hashCode(this.index));
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        return other instanceof LocalVariableEntry && equals((LocalVariableEntry) other);
+    }
+
+    public boolean equals(LocalVariableEntry other) {
+        return this.behaviorEntry.equals(other.behaviorEntry) && this.type.equals(other.type) && this.name.equals(other.name) && this.index == other.index;
+    }
+
+    @Override
+    public String toString() {
+        return this.behaviorEntry.toString() + "(" + this.index + ":" + this.name + ":" + this.type + ")";
+    }
+}
diff --git a/src/main/java/cuchaz/enigma/mapping/Translator.java b/src/main/java/cuchaz/enigma/mapping/Translator.java
index b05714c..6b636f2 100644
--- a/src/main/java/cuchaz/enigma/mapping/Translator.java
+++ b/src/main/java/cuchaz/enigma/mapping/Translator.java
@@ -58,6 +58,8 @@ public class Translator {
             return (T) translateEntry((ConstructorEntry) entry);
         } else if (entry instanceof ArgumentEntry) {
             return (T) translateEntry((ArgumentEntry) entry);
+        } else if (entry instanceof  LocalVariableEntry) {
+            return (T) translateEntry((LocalVariableEntry) entry);
         } else {
             throw new Error("Unknown entry type: " + entry.getClass().getName());
         }
@@ -74,11 +76,28 @@ public class Translator {
             return translate(entry);
         } else if (entry instanceof ArgumentEntry) {
             return translate((ArgumentEntry) entry);
+        } else if (entry instanceof  LocalVariableEntry) {
+            return translate((LocalVariableEntry) entry);
         } else {
             throw new Error("Unknown entry type: " + entry.getClass().getName());
         }
     }
 
+    public String translate(LocalVariableEntry in)
+    {
+        LocalVariableEntry translated = translateEntry(in);
+        if (translated.equals(in)) {
+            return null;
+        }
+        return translated.getName();
+    }
+
+    public LocalVariableEntry translateEntry(LocalVariableEntry in)
+    {
+        // TODO: Implement it
+        return in;
+    }
+
     public String translate(ClassEntry in) {
         ClassEntry translated = translateEntry(in);
         if (translated.equals(in)) {