match/unmatch button works

2 files changed, 47 insertions, 3 deletions

diff --git a/src/cuchaz/enigma/convert/ClassMatch.java b/src/cuchaz/enigma/convert/ClassMatch.java
index 9cecf70..eaaaa41 100644
--- a/src/cuchaz/enigma/convert/ClassMatch.java
+++ b/src/cuchaz/enigma/convert/ClassMatch.java
@@ -20,10 +20,16 @@ public class ClassMatch {
         }
         
         public ClassMatch(ClassEntry sourceClass, ClassEntry destClass) {
-                this.sourceClasses = Sets.newHashSet(sourceClass);
-                this.destClasses = Sets.newHashSet(destClass);
+                sourceClasses = Sets.newHashSet();
+                if (sourceClass != null) {
+                        sourceClasses.add(sourceClass);
+                }
+                destClasses = Sets.newHashSet();
+                if (destClass != null) {
+                        destClasses.add(destClass);
+                }
         }
-
+        
         public boolean isMatched() {
                 return sourceClasses.size() > 0 && destClasses.size() > 0;
         }
diff --git a/src/cuchaz/enigma/convert/Matches.java b/src/cuchaz/enigma/convert/Matches.java
index 5faa923..0b00b29 100644
--- a/src/cuchaz/enigma/convert/Matches.java
+++ b/src/cuchaz/enigma/convert/Matches.java
@@ -48,6 +48,22 @@ public class Matches implements Iterable<ClassMatch> {
                 m_matches.add(match);
                 indexMatch(match);
         }
+        
+        public void remove(ClassMatch match) {
+                for (ClassEntry sourceClass : match.sourceClasses) {
+                        m_matchesBySource.remove(sourceClass);
+                        m_uniqueMatches.remove(sourceClass);
+                        m_ambiguousMatchesBySource.remove(sourceClass);
+                        m_unmatchedSourceClasses.remove(sourceClass);
+                }
+                for (ClassEntry destClass : match.sourceClasses) {
+                        m_matchesByDest.remove(destClass);
+                        m_uniqueMatches.inverse().remove(destClass);
+                        m_ambiguousMatchesByDest.remove(destClass);
+                        m_unmatchedDestClasses.remove(destClass);
+                }
+                m_matches.remove(match);
+        }
 
         public int size() {
                 return m_matches.size();
@@ -112,4 +128,26 @@ public class Matches implements Iterable<ClassMatch> {
         public ClassMatch getMatchByDest(ClassEntry destClass) {
                 return m_matchesByDest.get(destClass);
         }
+        
+        public void removeSource(ClassEntry sourceClass) {
+                ClassMatch match = m_matchesBySource.get(sourceClass);
+                if (match != null) {
+                        remove(match);
+                        match.sourceClasses.remove(sourceClass);
+                        if (!match.sourceClasses.isEmpty() || !match.destClasses.isEmpty()) {
+                                add(match);
+                        }
+                }
+        }
+        
+        public void removeDest(ClassEntry destClass) {
+                ClassMatch match = m_matchesByDest.get(destClass);
+                if (match != null) {
+                        remove(match);
+                        match.destClasses.remove(destClass);
+                        if (!match.sourceClasses.isEmpty() || !match.destClasses.isEmpty()) {
+                                add(match);
+                        }
+                }
+        }
 }