/* * Copyright 2015 Red Hat, Inc. and/or its affiliates. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * */ package org.drools.core.util; import org.drools.core.util.index.TupleList; public class TupleRBTree<K extends Comparable< ? super K>> { public static final boolean VERIFY_RBTREE = false; private static final int INDENT_STEP = 4; public Node<K> root; public Node<K> nullNode; public void verifyProperties() { if ( VERIFY_RBTREE ) { verifyProperty1( root ); verifyProperty2( root ); // Property 3 is implicit verifyProperty4( root ); verifyProperty5( root ); } } private static void verifyProperty1(Node< ? > n) { assert nodeColor( n ) == Color.RED || nodeColor( n ) == Color.BLACK; if ( n == null ) return; verifyProperty1( n.left ); verifyProperty1( n.right ); } private static void verifyProperty2(Node< ? > root) { assert nodeColor( root ) == Color.BLACK; } private static Color nodeColor(Node< ? > n) { return n == null ? Color.BLACK : n.color; } private static void verifyProperty4(Node< ? > n) { if ( nodeColor( n ) == Color.RED ) { assert nodeColor( n.left ) == Color.BLACK; assert nodeColor( n.right ) == Color.BLACK; assert nodeColor( n.parent ) == Color.BLACK; } if ( n == null ) return; verifyProperty4( n.left ); verifyProperty4( n.right ); } private static void verifyProperty5(Node< ? > root) { verifyProperty5Helper( root, 0, -1 ); } private static int verifyProperty5Helper(Node< ? > n, int blackCount, int pathBlackCount) { if ( nodeColor( n ) == Color.BLACK ) { blackCount++; } if ( n == null ) { if ( pathBlackCount == -1 ) { pathBlackCount = blackCount; } else { assert blackCount == pathBlackCount; } return pathBlackCount; } pathBlackCount = verifyProperty5Helper( n.left, blackCount, pathBlackCount ); pathBlackCount = verifyProperty5Helper( n.right, blackCount, pathBlackCount ); return pathBlackCount; } public Node<K> lookup(K key) { if (key == null) { return nullNode; } Node<K> n = root; while ( n != null ) { int compResult = key.compareTo( n.key ); if ( compResult == 0 ) { return n; } else if ( compResult < 0 ) { n = n.left; } else { n = n.right; } } return n; } public enum Boundary { LOWER, UPPER }; public static class RangeFastIterator<K extends Comparable< ? super K>> implements FastIterator { private Node<K> upperBound; private Node<K> next; public RangeFastIterator(Node<K> lowerNearest, Node<K> upperNearest) { this.next = lowerNearest; this.upperBound = upperNearest; } public Entry next(Entry object) { Entry temp = next; next = checkUpperBound( recurse( next ) ); return temp; } public boolean isFullIterator() { return false; } private Node<K> recurse(Node<K> current) { if (current == null) { return null; } if (current.right != null) { Node<K> p = current.right; while (p.left != null) { p = p.left; } return p; } Node<K> p = current.parent; Node<K> ch = current; while (p != null && ch == p.right) { ch = p; p = p.parent; } return p; } public Node<K> checkUpperBound(Node<K> current) { if (upperBound == null) { return current; } return current == null || current.compareTo(upperBound) > 0 ? null : current; } } public boolean isEmpty() { return root == null; } public Node<K> first() { if (root == null) { return null; } Node<K> n = root; while (n.left != null) { n = n.left; } return n; } public Node<K> last() { if (root == null) { return null; } Node<K> n = root; while (n.right != null) { n = n.right; } return n; } public FastIterator fastIterator() { return root == null ? FastIterator.EMPTY : new RangeFastIterator( first(), null ); } @Override public String toString() { FastIterator iterator = fastIterator(); StringBuilder sb = new StringBuilder("["); boolean first = true; for (Entry entry = iterator.next(null); entry != null; entry = iterator.next(null)) { if (first) { first = false; } else { sb.append(", "); } sb.append(entry); } sb.append("]"); return sb.toString(); } public FastIterator range(K lowerBound, boolean testLowerEqual, K upperBound, boolean testUpperEqual) { Node<K> lowerNearest = findNearestNode( lowerBound, testLowerEqual, Boundary.LOWER ); Node<K> upperNearest = findNearestNode( upperBound, testUpperEqual, Boundary.UPPER ); if ( lowerNearest == null || upperNearest == null ) { return FastIterator.EMPTY; } if ( lowerNearest.key.compareTo( upperNearest.key ) > 0 ) { upperNearest = lowerNearest; } return new RangeFastIterator( lowerNearest, upperNearest ); } public void rangeUperBounded(K upperBound, boolean testUpperEqual) { Node<K> upperNearest = findNearestNode( upperBound, testUpperEqual, Boundary.UPPER ); } public void rangeLowerBounded(K upperBound, boolean testUpperEqual) { Node<K> upperNearest = findNearestNode( upperBound, testUpperEqual, Boundary.UPPER ); } public Node<K> findNearestNode(K key, boolean allowEqual, Boundary boundary) { if (key == null) { return allowEqual ? nullNode : null; } Node<K> nearest = null; Node<K> n = root; while ( n != null ) { int compResult = key.compareTo( n.key ); if ( allowEqual && compResult == 0 ) { return n; } boolean accepted = acceptNode(compResult, boundary); if ( acceptNode( compResult, boundary ) && ( nearest == null || acceptNode( n.key.compareTo( nearest.key ), boundary ) ) ) { nearest = n; } if ( compResult == 0 ) { n = boundary == Boundary.LOWER ? n.right : n.left; } else { n = accepted ^ boundary == Boundary.LOWER ? n.right : n.left; } } return nearest; } private boolean acceptNode(int compResult, Boundary boundary) { return compResult != 0 && ( compResult > 0 ^ boundary == Boundary.LOWER ); } private void rotateLeft(Node<K> n) { Node<K> r = n.right; replaceNode( n, r ); n.right = r.left; if ( r.left != null ) { r.left.parent = n; } r.left = n; n.parent = r; } private void rotateRight(Node<K> n) { Node<K> l = n.left; replaceNode( n, l ); n.left = l.right; if ( l.right != null ) { l.right.parent = n; } l.right = n; n.parent = l; } private void replaceNode(Node<K> oldn, Node<K> newn) { if ( oldn.parent == null ) { root = newn; } else { if ( oldn == oldn.parent.left ) oldn.parent.left = newn; else oldn.parent.right = newn; } if ( newn != null ) { newn.parent = oldn.parent; } } public Node<K> insert(K key) { if (key == null) { if (nullNode == null) { nullNode = new Node<K>( key ); } return nullNode; } Node<K> insertedNode; if ( root == null ) { insertedNode = new Node<K>( key ); root = insertedNode; } else { Node<K> n = root; while ( true ) { int compResult = key.compareTo( n.key ); if ( compResult == 0 ) { return n; } else if ( compResult < 0 ) { if ( n.left == null ) { insertedNode = new Node<K>( key ); n.left = insertedNode; break; } else { n = n.left; } } else { if ( n.right == null ) { insertedNode = new Node<K>( key ); n.right = insertedNode; break; } else { n = n.right; } } } insertedNode.parent = n; } insertCase1( insertedNode ); // verifyProperties(); return insertedNode; } private void insertCase1(Node<K> n) { if ( n.parent == null ) n.color = Color.BLACK; else insertCase2( n ); } private void insertCase2(Node<K> n) { if ( nodeColor( n.parent ) == Color.BLACK ) return; // Tree is still valid else insertCase3( n ); } void insertCase3(Node<K> n) { if ( nodeColor( n.uncle() ) == Color.RED ) { n.parent.color = Color.BLACK; n.uncle().color = Color.BLACK; n.grandparent().color = Color.RED; insertCase1( n.grandparent() ); } else { insertCase4( n ); } } void insertCase4(Node<K> n) { if ( n == n.parent.right && n.parent == n.grandparent().left ) { rotateLeft( n.parent ); n = n.left; } else if ( n == n.parent.left && n.parent == n.grandparent().right ) { rotateRight( n.parent ); n = n.right; } insertCase5(n); } void insertCase5(Node<K> n) { n.parent.color = Color.BLACK; n.grandparent().color = Color.RED; if ( n == n.parent.left && n.parent == n.grandparent().left ) { rotateRight( n.grandparent() ); } else { rotateLeft( n.grandparent() ); } } public void delete(K key) { Node<K> n = lookup(key); if ( n == null ) return; // Key not found, do nothing if ( n.left != null && n.right != null ) { // Copy key/value from predecessor and then delete it instead Node<K> pred = maximumNode( n.left ); pred.copyStateInto(n); n = pred; } Node<K> child = (n.right == null) ? n.left : n.right; if ( nodeColor( n ) == Color.BLACK ) { n.color = nodeColor( child ); deleteCase1( n ); } replaceNode( n, child ); if ( nodeColor( root ) == Color.RED ) { root.color = Color.BLACK; } if ( nodeColor( root ) == Color.RED ) { root.color = Color.BLACK; } // verifyProperties(); } private static <K extends Comparable< ? super K>, V> Node<K> maximumNode(Node<K> n) { while ( n.right != null ) { n = n.right; } return n; } private void deleteCase1(Node<K> n) { if ( n.parent == null ) return; else deleteCase2( n ); } private void deleteCase2(Node<K> n) { if ( nodeColor( n.sibling() ) == Color.RED ) { n.parent.color = Color.RED; n.sibling().color = Color.BLACK; if ( n == n.parent.left ) rotateLeft( n.parent ); else rotateRight( n.parent ); } deleteCase3( n ); } private void deleteCase3(Node<K> n) { if ( nodeColor( n.parent ) == Color.BLACK && nodeColor( n.sibling() ) == Color.BLACK && nodeColor( n.sibling().left ) == Color.BLACK && nodeColor( n.sibling().right ) == Color.BLACK ) { n.sibling().color = Color.RED; deleteCase1( n.parent ); } else deleteCase4( n ); } private void deleteCase4(Node<K> n) { if ( nodeColor( n.parent ) == Color.RED && nodeColor( n.sibling() ) == Color.BLACK && nodeColor( n.sibling().left ) == Color.BLACK && nodeColor( n.sibling().right ) == Color.BLACK ) { n.sibling().color = Color.RED; n.parent.color = Color.BLACK; } else deleteCase5( n ); } private void deleteCase5(Node<K> n) { if ( n == n.parent.left && nodeColor( n.sibling() ) == Color.BLACK && nodeColor( n.sibling().left ) == Color.RED && nodeColor( n.sibling().right ) == Color.BLACK ) { n.sibling().color = Color.RED; n.sibling().left.color = Color.BLACK; rotateRight( n.sibling() ); } else if ( n == n.parent.right && nodeColor( n.sibling() ) == Color.BLACK && nodeColor( n.sibling().right ) == Color.RED && nodeColor( n.sibling().left ) == Color.BLACK ) { n.sibling().color = Color.RED; n.sibling().right.color = Color.BLACK; rotateLeft( n.sibling() ); } deleteCase6( n ); } private void deleteCase6(Node<K> n) { n.sibling().color = nodeColor( n.parent ); n.parent.color = Color.BLACK; if ( n == n.parent.left ) { n.sibling().right.color = Color.BLACK; rotateLeft( n.parent ); } else { n.sibling().left.color = Color.BLACK; rotateRight( n.parent ); } } public void print() { printHelper( root, 0 ); } private static void printHelper(Node< ? > n, int indent) { if ( n == null ) { System.out.print( "<empty tree>" ); return; } if ( n.right != null ) { printHelper( n.right, indent + INDENT_STEP ); } for ( int i = 0; i < indent; i++ ) { System.out.print( " " ); } if ( n.color == Color.BLACK ) { System.out.println( n.key ); } else { System.out.println( "<" + n.key + ">" ); } if ( n.left != null ) { printHelper( n.left, indent + INDENT_STEP ); } } public enum Color { RED, BLACK } public static class Node<K extends Comparable< ? super K>> extends TupleList implements Entry<TupleList>, Comparable<Node<K>> { public K key; private Node<K> left; private Node<K> right; private Node<K> parent; private Color color = Color.RED; public Node(K key) { this.key = key; } public Node<K> grandparent() { return parent.parent; } public Node<K> sibling() { return this == parent.left ? parent.right :parent.left; } public Node<K> uncle() { return parent.sibling(); } public String toString() { return "Node key=" + key; } public void setNext(TupleList next) { // TODO Auto-generated method stub } public TupleList getNext() { // TODO Auto-generated method stub return null; } public int compareTo(Node<K> other) { return key.compareTo(other.key); } protected void copyStateInto(Node<K> other) { super.copyStateInto(other); other.key = key; } } }