SplayTreeMap() {
  _dummy = new SplayTreeNode<K, V>(null, null);
  _count = 0;
}

void clear() {
  _root = null;
  _count = 0;
}

bool containsKey(K key) {
  if (!isEmpty()) {
    splay_(key);
    if (_root.key.compareTo(key) == 0) return true;
  }
  return false;
}

bool containsValue(V value) {
  bool found = false;
  bool visit(SplayTreeNode node) {
    if (node === null) return false;
    if (node.value == value) return true;
    return visit(node.left) || visit(node.right);
  }
  return visit(_root);
}

Get the first key in the map. Returns null if the map is empty.

K firstKey() {
  if (_root === null) return null;
  SplayTreeNode<K, V> node = _root;
  while (node.left !== null) {
    node = node.left;
  }
  // Maybe implement a splay-method that can splay the minimum without
  // performing comparisons.
  splay_(node.key);
  return node.key;
}

Get the first key in the map that is strictly larger than key. Returns null if no key was not found.

K firstKeyAfter(K key) {
  splay_(key);
  K visit(SplayTreeNode node, K ifEmpty) {
    if (node === null) return ifEmpty;
    if (node.key.compareTo(key) > 0) {
      return visit(node.left, node.key);
    }
    if (node.key.compareTo(key) <= 0) {
      return visit(node.right, ifEmpty);
    }
  }
  return visit(_root, null);
}

void forEach(void f(K key, V value)) {
  List<SplayTreeNode<K, V>> list = new List<SplayTreeNode<K, V>>();
  SplayTreeNode<K, V> current = _root;
  while (current !== null) {
    if (current.left !== null) {
      list.add(current);
      current = current.left;
    } else {
      f(current.key, current.value);
      while (current.right === null) {
        if (list.isEmpty()) return;
        current = list.removeLast();
        f(current.key, current.value);
      }
      current = current.right;
    }
  }
}

Collection<K> getKeys() {
  List<K> list = new List<K>();
  forEach((K k, V v) { list.add(k); });
  return list;
}

Collection<V> getValues() {
  List<V> list = new List<V>();
  forEach((K k, V v) { list.add(v); });
  return list;
}

bool isEmpty() {
  // assert(!((_root === null) && (_count != 0)));
  // assert(!((_count == 0) && (_root !== null)));
  return (_root === null);
}

Get the last key in the map. Returns null if the map is empty.

K lastKey() {
  if (_root === null) return null;
  SplayTreeNode<K, V> node = _root;
  while (node.right !== null) {
    node = node.right;
  }
  // Maybe implement a splay-method that can splay the maximum without
  // performing comparisons.
  splay_(node.key);
  return node.key;
}

Get the last key in the map that is strictly smaller than key. Returns null if no key was not found.

K lastKeyBefore(K key) {
  splay_(key);
  K visit(SplayTreeNode node, K ifEmpty) {
    if (node === null) return ifEmpty;
    if (node.key.compareTo(key) >= 0) {
      return visit(node.left, ifEmpty);
    }
    if (node.key.compareTo(key) < 0) {
      return visit(node.right, node.key);
    }
  }
  return visit(_root, null);
}

int get length() {
  return _count;
}

void operator []=(K key, V value) {
  if (isEmpty()) {
    _count++;
    _root = new SplayTreeNode(key, value);
    return;
  }
  // Splay on the key to move the last node on the search path for
  // the key to the root of the tree.
  splay_(key);
  if (_root.key.compareTo(key) == 0) {
    _root.value = value;
    return;
  }
  SplayTreeNode<K, V> node = new SplayTreeNode(key, value);
  // assert(_count >= 0);
  _count++;
  if (key.compareTo(_root.key) > 0) {
    node.left = _root;
    node.right = _root.right;
    _root.right = null;
  } else {
    node.right = _root;
    node.left = _root.left;
    _root.left = null;
  }
  _root = node;
}

V operator [](K key) {
  if (!isEmpty()) {
    splay_(key);
    if (_root.key.compareTo(key) == 0) return _root.value;
  }
  return null;
}

V putIfAbsent(K key, V ifAbsent()) {
  if (containsKey(key)) return this[key];
  V value = ifAbsent();
  this[key] = value;
  return value;
}

V remove(K key) {
  if (isEmpty()) return null;
  splay_(key);
  if (_root.key.compareTo(key) != 0) return null;
  V value = _root.value;

  _count--;
  // assert(_count >= 0);
  if (_root.left === null) {
    _root = _root.right;
  } else {
    SplayTreeNode<K, V> right = _root.right;
    _root = _root.left;
    // Splay to make sure that the new root has an empty right child.
    splay_(key);
    // Insert the original right child as the right child of the new
    // root.
    _root.right = right;
  }
  return value;
}

Perform the splay operation for the given key. Moves the node with the given key to the top of the tree. If no node has the given key, the last node on the search path is moved to the top of the tree. This is the simplified top-down splaying algorithm from: "Self-adjusting Binary Search Trees" by Sleator and Tarjan.

void splay_(K key) {
  if (isEmpty()) return;

  // The right child of the dummy node will hold
  // the L tree of the algorithm.  The left child of the dummy node
  // will hold the R tree of the algorithm.  Using a dummy node, left
  // and right will always be nodes and we avoid special cases.
  SplayTreeNode<K, V> left = _dummy;
  SplayTreeNode<K, V> right = _dummy;
  SplayTreeNode<K, V> current = _root;
  while (true) {
    int comp = key.compareTo(current.key);
    if (comp < 0) {
      if (current.left === null) break;
      if (key.compareTo(current.left.key) < 0) {
        // Rotate right.
        SplayTreeNode<K, V> tmp = current.left;
        current.left = tmp.right;
        tmp.right = current;
        current = tmp;
        if (current.left === null) break;
      }
      // Link right.
      right.left = current;
      right = current;
      current = current.left;
    } else if (comp > 0) {
      if (current.right === null) break;
      if (key.compareTo(current.right.key) > 0) {
        // Rotate left.
        SplayTreeNode<K, V> tmp = current.right;
        current.right = tmp.left;
        tmp.left = current;
        current = tmp;
        if (current.right === null) break;
      }
      // Link left.
      left.right = current;
      left = current;
      current = current.right;
    } else {
      break;
    }
  }
  // Assemble.
  left.right = current.left;
  right.left = current.right;
  current.left = _dummy.right;
  current.right = _dummy.left;
  _root = current;

  _dummy.right = null;
  _dummy.left = null;
}

String toString() {
  return Maps.mapToString(this);
}