3 files changed, 238 insertions, 33 deletions

diff --git a/src/cuchaz/enigma/gui/ClassMatchingGui.java b/src/cuchaz/enigma/gui/ClassMatchingGui.java
index b674451..6943c3e 100644
--- a/src/cuchaz/enigma/gui/ClassMatchingGui.java
+++ b/src/cuchaz/enigma/gui/ClassMatchingGui.java
@@ -201,13 +201,10 @@ public class ClassMatchingGui {
                 
                 m_sourceClassLabel = new JLabel();
                 m_sourceClassLabel.setHorizontalAlignment(SwingConstants.RIGHT);
-                m_sourceClassLabel.setPreferredSize(new Dimension(300, 24));
                 m_destClassLabel = new JLabel();
                 m_destClassLabel.setHorizontalAlignment(SwingConstants.LEFT);
-                m_destClassLabel.setPreferredSize(new Dimension(300, 24));
                 
                 m_matchButton = new JButton();
-                m_matchButton.setPreferredSize(new Dimension(140, 24));
                 
                 m_advanceCheck = new JCheckBox("Advance to next likely match");
                 m_advanceCheck.addActionListener(new ActionListener() {
diff --git a/src/cuchaz/enigma/gui/CodeReader.java b/src/cuchaz/enigma/gui/CodeReader.java
index 05feb59..aa7e2db 100644
--- a/src/cuchaz/enigma/gui/CodeReader.java
+++ b/src/cuchaz/enigma/gui/CodeReader.java
@@ -7,12 +7,15 @@ import java.awt.event.ActionListener;
 import javax.swing.JEditorPane;
 import javax.swing.SwingUtilities;
 import javax.swing.Timer;
+import javax.swing.event.CaretEvent;
+import javax.swing.event.CaretListener;
 import javax.swing.text.BadLocationException;
 import javax.swing.text.Highlighter.HighlightPainter;
 
 import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 
 import cuchaz.enigma.Deobfuscator;
+import cuchaz.enigma.analysis.EntryReference;
 import cuchaz.enigma.analysis.SourceIndex;
 import cuchaz.enigma.analysis.Token;
 import cuchaz.enigma.mapping.ClassEntry;
@@ -26,8 +29,13 @@ public class CodeReader extends JEditorPane {
         
         private static final Object m_lock = new Object();
         
-        private SelectionHighlightPainter m_highlightPainter;
+        public static interface SelectionListener {
+                void onSelect(EntryReference<Entry,Entry> reference);
+        }
+        
+        private SelectionHighlightPainter m_selectionHighlightPainter;
         private SourceIndex m_sourceIndex;
+        private SelectionListener m_selectionListener;
         
         public CodeReader() {
                 
@@ -38,8 +46,28 @@ public class CodeReader extends JEditorPane {
                 DefaultSyntaxKit kit = (DefaultSyntaxKit)getEditorKit();
                 kit.toggleComponent(this, "de.sciss.syntaxpane.components.TokenMarker");
                 
-                m_highlightPainter = new SelectionHighlightPainter();
+                // hook events
+                addCaretListener(new CaretListener() {
+                        @Override
+                        public void caretUpdate(CaretEvent event) {
+                                if (m_selectionListener != null && m_sourceIndex != null) {
+                                        Token token = m_sourceIndex.getReferenceToken(event.getDot());
+                                        if (token != null) {
+                                                m_selectionListener.onSelect(m_sourceIndex.getDeobfReference(token));
+                                        } else {
+                                                m_selectionListener.onSelect(null);
+                                        }
+                                }
+                        }
+                });
+                
+                m_selectionHighlightPainter = new SelectionHighlightPainter();
                 m_sourceIndex = null;
+                m_selectionListener = null;
+        }
+        
+        public void setSelectionListener(SelectionListener val) {
+                m_selectionListener = val;
         }
         
         public void setCode(String code) {
@@ -49,6 +77,10 @@ public class CodeReader extends JEditorPane {
                 }
         }
         
+        public SourceIndex getSourceIndex() {
+                return m_sourceIndex;
+        }
+        
         public void decompileClass(ClassEntry classEntry, Deobfuscator deobfuscator) {
                 decompileClass(classEntry, deobfuscator, null);
         }
@@ -110,7 +142,7 @@ public class CodeReader extends JEditorPane {
         }
         
         public void navigateToToken(final Token token) {
-                navigateToToken(this, token, m_highlightPainter);
+                navigateToToken(this, token, m_selectionHighlightPainter);
         }
         
         // HACKHACK: someday we can update the main GUI to use this code reader
@@ -161,4 +193,22 @@ public class CodeReader extends JEditorPane {
                 });
                 timer.start();
         }
+        
+        public void setHighlightedTokens(Iterable<Token> tokens, HighlightPainter painter) {
+                for (Token token : tokens) {
+                        setHighlightedToken(token, painter);
+                }
+        }
+        
+        public void setHighlightedToken(Token token, HighlightPainter painter) {
+                try {
+                        getHighlighter().addHighlight(token.start, token.end, painter);
+                } catch (BadLocationException ex) {
+                        throw new IllegalArgumentException(ex);
+                }
+        }
+
+        public void clearHighlights() {
+                getHighlighter().removeAllHighlights();
+        }
 }
diff --git a/src/cuchaz/enigma/gui/FieldMatchingGui.java b/src/cuchaz/enigma/gui/FieldMatchingGui.java
index de9ba14..ef374c8 100644
--- a/src/cuchaz/enigma/gui/FieldMatchingGui.java
+++ b/src/cuchaz/enigma/gui/FieldMatchingGui.java
@@ -4,24 +4,34 @@ import java.awt.BorderLayout;
 import java.awt.Container;
 import java.awt.Dimension;
 import java.awt.FlowLayout;
-import java.util.Map;
+import java.awt.event.ActionEvent;
+import java.awt.event.ActionListener;
+import java.util.Collection;
+import java.util.Set;
 
 import javax.swing.BoxLayout;
+import javax.swing.JButton;
 import javax.swing.JFrame;
 import javax.swing.JLabel;
 import javax.swing.JPanel;
 import javax.swing.JScrollPane;
 import javax.swing.JSplitPane;
 import javax.swing.WindowConstants;
+import javax.swing.text.Highlighter.HighlightPainter;
+
+import com.google.common.collect.Sets;
 
 import cuchaz.enigma.Constants;
 import cuchaz.enigma.Deobfuscator;
+import cuchaz.enigma.analysis.EntryReference;
+import cuchaz.enigma.analysis.SourceIndex;
+import cuchaz.enigma.analysis.Token;
 import cuchaz.enigma.convert.ClassMatches;
 import cuchaz.enigma.convert.FieldMatches;
 import cuchaz.enigma.gui.ClassSelector.ClassSelectionListener;
 import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.Entry;
 import cuchaz.enigma.mapping.FieldEntry;
-import cuchaz.enigma.mapping.FieldMapping;
 import de.sciss.syntaxpane.DefaultSyntaxKit;
 
 
@@ -36,19 +46,28 @@ public class FieldMatchingGui {
         private ClassSelector m_sourceClasses;
         private CodeReader m_sourceReader;
         private CodeReader m_destReader;
+        private JButton m_matchButton;
+        private JLabel m_sourceLabel;
+        private JLabel m_destLabel;
+        private HighlightPainter m_unmatchedHighlightPainter;
+        private HighlightPainter m_matchedHighlightPainter;
 
         private ClassMatches m_classMatches;
         private FieldMatches m_fieldMatches;
-        private Map<FieldEntry,FieldMapping> m_droppedFieldMappings;
         private Deobfuscator m_sourceDeobfuscator;
         private Deobfuscator m_destDeobfuscator;
         private SaveListener m_saveListener;
+        private ClassEntry m_obfSourceClass;
+        private ClassEntry m_obfDestClass;
+        private FieldEntry m_obfSourceField;
+        private FieldEntry m_obfDestField;
+        private Set<FieldEntry> m_obfUnmatchedSourceFields;
+        private Set<FieldEntry> m_obfUnmatchedDestFields;
 
-        public FieldMatchingGui(ClassMatches classMatches, FieldMatches fieldMatches, Map<FieldEntry,FieldMapping> droppedFieldMappings, Deobfuscator sourceDeobfuscator, Deobfuscator destDeobfuscator) {
+        public FieldMatchingGui(ClassMatches classMatches, FieldMatches fieldMatches, Deobfuscator sourceDeobfuscator, Deobfuscator destDeobfuscator) {
                 
                 m_classMatches = classMatches;
                 m_fieldMatches = fieldMatches;
-                m_droppedFieldMappings = droppedFieldMappings;
                 m_sourceDeobfuscator = sourceDeobfuscator;
                 m_destDeobfuscator = destDeobfuscator;
                 
@@ -74,31 +93,57 @@ public class FieldMatchingGui {
                 JScrollPane sourceScroller = new JScrollPane(m_sourceClasses);
                 classesPanel.add(sourceScroller);
 
-                // init fields side
-                JPanel fieldsPanel = new JPanel();
-                fieldsPanel.setLayout(new BoxLayout(fieldsPanel, BoxLayout.PAGE_AXIS));
-                fieldsPanel.setPreferredSize(new Dimension(200, 0));
-                pane.add(fieldsPanel, BorderLayout.WEST);
-                fieldsPanel.add(new JLabel("Destination Fields"));
-                
                 // init readers
                 DefaultSyntaxKit.initKit();
                 m_sourceReader = new CodeReader();
+                m_sourceReader.setSelectionListener(new CodeReader.SelectionListener() {
+                        @Override
+                        public void onSelect(EntryReference<Entry,Entry> reference) {
+                                if (reference != null) {
+                                        onSelectSource(reference.entry);
+                                } else {
+                                        onSelectSource(null);
+                                }
+                        }
+                });
                 m_destReader = new CodeReader();
+                m_destReader.setSelectionListener(new CodeReader.SelectionListener() {
+                        @Override
+                        public void onSelect(EntryReference<Entry,Entry> reference) {
+                                if (reference != null) {
+                                        onSelectDest(reference.entry);
+                                } else {
+                                        onSelectDest(null);
+                                }
+                        }
+                });
 
                 // init all the splits
-                JSplitPane splitLeft = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, true, classesPanel, new JScrollPane(m_sourceReader));
+                JSplitPane splitRight = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, true, new JScrollPane(m_sourceReader), new JScrollPane(m_destReader));
+                splitRight.setResizeWeight(0.5); // resize 50:50
+                JSplitPane splitLeft = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, true, classesPanel, splitRight);
                 splitLeft.setResizeWeight(0); // let the right side take all the slack
-                JSplitPane splitRight = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, true, new JScrollPane(m_destReader), fieldsPanel);
-                splitRight.setResizeWeight(1); // let the left side take all the slack
-                JSplitPane splitCenter = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, true, splitLeft, splitRight);
-                splitCenter.setResizeWeight(0.5); // resize 50:50
-                pane.add(splitCenter, BorderLayout.CENTER);
-                splitCenter.resetToPreferredSizes();
+                pane.add(splitLeft, BorderLayout.CENTER);
+                splitLeft.resetToPreferredSizes();
                 
                 // init bottom panel
                 JPanel bottomPanel = new JPanel();
                 bottomPanel.setLayout(new FlowLayout());
+                pane.add(bottomPanel, BorderLayout.SOUTH);
+                
+                m_matchButton = new JButton("Match");
+                m_matchButton.addActionListener(new ActionListener() {
+                        @Override
+                        public void actionPerformed(ActionEvent event) {
+                                match();
+                        }
+                });
+                
+                m_sourceLabel = new JLabel();
+                bottomPanel.add(m_sourceLabel);
+                bottomPanel.add(m_matchButton);
+                m_destLabel = new JLabel();
+                bottomPanel.add(m_destLabel);
 
                 // show the frame
                 pane.doLayout();
@@ -106,26 +151,139 @@ public class FieldMatchingGui {
                 m_frame.setMinimumSize(new Dimension(640, 480));
                 m_frame.setVisible(true);
                 m_frame.setDefaultCloseOperation(WindowConstants.DISPOSE_ON_CLOSE);
+                
+                m_unmatchedHighlightPainter = new ObfuscatedHighlightPainter();
+                m_matchedHighlightPainter = new DeobfuscatedHighlightPainter();
 
                 // init state
                 m_saveListener = null;
-                m_fieldMatches.addUnmatchedSourceFields(droppedFieldMappings.keySet());
-                m_sourceClasses.setClasses(m_fieldMatches.getSourceClassesWithUnmatchedFields());
+                m_obfSourceClass = null;
+                m_obfDestClass = null;
+                m_obfSourceField = null;
+                m_obfDestField = null;
+                m_obfUnmatchedSourceFields = null;
+                m_obfUnmatchedDestFields = null;
+                updateSourceClasses();
+                updateMatchButton();
         }
 
         public void setSaveListener(SaveListener val) {
                 m_saveListener = val;
         }
         
+        private void updateSourceClasses() {
+                m_sourceClasses.setClasses(m_fieldMatches.getSourceClassesWithUnmatchedFields());
+                m_sourceClasses.expandAll();
+        }
+        
         protected void setSourceClass(ClassEntry obfSourceClass) {
                 
-                // get the matched dest class
-                final ClassEntry obfDestClass = m_classMatches.getUniqueMatches().get(obfSourceClass);
-                if (obfDestClass == null) {
-                        throw new Error("No matching dest class for source class: " + obfSourceClass);
+                m_obfSourceClass = obfSourceClass;
+                m_obfDestClass = m_classMatches.getUniqueMatches().get(obfSourceClass);
+                if (m_obfDestClass == null) {
+                        throw new Error("No matching dest class for source class: " + m_obfSourceClass);
                 }
+                
+                updateUnmatchedFields();
+                m_sourceReader.decompileClass(m_obfSourceClass, m_sourceDeobfuscator, new Runnable() {
+                        @Override
+                        public void run() {
+                                updateSourceHighlights();
+                        }
+                });
+                m_destReader.decompileClass(m_obfDestClass, m_destDeobfuscator, new Runnable() {
+                        @Override
+                        public void run() {
+                                updateDestHighlights();
+                        }
+                });
+        }
+        
+        private void updateUnmatchedFields() {
+                m_obfUnmatchedSourceFields = Sets.newHashSet(m_fieldMatches.getUnmatchedSourceFields(m_obfSourceClass));
+                m_obfUnmatchedDestFields = Sets.newHashSet();
+                for (FieldEntry destFieldEntry : m_destDeobfuscator.getJarIndex().getObfFieldEntries(m_obfDestClass)) {
+                        if (!m_fieldMatches.isDestMatched(destFieldEntry)) {
+                                m_obfUnmatchedDestFields.add(destFieldEntry);
+                        }
+                }
+        }
+
+        protected void updateSourceHighlights() {
+                highlightFields(m_sourceReader, m_sourceDeobfuscator, m_obfUnmatchedSourceFields, m_fieldMatches.matches().keySet());
+        }
+
+        protected void updateDestHighlights() {
+                highlightFields(m_destReader, m_destDeobfuscator, m_obfUnmatchedDestFields, m_fieldMatches.matches().values());
+        }
+        
+        private void highlightFields(CodeReader reader, Deobfuscator deobfuscator, Collection<FieldEntry> obfFieldEntries, Collection<FieldEntry> obfMatchedFieldEntries) {
+                reader.clearHighlights();
+                SourceIndex index = reader.getSourceIndex();
+                for (FieldEntry obfFieldEntry : obfFieldEntries) {
+                        FieldEntry deobfFieldEntry = deobfuscator.deobfuscateEntry(obfFieldEntry);
+                        Token token = index.getDeclarationToken(deobfFieldEntry);
+                        if (token == null) {
+                                System.err.println("WARNING: Can't find declaration token for " + deobfFieldEntry);
+                        } else {
+                                reader.setHighlightedToken(
+                                        token,
+                                        obfMatchedFieldEntries.contains(obfFieldEntry) ? m_matchedHighlightPainter : m_unmatchedHighlightPainter
+                                );
+                        }
+                }
+        }
+        
+        protected void onSelectSource(Entry entry) {
+                m_sourceLabel.setText("");
+                m_obfSourceField = null;
+                if (entry != null && entry instanceof FieldEntry) {
+                        FieldEntry fieldEntry = (FieldEntry)entry;
+                        FieldEntry obfFieldEntry = m_sourceDeobfuscator.obfuscateEntry(fieldEntry);
+                        if (m_obfUnmatchedSourceFields.contains(obfFieldEntry)) {
+                                m_obfSourceField = obfFieldEntry;
+                                m_sourceLabel.setText(fieldEntry.getName() + " " + fieldEntry.getType().toString());
+                        }
+                }
+                updateMatchButton();
+        }
 
-                m_sourceReader.decompileClass(obfSourceClass, m_sourceDeobfuscator);
-                m_destReader.decompileClass(obfDestClass, m_destDeobfuscator);
+        protected void onSelectDest(Entry entry) {
+                m_destLabel.setText("");
+                m_obfDestField = null;
+                if (entry != null && entry instanceof FieldEntry) {
+                        FieldEntry fieldEntry = (FieldEntry)entry;
+                        FieldEntry obfFieldEntry = m_destDeobfuscator.obfuscateEntry(fieldEntry);
+                        if (m_obfUnmatchedDestFields.contains(obfFieldEntry)) {
+                                m_obfDestField = obfFieldEntry;
+                                m_destLabel.setText(fieldEntry.getName() + " " + fieldEntry.getType().toString());
+                        }
+                }
+                updateMatchButton();
+        }
+                
+        private void updateMatchButton() {
+                m_matchButton.setEnabled(m_obfSourceField != null && m_obfDestField != null);
+        }
+        
+        protected void match() {
+                
+                // update the field matches
+                m_fieldMatches.makeMatch(m_obfSourceField, m_obfDestField);
+                save();
+
+                // update the ui
+                onSelectSource(null);
+                onSelectDest(null);
+                updateUnmatchedFields();
+                updateSourceHighlights();
+                updateDestHighlights();
+                updateSourceClasses();
+        }
+
+        private void save() {
+                if (m_saveListener != null) {
+                        m_saveListener.save(m_fieldMatches);
+                }
         }
 }