Resolve HashEntryTree#getSiblings building the full ancestor path

2 files changed, 26 insertions, 13 deletions

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 551fb1c8..fa9ed13d 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/tree/HashEntryTree.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/tree/HashEntryTree.java
@@ -24,7 +24,7 @@ public class HashEntryTree<T> implements EntryTree<T> {
 
         @Override
         public void insert(Entry<?> entry, T value) {
-                List<HashTreeNode<T>> path = computePath(entry);
+                List<HashTreeNode<T>> path = computePath(entry, true);
                 path.get(path.size() - 1).putValue(value);
                 if (value == null) {
                         removeDeadAlong(path);
@@ -34,7 +34,11 @@ public class HashEntryTree<T> implements EntryTree<T> {
         @Override
         @Nullable
         public T remove(Entry<?> entry) {
-                List<HashTreeNode<T>> path = computePath(entry);
+                List<HashTreeNode<T>> path = computePath(entry, false);
+                if (path.isEmpty()) {
+                        return null;
+                }
+
                 T value = path.get(path.size() - 1).removeValue();
 
                 removeDeadAlong(path);
@@ -68,7 +72,7 @@ public class HashEntryTree<T> implements EntryTree<T> {
 
         @Override
         public Collection<Entry<?>> getSiblings(Entry<?> entry) {
-                List<HashTreeNode<T>> path = computePath(entry);
+                List<HashTreeNode<T>> path = computePath(entry, false);
                 if (path.size() <= 1) {
                         return getSiblings(entry, root.keySet());
                 }
@@ -95,13 +99,13 @@ public class HashEntryTree<T> implements EntryTree<T> {
                         if (node == null) {
                                 return null;
                         }
-                        node = node.getChild(parentChain.get(i), false);
+                        node = node.getChild(parentChain.get(i));
                 }
 
                 return node;
         }
 
-        private List<HashTreeNode<T>> computePath(Entry<?> target) {
+        private List<HashTreeNode<T>> computePath(Entry<?> target, boolean make) {
                 List<Entry<?>> ancestry = target.getAncestry();
                 if (ancestry.isEmpty()) {
                         return Collections.emptyList();
@@ -110,11 +114,17 @@ public class HashEntryTree<T> implements EntryTree<T> {
                 List<HashTreeNode<T>> path = new ArrayList<>(ancestry.size());
 
                 Entry<?> rootEntry = ancestry.get(0);
-                HashTreeNode<T> node = root.computeIfAbsent(rootEntry, HashTreeNode::new);
+                HashTreeNode<T> node = make ? root.computeIfAbsent(rootEntry, HashTreeNode::new) : root.get(rootEntry);
                 path.add(node);
 
                 for (int i = 1; i < ancestry.size(); i++) {
-                        node = node.getChild(ancestry.get(i), true);
+                        if (node == null) {
+                                return Collections.emptyList();
+                        }
+
+                        Entry<?> ancestor = ancestry.get(i);
+                        node = make ? node.computeChild(ancestor) : node.getChild(ancestor);
+
                         path.add(node);
                 }
 
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/tree/HashTreeNode.java b/src/main/java/cuchaz/enigma/translation/mapping/tree/HashTreeNode.java
index 90e91647..0a990bd5 100644
--- a/src/main/java/cuchaz/enigma/translation/mapping/tree/HashTreeNode.java
+++ b/src/main/java/cuchaz/enigma/translation/mapping/tree/HashTreeNode.java
@@ -2,6 +2,7 @@ package cuchaz.enigma.translation.mapping.tree;
 
 import cuchaz.enigma.translation.representation.entry.Entry;
 
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.util.Collection;
 import java.util.HashMap;
@@ -27,12 +28,14 @@ public class HashTreeNode<T> implements EntryTreeNode<T>, Iterable<HashTreeNode<
                 return value;
         }
 
-        HashTreeNode<T> getChild(Entry<?> entry, boolean create) {
-                if (create) {
-                        return children.computeIfAbsent(entry, HashTreeNode::new);
-                } else {
-                        return children.get(entry);
-                }
+        @Nullable
+        HashTreeNode<T> getChild(Entry<?> entry) {
+                return children.get(entry);
+        }
+
+        @Nonnull
+        HashTreeNode<T> computeChild(Entry<?> entry) {
+                return children.computeIfAbsent(entry, HashTreeNode::new);
         }
 
         void remove(Entry<?> entry) {