BUG-7464: Add dedicated key set classes

[yangtools.git] / third-party / triemap / src / main / java / org / opendaylight / yangtools / triemap / TrieMap.java

diff --git a/third-party/triemap/src/main/java/org/opendaylight/yangtools/triemap/TrieMap.java b/third-party/triemap/src/main/java/org/opendaylight/yangtools/triemap/TrieMap.java
index b93fa4a367fea9b9043e0a4adcbb790a176e2c5b..a7a92ee1bc1687ac1a6209dcb8284f1908f76408 100644 (file)
--- a/third-party/triemap/src/main/java/org/opendaylight/yangtools/triemap/TrieMap.java
+++ b/third-party/triemap/src/main/java/org/opendaylight/yangtools/triemap/TrieMap.java
@@ -22,7 +22,6 @@ import com.google.common.annotations.Beta;
 import java.io.ObjectStreamException;
 import java.io.Serializable;
 import java.util.AbstractMap;
-import java.util.Iterator;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
@@ -45,6 +44,7 @@ public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements Concurr
     private final Equivalence<? super K> equiv;
 
     private AbstractEntrySet<K, V> entrySet;
+    private AbstractKeySet<K> keySet;
 
     TrieMap(final Equivalence<? super K> equiv) {
         this.equiv = equiv;
@@ -100,6 +100,15 @@ public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements Concurr
         return ret;
     }
 
+    @Override
+    public final Set<K> keySet() {
+        AbstractKeySet<K> ret = keySet;
+        if (ret == null) {
+            keySet = ret = createKeySet();
+        }
+        return ret;
+    }
+
     @Override
     public final V get(final Object key) {
         @SuppressWarnings("unchecked")
@@ -135,6 +144,8 @@ public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements Concurr
 
     abstract AbstractEntrySet<K, V> createEntrySet();
 
+    abstract AbstractKeySet<K> createKeySet();
+
     abstract boolean isReadOnly();
 
     abstract INode<K, V> RDCSS_READ_ROOT(boolean abort);
@@ -149,7 +160,7 @@ public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements Concurr
      *
      * @return
      */
-    abstract Iterator<Entry<K, V>> iterator();
+    abstract AbstractIterator<K, V> iterator();
 
     /* internal methods provided for subclasses */
 
@@ -159,8 +170,8 @@ public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements Concurr
      *
      * @return
      */
-    final Iterator<Entry<K, V>> immutableIterator() {
-        return new TrieMapReadOnlyIterator<>(0, immutableSnapshot());
+    final ImmutableIterator<K, V> immutableIterator() {
+        return new ImmutableIterator<>(immutableSnapshot());
     }
 
     @SuppressWarnings("null")