001/**
002 *
003 * Licensed to the Apache Software Foundation (ASF) under one
004 * or more contributor license agreements.  See the NOTICE file
005 * distributed with this work for additional information
006 * regarding copyright ownership.  The ASF licenses this file
007 * to you under the Apache License, Version 2.0 (the
008 * "License"); you may not use this file except in compliance
009 * with the License.  You may obtain a copy of the License at
010 *
011 *     http://www.apache.org/licenses/LICENSE-2.0
012 *
013 * Unless required by applicable law or agreed to in writing, software
014 * distributed under the License is distributed on an "AS IS" BASIS,
015 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
016 * See the License for the specific language governing permissions and
017 * limitations under the License.
018 */
019
020package org.apache.hadoop.hbase.util;
021
022import java.util.Iterator;
023import java.util.concurrent.atomic.AtomicBoolean;
024
025import org.apache.yetus.audience.InterfaceAudience;
026import org.apache.yetus.audience.InterfaceStability;
027
028/**
029 * Helper class that allows to create and manipulate an AvlTree.
030 * The main utility is in cases where over time we have a lot of add/remove of the same object,
031 * and we want to avoid all the allocations/deallocations of the "node" objects that the
032 * java containers will create.
033 */
034@InterfaceAudience.Private
035@InterfaceStability.Evolving
036public final class AvlUtil {
037  private AvlUtil() {}
038
039  /**
040   * This class represent a node that will be used in an AvlTree.
041   * Instead of creating another object for the tree node,
042   * like the TreeMap and the other java contains, here the node can be extended
043   * and the content can be embedded directly in the node itself.
044   * This is useful in cases where over time we have a lot of add/remove of the same object.
045   */
046  @InterfaceAudience.Private
047  public static abstract class AvlNode<TNode extends AvlNode> {
048    protected TNode avlLeft;
049    protected TNode avlRight;
050    protected int avlHeight;
051
052    public abstract int compareTo(TNode other);
053  }
054
055  /**
056   * This class extends the AvlNode and adds two links that will be used in conjunction
057   * with the AvlIterableList class.
058   * This is useful in situations where your node must be in a map to have a quick lookup by key,
059   * but it also require to be in something like a list/queue.
060   * This is useful in cases where over time we have a lot of add/remove of the same object.
061   */
062  @InterfaceAudience.Private
063  public static abstract class AvlLinkedNode<TNode extends AvlLinkedNode> extends AvlNode<TNode> {
064    protected TNode iterNext = null;
065    protected TNode iterPrev = null;
066  }
067
068  @InterfaceAudience.Private
069  public interface AvlInsertOrReplace<TNode extends AvlNode> {
070    TNode insert(Object searchKey);
071    TNode replace(Object searchKey, TNode prevNode);
072  }
073
074  /**
075   * The AvlTree allows to lookup an object using a custom key.
076   * e.g. the java Map allows only to lookup by key using the Comparator
077   * specified in the constructor.
078   * In this case you can pass a specific comparator for every needs.
079   */
080  @InterfaceAudience.Private
081  public static interface AvlKeyComparator<TNode extends AvlNode> {
082    int compareKey(TNode node, Object key);
083  }
084
085  /**
086   * Visitor that allows to traverse a set of AvlNodes.
087   * If you don't like the callback style of the visitor you can always use the AvlTreeIterator.
088   */
089  @InterfaceAudience.Private
090  public static interface AvlNodeVisitor<TNode extends AvlNode> {
091    /**
092     * @param node the node that we are currently visiting
093     * @return false to stop the iteration. true to continue.
094     */
095    boolean visitNode(TNode node);
096  }
097
098  /**
099   * Helper class that allows to create and manipulate an AVL Tree
100   */
101  @InterfaceAudience.Private
102  public static class AvlTree {
103    /**
104     * @param root the current root of the tree
105     * @param key the key for the node we are trying to find
106     * @param keyComparator the comparator to use to match node and key
107     * @return the node that matches the specified key or null in case of node not found.
108     */
109    public static <TNode extends AvlNode> TNode get(TNode root, final Object key,
110        final AvlKeyComparator<TNode> keyComparator) {
111      while (root != null) {
112        int cmp = keyComparator.compareKey(root, key);
113        if (cmp > 0) {
114          root = (TNode)root.avlLeft;
115        } else if (cmp < 0) {
116          root = (TNode)root.avlRight;
117        } else {
118          return (TNode)root;
119        }
120      }
121      return null;
122    }
123
124    /**
125     * @param root the current root of the tree
126     * @return the first (min) node of the tree
127     */
128    public static <TNode extends AvlNode> TNode getFirst(TNode root) {
129      if (root != null) {
130        while (root.avlLeft != null) {
131          root = (TNode)root.avlLeft;
132        }
133      }
134      return root;
135    }
136
137    /**
138     * @param root the current root of the tree
139     * @return the last (max) node of the tree
140     */
141    public static <TNode extends AvlNode> TNode getLast(TNode root) {
142      if (root != null) {
143        while (root.avlRight != null) {
144          root = (TNode)root.avlRight;
145        }
146      }
147      return root;
148    }
149
150    /**
151     * Insert a node into the tree. It uses the AvlNode.compareTo() for ordering.
152     * NOTE: The node must not be already in the tree.
153     * @param root the current root of the tree
154     * @param node the node to insert
155     * @return the new root of the tree
156     */
157    public static <TNode extends AvlNode> TNode insert(TNode root, TNode node) {
158      if (root == null) return node;
159
160      int cmp = node.compareTo(root);
161      assert cmp != 0 : "node already inserted: " + root;
162      if (cmp < 0) {
163        root.avlLeft = insert(root.avlLeft, node);
164      } else {
165        root.avlRight = insert(root.avlRight, node);
166      }
167      return balance(root);
168    }
169
170    /**
171     * Insert a node into the tree.
172     * This is useful when you want to create a new node or replace the content
173     * depending if the node already exists or not.
174     * Using AvlInsertOrReplace class you can return the node to add/replace.
175     *
176     * @param root the current root of the tree
177     * @param key the key for the node we are trying to insert
178     * @param keyComparator the comparator to use to match node and key
179     * @param insertOrReplace the class to use to insert or replace the node
180     * @return the new root of the tree
181     */
182    public static <TNode extends AvlNode> TNode insert(TNode root, Object key,
183        final AvlKeyComparator<TNode> keyComparator,
184        final AvlInsertOrReplace<TNode> insertOrReplace) {
185      if (root == null) {
186        return insertOrReplace.insert(key);
187      }
188
189      int cmp = keyComparator.compareKey(root, key);
190      if (cmp < 0) {
191        root.avlLeft = insert((TNode)root.avlLeft, key, keyComparator, insertOrReplace);
192      } else if (cmp > 0) {
193        root.avlRight = insert((TNode)root.avlRight, key, keyComparator, insertOrReplace);
194      } else {
195        TNode left = (TNode)root.avlLeft;
196        TNode right = (TNode)root.avlRight;
197        root = insertOrReplace.replace(key, root);
198        root.avlLeft = left;
199        root.avlRight = right;
200        return root;
201      }
202      return balance(root);
203    }
204
205    private static <TNode extends AvlNode> TNode removeMin(TNode p) {
206      if (p.avlLeft == null)
207        return (TNode)p.avlRight;
208      p.avlLeft = removeMin(p.avlLeft);
209      return balance(p);
210    }
211
212    /**
213     * Removes the node matching the specified key from the tree
214     * @param root the current root of the tree
215     * @param key the key for the node we are trying to find
216     * @param keyComparator the comparator to use to match node and key
217     * @return the new root of the tree
218     */
219    public static <TNode extends AvlNode> TNode remove(TNode root, Object key,
220        final AvlKeyComparator<TNode> keyComparator) {
221      return remove(root, key, keyComparator, null);
222    }
223
224    /**
225     * Removes the node matching the specified key from the tree
226     * @param root the current root of the tree
227     * @param key the key for the node we are trying to find
228     * @param keyComparator the comparator to use to match node and key
229     * @param removed will be set to true if the node was found and removed, otherwise false
230     * @return the new root of the tree
231     */
232    public static <TNode extends AvlNode> TNode remove(TNode root, Object key,
233        final AvlKeyComparator<TNode> keyComparator, final AtomicBoolean removed) {
234      if (root == null) return null;
235
236      int cmp = keyComparator.compareKey(root, key);
237      if (cmp == 0) {
238        if (removed != null) removed.set(true);
239
240        TNode q = (TNode)root.avlLeft;
241        TNode r = (TNode)root.avlRight;
242        if (r == null) return q;
243        TNode min = getFirst(r);
244        min.avlRight = removeMin(r);
245        min.avlLeft = q;
246        return balance(min);
247      } else if (cmp > 0) {
248        root.avlLeft = remove((TNode)root.avlLeft, key, keyComparator);
249      } else /* if (cmp < 0) */ {
250        root.avlRight = remove((TNode)root.avlRight, key, keyComparator);
251      }
252      return balance(root);
253    }
254
255    /**
256     * Visit each node of the tree
257     * @param root the current root of the tree
258     * @param visitor the AvlNodeVisitor instance
259     */
260    public static <TNode extends AvlNode> void visit(final TNode root,
261        final AvlNodeVisitor<TNode> visitor) {
262      if (root == null) return;
263
264      final AvlTreeIterator<TNode> iterator = new AvlTreeIterator<>(root);
265      boolean visitNext = true;
266      while (visitNext && iterator.hasNext()) {
267        visitNext = visitor.visitNode(iterator.next());
268      }
269    }
270
271    private static <TNode extends AvlNode> TNode balance(TNode p) {
272      fixHeight(p);
273      int balance = balanceFactor(p);
274      if (balance == 2) {
275        if (balanceFactor(p.avlRight) < 0) {
276          p.avlRight = rotateRight(p.avlRight);
277        }
278        return rotateLeft(p);
279      } else if (balance == -2) {
280        if (balanceFactor(p.avlLeft) > 0) {
281          p.avlLeft = rotateLeft(p.avlLeft);
282        }
283        return rotateRight(p);
284      }
285      return p;
286    }
287
288    private static <TNode extends AvlNode> TNode rotateRight(TNode p) {
289      TNode q = (TNode)p.avlLeft;
290      p.avlLeft = q.avlRight;
291      q.avlRight = p;
292      fixHeight(p);
293      fixHeight(q);
294      return q;
295    }
296
297    private static <TNode extends AvlNode> TNode rotateLeft(TNode q) {
298      TNode p = (TNode)q.avlRight;
299      q.avlRight = p.avlLeft;
300      p.avlLeft = q;
301      fixHeight(q);
302      fixHeight(p);
303      return p;
304    }
305
306    private static <TNode extends AvlNode> void fixHeight(TNode node) {
307      final int heightLeft = height(node.avlLeft);
308      final int heightRight = height(node.avlRight);
309      node.avlHeight = 1 + Math.max(heightLeft, heightRight);
310    }
311
312    private static <TNode extends AvlNode> int height(TNode node) {
313      return node != null ? node.avlHeight : 0;
314    }
315
316    private static <TNode extends AvlNode> int balanceFactor(TNode node) {
317      return height(node.avlRight) - height(node.avlLeft);
318    }
319  }
320
321  /**
322   * Iterator for the AvlTree
323   */
324  @InterfaceAudience.Private
325  public static class AvlTreeIterator<TNode extends AvlNode> implements Iterator<TNode> {
326    private final Object[] stack = new Object[64];
327
328    private TNode current = null;
329    private int height = 0;
330
331    public AvlTreeIterator() {
332    }
333
334    /**
335     * Create the iterator starting from the first (min) node of the tree
336     */
337    public AvlTreeIterator(final TNode root) {
338      seekFirst(root);
339    }
340
341    /**
342     * Create the iterator starting from the specified key
343     * @param root the current root of the tree
344     * @param key the key for the node we are trying to find
345     * @param keyComparator the comparator to use to match node and key
346     */
347    public AvlTreeIterator(final TNode root, final Object key,
348        final AvlKeyComparator<TNode> keyComparator) {
349      seekTo(root, key, keyComparator);
350    }
351
352    @Override
353    public boolean hasNext() {
354      return current != null;
355    }
356
357    @Override
358    public TNode next() {
359      final TNode node = this.current;
360      seekNext();
361      return node;
362    }
363
364    @Override
365    public void remove() {
366      throw new UnsupportedOperationException();
367    }
368
369    /**
370     * Reset the iterator, and seeks to the first (min) node of the tree
371     * @param root the current root of the tree
372     */
373    public void seekFirst(final TNode root) {
374      current = root;
375      height = 0;
376      if (root != null) {
377        while (current.avlLeft != null) {
378          stack[height++] = current;
379          current = (TNode)current.avlLeft;
380        }
381      }
382    }
383
384    /**
385     * Reset the iterator, and seeks to the specified key
386     * @param root the current root of the tree
387     * @param key the key for the node we are trying to find
388     * @param keyComparator the comparator to use to match node and key
389     */
390    public void seekTo(final TNode root, final Object key,
391        final AvlKeyComparator<TNode> keyComparator) {
392      current = null;
393      height = 0;
394
395      TNode node = root;
396      while (node != null) {
397        if (keyComparator.compareKey(node, key) >= 0) {
398          if (node.avlLeft != null) {
399            stack[height++] = node;
400            node = (TNode)node.avlLeft;
401          } else {
402            current = node;
403            return;
404          }
405        } else {
406          if (node.avlRight != null) {
407            stack[height++] = node;
408            node = (TNode)node.avlRight;
409          } else {
410            if (height > 0) {
411              TNode parent = (TNode)stack[--height];
412              while (node == parent.avlRight) {
413                if (height == 0) {
414                  current = null;
415                  return;
416                }
417                node = parent;
418                parent = (TNode)stack[--height];
419              }
420              current = parent;
421              return;
422            }
423            current = null;
424            return;
425          }
426        }
427      }
428    }
429
430    private void seekNext() {
431      if (current == null) return;
432      if (current.avlRight != null) {
433        stack[height++] = current;
434        current = (TNode)current.avlRight;
435        while (current.avlLeft != null) {
436          stack[height++] = current;
437          current = (TNode)current.avlLeft;
438        }
439      } else {
440        TNode node;
441        do {
442          if (height == 0) {
443            current = null;
444            return;
445          }
446          node = current;
447          current = (TNode)stack[--height];
448        } while (current.avlRight == node);
449      }
450    }
451  }
452
453  /**
454   * Helper class that allows to create and manipulate a linked list of AvlLinkedNodes
455   */
456  @InterfaceAudience.Private
457  public static class AvlIterableList {
458    /**
459     * @param node the current node
460     * @return the successor of the current node
461     */
462    public static <TNode extends AvlLinkedNode> TNode readNext(TNode node) {
463      return (TNode)node.iterNext;
464    }
465
466    /**
467     * @param node the current node
468     * @return the predecessor of the current node
469     */
470    public static <TNode extends AvlLinkedNode> TNode readPrev(TNode node) {
471      return (TNode)node.iterPrev;
472    }
473
474    /**
475     * @param head the head of the linked list
476     * @param node the node to add to the front of the list
477     * @return the new head of the list
478     */
479    public static <TNode extends AvlLinkedNode> TNode prepend(TNode head, TNode node) {
480      assert !isLinked(node) : node + " is already linked";
481      if (head != null) {
482        TNode tail = (TNode)head.iterPrev;
483        tail.iterNext = node;
484        head.iterPrev = node;
485        node.iterNext = head;
486        node.iterPrev = tail;
487      } else {
488        node.iterNext = node;
489        node.iterPrev = node;
490      }
491      return node;
492    }
493
494    /**
495     * @param head the head of the linked list
496     * @param node the node to add to the tail of the list
497     * @return the new head of the list
498     */
499    public static <TNode extends AvlLinkedNode> TNode append(TNode head, TNode node) {
500      assert !isLinked(node) : node + " is already linked";
501      if (head != null) {
502        TNode tail = (TNode)head.iterPrev;
503        tail.iterNext = node;
504        node.iterNext = head;
505        node.iterPrev = tail;
506        head.iterPrev = node;
507        return head;
508      }
509      node.iterNext = node;
510      node.iterPrev = node;
511      return node;
512    }
513
514    /**
515     * @param head the head of the current linked list
516     * @param otherHead the head of the list to append to the current list
517     * @return the new head of the current list
518     */
519    public static <TNode extends AvlLinkedNode> TNode appendList(TNode head, TNode otherHead) {
520      if (head == null) return otherHead;
521      if (otherHead == null) return head;
522
523      TNode tail = (TNode)head.iterPrev;
524      TNode otherTail = (TNode)otherHead.iterPrev;
525      tail.iterNext = otherHead;
526      otherHead.iterPrev = tail;
527      otherTail.iterNext = head;
528      head.iterPrev = otherTail;
529      return head;
530    }
531
532    /**
533     * @param head the head of the linked list
534     * @param node the node to remove from the list
535     * @return the new head of the list
536     */
537    public static <TNode extends AvlLinkedNode> TNode remove(TNode head, TNode node) {
538      assert isLinked(node) : node + " is not linked";
539      if (node != node.iterNext) {
540        node.iterPrev.iterNext = node.iterNext;
541        node.iterNext.iterPrev = node.iterPrev;
542        head = (head == node) ? (TNode)node.iterNext : head;
543      } else {
544        head = null;
545      }
546      node.iterNext = null;
547      node.iterPrev = null;
548      return head;
549    }
550
551    /**
552     * @param head the head of the linked list
553     * @param base the node which we want to add the {@code node} before it
554     * @param node the node which we want to add it before the {@code base} node
555     */
556    public static <TNode extends AvlLinkedNode> TNode prepend(TNode head, TNode base, TNode node) {
557      assert !isLinked(node) : node + " is already linked";
558      node.iterNext = base;
559      node.iterPrev = base.iterPrev;
560      base.iterPrev.iterNext = node;
561      base.iterPrev = node;
562      return head == base ? node : head;
563    }
564
565    /**
566     * @param node the node to check
567     * @return true if the node is linked to a list, false otherwise
568     */
569    public static <TNode extends AvlLinkedNode> boolean isLinked(TNode node) {
570      return node.iterPrev != null && node.iterNext != null;
571    }
572  }
573}