// based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html
// TODO: implement NavigableSet

// RAM use in a CompressedOOPS JVM (bytes per element):
//   ~12   when elements are inserted in sorted order
//   ~13.7 when elements are inserted in random order
//
// (Obviously more for very small sets.)

sclass HyperCompactTreeSet<A> extends AbstractSet<A> {

  // A symbol table implemented using a left-leaning red-black BST.
  // This is the 2-3 version.
 
  // Note: We sometimes cast the nullSentinel to A
  // which is technically incorrect but ok because of type erasure.
  // May want to fix just to be clean on the source level too.
  
  private static final boolean RED   = true;
  private static final boolean BLACK = false;
  
  // replacement for null elements
  private static final new O nullSentinel;

  private Node<A> root;     // root of the BST
  int size; // size of tree set

  // BST helper node data type
  abstract sclass Node<A> {
    A val;         // associated data
    
    Node<A> left() { null; } // get left subtree
    abstract Node<A> setLeft(Node<A> left); // set left subtree - return potentially replaced node
    
    Node<A> right() { null; } // get right subtree
    abstract Node<A> setRight(Node<A> right); // set right subtree - return potentially replaced node
    
    abstract bool color();
    abstract Node<A> convertToBlack();
    abstract Node<A> convertToRed();
    abstract Node<A> invertColor();
    Node<A> convertToColor(bool color) { ret color == RED ? convertToRed() : convertToBlack(); }
    
    abstract bool isLeaf();
  }
  
  // This represents a common case near a red leaf - a black node
  // containing a red leaf in the left slot and null in the right slot.
  // Combined with the direct storage of black leaf children in NonLeaf,
  // this is all we need to get rid of all the leaf overhead. Yay!
  sclass SpecialNode<A> extends Node<A> {
    A leftVal;
    
    *(A *leftVal, A *val) {}
    
    bool color() { ret BLACK; }
    Node<A> convertToBlack() { this; }
    Node<A> convertToRed() { ret newNode(RED, val, left(), right()); }
    Node invertColor() { ret convertToRed(); }
    
    Node<A> left() {
      ret newLeaf(RED, leftVal);
    }
    
    Node<A> setLeft(Node<A> left) {
      // Can we keep the optimized representation? (Probably this
      // is never going to be true.)
      if (left != null && left.isLeaf() && left.color() == RED)
        ret this with leftVal = left.val;
      else
        ret newNode(BLACK, val, left, right());
    }
    
    Node<A> right() {
      null;
    }
    
    Node<A> setRight(Node<A> right) {
      if (right == null) this;
      ret newNode(color(), val, left(), right);
    }
    
    bool isLeaf() { false; }
  }
  
  asclass NonLeaf<A> extends Node<A> {
    // either a Node or a (sentinelled) direct user value
    O left, right;
    
    // color of leaf if left is a user value
    bool defaultLeftLeafColor() { ret BLACK; }
    
    // color of leaf if right is a user value
    bool defaultRightLeafColor() { ret BLACK; }

    Node<A> left() {
      ret left == null ? null
        : left instanceof Node ? (Node) left
        : newLeaf(defaultLeftLeafColor(), (A) left);
    }
    
    void setLeft_noMorph(Node<A> left) {
      this.left = left != null && left.isLeaf() && left.color() == defaultLeftLeafColor() ? left.val : left;
    }
    
    void setRight_noMorph(Node<A> right) {
      this.right = right != null && right.isLeaf() && right.color() == defaultRightLeafColor() ? right.val : right;
    }
    
    Node<A> setLeft(Node<A> left) {
      ifndef HyperCompactTreeSet_disableSpecialNodes
      if (color() == BLACK && right == null && left != null && left.isLeaf() && left.color() == RED)
        ret new SpecialNode(left.val, val);
      endifndef
      
      setLeft_noMorph(left);
      if (left == null && right() == null) ret newLeaf(color(), val);
      this;
    }
    
    Node<A> right() {
      ret right == null ? null
        : right instanceof Node ? (Node) right
        : newLeaf(defaultRightLeafColor(), (A) right);
    }
    
    Node<A> setRight(Node<A> right) {
      // Setting right to null may produce either a leaf or a
      // special node, so we just go through newNode.
      if (right == null && this.right != null)
        ret newNode(color(), val, left(), null);
        
      // New right is not null, so we compress (if possible) and store it
      setRight_noMorph(right);
      this;
    }
    
    bool isLeaf() { false; }
  }
    
  sclass BlackNode<A> extends NonLeaf<A> {
    *(A *val) {}
    
    bool color() { ret BLACK; }
    Node<A> convertToBlack() { this; }
    Node<A> convertToRed() { ret newNode(RED, val, left(), right()); }
    Node invertColor() { ret convertToRed(); }
  }

  sclass RedNode<A> extends NonLeaf<A> {
    *(A *val) {}

    bool color() { ret RED; }
    Node<A> convertToBlack() { ret newNode(BLACK, val, left(), right()); }
    Node<A> convertToRed() { this; }
    Node invertColor() { ret convertToBlack(); }
  }
  
  asclass Leaf<A> extends Node<A> {
    bool isLeaf() { true; }
    
    Node<A> setLeft(Node<A> left) {
      ret left == null ? this : newNode(color(), val, left, null);
    }
    
    Node<A> setRight(Node<A> right) {
      ret right == null ? this : newNode(color(), val, null, right);
    }
  }
  
  sclass BlackLeaf<A> extends Leaf<A> {
    *(A *val) {}

    bool color() { ret BLACK; }
    Node<A> convertToBlack() { this; }
    Node<A> convertToRed() { ret new RedLeaf(val); }
    Node invertColor() { ret convertToRed(); }
  }

  sclass RedLeaf<A> extends Leaf<A> {
    *(A *val) {}

    bool color() { ret RED; }
    Node<A> convertToBlack() { ret new BlackLeaf(val); }
    Node<A> convertToRed() { this; }
    Node invertColor() { ret convertToBlack(); }
  }
  
  *() {}
  *(Cl<? extends A> cl) { addAll(cl); }
  
  private static O deSentinel(O o) {
    ret o == nullSentinel ? null : o;
  }
  
  private static O sentinel(O o) {
    ret o == null ? nullSentinel : o;
  }

  // returns false on null (algorithm needs this)
  static bool <A> isRed(Node<A> x) {
    ret x != null && x.color() == RED;
  }

  static <A> Node<A> newLeaf(bool color, A val) {
    ret color == RED ? new RedLeaf(val) : new BlackLeaf(val);
  }

  static <A> Node<A> newNode(bool color, A val, Node<A> left, Node<A> right) {
    // Make leaf (always a temporary object now)
    if (left == null && right == null)
      ret newLeaf(color, val);
      
    ifndef HyperCompactTreeSet_disableSpecialNodes
      // Make special node
      if (color == BLACK
        && right == null
        && left != null && left.isLeaf() && left.color() == RED)
        ret new SpecialNode(left.val, val);
    endifndef
      
    // Make normal non-leaf
    NonLeaf node = color == RED ? new RedNode(val) : new BlackNode(val);
    node.setLeft_noMorph(left);
    node.setRight_noMorph(right);
    ret node;
  }

  public int size() {
    ret size;
  }

  public bool isEmpty() {
    ret root == null;
  }

  public bool add(A val) {
    val = (A) sentinel(val);
    int oldSize = size;
    root = put(root, val);
    root = root.convertToBlack();
    ifdef CompactTreeSet_debug assertTrue(check()); endifdef
    ret size > oldSize;
  }

  // insert the value in the subtree rooted at h
  private Node<A> put(Node<A> h, A val) { 
    if (h == null) { ++size; ret new RedLeaf(val); }

    int cmp = compare_deSentinel(val, h.val);
    if      (cmp < 0) h = h.setLeft(put(h.left(),  val)); 
    else if (cmp > 0) h = h.setRight(put(h.right(), val)) ;
    else              { /*h.val = val;*/ } // no overwriting

    // fix-up any right-leaning links
    if (isRed(h.right()) && !isRed(h.left()))      h = rotateLeft(h);
    if (isRed(h.left())  &&  isRed(h.left().left())) h = rotateRight(h);
    if (isRed(h.left())  &&  isRed(h.right()))     h = flipColors(h);

    ret h;
  }
  
  final int compare_deSentinel(A a, A b) {
    ret compare((A) deSentinel(a), (A) deSentinel(b));
  }
  
  // override me if you wish
  int compare(A a, A b) {
    ret cmp(a, b);
  }

  public bool remove(O key) { 
    if (!contains(key)) false;

    key = sentinel(key);
    
    // if both children of root are black, set root to red
    if (!isRed(root.left()) && !isRed(root.right()))
        root = root.convertToRed();

    root = delete(root, (A) key);
    if (!isEmpty()) root = root.convertToBlack();
    // assert check();
    true;
  }

  // delete the key-value pair with the given key rooted at h
  private Node delete(Node<A> h, A key) { 
    // assert get(h, key) != null;

    if (compare_deSentinel(key, h.val) < 0)  {
      if (!isRed(h.left()) && !isRed(h.left().left()))
          h = moveRedLeft(h);
      h = h.setLeft(delete(h.left(), key));
    }
    else {
      if (isRed(h.left()))
          h = rotateRight(h);
      if (compare_deSentinel(key, h.val) == 0 && (h.right() == null)) {
          --size; null;
      } if (!isRed(h.right()) && !isRed(h.right().left()))
          h = moveRedRight(h);
      if (compare_deSentinel(key, h.val) == 0) {
          --size;
          Node<A> x = min(h.right());
          h.val = x.val;
          // h.val = get(h.right(), min(h.right()).val);
          // h.val = min(h.right()).val;
          h = h.setRight(deleteMin(h.right()));
      }
      else h = h.setRight(delete(h.right(), key));
    }
    return balance(h);
  }

  // make a left-leaning link lean to the right
  private Node<A> rotateRight(Node<A> h) {
      // assert (h != null) && isRed(h.left());
      Node<A> x = h.left();
      h = h.setLeft(x.right());
      x = x.setRight(h);
      x = x.convertToColor(x.right().color());
      x = x.setRight(x.right().convertToRed());
      ret x;
  }

  // make a right-leaning link lean to the left
  private Node<A> rotateLeft(Node<A> h) {
      // assert (h != null) && isRed(h.right());
      Node<A> x = h.right();
      h = h.setRight(x.left());
      x = x.setLeft(h);
      x = x.convertToColor(x.left().color());
      x = x.setLeft(x.left().convertToRed());
      ret x;
  }

  // flip the colors of a node and its two children
  private Node<A> flipColors(Node<A> h) {
      // h must have opposite color of its two children
      // assert (h != null) && (h.left() != null) && (h.right() != null);
      // assert (!isRed(h) &&  isRed(h.left()) &&  isRed(h.right()))
      //    || (isRed(h)  && !isRed(h.left()) && !isRed(h.right()));
      h = h.setLeft(h.left().invertColor());
      h = h.setRight(h.right().invertColor());
      ret h.invertColor();
  }

  // Assuming that h is red and both h.left() and h.left().left()
  // are black, make h.left() or one of its children red.
  private Node<A> moveRedLeft(Node<A> h) {
      // assert (h != null);
      // assert isRed(h) && !isRed(h.left()) && !isRed(h.left().left());

      h = flipColors(h);
      if (isRed(h.right().left())) { 
          h = h.setRight(rotateRight(h.right()));
          h = rotateLeft(h);
          h = flipColors(h);
      }
      ret h;
  }

  // Assuming that h is red and both h.right() and h.right().left()
  // are black, make h.right() or one of its children red.
  private Node<A> moveRedRight(Node<A> h) {
      // assert (h != null);
      // assert isRed(h) && !isRed(h.right()) && !isRed(h.right().left());
      h = flipColors(h);
      if (isRed(h.left().left())) { 
          h = rotateRight(h);
          h = flipColors(h);
      }
      ret h;
  }

  // restore red-black tree invariant
  private Node<A> balance(Node<A> h) {
      // assert (h != null);

      if (isRed(h.right()))                      h = rotateLeft(h);
      if (isRed(h.left()) && isRed(h.left().left())) h = rotateRight(h);
      if (isRed(h.left()) && isRed(h.right()))     h = flipColors(h);

      ret h;
  }


  /**
   * Returns the height of the BST (for debugging).
   * @return the height of the BST (a 1-node tree has height 0)
   */
  public int height() {
      ret height(root);
  }
  
  private int height(Node<A> x) {
      if (x == null) return -1;
      return 1 + Math.max(height(x.left()), height(x.right()));
  }
  
  public bool contains(O val) {
    ret find(root, (A) sentinel(val)) != null;
  }
  
  public A find(A probeVal) {
    probeVal = (A) sentinel(probeVal);
    Node<A> n = find(root, probeVal);
    ret n == null ? null : n.val;
  }

  // value associated with the given key in subtree rooted at x; null if no such key
  private A get(Node<A> x, A key) {
    x = find(x, key);
    ret x == null ? null : x.val;
  }

  Node<A> find(Node<A> x, A key) {
    while (x != null) {
      int cmp = compare_deSentinel(key, x.val);
      if      (cmp < 0) x = x.left();
      else if (cmp > 0) x = x.right();
      else              ret x;
    }
    null;
  }

  private boolean check() {
      if (!is23())             println("Not a 2-3 tree");
      if (!isBalanced())       println("Not balanced");
      return is23() && isBalanced();
  }

  // Does the tree have no red right links, and at most one (left)
  // red links in a row on any path?
  private boolean is23() { return is23(root); }
  private boolean is23(Node<A> x) {
      if (x == null) true;
      if (isRed(x.right())) false;
      if (x != root && isRed(x) && isRed(x.left())) false;
      ret is23(x.left()) && is23(x.right());
  } 

  // do all paths from root to leaf have same number of black edges?
  private bool isBalanced() { 
      int black = 0;     // number of black links on path from root to min
      Node<A> x = root;
      while (x != null) {
          if (!isRed(x)) black++;
          x = x.left();
      }
      ret isBalanced(root, black);
  }

  // does every path from the root to a leaf have the given number of black links?
  private boolean isBalanced(Node<A> x, int black) {
      if (x == null) return black == 0;
      if (!isRed(x)) black--;
      return isBalanced(x.left(), black) && isBalanced(x.right(), black);
  }
  
  public void clear() { root = null; size = 0; }
  
  // the smallest key in subtree rooted at x; null if no such key
  private Node<A> min(Node<A> x) { 
    // assert x != null;
    while (x.left() != null) x = x.left();
    ret x;
  }
  
  private Node<A> deleteMin(Node<A> h) { 
    if (h.left() == null)
      return null;

    if (!isRed(h.left()) && !isRed(h.left().left()))
      h = moveRedLeft(h);

    h = h.setLeft(deleteMin(h.left()));
    ret balance(h);
  }
  
  public Iterator<A> iterator() {
    ret new MyIterator;
  }

  class MyIterator extends ItIt<A> {
    new L<Node<A>> path;
    
    *() {
      fetch(root);
    }
    
    void fetch(Node<A> node) {
      while (node != null) {
        path.add(node);
        node = node.left();
      }
    }
    
    public bool hasNext() { ret !path.isEmpty(); }

    public A next() {
      if (path.isEmpty()) fail("no more elements");
      Node<A> node = popLast(path);
      // last node is always a leaf, so left is null
      // so proceed to fetch right branch
      fetch(node.right());
      ret (A) deSentinel(node.val);
    }
  }
  
  // Returns the smallest key in the symbol table greater than or equal to {@code key}.
  public A ceiling(A key) {
    key = (A) sentinel(key);
    Node<A> x = ceiling(root, key);
    ret x == null ? null : x.val;
  }

  // the smallest key in the subtree rooted at x greater than or equal to the given key
  Node<A> ceiling(Node<A> x, A key) {  
    if (x == null) null;
    int cmp = compare_deSentinel(key, x.val);
    if (cmp == 0) ret x;
    if (cmp > 0)  ret ceiling(x.right(), key);
    Node<A> t = ceiling(x.left(), key);
    if (t != null) ret t; 
    else           ret x;
  }
  
  public A floor(A key) {
    key = (A) sentinel(key);
    Node<A> x = floor(root, key);
    ret x == null ? null : x.val;
  }    

  // the largest key in the subtree rooted at x less than or equal to the given key
  Node<A> floor(Node<A> x, A key) {
    if (x == null) null;
    int cmp = compare_deSentinel(key, x.val);
    if (cmp == 0) ret x;
    if (cmp < 0)  ret floor(x.left(), key);
    Node<A> t = floor(x.right(), key);
    if (t != null) ret t; 
    else           ret x;
  }
  
  void testInternalStructure() {
    ifndef HyperCompactTreeSet_disableSpecialNodes
      // one leaf object (root) is allowed - could even optimize that
      assertTrue(countLeafObjects() <= 1);
    endifndef
  }
  
  // count leaf objects we didn't optimize away
  int countLeafObjects(Node node default root) {
    if (node instanceof Leaf) ret 1;
    if (node cast NonLeaf)
      ret countLeafObjects(optCast Node(node.left))
        + countLeafObjects(optCast Node(node.right));
    ret 0;
  }
  
  Cl<NonLeaf> unoptimizedNodes() {
    new L<NonLeaf> out;
    findUnoptimizedNodes(out);
    ret out;
  }
  
  void findUnoptimizedNodes(Node node default root, L<NonLeaf> out) {
    if (node == null) ret;
    if (node cast NonLeaf) {
      if (isUnoptimizedNode(node)) out.add(node);
      findUnoptimizedNodes(optCast Node(node.left), out);
      findUnoptimizedNodes(optCast Node(node.right), out);
    }
  }
  
  bool isUnoptimizedNode(Node node) {
    if (node cast NonLeaf)
      ret node.left instanceof Leaf
        || node.right instanceof Leaf;
    false;
  }
  
  // Compact me (minimize space use). Returns a compacted copy.
  // This only has an effect when elements were inserted in non-sorted
  // or and/or elements were removed.
  selfType compact() {
    ret new HyperCompactTreeSet(this);
  }
}

Began life as a copy of #1032628