Fix TreeMap iterator

Fix #519
This commit is contained in:
Alexey Andreev 2021-03-06 21:08:51 +03:00
parent ddf194af31
commit 8b4f401bcb
2 changed files with 55 additions and 27 deletions

View File

@ -27,7 +27,7 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
int height = 1; int height = 1;
int size = 1; int size = 1;
public TreeNode(K key) { TreeNode(K key) {
super(key, null); super(key, null);
} }
@ -548,7 +548,7 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
return copy; return copy;
} }
private static class EntrySet<K, V> extends TAbstractSet<Entry<K, V>> { static class EntrySet<K, V> extends TAbstractSet<Entry<K, V>> {
private int modCount = -1; private int modCount = -1;
private TTreeMap<K, V> owner; private TTreeMap<K, V> owner;
private K from; private K from;
@ -560,7 +560,7 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
private int cachedSize; private int cachedSize;
private boolean reverse; private boolean reverse;
public EntrySet(TTreeMap<K, V> owner, K from, boolean fromIncluded, boolean fromChecked, EntrySet(TTreeMap<K, V> owner, K from, boolean fromIncluded, boolean fromChecked,
K to, boolean toIncluded, boolean toChecked, boolean reverse) { K to, boolean toIncluded, boolean toChecked, boolean reverse) {
this.owner = owner; this.owner = owner;
this.from = from; this.from = from;
@ -615,13 +615,7 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
} else { } else {
fromPath = owner.pathToFirst(false); fromPath = owner.pathToFirst(false);
} }
TreeNode<K, V> toEntry; return new EntryIterator<>(owner, fromPath, to, toChecked, toIncluded, false);
if (toChecked) {
toEntry = toIncluded ? owner.findExactOrNext(to, true) : owner.findNext(to, true);
} else {
toEntry = owner.firstNode(true);
}
return new EntryIterator<>(owner, fromPath, toEntry, false);
} }
private TIterator<Entry<K, V>> descendingIterator() { private TIterator<Entry<K, V>> descendingIterator() {
@ -631,13 +625,7 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
} else { } else {
toPath = owner.pathToFirst(true); toPath = owner.pathToFirst(true);
} }
TreeNode<K, V> fromEntry; return new EntryIterator<>(owner, toPath, from, fromIncluded, fromChecked, true);
if (fromChecked) {
fromEntry = fromIncluded ? owner.findExactOrNext(to, false) : owner.findNext(to, false);
} else {
fromEntry = owner.firstNode(false);
}
return new EntryIterator<>(owner, toPath, fromEntry, true);
} }
@Override @Override
@ -671,22 +659,28 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
} }
} }
private static class EntryIterator<K, V> implements TIterator<Entry<K, V>> { static class EntryIterator<K, V> implements TIterator<Entry<K, V>> {
private int modCount; private int modCount;
private TTreeMap<K, V> owner; private TTreeMap<K, V> owner;
private TreeNode<K, V>[] path; private TreeNode<K, V>[] path;
private TreeNode<K, V> last; private TreeNode<K, V> last;
private TreeNode<K, V> to; private K to;
private boolean toChecked;
private boolean toIncluded;
private int depth; private int depth;
private boolean reverse; private boolean reverse;
public EntryIterator(TTreeMap<K, V> owner, TreeNode<K, V>[] path, TreeNode<K, V> to, boolean reverse) { EntryIterator(TTreeMap<K, V> owner, TreeNode<K, V>[] path, K to, boolean toChecked, boolean toIncluded,
boolean reverse) {
this.owner = owner; this.owner = owner;
modCount = owner.modCount; modCount = owner.modCount;
this.path = TArrays.copyOf(path, owner.root == null ? 0 : owner.root.height); this.path = TArrays.copyOf(path, owner.root == null ? 0 : owner.root.height);
depth = path.length; depth = path.length;
this.to = to; this.to = to;
this.toChecked = toChecked;
this.toIncluded = toIncluded;
this.reverse = reverse; this.reverse = reverse;
checkFinished();
} }
@Override @Override
@ -712,10 +706,28 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
node = node.forward(reverse); node = node.forward(reverse);
} }
} }
if (last == to) {
checkFinished();
return last;
}
private void checkFinished() {
if (!toChecked || depth == 0) {
return;
}
int cmp = owner.comparator.compare(path[depth - 1].getKey(), to);
if (reverse) {
cmp = -cmp;
}
if (toIncluded) {
if (cmp > 0) {
depth = 0; depth = 0;
} }
return last; } else {
if (cmp >= 0) {
depth = 0;
}
}
} }
@Override @Override
@ -732,7 +744,7 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
} }
} }
private static class MapView<K, V> extends TAbstractMap<K, V> implements TNavigableMap<K, V>, TSerializable { static class MapView<K, V> extends TAbstractMap<K, V> implements TNavigableMap<K, V>, TSerializable {
private int modCount = -1; private int modCount = -1;
private int cachedSize; private int cachedSize;
private TTreeMap<K, V> owner; private TTreeMap<K, V> owner;
@ -746,7 +758,7 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
private boolean reverse; private boolean reverse;
private NavigableKeySet<K, V> cachedNavigableKeySet; private NavigableKeySet<K, V> cachedNavigableKeySet;
public MapView(TTreeMap<K, V> owner, K from, boolean fromIncluded, boolean fromChecked, MapView(TTreeMap<K, V> owner, K from, boolean fromIncluded, boolean fromChecked,
K to, boolean toIncluded, boolean toChecked, boolean reverse) { K to, boolean toIncluded, boolean toChecked, boolean reverse) {
this.owner = owner; this.owner = owner;
this.from = from; this.from = from;
@ -1079,10 +1091,10 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
} }
} }
private static class NavigableKeySet<K, V> extends TAbstractSet<K> implements TNavigableSet<K> { static class NavigableKeySet<K, V> extends TAbstractSet<K> implements TNavigableSet<K> {
private TNavigableMap<K, V> map; private TNavigableMap<K, V> map;
public NavigableKeySet(TNavigableMap<K, V> map) { NavigableKeySet(TNavigableMap<K, V> map) {
this.map = map; this.map = map;
} }

View File

@ -694,4 +694,20 @@ public class TreeMapTest {
assertEquals(101, tm.size()); assertEquals(101, tm.size());
} }
} }
@Test
public void submap() {
TreeMap<Integer, Integer> map = new TreeMap<>();
map.put(1, 1);
map.put(3, 15);
map.put(4, 20);
map.put(6, 13);
map.put(10, 119);
assertEquals("{}", map.subMap(0, 0).toString());
assertEquals("{}", map.subMap(7, 9).toString());
assertEquals("{3=15, 4=20, 6=13}", map.subMap(3, 9).toString());
assertEquals("{10=119}", map.subMap(10, 29).toString());
assertEquals("{}", map.subMap(29, 100).toString());
}
} }