add key-by-key tests

fix adoption not propagating upward

Author: supersaiyansubtlety

enigma/src/main/java/org/quiltmc/enigma/util/multi_trie/MutableMapNode.java

@@ -74,7 +74,7 @@ public boolean clearLeaves() {
 	/**
 	 * Removes the branch node associated with the passed {@code key} if that node is empty.
 	 *
-	 * @implNote This should only be passed one of <em>this</em> node's branches.
+	 * @implNote This should only be passed one of <em>this</em> node's {@linkplain #getBranches() branches}.
 	 *
 	 * @return {@code true} if the branch was pruned, or {@code false otherwise}
 	 */
@@ -88,6 +88,26 @@ protected boolean pruneIfEmpty(Branch<K, V, B> branch) {
 		}
 	}
 
+	/**
+	 * If the passed {@code branch} is an {@linkplain #orphans orphan},
+	 * removes it from {@linkplain #orphans} and puts it in {@linkplain #getBranches() branches}.
+	 *
+	 * @implNote only non-empty branches should be passed to this method;
+	 * it's called when a node may have changed from empty to non-empty
+	 *
+	 * @return {@code true} if the passed {@code branch} was an orphan, or {@code false} otherwise
+	 */
+	protected boolean tryAdopt(Branch<K, V, B> branch) {
+		final boolean wasOrphan = this.orphans.remove(branch.getKey()) != null;
+		if (wasOrphan) {
+			this.getBranches().put(branch.getKey(), branch.getSelf());
+
+			return true;
+		} else {
+			return false;
+		}
+	}
+
 	/**
 	 * @return a new, empty branch node instance
 	 */
@@ -128,10 +148,7 @@ protected abstract static class Branch<K, V, B extends Branch<K, V, B>> extends
 		@Override
 		public void put(V value) {
 			super.put(value);
-			final boolean wasOrphan = this.getParent().orphans.remove(this.getKey()) != null;
-			if (wasOrphan) {
-				this.getParent().getBranches().put(this.getKey(), this.getSelf());
-			}
+			this.getParent().tryAdopt(this);
 		}
 
 		@Override
@@ -166,5 +183,16 @@ protected boolean pruneIfEmpty(Branch<K, V, B> branch) {
 				return false;
 			}
 		}
+
+		@Override
+		protected boolean tryAdopt(Branch<K, V, B> branch) {
+			if (super.tryAdopt(branch)) {
+				this.getParent().tryAdopt(this);
+
+				return true;
+			} else {
+				return false;
+			}
+		}
 	}
 }

enigma/src/test/java/org/quiltmc/enigma/util/multi_trie/CompositeStringMultiTrieTest.java

@@ -15,12 +15,65 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-// TODO key-by-key access tests
 public class CompositeStringMultiTrieTest {
 	private static final String VALUES = "values";
 	private static final String LEAVES = "leaves";
 	private static final String BRANCHES = "branches";
 
+	private static final String KEY_BY_KEY_SUBJECT = "key-by-key subject";
+
+	// test key-by-key put's orphan logic
+	@Test
+	void testPutKeyByKeyRootDown() {
+		final CompositeStringMultiTrie<Integer> trie = CompositeStringMultiTrie.createHashed();
+
+		for (int depth = 0; depth < KEY_BY_KEY_SUBJECT.length(); depth++) {
+			MutableMapNode<Character, Integer, ?> node = trie.getRoot();
+			for (int iKey = 0; iKey <= depth; iKey++) {
+				node = node.next(KEY_BY_KEY_SUBJECT.charAt(iKey));
+			}
+
+			node.put(depth);
+
+			assertOneLeaf(node);
+
+			assertTrieSize(trie, depth + 1);
+		}
+	}
+
+	@Test
+	void testPutKeyByKeyStemUp() {
+		final CompositeStringMultiTrie<Integer> trie = CompositeStringMultiTrie.createHashed();
+
+		for (int depth = KEY_BY_KEY_SUBJECT.length() - 1; depth >= 0; depth--) {
+			MutableMapNode<Character, Integer, ?> node = trie.getRoot();
+			for (int iKey = 0; iKey <= depth; iKey++) {
+				node = node.next(KEY_BY_KEY_SUBJECT.charAt(iKey));
+			}
+
+			node.put(depth);
+
+			assertOneLeaf(node);
+
+			assertTrieSize(trie, KEY_BY_KEY_SUBJECT.length() - depth);
+		}
+	}
+
+	private static void assertOneLeaf(MutableMapNode<Character, Integer, ?> node) {
+		assertEquals(
+			1, node.streamLeaves().count(),
+			() -> "Expected node to have only one leaf, but had the following: " + node.streamLeaves().toList()
+		);
+	}
+
+	private static void assertTrieSize(CompositeStringMultiTrie<Integer> trie, int expectedSize) {
+		assertEquals(
+			expectedSize, trie.getSize(),
+			() -> "Expected node to have %s values, but had the following: %s"
+				.formatted(expectedSize, trie.getRoot().streamValues().toList())
+		);
+	}
+
 	@Test
 	void testPut() {
 		final CompositeStringMultiTrie<Association> trie = Association.createAndPopulateTrie();