From b34305a235cdaaeadecad874931f2f2b62971711 Mon Sep 17 00:00:00 2001
From: gegy1000
Date: Sat, 16 Feb 2019 17:15:14 +0200
Subject: Resolve HashEntryTree#getSiblings building the full ancestor path

---
 .../translation/mapping/tree/HashEntryTree.java    | 24 +++++++++++++++-------
 .../translation/mapping/tree/HashTreeNode.java     | 15 ++++++++------
 2 files changed, 26 insertions(+), 13 deletions(-)

(limited to 'src/main/java')

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) {
-- 
cgit v1.2.3