Fix #110 and remap indices with matched bridge method names

15 files changed, 124 insertions, 105 deletions

diff --git a/src/main/java/cuchaz/enigma/translation/mapping/EntryMap.java b/src/main/java/cuchaz/enigma/translation/mapping/EntryMap.java
index 6af4846..e1a3253 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/EntryMap.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/EntryMap.java
@@ -3,7 +3,7 @@ package cuchaz.enigma.translation.mapping;
 import cuchaz.enigma.translation.representation.entry.Entry;
 
 import javax.annotation.Nullable;
-import java.util.Collection;
+import java.util.stream.Stream;
 
 public interface EntryMap<T> {
         void insert(Entry<?> entry, T value);
@@ -18,7 +18,7 @@ public interface EntryMap<T> {
                 return get(entry) != null;
         }
 
-        Collection<Entry<?>> getAllEntries();
+        Stream<Entry<?>> getAllEntries();
 
         boolean isEmpty();
 }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/EntryMapping.java b/src/main/java/cuchaz/enigma/translation/mapping/EntryMapping.java
index f11cdef..b74cc0b 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/EntryMapping.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/EntryMapping.java
@@ -27,4 +27,21 @@ public class EntryMapping {
                 }
                 return accessModifier;
         }
+
+        @Override
+        public boolean equals(Object obj) {
+                if (obj == this) return true;
+
+                if (obj instanceof EntryMapping) {
+                        EntryMapping mapping = (EntryMapping) obj;
+                        return mapping.targetName.equals(targetName) && mapping.accessModifier.equals(accessModifier);
+                }
+
+                return false;
+        }
+
+        @Override
+        public int hashCode() {
+                return targetName.hashCode() + accessModifier.hashCode() * 31;
+        }
 }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/EntryRemapper.java b/src/main/java/cuchaz/enigma/translation/mapping/EntryRemapper.java
index fa421b9..ed9f820 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/EntryRemapper.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/EntryRemapper.java
@@ -11,6 +11,7 @@ import cuchaz.enigma.translation.representation.entry.Entry;
 
 import javax.annotation.Nullable;
 import java.util.Collection;
