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