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 0aacd4f5f71109d6b8285d8ee2fe1e3c3dd547ad..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;
@@ -42,12 +41,11 @@ import java.util.concurrent.ConcurrentMap;
 public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements ConcurrentMap<K,V>, Serializable {
     private static final long serialVersionUID = 1L;
 
-    /**
-     * EntrySet
-     */
-    private final EntrySet<K, V> entrySet = new EntrySet<>(this);
     private final Equivalence<? super K> equiv;
 
+    private AbstractEntrySet<K, V> entrySet;
+    private AbstractKeySet<K> keySet;
+
     TrieMap(final Equivalence<? super K> equiv) {
         this.equiv = equiv;
     }
@@ -95,7 +93,20 @@ public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements Concurr
 
     @Override
     public final Set<Entry<K, V>> entrySet() {
-        return entrySet;
+        AbstractEntrySet<K, V> ret = entrySet;
+        if (ret == null) {
+            entrySet = ret = createEntrySet();
+        }
+        return ret;
+    }
+
+    @Override
+    public final Set<K> keySet() {
+        AbstractKeySet<K> ret = keySet;
+        if (ret == null) {
+            keySet = ret = createKeySet();
+        }
+        return ret;
     }
 
     @Override
@@ -131,12 +142,38 @@ public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements Concurr
 
     /* internal methods implemented by subclasses */
 
+    abstract AbstractEntrySet<K, V> createEntrySet();
+
+    abstract AbstractKeySet<K> createKeySet();
+
     abstract boolean isReadOnly();
 
     abstract INode<K, V> RDCSS_READ_ROOT(boolean abort);
 
+    /**
+     * Return an iterator over a TrieMap.
+     *
+     * If this is a read-only snapshot, it would return a read-only iterator.
+     *
+     * If it is the original TrieMap or a non-readonly snapshot, it would return
+     * an iterator that would allow for updates.
+     *
+     * @return
+     */
+    abstract AbstractIterator<K, V> iterator();
+
     /* internal methods provided for subclasses */
 
+    /**
+     * Return an iterator over a TrieMap.
+     * This is a read-only iterator.
+     *
+     * @return
+     */
+    final ImmutableIterator<K, V> immutableIterator() {
+        return new ImmutableIterator<>(immutableSnapshot());
+    }
+
     @SuppressWarnings("null")
     static <V> V toNullable(final Optional<V> opt) {
         return opt.orElse(null);
@@ -185,29 +222,4 @@ public abstract class TrieMap<K, V> extends AbstractMap<K, V> implements Concurr
 
         return (V) res;
     }
-
-    /**
-     * Return an iterator over a TrieMap.
-     *
-     * If this is a read-only snapshot, it would return a read-only iterator.
-     *
-     * If it is the original TrieMap or a non-readonly snapshot, it would return
-     * an iterator that would allow for updates.
-     *
-     * @return
-     */
-    final Iterator<Entry<K, V>> iterator() {
-        // FIXME: it would be nice to have a ReadWriteTrieMap with read-only iterator
-        return isReadOnly() ? new TrieMapReadOnlyIterator<>(0, this) : new TrieMapIterator<>(0, this);
-    }
-
-    /**
-     * Return an iterator over a TrieMap.
-     * This is a read-only iterator.
-     *
-     * @return
-     */
-    final Iterator<Entry<K, V>> readOnlyIterator() {
-        return new TrieMapReadOnlyIterator<>(0, immutableSnapshot());
-    }
 }