+import java.util.stream.Stream;
 
 public class EntryRemapper {
         private final DeltaTrackingTree<EntryMapping> obfToDeobf;
@@ -69,14 +70,10 @@ public class EntryRemapper {
                 return deobfuscator;
         }
 
-        public Collection<Entry<?>> getObfEntries() {
+        public Stream<Entry<?>> getObfEntries() {
                 return obfToDeobf.getAllEntries();
         }
 
-        public Collection<Entry<?>> getObfRootEntries() {
-                return obfToDeobf.getRootEntries();
-        }
-
         public Collection<Entry<?>> getObfChildren(Entry<?> obfuscatedEntry) {
                 return obfToDeobf.getChildren(obfuscatedEntry);
         }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/IndexEntryResolver.java b/src/main/java/cuchaz/enigma/translation/mapping/IndexEntryResolver.java
index 1f2290a..6c8ed76 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/IndexEntryResolver.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/IndexEntryResolver.java
@@ -91,6 +91,8 @@ public class IndexEntryResolver implements EntryResolver {
                                 Set<Entry<ClassEntry>> resolvedBridge = resolveChildEntry(bridgeMethod, strategy);
                                 if (!resolvedBridge.isEmpty()) {
                                         return resolvedBridge;
+                                } else {
+                                        return Collections.singleton(bridgeMethod);
                                 }
                         }
                 }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/MappingDelta.java b/src/main/java/cuchaz/enigma/translation/mapping/MappingDelta.java
index 9f1f468..1407bb6 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/MappingDelta.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/MappingDelta.java
@@ -2,55 +2,53 @@ package cuchaz.enigma.translation.mapping;
 
 import cuchaz.enigma.translation.Translatable;
 import cuchaz.enigma.translation.Translator;
-import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.translation.mapping.tree.EntryTreeNode;
+import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
 import cuchaz.enigma.translation.representation.entry.Entry;
 
+import java.util.stream.Stream;
+
 public class MappingDelta<T> implements Translatable {
         public static final Object PLACEHOLDER = new Object();
 
         private final EntryTree<T> baseMappings;
 
-        private final EntryTree<Object> additions;
-        private final EntryTree<Object> deletions;
+        private final EntryTree<Object> changes;
 
-        public MappingDelta(EntryTree<T> baseMappings, EntryTree<Object> additions, EntryTree<Object> deletions) {
+        public MappingDelta(EntryTree<T> baseMappings, EntryTree<Object> changes) {
                 this.baseMappings = baseMappings;
-                this.additions = additions;
-                this.deletions = deletions;
+                this.changes = changes;
         }
 
         public MappingDelta(EntryTree<T> baseMappings) {
-                this(baseMappings, new HashEntryTree<>(), new HashEntryTree<>());
+                this(baseMappings, new HashEntryTree<>());
         }
 
         public static <T> MappingDelta<T> added(EntryTree<T> mappings) {
-                EntryTree<Object> additions = new HashEntryTree<>();
-                for (Entry<?> entry : mappings.getAllEntries()) {
-                        additions.insert(entry, PLACEHOLDER);
-                }
+                EntryTree<Object> changes = new HashEntryTree<>();
+                mappings.getAllEntries().forEach(entry -> changes.insert(entry, PLACEHOLDER));
 
-                return new MappingDelta<>(new HashEntryTree<>(), additions, new HashEntryTree<>());
+                return new MappingDelta<>(new HashEntryTree<>(), changes);
         }
 
         public EntryTree<T> getBaseMappings() {
                 return baseMappings;
         }
 
-        public EntryTree<?> getAdditions() {
-                return additions;
+        public EntryTree<?> getChanges() {
+                return changes;
         }
 
-        public EntryTree<?> getDeletions() {
-                return deletions;
+        public Stream<Entry<?>> getChangedRoots() {
+                return changes.getRootNodes().map(EntryTreeNode::getEntry);
         }
 
         @Override
         public MappingDelta<T> translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings) {
                 return new MappingDelta<>(
                                 translator.translate(baseMappings),
-                                translator.translate(additions),
-                                translator.translate(deletions)
+                                translator.translate(changes)
                 );
         }
 }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/MappingsChecker.java b/src/main/java/cuchaz/enigma/translation/mapping/MappingsChecker.java
index 77d75ec..af9e63a 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/MappingsChecker.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/MappingsChecker.java
@@ -11,6 +11,7 @@
 
 package cuchaz.enigma.translation.mapping;
 
+import cuchaz.enigma.ProgressListener;
 import cuchaz.enigma.analysis.index.JarIndex;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
 import cuchaz.enigma.translation.mapping.tree.EntryTreeNode;
@@ -22,6 +23,7 @@ import cuchaz.enigma.translation.representation.entry.MethodEntry;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 public class MappingsChecker {
         private final JarIndex index;
@@ -32,14 +34,19 @@ public class MappingsChecker {
                 this.mappings = mappings;
         }
 
-        public Dropped dropBrokenMappings() {
+        public Dropped dropBrokenMappings(ProgressListener progress) {
                 Dropped dropped = new Dropped();
 
-                Collection<Entry<?>> obfEntries = mappings.getAllEntries();
+                Collection<Entry<?>> obfEntries = mappings.getAllEntries()
+                                .filter(e -> e instanceof ClassEntry || e instanceof MethodEntry || e instanceof FieldEntry)
+                                .collect(Collectors.toList());
+
+                progress.init(obfEntries.size(), "Checking for dropped mappings");
+
+                int steps = 0;
                 for (Entry<?> entry : obfEntries) {
-                        if (entry instanceof ClassEntry || entry instanceof MethodEntry || entry instanceof FieldEntry) {
-                                tryDropEntry(dropped, entry);
-                        }
+                        progress.step(steps++, entry.toString());
+                        tryDropEntry(dropped, entry);
                 }
 
                 dropped.apply(mappings);
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/serde/EnigmaMappingsReader.java b/src/main/java/cuchaz/enigma/translation/mapping/serde/EnigmaMappingsReader.java
index d36bc0b..3f61325 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/serde/EnigmaMappingsReader.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/serde/EnigmaMappingsReader.java
@@ -1,12 +1,13 @@
 package cuchaz.enigma.translation.mapping.serde;
 
 import com.google.common.base.Charsets;
+import cuchaz.enigma.ProgressListener;
 import cuchaz.enigma.throwables.MappingParseException;
 import cuchaz.enigma.translation.mapping.AccessModifier;
 import cuchaz.enigma.translation.mapping.EntryMapping;
 import cuchaz.enigma.translation.mapping.MappingPair;
-import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
 import cuchaz.enigma.translation.representation.MethodDescriptor;
 import cuchaz.enigma.translation.representation.TypeDescriptor;
 import cuchaz.enigma.translation.representation.entry.*;
@@ -24,22 +25,32 @@ import java.util.stream.Collectors;
 public enum EnigmaMappingsReader implements MappingsReader {
         FILE {
                 @Override
-                public EntryTree<EntryMapping> read(Path path) throws IOException, MappingParseException {
+                public EntryTree<EntryMapping> read(Path path, ProgressListener progress) throws IOException, MappingParseException {
+                        progress.init(1, "Loading mapping file");
+
                         EntryTree<EntryMapping> mappings = new HashEntryTree<>();
                         readFile(path, mappings);
+
+                        progress.step(1, "Done!");
+
                         return mappings;
                 }
         },
         DIRECTORY {
                 @Override
-                public EntryTree<EntryMapping> read(Path path) throws IOException, MappingParseException {
+                public EntryTree<EntryMapping> read(Path root, ProgressListener progress) throws IOException, MappingParseException {
                         EntryTree<EntryMapping> mappings = new HashEntryTree<>();
 
-                        List<Path> files = Files.walk(path)
+                        List<Path> files = Files.walk(root)
                                         .filter(f -> !Files.isDirectory(f))
                                         .filter(f -> f.toString().endsWith(".mapping"))
                                         .collect(Collectors.toList());
+
+                        progress.init(files.size(), "Loading mapping files");
+                        int step = 0;
+
                         for (Path file : files) {
+                                progress.step(step++, root.relativize(file).toString());
                                 if (Files.isHidden(file)) {
                                         continue;
                                 }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/serde/EnigmaMappingsWriter.java b/src/main/java/cuchaz/enigma/translation/mapping/serde/EnigmaMappingsWriter.java
index 8c55fa9..618bb0a 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/serde/EnigmaMappingsWriter.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/serde/EnigmaMappingsWriter.java
@@ -28,6 +28,7 @@ import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
@@ -35,7 +36,7 @@ public enum EnigmaMappingsWriter implements MappingsWriter {
         FILE {
                 @Override
                 public void write(EntryTree<EntryMapping> mappings, MappingDelta<EntryMapping> delta, Path path, ProgressListener progress) {
-                        Collection<ClassEntry> classes = mappings.getRootEntries().stream()
+                        Collection<ClassEntry> classes = mappings.getRootNodes()
                                         .filter(entry -> entry instanceof ClassEntry)
                                         .map(entry -> (ClassEntry) entry)
                                         .collect(Collectors.toList());
@@ -56,19 +57,19 @@ public enum EnigmaMappingsWriter implements MappingsWriter {
         DIRECTORY {
                 @Override
                 public void write(EntryTree<EntryMapping> mappings, MappingDelta<EntryMapping> delta, Path path, ProgressListener progress) {
-                        applyDeletions(delta.getBaseMappings(), delta.getDeletions(), path);
-
-                        Collection<ClassEntry> classes = delta.getAdditions().getRootEntries().stream()
+                        Collection<ClassEntry> changedClasses = delta.getChangedRoots()
                                         .filter(entry -> entry instanceof ClassEntry)
                                         .map(entry -> (ClassEntry) entry)
                                         .collect(Collectors.toList());
 
-                        progress.init(classes.size(), "Writing classes");
+                        applyDeletions(path, changedClasses, mappings, delta.getBaseMappings());
+
+                        progress.init(changedClasses.size(), "Writing classes");
 
-                        Translator translator = new MappingTranslator(mappings, VoidEntryResolver.INSTANCE);
                         AtomicInteger steps = new AtomicInteger();
 
-                        classes.parallelStream().forEach(classEntry -> {
+                        Translator translator = new MappingTranslator(mappings, VoidEntryResolver.INSTANCE);
+                        changedClasses.parallelStream().forEach(classEntry -> {
                                 progress.step(steps.getAndIncrement(), classEntry.getFullName());
 
                                 try {
@@ -86,12 +87,12 @@ public enum EnigmaMappingsWriter implements MappingsWriter {
                         });
                 }
 
-                private void applyDeletions(EntryTree<EntryMapping> baseMappings, EntryTree<?> deletions, Path root) {
-                        Translator oldMappingTranslator = new MappingTranslator(baseMappings, VoidEntryResolver.INSTANCE);
+                private void applyDeletions(Path root, Collection<ClassEntry> changedClasses, EntryTree<EntryMapping> mappings, EntryTree<EntryMapping> oldMappings) {
+                        Translator oldMappingTranslator = new MappingTranslator(oldMappings, VoidEntryResolver.INSTANCE);
 
-                        Collection<ClassEntry> deletedClasses = deletions.getRootEntries().stream()
-                                        .filter(e -> e instanceof ClassEntry)
-                                        .map(e -> oldMappingTranslator.translate((ClassEntry) e))
+                        Collection<ClassEntry> deletedClasses = changedClasses.stream()
+                                        .filter(e -> !Objects.equals(oldMappings.get(e), mappings.get(e)))
+                                        .map(oldMappingTranslator::translate)
                                         .collect(Collectors.toList());
 
                         for (ClassEntry classEntry : deletedClasses) {
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java b/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
index 622a0e1..2528352 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
@@ -35,11 +35,11 @@ public enum MappingFormat {
                 writer.write(mappings, delta, path, progressListener);
         }
 
-        public EntryTree<EntryMapping> read(Path path) throws IOException, MappingParseException {
+        public EntryTree<EntryMapping> read(Path path, ProgressListener progressListener) throws IOException, MappingParseException {
                 if (reader == null) {
                         throw new IllegalStateException(name() + " does not support reading");
                 }
-                return reader.read(path);
+                return reader.read(path, progressListener);
         }
 
         @Nullable
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingsReader.java b/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingsReader.java
index f239ee6..af0933f 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingsReader.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingsReader.java
@@ -1,5 +1,6 @@
 package cuchaz.enigma.translation.mapping.serde;
 
+import cuchaz.enigma.ProgressListener;
 import cuchaz.enigma.throwables.MappingParseException;
 import cuchaz.enigma.translation.mapping.EntryMapping;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
@@ -8,5 +9,5 @@ import java.io.IOException;
 import java.nio.file.Path;
 
 public interface MappingsReader {
-        EntryTree<EntryMapping> read(Path path) throws MappingParseException, IOException;
+        EntryTree<EntryMapping> read(Path path, ProgressListener progress) throws MappingParseException, IOException;
 }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/serde/TinyMappingsReader.java b/src/main/java/cuchaz/enigma/translation/mapping/serde/TinyMappingsReader.java
index e0afc3e..69d8235 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/serde/TinyMappingsReader.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/serde/TinyMappingsReader.java
@@ -1,11 +1,12 @@
 package cuchaz.enigma.translation.mapping.serde;
 
 import com.google.common.base.Charsets;
+import cuchaz.enigma.ProgressListener;
 import cuchaz.enigma.throwables.MappingParseException;
 import cuchaz.enigma.translation.mapping.EntryMapping;
 import cuchaz.enigma.translation.mapping.MappingPair;
-import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
 import cuchaz.enigma.translation.representation.MethodDescriptor;
 import cuchaz.enigma.translation.representation.TypeDescriptor;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
@@ -22,15 +23,19 @@ public enum TinyMappingsReader implements MappingsReader {
         INSTANCE;
 
         @Override
-        public EntryTree<EntryMapping> read(Path path) throws IOException, MappingParseException {
-                return read(path, Files.readAllLines(path, Charsets.UTF_8));
+        public EntryTree<EntryMapping> read(Path path, ProgressListener progress) throws IOException, MappingParseException {
+                return read(path, Files.readAllLines(path, Charsets.UTF_8), progress);
         }
 
-        private EntryTree<EntryMapping> read(Path path, List<String> lines) throws MappingParseException {
+        private EntryTree<EntryMapping> read(Path path, List<String> lines, ProgressListener progress) throws MappingParseException {
                 EntryTree<EntryMapping> mappings = new HashEntryTree<>();
                 lines.remove(0);
 
+                progress.init(lines.size(), "Loading mapping file");
+
                 for (int lineNumber = 0; lineNumber < lines.size(); lineNumber++) {
+                        progress.step(lineNumber, "");
+
                         String line = lines.get(lineNumber);
 
                         try {
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/tree/DeltaTrackingTree.java b/src/main/java/cuchaz/enigma/translation/mapping/tree/DeltaTrackingTree.java
index 36be5e1..255fa5f 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/tree/DeltaTrackingTree.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/tree/DeltaTrackingTree.java
@@ -10,13 +10,13 @@ import cuchaz.enigma.translation.representation.entry.Entry;
 import javax.annotation.Nullable;
 import java.util.Collection;
 import java.util.Iterator;
+import java.util.stream.Stream;
 
 public class DeltaTrackingTree<T> implements EntryTree<T> {
         private final EntryTree<T> delegate;
 
         private EntryTree<T> deltaReference;
-        private EntryTree<Object> additions = new HashEntryTree<>();
-        private EntryTree<Object> deletions = new HashEntryTree<>();
+        private EntryTree<Object> changes = new HashEntryTree<>();
 
         public DeltaTrackingTree(EntryTree<T> delegate) {
                 this.delegate = delegate;
@@ -29,30 +29,19 @@ public class DeltaTrackingTree<T> implements EntryTree<T> {
 
         @Override
         public void insert(Entry<?> entry, T value) {
-                if (value != null) {
-                        trackAddition(entry);
-                } else {
-                        trackDeletion(entry);
-                }
+                trackChange(entry);
                 delegate.insert(entry, value);
         }
 
         @Nullable
         @Override
         public T remove(Entry<?> entry) {
-                T value = delegate.remove(entry);
-                trackDeletion(entry);
-                return value;
+                trackChange(entry);
+                return delegate.remove(entry);
         }
 
-        public void trackAddition(Entry<?> entry) {
-                deletions.insert(entry, MappingDelta.PLACEHOLDER);
-                additions.insert(entry, MappingDelta.PLACEHOLDER);
-        }
-
-        public void trackDeletion(Entry<?> entry) {
-                additions.remove(entry);
-                deletions.insert(entry, MappingDelta.PLACEHOLDER);
+        public void trackChange(Entry<?> entry) {
+                changes.insert(entry, MappingDelta.PLACEHOLDER);
         }
 
         @Nullable
@@ -78,25 +67,19 @@ public class DeltaTrackingTree<T> implements EntryTree<T> {
         }
 
         @Override
-        public Collection<EntryTreeNode<T>> getAllNodes() {
-                return delegate.getAllNodes();
-        }
-
-        @Override
-        public Collection<Entry<?>> getRootEntries() {
-                return delegate.getRootEntries();
+        public Stream<EntryTreeNode<T>> getRootNodes() {
+                return delegate.getRootNodes();
         }
 
         @Override
         public DeltaTrackingTree<T> translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings) {
                 DeltaTrackingTree<T> translatedTree = new DeltaTrackingTree<>(delegate.translate(translator, resolver, mappings));
-                translatedTree.additions = additions.translate(translator, resolver, mappings);
-                translatedTree.deletions = deletions.translate(translator, resolver, mappings);
+                translatedTree.changes = changes.translate(translator, resolver, mappings);
                 return translatedTree;
         }
 
         @Override
-        public Collection<Entry<?>> getAllEntries() {
+        public Stream<Entry<?>> getAllEntries() {
                 return delegate.getAllEntries();
         }
 
@@ -111,18 +94,17 @@ public class DeltaTrackingTree<T> implements EntryTree<T> {
         }
 
         public MappingDelta<T> takeDelta() {
-                MappingDelta<T> delta = new MappingDelta<>(deltaReference, additions, deletions);
+                MappingDelta<T> delta = new MappingDelta<>(deltaReference, changes);
                 resetDelta();
                 return delta;
         }
 
         private void resetDelta() {
                 deltaReference = new HashEntryTree<>(delegate);
-                additions = new HashEntryTree<>();
-                deletions = new HashEntryTree<>();
+                changes = new HashEntryTree<>();
         }
 
         public boolean isDirty() {
-                return !additions.isEmpty() || !deletions.isEmpty();
+                return !changes.isEmpty();
         }
 }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/tree/EntryTree.java b/src/main/java/cuchaz/enigma/translation/mapping/tree/EntryTree.java
index 4f341f4..daaefcc 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/tree/EntryTree.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/tree/EntryTree.java
@@ -9,6 +9,7 @@ import cuchaz.enigma.translation.representation.entry.Entry;
 
 import javax.annotation.Nullable;
 import java.util.Collection;
+import java.util.stream.Stream;
 
 public interface EntryTree<T> extends EntryMap<T>, Iterable<EntryTreeNode<T>>, Translatable {
         Collection<Entry<?>> getChildren(Entry<?> entry);
@@ -18,9 +19,7 @@ public interface EntryTree<T> extends EntryMap<T>, Iterable<EntryTreeNode<T>>, T
         @Nullable
         EntryTreeNode<T> findNode(Entry<?> entry);
 
-        Collection<EntryTreeNode<T>> getAllNodes();
-
-        Collection<Entry<?>> getRootEntries();
+        Stream<EntryTreeNode<T>> getRootNodes();
 
         @Override
         EntryTree<T> translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings);
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/tree/EntryTreeNode.java b/src/main/java/cuchaz/enigma/translation/mapping/tree/EntryTreeNode.java
index 734b60c..affcd50 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/tree/EntryTreeNode.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/tree/EntryTreeNode.java
@@ -33,4 +33,8 @@ public interface EntryTreeNode<T> {
                                 .map(EntryTreeNode::getEntry)
                                 .collect(Collectors.toList());
         }
+
+        default boolean hasValue() {
+                return getValue() != null;
+        }
 }
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/tree/HashEntryTree.java b/src/main/java/cuchaz/enigma/translation/mapping/tree/HashEntryTree.java
index c5fc473..bb21de6 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/tree/HashEntryTree.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/tree/HashEntryTree.java
@@ -8,7 +8,9 @@ import cuchaz.enigma.translation.representation.entry.Entry;
 
 import javax.annotation.Nullable;
 import java.util.*;
-import java.util.stream.Collectors;
+import java.util.function.Function;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 public class HashEntryTree<T> implements EntryTree<T> {
         private final Map<Entry<?>, HashTreeNode<T>> root = new HashMap<>();
@@ -17,7 +19,7 @@ public class HashEntryTree<T> implements EntryTree<T> {
         }
 
         public HashEntryTree(EntryTree<T> tree) {
-                for (EntryTreeNode<T> node : tree.getAllNodes()) {
+                for (EntryTreeNode<T> node : tree) {
                         insert(node.getEntry(), node.getValue());
                 }
         }
@@ -151,31 +153,24 @@ public class HashEntryTree<T> implements EntryTree<T> {
         }
 
         @Override
-        @SuppressWarnings("unchecked")
         public Iterator<EntryTreeNode<T>> iterator() {
-                Collection<EntryTreeNode<T>> values = (Collection) root.values();
-                return values.iterator();
-        }
-
-        @Override
-        public Collection<EntryTreeNode<T>> getAllNodes() {
                 Collection<EntryTreeNode<T>> nodes = new ArrayList<>();
                 for (EntryTreeNode<T> node : root.values()) {
                         nodes.addAll(node.getNodesRecursively());
                 }
-                return nodes;
+                return nodes.iterator();
         }
 
         @Override
-        public Collection<Entry<?>> getAllEntries() {
-                return getAllNodes().stream()
-                                .map(EntryTreeNode::getEntry)
-                                .collect(Collectors.toList());
+        public Stream<Entry<?>> getAllEntries() {
+                return StreamSupport.stream(spliterator(), false)
+                                .filter(EntryTreeNode::hasValue)
+                                .map(EntryTreeNode::getEntry);
         }
 
         @Override
-        public Collection<Entry<?>> getRootEntries() {
-                return root.keySet();
+        public Stream<EntryTreeNode<T>> getRootNodes() {
+                return root.values().stream().map(Function.identity());
         }
 
         @Override
@@ -186,7 +181,7 @@ public class HashEntryTree<T> implements EntryTree<T> {
         @Override
         public HashEntryTree<T> translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings) {
                 HashEntryTree<T> translatedTree = new HashEntryTree<>();
-                for (EntryTreeNode<T> node : getAllNodes()) {
+                for (EntryTreeNode<T> node : this) {
                         translatedTree.insert(translator.translate(node.getEntry()), node.getValue());
                 }
                 return translatedTree;