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     * Visitor
088     * @param node the node that we are currently visiting
089     * @return false to stop the iteration. true to continue.
090     */
091    boolean visitNode(TNode node);
092  }
093
094  /**
095   * Helper class that allows to create and manipulate an AVL Tree
096   */
097  @InterfaceAudience.Private
098  public static class AvlTree {
099    /**
100     * Return the node that matches the specified key or null in case of node not found.
101     * @param root          the current root of the tree
102     * @param key           the key for the node we are trying to find
103     * @param keyComparator the comparator to use to match node and key
104     * @return the node that matches the specified key or null in case of node not found.
105     */
106    public static <TNode extends AvlNode> TNode get(TNode root, final Object key,
107      final AvlKeyComparator<TNode> keyComparator) {
108      while (root != null) {
109        int cmp = keyComparator.compareKey(root, key);
110        if (cmp > 0) {
111          root = (TNode) root.avlLeft;
112        } else if (cmp < 0) {
113          root = (TNode) root.avlRight;
114        } else {
115          return (TNode) root;
116        }
117      }
118      return null;
119    }
120
121    /**
122     * Return the first node of the tree.
123     * @param root the current root of the tree
124     * @return the first (min) node of the tree
125     */
126    public static <TNode extends AvlNode> TNode getFirst(TNode root) {
127      if (root != null) {
128        while (root.avlLeft != null) {
129          root = (TNode) root.avlLeft;
130        }
131      }
132      return root;
133    }
134
135    /**
136     * Return the last node of the tree.
137     * @param root the current root of the tree
138     * @return the last (max) node of the tree
139     */
140    public static <TNode extends AvlNode> TNode getLast(TNode root) {
141      if (root != null) {
142        while (root.avlRight != null) {
143          root = (TNode) root.avlRight;
144        }
145      }
146      return root;
147    }
148
149    /**
150     * Insert a node into the tree. It uses the AvlNode.compareTo() for ordering. NOTE: The node
151     * must not be already in the tree.
152     * @param root the current root of the tree
153     * @param node the node to insert
154     * @return the new root of the tree
155     */
156    public static <TNode extends AvlNode> TNode insert(TNode root, TNode node) {
157      if (root == null) return node;
158
159      int cmp = node.compareTo(root);
160      assert cmp != 0 : "node already inserted: " + root;
161      if (cmp < 0) {
162        root.avlLeft = insert(root.avlLeft, node);
163      } else {
164        root.avlRight = insert(root.avlRight, node);
165      }
166      return balance(root);
167    }
168
169    /**
170     * Insert a node into the tree. This is useful when you want to create a new node or replace the
171     * content depending if the node already exists or not. Using AvlInsertOrReplace class you can
172     * return the node to add/replace.
173     * @param root            the current root of the tree
174     * @param key             the key for the node we are trying to insert
175     * @param keyComparator   the comparator to use to match node and key
176     * @param insertOrReplace the class to use to insert or replace the node
177     * @return the new root of the tree
178     */
179    public static <TNode extends AvlNode> TNode insert(TNode root, Object key,
180      final AvlKeyComparator<TNode> keyComparator,
181      final AvlInsertOrReplace<TNode> insertOrReplace) {
182      if (root == null) {
183        return insertOrReplace.insert(key);
184      }
185
186      int cmp = keyComparator.compareKey(root, key);
187      if (cmp < 0) {
188        root.avlLeft = insert((TNode) root.avlLeft, key, keyComparator, insertOrReplace);
189      } else if (cmp > 0) {
190        root.avlRight = insert((TNode) root.avlRight, key, keyComparator, insertOrReplace);
191      } else {
192        TNode left = (TNode) root.avlLeft;
193        TNode right = (TNode) root.avlRight;
194        root = insertOrReplace.replace(key, root);
195        root.avlLeft = left;
196        root.avlRight = right;
197        return root;
198      }
199      return balance(root);
200    }
201
202    private static <TNode extends AvlNode> TNode removeMin(TNode p) {
203      if (p.avlLeft == null) return (TNode) p.avlRight;
204      p.avlLeft = removeMin(p.avlLeft);
205      return balance(p);
206    }
207
208    /**
209     * Removes the node matching the specified key from the tree
210     * @param root          the current root of the tree
211     * @param key           the key for the node we are trying to find
212     * @param keyComparator the comparator to use to match node and key
213     * @return the new root of the tree
214     */
215    public static <TNode extends AvlNode> TNode remove(TNode root, Object key,
216      final AvlKeyComparator<TNode> keyComparator) {
217      return remove(root, key, keyComparator, null);
218    }
219
220    /**
221     * Removes the node matching the specified key from the tree
222     * @param root          the current root of the tree
223     * @param key           the key for the node we are trying to find
224     * @param keyComparator the comparator to use to match node and key
225     * @param removed       will be set to true if the node was found and removed, otherwise false
226     * @return the new root of the tree
227     */
228    public static <TNode extends AvlNode> TNode remove(TNode root, Object key,
229      final AvlKeyComparator<TNode> keyComparator, final AtomicBoolean removed) {
230      if (root == null) return null;
231
232      int cmp = keyComparator.compareKey(root, key);
233      if (cmp == 0) {
234        if (removed != null) removed.set(true);
235
236        TNode q = (TNode) root.avlLeft;
237        TNode r = (TNode) root.avlRight;
238        if (r == null) return q;
239        TNode min = getFirst(r);
240        min.avlRight = removeMin(r);
241        min.avlLeft = q;
242        return balance(min);
243      } else if (cmp > 0) {
244        root.avlLeft = remove((TNode) root.avlLeft, key, keyComparator);
245      } else /* if (cmp < 0) */ {
246        root.avlRight = remove((TNode) root.avlRight, key, keyComparator);
247      }
248      return balance(root);
249    }
250
251    /**
252     * Visit each node of the tree
253     * @param root    the current root of the tree
254     * @param visitor the AvlNodeVisitor instance
255     */
256    public static <TNode extends AvlNode> void visit(final TNode root,
257      final AvlNodeVisitor<TNode> visitor) {
258      if (root == null) return;
259
260      final AvlTreeIterator<TNode> iterator = new AvlTreeIterator<>(root);
261      boolean visitNext = true;
262      while (visitNext && iterator.hasNext()) {
263        visitNext = visitor.visitNode(iterator.next());
264      }
265    }
266
267    private static <TNode extends AvlNode> TNode balance(TNode p) {
268      fixHeight(p);
269      int balance = balanceFactor(p);
270      if (balance == 2) {
271        if (balanceFactor(p.avlRight) < 0) {
272          p.avlRight = rotateRight(p.avlRight);
273        }
274        return rotateLeft(p);
275      } else if (balance == -2) {
276        if (balanceFactor(p.avlLeft) > 0) {
277          p.avlLeft = rotateLeft(p.avlLeft);
278        }
279        return rotateRight(p);
280      }
281      return p;
282    }
283
284    private static <TNode extends AvlNode> TNode rotateRight(TNode p) {
285      TNode q = (TNode) p.avlLeft;
286      p.avlLeft = q.avlRight;
287      q.avlRight = p;
288      fixHeight(p);
289      fixHeight(q);
290      return q;
291    }
292
293    private static <TNode extends AvlNode> TNode rotateLeft(TNode q) {
294      TNode p = (TNode) q.avlRight;
295      q.avlRight = p.avlLeft;
296      p.avlLeft = q;
297      fixHeight(q);
298      fixHeight(p);
299      return p;
300    }
301
302    private static <TNode extends AvlNode> void fixHeight(TNode node) {
303      final int heightLeft = height(node.avlLeft);
304      final int heightRight = height(node.avlRight);
305      node.avlHeight = 1 + Math.max(heightLeft, heightRight);
306    }
307
308    private static <TNode extends AvlNode> int height(TNode node) {
309      return node != null ? node.avlHeight : 0;
310    }
311
312    private static <TNode extends AvlNode> int balanceFactor(TNode node) {
313      return height(node.avlRight) - height(node.avlLeft);
314    }
315  }
316
317  /**
318   * Iterator for the AvlTree
319   */
320  @InterfaceAudience.Private
321  public static class AvlTreeIterator<TNode extends AvlNode> implements Iterator<TNode> {
322    private final Object[] stack = new Object[64];
323
324    private TNode current = null;
325    private int height = 0;
326
327    public AvlTreeIterator() {
328    }
329
330    /**
331     * Create the iterator starting from the first (min) node of the tree
332     * @param root the current root of the tree
333     */
334    public AvlTreeIterator(final TNode root) {
335      seekFirst(root);
336    }
337
338    /**
339     * Create the iterator starting from the specified key
340     * @param root          the current root of the tree
341     * @param key           the key for the node we are trying to find
342     * @param keyComparator the comparator to use to match node and key
343     */
344    public AvlTreeIterator(final TNode root, final Object key,
345      final AvlKeyComparator<TNode> keyComparator) {
346      seekTo(root, key, keyComparator);
347    }
348
349    @Override
350    public boolean hasNext() {
351      return current != null;
352    }
353
354    @Override
355    public TNode next() {
356      final TNode node = this.current;
357      seekNext();
358      return node;
359    }
360
361    @Override
362    public void remove() {
363      throw new UnsupportedOperationException();
364    }
365
366    /**
367     * Reset the iterator, and seeks to the first (min) node of the tree
368     * @param root the current root of the tree
369     */
370    public void seekFirst(final TNode root) {
371      current = root;
372      height = 0;
373      if (root != null) {
374        while (current.avlLeft != null) {
375          stack[height++] = current;
376          current = (TNode) current.avlLeft;
377        }
378      }
379    }
380
381    /**
382     * Reset the iterator, and seeks to the specified key
383     * @param root          the current root of the tree
384     * @param key           the key for the node we are trying to find
385     * @param keyComparator the comparator to use to match node and key
386     */
387    public void seekTo(final TNode root, final Object key,
388      final AvlKeyComparator<TNode> keyComparator) {
389      current = null;
390      height = 0;
391
392      TNode node = root;
393      while (node != null) {
394        if (keyComparator.compareKey(node, key) >= 0) {
395          if (node.avlLeft != null) {
396            stack[height++] = node;
397            node = (TNode) node.avlLeft;
398          } else {
399            current = node;
400            return;
401          }
402        } else {
403          if (node.avlRight != null) {
404            stack[height++] = node;
405            node = (TNode) node.avlRight;
406          } else {
407            if (height > 0) {
408              TNode parent = (TNode) stack[--height];
409              while (node == parent.avlRight) {
410                if (height == 0) {
411                  current = null;
412                  return;
413                }
414                node = parent;
415                parent = (TNode) stack[--height];
416              }
417              current = parent;
418              return;
419            }
420            current = null;
421            return;
422          }
423        }
424      }
425    }
426
427    private void seekNext() {
428      if (current == null) return;
429      if (current.avlRight != null) {
430        stack[height++] = current;
431        current = (TNode) current.avlRight;
432        while (current.avlLeft != null) {
433          stack[height++] = current;
434          current = (TNode) current.avlLeft;
435        }
436      } else {
437        TNode node;
438        do {
439          if (height == 0) {
440            current = null;
441            return;
442          }
443          node = current;
444          current = (TNode) stack[--height];
445        } while (current.avlRight == node);
446      }
447    }
448  }
449
450  /**
451   * Helper class that allows to create and manipulate a linked list of AvlLinkedNodes
452   */
453  @InterfaceAudience.Private
454  public static class AvlIterableList {
455    /**
456     * Return the successor of the current node
457     * @param node the current node
458     * @return the successor of the current node
459     */
460    public static <TNode extends AvlLinkedNode> TNode readNext(TNode node) {
461      return (TNode) node.iterNext;
462    }
463
464    /**
465     * Return the predecessor of the current node
466     * @param node the current node
467     * @return the predecessor of the current node
468     */
469    public static <TNode extends AvlLinkedNode> TNode readPrev(TNode node) {
470      return (TNode) node.iterPrev;
471    }
472
473    /**
474     * Prepend a node to the tree
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     * Append a node to the tree
496     * @param head the head of the linked list
497     * @param node the node to add to the tail of the list
498     * @return the new head of the list
499     */
500    public static <TNode extends AvlLinkedNode> TNode append(TNode head, TNode node) {
501      assert !isLinked(node) : node + " is already linked";
502      if (head != null) {
503        TNode tail = (TNode) head.iterPrev;
504        tail.iterNext = node;
505        node.iterNext = head;
506        node.iterPrev = tail;
507        head.iterPrev = node;
508        return head;
509      }
510      node.iterNext = node;
511      node.iterPrev = node;
512      return node;
513    }
514
515    /**
516     * Append a list of nodes to the tree
517     * @param head      the head of the current linked list
518     * @param otherHead the head of the list to append to the current list
519     * @return the new head of the current list
520     */
521    public static <TNode extends AvlLinkedNode> TNode appendList(TNode head, TNode otherHead) {
522      if (head == null) return otherHead;
523      if (otherHead == null) return head;
524
525      TNode tail = (TNode) head.iterPrev;
526      TNode otherTail = (TNode) otherHead.iterPrev;
527      tail.iterNext = otherHead;
528      otherHead.iterPrev = tail;
529      otherTail.iterNext = head;
530      head.iterPrev = otherTail;
531      return head;
532    }
533
534    /**
535     * Remove a node from the tree
536     * @param head the head of the linked list
537     * @param node the node to remove from the list
538     * @return the new head of the list
539     */
540    public static <TNode extends AvlLinkedNode> TNode remove(TNode head, TNode node) {
541      assert isLinked(node) : node + " is not linked";
542      if (node != node.iterNext) {
543        node.iterPrev.iterNext = node.iterNext;
544        node.iterNext.iterPrev = node.iterPrev;
545        head = (head == node) ? (TNode) node.iterNext : head;
546      } else {
547        head = null;
548      }
549      node.iterNext = null;
550      node.iterPrev = null;
551      return head;
552    }
553
554    /**
555     * Prepend a node to the tree before a specific node
556     * @param head the head of the linked list
557     * @param base the node which we want to add the {@code node} before it
558     * @param node the node which we want to add it before the {@code base} node
559     */
560    public static <TNode extends AvlLinkedNode> TNode prepend(TNode head, TNode base, TNode node) {
561      assert !isLinked(node) : node + " is already linked";
562      node.iterNext = base;
563      node.iterPrev = base.iterPrev;
564      base.iterPrev.iterNext = node;
565      base.iterPrev = node;
566      return head == base ? node : head;
567    }
568
569    /** Return true if the node is linked to a list, false otherwise */
570    public static <TNode extends AvlLinkedNode> boolean isLinked(TNode node) {
571      return node.iterPrev != null && node.iterNext != null;
572    }
573  }
574}