diff --git a/luna/src/main/java/org/classdump/luna/util/TraversableHashMap.java b/luna/src/main/java/org/classdump/luna/util/TraversableHashMap.java index d0e40980..2542aba5 100644 --- a/luna/src/main/java/org/classdump/luna/util/TraversableHashMap.java +++ b/luna/src/main/java/org/classdump/luna/util/TraversableHashMap.java @@ -46,6 +46,9 @@ public class TraversableHashMap implements Map { private final Set> entrySet; private K firstKey; private K lastKey; + // replicate 'next' behavior of reference implementation, where removing key from table does not invalidate it immediately + private K lingeringKey; + private K lingeringSuccessor; /** * Constructs a new empty map. @@ -133,10 +136,12 @@ public class TraversableHashMap implements Map { Entry e = entries.remove(key); if (e != null) { - K prevKey = e.getPreviousKey(); K nextKey = e.getNextKey(); + lingeringKey = (K) key; + lingeringSuccessor = nextKey; + if (prevKey != null) { entries.get(prevKey).setNextKey(nextKey); } else { @@ -202,6 +207,9 @@ public class TraversableHashMap implements Map { Objects.requireNonNull(key); Entry e = entries.get(key); if (e == null) { + if (key.equals(lingeringKey)) + return lingeringSuccessor; + throw new NoSuchElementException(key.toString()); } return e.getNextKey();