1 /**
2  *  Copyright Terracotta, Inc.
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 net.sf.ehcache.store.disk.ods;
17
18 import java.util.AbstractSet;
19 import java.util.Comparator;
20 import java.util.Iterator;
21 import java.util.SortedSet;
22 import java.util.Stack;
23
24 /**
25  * A AA-Tree based SortedSet implementation
26  *
27  * @param <T> type of values stored
28  *
29  * @author Chris Dennis
30  */
31 public class AATreeSet<T extends Comparable> extends AbstractSet<T> implements SortedSet<T> {
32
33   private static final Node<?> TERMINAL = new TerminalNode();
34
35   private Node<T> root = terminal();
36
37   private int     size;
38   private boolean mutated;
39
40   private Node<T> item = terminal();
41   private Node<T> heir = terminal();
42   private T       removed;
43
44   @Override
45   public boolean add(T o) {
46     try {
47       root = insert(root, o);
48       if (mutated) {
49         size++;
50       }
51       return mutated;
52     } finally {
53       mutated = false;
54     }
55   }
56
57   @Override
58   public boolean remove(Object o) {
59     try {
60       root = remove(root, (T) o);
61       if (mutated) {
62         size--;
63       }
64       return mutated;
65     } finally {
66       heir = terminal();
67       item = terminal();
68       mutated = false;
69       removed = null;
70     }
71   }
72
73   /**
74    * Remove the node matching this object and return it.
75    */
76   public T removeAndReturn(Object o) {
77     try {
78       root = remove(root, (T) o);
79       if (mutated) {
80         size--;
81       }
82       return removed;
83     } finally {
84       heir = terminal();
85       item = terminal();
86       mutated = false;
87       removed = null;
88     }
89   }
90
91   @Override
92   public void clear() {
93     root = terminal();
94     size = 0;
95   }
96
97   @Override
98   public Iterator<T> iterator() {
99     return new TreeIterator();
100   }
101
102   @Override
103   public int size() {
104     return size;
105   }
106
107   @Override
108   public boolean isEmpty() {
109     return root == terminal();
110   }
111
112   /**
113    * {@inheritDoc}
114    */
115   public Comparator<? super T> comparator() {
116     return null;
117   }
118
119   /**
120    * {@inheritDoc}
121    */
122   public SortedSet<T> subSet(T fromElement, T toElement) {
123     return new SubSet(fromElement, toElement);
124   }
125
126   /**
127    * {@inheritDoc}
128    */
129   public SortedSet<T> headSet(T toElement) {
130     return new SubSet(null, toElement);
131   }
132
133   /**
134    * {@inheritDoc}
135    */
136   public SortedSet<T> tailSet(T fromElement) {
137     return new SubSet(fromElement, null);
138   }
139
140   /**
141    * {@inheritDoc}
142    */
143   public T first() {
144     Node<T> leftMost = root;
145     while (leftMost.getLeft() != terminal()) {
146       leftMost = leftMost.getLeft();
147     }
148     return leftMost.getPayload();
149   }
150
151   /**
152    * {@inheritDoc}
153    */
154   public T last() {
155     Node<T> rightMost = root;
156     while (rightMost.getRight() != terminal()) {
157       rightMost = rightMost.getRight();
158     }
159     return rightMost.getPayload();
160   }
161
162   /**
163    * Find the node within this tree equal to the probe node.
164    */
165   public T find(Object probe) {
166     return find(root, (T) probe).getPayload();
167   }
168
169   private Node<T> terminal() {
170     return (Node<T>) TERMINAL;
171   }
172
173   /**
174    * Returns the root node of this tree.
175    */
176   protected final Node<T> getRoot() {
177       return root;
178   }
179
180   private Node<T> find(Node<T> top, T probe) {
181     if (top == terminal()) {
182       return top;
183     } else {
184       int direction = top.compareTo(probe);
185       if (direction > 0) {
186         return find(top.getLeft(), probe);
187       } else if (direction < 0) {
188         return find(top.getRight(), probe);
189       } else {
190         return top;
191       }
192     }
193   }
194
195   private Node<T> insert(Node<T> top, T data) {
196     if (top == terminal()) {
197       mutated = true;
198       return createNode(data);
199     } else {
200       int direction = top.compareTo(data);
201       if (direction > 0) {
202         top.setLeft(insert(top.getLeft(), data));
203       } else if (direction < 0) {
204         top.setRight(insert(top.getRight(), data));
205       } else {
206         return top;
207       }
208       top = skew(top);
209       top = split(top);
210       return top;
211     }
212   }
213
214   private Node<T> createNode(T data) {
215     if (data instanceof Node<?>) {
216       return (Node<T>) data;
217     } else {
218       return new TreeNode<T>(data);
219     }
220   }
221
222   private Node<T> remove(Node<T> top, T data) {
223     if (top != terminal()) {
224       int direction = top.compareTo(data);
225
226       heir = top;
227       if (direction > 0) {
228         top.setLeft(remove(top.getLeft(), data));
229       } else {
230         item = top;
231         top.setRight(remove(top.getRight(), data));
232       }
233
234       if (top == heir) {
235         if (item != terminal() && item.compareTo(data) == 0) {
236           mutated = true;
237           item.swapPayload(top);
238           removed = top.getPayload();
239           top = top.getRight();
240         }
241       } else {
242         if (top.getLeft().getLevel() < top.getLevel() - 1 || top.getRight().getLevel() < top.getLevel() - 1) {
243           if (top.getRight().getLevel() > top.decrementLevel()) {
244             top.getRight().setLevel(top.getLevel());
245           }
246
247           top = skew(top);
248           top.setRight(skew(top.getRight()));
249           top.getRight().setRight(skew(top.getRight().getRight()));
250           top = split(top);
251           top.setRight(split(top.getRight()));
252         }
253       }
254     }
255     return top;
256   }
257
258   private static <T> Node<T> skew(Node<T> top) {
259     if (top.getLeft().getLevel() == top.getLevel() && top.getLevel() != 0) {
260       Node<T> save = top.getLeft();
261       top.setLeft(save.getRight());
262       save.setRight(top);
263       top = save;
264     }
265
266     return top;
267   }
268
269   private static <T> Node<T> split(Node<T> top) {
270     if (top.getRight().getRight().getLevel() == top.getLevel() && top.getLevel() != 0) {
271       Node<T> save = top.getRight();
272       top.setRight(save.getLeft());
273       save.setLeft(top);
274       top = save;
275       top.incrementLevel();
276     }
277
278     return top;
279   }
280
281   /**
282    * Interface implemented by nodes within this tree.
283    */
284   public static interface Node<E> {
285
286     /**
287      * Compare this node to the supplied 'data' object.
288      */
289     public int compareTo(E data);
290
291     /**
292      * Set this node's left child.
293      */
294     public void setLeft(Node<E> node);
295
296     /**
297      * Set this node's right child.
298      */
299     public void setRight(Node<E> node);
300
301     /**
302      * Get this node's left child.
303      */
304     public Node<E> getLeft();
305
306     /**
307      * Get this node's right child.
308      */
309     public Node<E> getRight();
310
311     /**
312      * Get this node's level.
313      */
314     public int getLevel();
315
316     /**
317      * Set this node's level.
318      */
319     public void setLevel(int value);
320
321     /**
322      * Decrement and then return this node's new level.
323      */
324     public int decrementLevel();
325
326     /**
327      * Increment and then return this node's new level.
328      */
329     public int incrementLevel();
330
331     /**
332      * Swap the payload objects between this node and the supplied node.
333      */
334     public void swapPayload(Node<E> with);
335
336     /**
337      * Return the 'value' object held within this node.
338      */
339     public E getPayload();
340   }
341
342   /**
343    * Abstract node implementation that can be extended with a custom payload.
344    */
345   public abstract static class AbstractTreeNode<E> implements Node<E> {
346
347     private Node<E> left;
348     private Node<E> right;
349     private int     level;
350
351     /**
352      * Constructs an initial (leaf) node.
353      */
354     public AbstractTreeNode() {
355       this(1);
356     }
357
358     private AbstractTreeNode(int level) {
359       this.left = (Node<E>) TERMINAL;
360       this.right = (Node<E>) TERMINAL;
361       this.level = level;
362     }
363
364     /**
365      * {@inheritDoc}
366      */
367     public void setLeft(Node<E> node) {
368       left = node;
369     }
370
371     /**
372      * {@inheritDoc}
373      */
374     public void setRight(Node<E> node) {
375       right = node;
376     }
377
378     /**
379      * {@inheritDoc}
380      */
381     public Node<E> getLeft() {
382       return left;
383     }
384
385     /**
386      * {@inheritDoc}
387      */
388     public Node<E> getRight() {
389       return right;
390     }
391
392     /**
393      * {@inheritDoc}
394      */
395     public int getLevel() {
396       return level;
397     }
398
399     /**
400      * {@inheritDoc}
401      */
402     public void setLevel(int value) {
403       level = value;
404     }
405
406     /**
407      * {@inheritDoc}
408      */
409     public int decrementLevel() {
410       return --level;
411     }
412
413     /**
414      * {@inheritDoc}
415      */
416     public int incrementLevel() {
417       return ++level;
418     }
419   }
420
421   /**
422    * Default Comparable wrapping node implementation.
423    */
424   private static final class TreeNode<E extends Comparable> extends AbstractTreeNode<E> {
425
426     private E payload;
427
428     public TreeNode(E payload) {
429       super();
430       this.payload = payload;
431     }
432
433     public int compareTo(E data) {
434       return payload.compareTo(data);
435     }
436
437     public void swapPayload(Node<E> node) {
438       if (node instanceof TreeNode<?>) {
439         TreeNode<E> treeNode = (TreeNode<E>) node;
440         E temp = treeNode.payload;
441         treeNode.payload = this.payload;
442         this.payload = temp;
443       } else {
444         throw new IllegalArgumentException();
445       }
446     }
447
448     public E getPayload() {
449       return payload;
450     }
451   }
452
453   /**
454    * Node implementation that is used to mark the terminal (leaf) nodes of the tree.
455    */
456   private static final class TerminalNode<E> extends AbstractTreeNode<E> {
457
458     TerminalNode() {
459       super(0);
460       super.setLeft(this);
461       super.setRight(this);
462     }
463
464     public int compareTo(E data) {
465       return 0;
466     }
467
468     @Override
469     public void setLeft(Node<E> right) {
470         if (right != TERMINAL) {
471             throw new AssertionError();
472         }
473     }
474
475     @Override
476     public void setRight(Node<E> left) {
477         if (left != TERMINAL) {
478             throw new AssertionError();
479         }
480     }
481
482     @Override
483     public void setLevel(int value) {
484       throw new AssertionError();
485     }
486
487     @Override
488     public int decrementLevel() {
489       throw new AssertionError();
490     }
491
492     @Override
493     public int incrementLevel() {
494       throw new AssertionError();
495     }
496
497     public void swapPayload(Node<E> payload) {
498       throw new AssertionError();
499     }
500
501     public E getPayload() {
502       return null;
503     }
504   }
505
506   /**
507    * SortedSet representation of a range of the tree.
508    */
509   private class SubSet extends AbstractSet<T> implements SortedSet<T> {
510
511     private final T start;
512     private final T end;
513
514     SubSet(T start, T end) {
515       this.start = start;
516       this.end = end;
517     }
518
519     @Override
520     public boolean add(T o) {
521       if (inRange(o)) {
522         return AATreeSet.this.add(o);
523       } else {
524         throw new IllegalArgumentException();
525       }
526     }
527
528     @Override
529     public boolean remove(Object o) {
530       if (inRange((T) o)) {
531         return remove(o);
532       } else {
533         return false;
534       }
535     }
536
537     @Override
538     public void clear() {
539       throw new UnsupportedOperationException();
540     }
541
542     @Override
543     public Iterator<T> iterator() {
544       if (end == null) {
545         return new SubTreeIterator(start, end);
546       } else {
547         throw new UnsupportedOperationException();
548       }
549     }
550
551     @Override
552     public int size() {
553       throw new UnsupportedOperationException();
554     }
555
556     @Override
557     public boolean isEmpty() {
558       return !iterator().hasNext();
559     }
560
561     public Comparator<? super T> comparator() {
562       return null;
563     }
564
565     public SortedSet<T> subSet(T fromElement, T toElement) {
566       if (inRangeInclusive(fromElement) && inRangeInclusive(toElement)) {
567         return new SubSet(fromElement, toElement);
568       } else {
569         throw new IllegalArgumentException();
570       }
571     }
572
573     public SortedSet<T> headSet(T toElement) {
574       if (inRangeInclusive(toElement)) {
575         return new SubSet(start, toElement);
576       } else {
577         throw new IllegalArgumentException();
578       }
579     }
580
581     public SortedSet<T> tailSet(T fromElement) {
582       if (inRangeInclusive(fromElement)) {
583         return new SubSet(fromElement, end);
584       } else {
585         throw new IllegalArgumentException();
586       }
587     }
588
589     public T first() {
590       if (start == null) {
591         return AATreeSet.this.first();
592       } else {
593         throw new UnsupportedOperationException();
594       }
595     }
596
597     public T last() {
598       if (end == null) {
599         return AATreeSet.this.last();
600       } else {
601         throw new UnsupportedOperationException();
602       }
603     }
604
605     private boolean inRange(T value) {
606       return (start == null || start.compareTo(value) <= 0) && (end == null || end.compareTo(value) > 0);
607     }
608
609     private boolean inRangeInclusive(T value) {
610       return (start == null || start.compareTo(value) <= 0) && (end == null || end.compareTo(value) >= 0);
611     }
612   }
613
614   /**
615    * Iterator that iterates over the tree.
616    */
617   private class TreeIterator implements Iterator<T> {
618
619     private final Stack<Node<T>> path = new Stack<Node<T>>();
620     private Node<T>                        next;
621
622     TreeIterator() {
623       path.push(terminal());
624       Node<T> leftMost = root;
625       while (leftMost.getLeft() != terminal()) {
626         path.push(leftMost);
627         leftMost = leftMost.getLeft();
628       }
629       next = leftMost;
630     }
631
632     TreeIterator(T start) {
633       path.push(terminal());
634       Node<T> current = root;
635       while (true) {
636         int direction = current.compareTo(start);
637         if (direction > 0) {
638           if (current.getLeft() == terminal()) {
639             next = current;
640             break;
641           } else {
642             path.push(current);
643             current = current.getLeft();
644           }
645         } else if (direction < 0) {
646           if (current.getRight() == terminal()) {
647             next = path.pop();
648             break;
649           } else {
650             current = current.getRight();
651           }
652         } else {
653           next = current;
654           break;
655         }
656       }
657     }
658
659     public boolean hasNext() {
660       return next != terminal();
661     }
662
663     public T next() {
664       Node<T> current = next;
665       advance();
666       return current.getPayload();
667     }
668
669     private void advance() {
670       Node<T> successor = next.getRight();
671       if (successor != terminal()) {
672         while (successor.getLeft() != terminal()) {
673           path.push(successor);
674           successor = successor.getLeft();
675         }
676         next = successor;
677       } else {
678         next = path.pop();
679       }
680     }
681
682     public void remove() {
683       throw new UnsupportedOperationException();
684     }
685
686   }
687
688   /**
689    * Iterator that iterates over a subset of the tree.
690    */
691   private class SubTreeIterator extends TreeIterator {
692     SubTreeIterator(T start, T end) {
693       super(start);
694     }
695   }
696 }
697