8a2b674ef90642c8cbf67c9af59fcaf4dd5b6d41
[yangtools.git] / third-party / triemap / src / main / java / org / opendaylight / yangtools / triemap / TrieMap.java
1 /*
2  * (C) Copyright 2016 Pantheon Technologies, s.r.o. and others.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package org.opendaylight.yangtools.triemap;
17
18 import com.google.common.base.Preconditions;
19 import java.io.IOException;
20 import java.io.ObjectInputStream;
21 import java.io.ObjectOutputStream;
22 import java.io.Serializable;
23 import java.lang.reflect.Field;
24 import java.util.AbstractMap;
25 import java.util.AbstractSet;
26 import java.util.ArrayList;
27 import java.util.Arrays;
28 import java.util.Iterator;
29 import java.util.List;
30 import java.util.Map;
31 import java.util.Optional;
32 import java.util.Set;
33 import java.util.concurrent.ConcurrentMap;
34 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
35
36 /***
37  * This is a port of Scala's TrieMap class from the Scala Collections library. This implementation does not support
38  * null keys nor null values.
39  *
40  * @author Roman Levenstein <romixlev@gmail.com>
41  *
42  * @param <K>
43  * @param <V>
44  */
45 @SuppressWarnings({"unchecked", "rawtypes", "unused"})
46 public final class TrieMap<K, V> extends AbstractMap<K, V> implements ConcurrentMap<K,V>, Serializable {
47     private static final AtomicReferenceFieldUpdater<TrieMap, Object> ROOT_UPDATER =
48             AtomicReferenceFieldUpdater.newUpdater(TrieMap.class, Object.class, "root");
49     private static final long serialVersionUID = 1L;
50     private static final Field READONLY_FIELD;
51
52     static {
53         final Field f;
54         try {
55             f = TrieMap.class.getDeclaredField("readOnly");
56         } catch (NoSuchFieldException e) {
57             throw new ExceptionInInitializerError(e);
58         } catch (SecurityException e) {
59             throw new ExceptionInInitializerError(e);
60         }
61         f.setAccessible(true);
62         READONLY_FIELD = f;
63     }
64
65     /**
66      * EntrySet
67      */
68     private transient final EntrySet entrySet = new EntrySet ();
69
70     private final Equivalence<? super K> equiv;
71
72     private transient volatile Object root;
73     private final transient boolean readOnly;
74
75     TrieMap(final INode<K, V> r, final Equivalence<? super K> equiv, final boolean readOnly) {
76         this.root = r;
77         this.equiv = equiv;
78         this.readOnly = readOnly;
79     }
80
81     public TrieMap() {
82         this(newRootNode(), Equivalence.equals(), false);
83     }
84
85     /* internal methods */
86
87     private static <K,V> INode<K, V> newRootNode() {
88         final Gen gen = new Gen();
89         return new INode<>(gen, new CNode<>(gen));
90     }
91
92     final boolean CAS_ROOT(final Object ov, final Object nv) {
93         Preconditions.checkState(!readOnly, "Attempted to modify a read-only snapshot");
94         return ROOT_UPDATER.compareAndSet (this, ov, nv);
95     }
96
97     // FIXME: abort = false by default
98     final INode<K, V> readRoot(final boolean abort) {
99         return RDCSS_READ_ROOT(abort);
100     }
101
102     final INode<K, V> readRoot() {
103         return RDCSS_READ_ROOT(false);
104     }
105
106     final INode<K, V> RDCSS_READ_ROOT() {
107         return RDCSS_READ_ROOT(false);
108     }
109
110     final INode<K, V> RDCSS_READ_ROOT(final boolean abort) {
111         final Object r = /* READ */root;
112         if (r instanceof INode) {
113             return (INode<K, V>) r;
114         } else if (r instanceof RDCSS_Descriptor) {
115             return RDCSS_Complete (abort);
116         } else {
117             throw new IllegalStateException("Unhandled root " + r);
118         }
119     }
120
121     private INode<K, V> RDCSS_Complete(final boolean abort) {
122         while (true) {
123             final Object v = /* READ */root;
124             if (v instanceof INode) {
125                 return (INode<K, V>) v;
126             }
127
128             if (v instanceof RDCSS_Descriptor) {
129                 final RDCSS_Descriptor<K, V> desc = (RDCSS_Descriptor<K, V>) v;
130                 final INode<K, V> ov = desc.old;
131                 final MainNode<K, V> exp = desc.expectedmain;
132                 final INode<K, V> nv = desc.nv;
133
134                 if (abort) {
135                     if (CAS_ROOT(desc, ov)) {
136                         return ov;
137                     }
138
139                     // Tail recursion: return RDCSS_Complete(abort);
140                     continue;
141                 }
142
143                 final MainNode<K, V> oldmain = ov.gcasRead(this);
144                 if (oldmain == exp) {
145                     if (CAS_ROOT(desc, nv)) {
146                         desc.committed = true;
147                         return nv;
148                     }
149
150                     // Tail recursion: return RDCSS_Complete(abort);
151                     continue;
152                 }
153
154                 if (CAS_ROOT(desc, ov)) {
155                     return ov;
156                 }
157
158                 // Tail recursion: return RDCSS_Complete(abort);
159                 continue;
160             }
161
162             throw new IllegalStateException("Unexpected root " + v);
163         }
164     }
165
166     private boolean RDCSS_ROOT(final INode<K, V> ov, final MainNode<K, V> expectedmain, final INode<K, V> nv) {
167         final RDCSS_Descriptor<K, V> desc = new RDCSS_Descriptor<> (ov, expectedmain, nv);
168         if (CAS_ROOT(ov, desc)) {
169             RDCSS_Complete(false);
170             return /* READ */desc.committed;
171         }
172
173         return false;
174     }
175
176     private void inserthc(final K k, final int hc, final V v) {
177         while (true) {
178             final INode<K, V> r = RDCSS_READ_ROOT();
179             if (r.rec_insert(k, v, hc, 0, null, r.gen, this)) {
180                 // Successful, we are done
181                 return;
182             }
183
184             // Tail recursion: inserthc(k, hc, v);
185         }
186     }
187
188     private Optional<V> insertifhc(final K k, final int hc, final V v, final Object cond) {
189         while (true) {
190             final INode<K, V> r = RDCSS_READ_ROOT();
191             final Optional<V> ret = r.rec_insertif(k, v, hc, cond, 0, null, r.gen, this);
192             if (ret != null) {
193                 return ret;
194             }
195
196             // Tail recursion: return insertifhc(k, hc, v, cond);
197         }
198     }
199
200     private Object lookuphc(final K k, final int hc) {
201         while (true) {
202             final INode<K, V> r = RDCSS_READ_ROOT ();
203             final Object res = r.rec_lookup(k, hc, 0, null, r.gen, this);
204             if (!INode.RESTART.equals(res)) {
205                 return res;
206             }
207
208             // Tail recursion: lookuphc(k, hc)
209         }
210     }
211
212     private Optional<V> removehc(final K k, final V v, final int hc) {
213         while (true) {
214             final INode<K, V> r = RDCSS_READ_ROOT();
215             final Optional<V> res = r.rec_remove(k, v, hc, 0, null, r.gen, this);
216             if (res != null) {
217                 return res;
218             }
219
220             // Tail recursion: return removehc(k, v, hc);
221         }
222     }
223
224     /**
225      * Ensure this instance is read-write, throw UnsupportedOperationException
226      * otherwise. Used by Map-type methods for quick check.
227      */
228     private void ensureReadWrite() {
229         if (readOnly) {
230             throw new UnsupportedOperationException("Attempted to modify a read-only view");
231         }
232     }
233
234     boolean isReadOnly() {
235         return readOnly;
236     }
237
238     /* public methods */
239
240     /**
241      * Returns a snapshot of this TrieMap. This operation is lock-free and
242      * linearizable.
243      *
244      * The snapshot is lazily updated - the first time some branch in the
245      * snapshot or this TrieMap are accessed, they are rewritten. This means
246      * that the work of rebuilding both the snapshot and this TrieMap is
247      * distributed across all the threads doing updates or accesses subsequent
248      * to the snapshot creation.
249      */
250     public TrieMap<K, V> snapshot() {
251         while (true) {
252             final INode<K, V> r = RDCSS_READ_ROOT();
253             final MainNode<K, V> expmain = r.gcasRead(this);
254             if (RDCSS_ROOT(r, expmain, r.copyToGen(new Gen(), this))) {
255                 return new TrieMap<> (r.copyToGen(new Gen(), this), equiv, readOnly);
256             }
257
258             // Tail recursion: return snapshot();
259         }
260     }
261
262     /**
263      * Returns a read-only snapshot of this TrieMap. This operation is lock-free
264      * and linearizable.
265      *
266      * The snapshot is lazily updated - the first time some branch of this
267      * TrieMap are accessed, it is rewritten. The work of creating the snapshot
268      * is thus distributed across subsequent updates and accesses on this
269      * TrieMap by all threads. Note that the snapshot itself is never rewritten
270      * unlike when calling the `snapshot` method, but the obtained snapshot
271      * cannot be modified.
272      *
273      * This method is used by other methods such as `size` and `iterator`.
274      */
275     public TrieMap<K, V> readOnlySnapshot() {
276         // Is it a snapshot of a read-only snapshot?
277         if (readOnly) {
278             return this;
279         }
280
281         while (true) {
282             final INode<K, V> r = RDCSS_READ_ROOT();
283             final MainNode<K, V> expmain = r.gcasRead(this);
284             if (RDCSS_ROOT(r, expmain, r.copyToGen (new Gen(), this))) {
285                 return new TrieMap<>(r, equiv, true);
286             }
287
288             // Tail recursion: return readOnlySnapshot();
289         }
290     }
291
292     @Override
293     public void clear() {
294         while (true) {
295             final INode<K, V> r = RDCSS_READ_ROOT();
296             if (RDCSS_ROOT(r, r.gcasRead(this), newRootNode())) {
297                 return;
298             }
299         }
300     }
301
302     int computeHash(final K k) {
303         return equiv.hash(k);
304     }
305
306     boolean equal(final K k1, final K k2) {
307         return equiv.equivalent(k1, k2);
308     }
309
310     V lookup(final K k) {
311         return (V) lookuphc(k, computeHash(k));
312     }
313
314     @Override
315     public V get(final Object k) {
316         return lookup((K)k);
317     }
318
319     @Override
320     public V put(final K key, final V value) {
321         ensureReadWrite();
322         final int hc = computeHash(key);
323         return insertifhc (key, hc, value, null).orElse(null);
324     }
325
326     TrieMap<K, V> add(final K k, final V v) {
327         final int hc = computeHash (k);
328         inserthc (k, hc, v);
329         return this;
330     }
331
332     @Override
333     public V remove(final Object k) {
334         ensureReadWrite();
335         final int hc = computeHash ((K)k);
336         return removehc ((K)k, (V) null, hc).orElse(null);
337     }
338
339     @Override
340     public V putIfAbsent(final K k, final V v) {
341         ensureReadWrite();
342         final int hc = computeHash (k);
343         return insertifhc (k, hc, v, INode.KEY_ABSENT).orElse(null);
344     }
345
346     @Override
347     public boolean remove(final Object k, final Object v) {
348         ensureReadWrite();
349         final int hc = computeHash ((K)k);
350         return removehc((K)k, (V)v, hc).isPresent();
351     }
352
353     @Override
354     public boolean replace(final K k, final V oldvalue, final V newvalue) {
355         ensureReadWrite();
356         final int hc = computeHash (k);
357         return insertifhc (k, hc, newvalue, oldvalue).isPresent();
358     }
359
360     @Override
361     public V replace(final K k, final V v) {
362         ensureReadWrite();
363         final int hc = computeHash (k);
364         return insertifhc (k, hc, v, INode.KEY_PRESENT).orElse(null);
365     }
366
367     /***
368      * Return an iterator over a TrieMap.
369      *
370      * If this is a read-only snapshot, it would return a read-only iterator.
371      *
372      * If it is the original TrieMap or a non-readonly snapshot, it would return
373      * an iterator that would allow for updates.
374      *
375      * @return
376      */
377     Iterator<Entry<K, V>> iterator() {
378         return readOnly ? new TrieMapReadOnlyIterator<>(0, this) : new TrieMapIterator<>(0, this);
379     }
380
381     /***
382      * Return an iterator over a TrieMap.
383      * This is a read-only iterator.
384      *
385      * @return
386      */
387     Iterator<Entry<K, V>> readOnlyIterator() {
388         return new TrieMapReadOnlyIterator<>(0, readOnly ? this : readOnlySnapshot());
389     }
390
391     private int cachedSize() {
392         INode<K, V> r = RDCSS_READ_ROOT ();
393         return r.cachedSize (this);
394     }
395
396     @Override
397     public int size() {
398         return readOnly ? cachedSize() : readOnlySnapshot().size();
399     }
400
401     @Override
402     public boolean containsKey(final Object key) {
403         return lookup((K) key) != null;
404     }
405
406     @Override
407     public Set<Entry<K, V>> entrySet() {
408         return entrySet;
409     }
410
411     private static final class RDCSS_Descriptor<K, V> {
412         INode<K, V> old;
413         MainNode<K, V> expectedmain;
414         INode<K, V> nv;
415         volatile boolean committed = false;
416
417         RDCSS_Descriptor (final INode<K, V> old, final MainNode<K, V> expectedmain, final INode<K, V> nv) {
418             this.old = old;
419             this.expectedmain = expectedmain;
420             this.nv = nv;
421         }
422     }
423
424     /***
425      * This iterator is a read-only one and does not allow for any update
426      * operations on the underlying data structure.
427      *
428      * @param <K>
429      * @param <V>
430      */
431     private static final class TrieMapReadOnlyIterator<K, V> extends TrieMapIterator<K, V> {
432         TrieMapReadOnlyIterator (final int level, final TrieMap<K, V> ct, final boolean mustInit) {
433             super (level, ct, mustInit);
434         }
435
436         TrieMapReadOnlyIterator (final int level, final TrieMap<K, V> ct) {
437             this (level, ct, true);
438         }
439         @Override
440         void initialize () {
441             assert (ct.isReadOnly ());
442             super.initialize ();
443         }
444
445         @Override
446         public void remove () {
447             throw new UnsupportedOperationException ("Operation not supported for read-only iterators");
448         }
449
450         @Override
451         Entry<K, V> nextEntry(final Entry<K, V> rr) {
452             // Return non-updatable entry
453             return rr;
454         }
455     }
456
457     private static class TrieMapIterator<K, V> implements Iterator<Entry<K, V>> {
458         private int level;
459         protected TrieMap<K, V> ct;
460         private final boolean mustInit;
461         private final BasicNode[][] stack = new BasicNode[7][];
462         private final int[] stackpos = new int[7];
463         private int depth = -1;
464         private Iterator<Entry<K, V>> subiter = null;
465         private KVNode<K, V> current = null;
466         private Entry<K, V> lastReturned = null;
467
468         TrieMapIterator (final int level, final TrieMap<K, V> ct, final boolean mustInit) {
469             this.level = level;
470             this.ct = ct;
471             this.mustInit = mustInit;
472             if (this.mustInit) {
473                 initialize ();
474             }
475         }
476
477         TrieMapIterator (final int level, final TrieMap<K, V> ct) {
478             this (level, ct, true);
479         }
480
481
482         @Override
483         public boolean hasNext () {
484             return (current != null) || (subiter != null);
485         }
486
487         @Override
488         public Entry<K, V> next () {
489             if (hasNext ()) {
490                 Entry<K, V> r = null;
491                 if (subiter != null) {
492                     r = subiter.next ();
493                     checkSubiter ();
494                 } else {
495                     r = current.kvPair ();
496                     advance ();
497                 }
498
499                 lastReturned = r;
500                 if (r != null) {
501                     final Entry<K, V> rr = r;
502                     return nextEntry(rr);
503                 }
504                 return r;
505             } else {
506                 // return Iterator.empty ().next ();
507                 return null;
508             }
509         }
510
511         Entry<K, V> nextEntry(final Entry<K, V> rr) {
512             return new Entry<K, V>() {
513                 private V updated = null;
514
515                 @Override
516                 public K getKey () {
517                     return rr.getKey ();
518                 }
519
520                 @Override
521                 public V getValue () {
522                     return (updated == null)?rr.getValue (): updated;
523                 }
524
525                 @Override
526                 public V setValue (final V value) {
527                     updated = value;
528                     return ct.replace (getKey (), value);
529                 }
530             };
531         }
532
533         private void readin (final INode<K, V> in) {
534             MainNode<K, V> m = in.gcasRead (ct);
535             if (m instanceof CNode) {
536                 CNode<K, V> cn = (CNode<K, V>) m;
537                 depth += 1;
538                 stack [depth] = cn.array;
539                 stackpos [depth] = -1;
540                 advance ();
541             } else if (m instanceof TNode) {
542                 current = (TNode<K, V>) m;
543             } else if (m instanceof LNode) {
544                 subiter = ((LNode<K, V>) m).iterator();
545                 checkSubiter ();
546             } else if (m == null) {
547                 current = null;
548             }
549         }
550
551         // @inline
552         private void checkSubiter () {
553             if (!subiter.hasNext ()) {
554                 subiter = null;
555                 advance ();
556             }
557         }
558
559         // @inline
560         void initialize () {
561 //            assert (ct.isReadOnly ());
562             INode<K, V> r = ct.RDCSS_READ_ROOT ();
563             readin (r);
564         }
565
566         void advance () {
567             if (depth >= 0) {
568                 int npos = stackpos [depth] + 1;
569                 if (npos < stack [depth].length) {
570                     stackpos [depth] = npos;
571                     BasicNode elem = stack [depth] [npos];
572                     if (elem instanceof SNode) {
573                         current = (SNode<K, V>) elem;
574                     } else if (elem instanceof INode) {
575                         readin ((INode<K, V>) elem);
576                     }
577                 } else {
578                     depth -= 1;
579                     advance ();
580                 }
581             } else {
582                 current = null;
583             }
584         }
585
586         protected TrieMapIterator<K, V> newIterator (final int _lev, final TrieMap<K, V> _ct, final boolean _mustInit) {
587             return new TrieMapIterator<> (_lev, _ct, _mustInit);
588         }
589
590         protected void dupTo (final TrieMapIterator<K, V> it) {
591             it.level = this.level;
592             it.ct = this.ct;
593             it.depth = this.depth;
594             it.current = this.current;
595
596             // these need a deep copy
597             System.arraycopy (this.stack, 0, it.stack, 0, 7);
598             System.arraycopy (this.stackpos, 0, it.stackpos, 0, 7);
599
600             // this one needs to be evaluated
601             if (this.subiter == null) {
602                 it.subiter = null;
603             } else {
604                 List<Entry<K, V>> lst = toList (this.subiter);
605                 this.subiter = lst.iterator ();
606                 it.subiter = lst.iterator ();
607             }
608         }
609
610         // /** Returns a sequence of iterators over subsets of this iterator.
611         // * It's used to ease the implementation of splitters for a parallel
612         // version of the TrieMap.
613         // */
614         // protected def subdivide(): Seq[Iterator[(K, V)]] = if (subiter ne
615         // null) {
616         // // the case where an LNode is being iterated
617         // val it = subiter
618         // subiter = null
619         // advance()
620         // this.level += 1
621         // Seq(it, this)
622         // } else if (depth == -1) {
623         // this.level += 1
624         // Seq(this)
625         // } else {
626         // var d = 0
627         // while (d <= depth) {
628         // val rem = stack(d).length - 1 - stackpos(d)
629         // if (rem > 0) {
630         // val (arr1, arr2) = stack(d).drop(stackpos(d) + 1).splitAt(rem / 2)
631         // stack(d) = arr1
632         // stackpos(d) = -1
633         // val it = newIterator(level + 1, ct, false)
634         // it.stack(0) = arr2
635         // it.stackpos(0) = -1
636         // it.depth = 0
637         // it.advance() // <-- fix it
638         // this.level += 1
639         // return Seq(this, it)
640         // }
641         // d += 1
642         // }
643         // this.level += 1
644         // Seq(this)
645         // }
646
647         private List<Entry<K, V>> toList (final Iterator<Entry<K, V>> it) {
648             ArrayList<Entry<K, V>> list = new ArrayList<> ();
649             while (it.hasNext ()) {
650                 list.add (it.next ());
651             }
652             return list;
653         }
654
655         void printDebug () {
656             System.out.println ("ctrie iterator");
657             System.out.println (Arrays.toString (stackpos));
658             System.out.println ("depth: " + depth);
659             System.out.println ("curr.: " + current);
660             // System.out.println(stack.mkString("\n"));
661         }
662
663         @Override
664         public void remove () {
665             if (lastReturned != null) {
666                 ct.remove (lastReturned.getKey ());
667                 lastReturned = null;
668             } else {
669                 throw new IllegalStateException();
670             }
671         }
672
673     }
674
675     /***
676      * Support for EntrySet operations required by the Map interface
677      */
678     private final class EntrySet extends AbstractSet<Entry<K, V>> {
679
680         @Override
681         public Iterator<Entry<K, V>> iterator () {
682             return TrieMap.this.iterator ();
683         }
684
685         @Override
686         public final boolean contains (final Object o) {
687             if (!(o instanceof Entry)) {
688                 return false;
689             }
690             final Entry<K, V> e = (Entry<K, V>) o;
691             final K k = e.getKey ();
692             final V v = lookup (k);
693             return v != null;
694         }
695
696         @Override
697         public final boolean remove (final Object o) {
698             if (!(o instanceof Entry)) {
699                 return false;
700             }
701             final Entry<K, V> e = (Entry<K, V>) o;
702             final K k = e.getKey();
703             return null != TrieMap.this.remove(k);
704         }
705
706         @Override
707         public final int size () {
708             int size = 0;
709             for (final Iterator<?> i = iterator (); i.hasNext (); i.next ()) {
710                 size++;
711             }
712             return size;
713         }
714
715         @Override
716         public final void clear () {
717             TrieMap.this.clear ();
718         }
719     }
720
721     private void readObject(final ObjectInputStream inputStream) throws IOException, ClassNotFoundException {
722         inputStream.defaultReadObject();
723         this.root = newRootNode();
724
725         final boolean ro = inputStream.readBoolean();
726         final int size = inputStream.readInt();
727         for (int i = 0; i < size; ++i) {
728             final K key = (K)inputStream.readObject();
729             final V value = (V)inputStream.readObject();
730             add(key, value);
731         }
732
733         // Propagate the read-only bit
734         try {
735             READONLY_FIELD.setBoolean(this, ro);
736         } catch (IllegalAccessException e) {
737             throw new IOException("Failed to set read-only flag", e);
738         }
739     }
740
741     private void writeObject(final ObjectOutputStream outputStream) throws IOException {
742         outputStream.defaultWriteObject();
743
744         final Map<K, V> ro = readOnlySnapshot();
745         outputStream.writeBoolean(readOnly);
746         outputStream.writeInt(ro.size());
747
748         for (Entry<K, V> e : ro.entrySet()) {
749             outputStream.writeObject(e.getKey());
750             outputStream.writeObject(e.getValue());
751         }
752     }
753 }