io.takamaka.code.util.StorageTreeIntMap Maven / Gradle / Ivy
Show all versions of io-takamaka-code Show documentation
/*
Copyright 2021 Fausto Spoto
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
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 io.takamaka.code.util;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.function.IntFunction;
import java.util.function.Supplier;
import java.util.function.UnaryOperator;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;
import io.takamaka.code.lang.Exported;
import io.takamaka.code.lang.Storage;
import io.takamaka.code.lang.View;
/**
* A map from integer keys to (possibly {@code null}) storage values,
* that can be kept in storage. By iterating on this object, one gets
* the key/value pairs of the map, in increasing key order.
*
* This code is derived from Sedgewick and Wayne's code for
* red-black trees, with some adaptation. It implements an associative
* map from keys to values. The map can be kept in storage.
* Values must have types allowed in storage. Keys are kept in increasing order.
*
* This class represents an ordered symbol table of generic key-value pairs.
* It supports the usual put, get, contains,
* remove, size, and is-empty methods.
* It also provides ordered methods for finding the minimum,
* maximum, floor, and ceiling.
* A symbol table implements the associative array abstraction:
* when associating a value with a key that is already in the symbol table,
* the convention is to replace the old value with the new value.
*
* This implementation uses a left-leaning red-black BST.
* The put, contains, delete, minimum,
* maximum, ceiling, and floor operations each take
* logarithmic time in the worst case, if the tree becomes unbalanced.
* The size, and is-empty operations take constant time.
* Construction takes constant time.
*
* For additional documentation, see Section 3.3 of
* Algorithms, 4th Edition by Robert Sedgewick and Kevin Wayne.
*
* @author Robert Sedgewick
* @author Kevin Wayne
* @param the type of the values
*/
public class StorageTreeIntMap extends Storage implements StorageIntMap {
/**
* The root of the tree.
*/
private Node root;
/**
* Builds an empty map.
*/
public StorageTreeIntMap() {}
/**
* Creates a map initialized to the same bindings as the given parent map.
*
* @param parent the parent map
*/
public StorageTreeIntMap(Map parent) {
parent.forEach(this::put);
}
/**
* Yields a snapshot of the given map.
*
* @param parent the map
*/
private StorageTreeIntMap(StorageTreeIntMap parent) {
this.root = parent.root;
}
private void mkRootBlack() {
if (isRed(root))
root = Node.mkBlack(root.key, root.value, root.size, root.left, root.right);
}
private void mkRootRed() {
if (isBlack(root))
root = Node.mkRed(root.key, root.value, root.size, root.left, root.right);
}
/**
* A node of the binary search tree that implements the map.
*/
private abstract static class Node extends Storage implements Entry {
protected final int key;
protected final V value; // possibly null
protected final Node left, right;
/**
* Count of the subtree nodes.
*/
protected final int size;
private Node(int key, V value, int size, Node left, Node right) {
this.key = key;
this.value = value;
this.size = size;
this.left = left;
this.right = right;
}
protected static Node mkBlack(int key, V value, int size, Node left, Node right) {
return new BlackNode<>(key, value, size, left, right);
}
protected static Node mkRed(int key, V value, int size, Node left, Node right) {
return new RedNode<>(key, value, size, left, right);
}
protected static Node mkRed(int key, V value) {
return new RedNode<>(key, value, 1, null, null);
}
@Override
public int getKey() {
return key;
}
@Override
public V getValue() {
return value;
}
@Override
public int hashCode() { // unused, but needed to satisfy white-listing for addition of Nodes inside Java collections
return 42;
}
protected abstract Node setValue(V value);
protected abstract Node setLeft(Node left);
protected abstract Node setRight(Node right);
protected abstract Node rotateRight();
protected abstract Node rotateLeft();
protected abstract Node flipColors();
protected abstract Node fixSize();
protected abstract Node flipColor();
private Node moveRedLeft() {
// assert isRed(this) && isBlack(left) && isBlack(left.left);
Node h = flipColors();
return isRed(h.right.left) ? h.setRight(h.right.rotateRight()).rotateLeft().flipColors() : h;
}
private Node moveRedRight() {
// assert isRed(this) && isBlack(right) && isBlack(right.left);
Node h = flipColors();
return isRed(h.left.left) ? h.rotateRight().flipColors() : h;
}
// restore red-black tree invariant
private Node balance() {
Node h = this;
if (isRed(h.right)) h = h.rotateLeft();
if (isRed(h.left) && isRed(h.left.left)) h = h.rotateRight();
if (isRed(h.left) && isRed(h.right)) h = h.flipColors();
return h.fixSize();
}
}
private static class RedNode extends Node {
private RedNode(int key, V value, int size, Node left, Node right) {
super(key, value, size, left, right);
}
@Override
protected Node fixSize() {
return mkRed(key, value, size(left) + size(right) + 1, left, right);
}
@Override
protected Node flipColor() {
return mkBlack(key, value, size, left, right);
}
@Override
protected Node rotateLeft() {
final Node x = right;
Node newThis = mkRed(key, value, size(x.left) + size(left) + 1, left, x.left);
return mkRed(x.key, x.value, size, newThis, x.right);
}
@Override
protected Node rotateRight() {
// assert isRed(left);
final Node x = left;
Node newThis = mkRed(key, value, size(x.right) + size(right) + 1, x.right, right);
return mkRed(x.key, x.value, size, x.left, newThis);
}
@Override
protected Node setValue(V value) {
return mkRed(key, value, size, left, right);
}
@Override
protected Node setLeft(Node left) {
return mkRed(key, value, size, left, right);
}
@Override
protected Node setRight(Node right) {
return mkRed(key, value, size, left, right);
}
@Override
protected Node flipColors() {
// h must have opposite color of its two children
// assert (h != null) && (h.left != null) && (h.right != null);
// assert (isBlack(h) && isRed(h.left) && isRed(h.right))
// || (isRed(h) && isBlack(h.left) && isBlack(h.right));
return mkBlack(key, value, size, left.flipColor(), right.flipColor());
}
}
private static class BlackNode extends Node {
private BlackNode(int key, V value, int size, Node left, Node right) {
super(key, value, size, left, right);
}
@Override
protected Node fixSize() {
return mkBlack(key, value, size(left) + size(right) + 1, left, right);
}
@Override
protected Node flipColor() {
return mkRed(key, value, size, left, right);
}
@Override
protected Node rotateLeft() {
final Node x = right;
Node newThis = mkRed(key, value, size(x.left) + size(left) + 1, left, x.left);
return mkBlack(x.key, x.value, size, newThis, x.right);
}
@Override
protected Node rotateRight() {
final Node x = left;
Node newThis = mkRed(key, value, size(x.right) + size(right) + 1, x.right, right);
return mkBlack(x.key, x.value, size, x.left, newThis);
}
@Override
protected Node setValue(V value) {
return mkBlack(key, value, size, left, right);
}
@Override
protected Node setLeft(Node left) {
return mkBlack(key, value, size, left, right);
}
@Override
protected Node setRight(Node right) {
return mkBlack(key, value, size, left, right);
}
@Override
protected Node flipColors() {
// this must have opposite color of its two children
// assert (left != null) && (right != null);
// assert ( isBlack(this) && isRed(left) && isRed(right))
// || (isRed(this) && isBlack(left) && isBlack(right));
return mkRed(key, value, size, left.flipColor(), right.flipColor());
}
}
/**
* Determines if the given node is red.
*
* @param x the node
* @return true if and only if {@code x} is red
*/
private static boolean isRed(Node x) {
return x instanceof RedNode;
}
/**
* Determines if the given node is black.
*
* @param x the node
* @return true if and only if {@code x} is black
*/
private static boolean isBlack(Node x) {
return x == null || x instanceof BlackNode;
}
/**
* Yields the number of nodes in the subtree rooted at x.
*
* @param x the root of the subtree
* @return the number of nodes. Yields 0 if {@code x} is {@code null}
*/
private static int size(Node x) {
return x == null ? 0 : x.size;
}
@Override
public @View int size() {
return size(root);
}
@Override
public @View boolean isEmpty() {
return root == null;
}
private static int compareTo(int key1, int key2) {
return key1 - key2;
}
@Override
public @View V get(int key) {
return get(root, key);
}
/**
* Yields the value associated with the given key in subtree rooted at x;
*
* @param x the root of the subtree
* @param key the key
* @return the value. Yields {@code null} if the key is not found
*/
private static V get(Node x, int key) {
while (x != null) {
int cmp = compareTo(key, x.key);
if (cmp < 0) x = x.left;
else if (cmp > 0) x = x.right;
else return x.value;
}
return null;
}
@Override
public @View V getOrDefault(int key, V _default) {
return getOrDefault(root, key, _default);
}
private static V getOrDefault(Node x, int key, V _default) {
while (x != null) {
int cmp = compareTo(key, x.key);
if (cmp < 0) x = x.left;
else if (cmp > 0) x = x.right;
else return x.value;
}
return _default;
}
@Override
public V getOrDefault(int key, Supplier _default) {
return getOrDefault(root, key, _default);
}
// value associated with the given key in subtree rooted at x; uses supplier if no such key is found
private static V getOrDefault(Node x, int key, Supplier _default) {
while (x != null) {
int cmp = compareTo(key, x.key);
if (cmp < 0) x = x.left;
else if (cmp > 0) x = x.right;
else return x.value;
}
return _default.get();
}
@Override
public @View boolean containsKey(int key) {
return containsKey(root, key);
}
/**
* Checks if the given key is contained in the subtree rooted at x.
*
* @param x the root of the subtree
* @param key the key
* @return true if and only if that condition holds
*/
private static boolean containsKey(Node x, int key) {
while (x != null) {
int cmp = compareTo(key, x.key);
if (cmp < 0) x = x.left;
else if (cmp > 0) x = x.right;
else return true;
}
return false;
}
@Override
public void put(int key, V value) {
root = put(root, key, value);
mkRootBlack();
}
// insert the key-value pair in the subtree rooted at h
private static Node put(Node h, int key, V value) {
if (h == null) return Node.mkRed(key, value);
int cmp = compareTo(key, h.key);
if (cmp < 0) h = h.setLeft(put(h.left, key, value));
else if (cmp > 0) h = h.setRight(put(h.right, key, value));
else h = h.setValue(value);
// fix-up any right-leaning links
if (isRed(h.right) && isBlack(h.left)) h = h.rotateLeft();
if (isRed(h.left) && isRed(h.left.left)) h = h.rotateRight();
if (isRed(h.left) && isRed(h.right)) h = h.flipColors();
return h.fixSize();
}
@Override
public void removeMin() {
if (isEmpty()) throw new NoSuchElementException();
// if both children of root are black, set root to red
if (isBlack(root.left) && isBlack(root.right))
mkRootRed();
root = removeMin(root);
if (!isEmpty()) mkRootBlack();
}
// delete the key-value pair with the minimum key rooted at h
private static Node removeMin(Node h) {
if (h.left == null)
return null;
if (isBlack(h.left) && isBlack(h.left.left))
h = h.moveRedLeft();
return h.setLeft(removeMin(h.left)).balance();
}
@Override
public void removeMax() {
if (isEmpty()) throw new NoSuchElementException();
// if both children of root are black, set root to red
if (isBlack(root.left) && isBlack(root.right))
mkRootRed();
root = removeMax(root);
if (!isEmpty()) mkRootBlack();
}
// delete the key-value pair with the maximum key rooted at h
private static Node removeMax(Node h) {
if (isRed(h.left))
h = h.rotateRight();
if (h.right == null)
return null;
if (isBlack(h.right) && isBlack(h.right.left))
h = h.moveRedRight();
return h.setRight(removeMax(h.right)).balance();
}
@Override
public void remove(int key) {
if (containsKey(key)) {
// if both children of root are black, set root to red
if (isBlack(root.left) && isBlack(root.right))
mkRootRed();
root = remove(root, key);
if (!isEmpty()) mkRootBlack();
}
}
// delete the key-value pair with the given key rooted at h
private static Node remove(Node h, int key) {
if (compareTo(key, h.key) < 0) {
if (isBlack(h.left) && isBlack(h.left.left))
h = h.moveRedLeft();
h = h.setLeft(remove(h.left, key));
}
else {
if (isRed(h.left))
h = h.rotateRight();
if (compareTo(key, h.key) == 0 && (h.right == null))
return null;
if (isBlack(h.right) && isBlack(h.right.left))
h = h.moveRedRight();
if (compareTo(key, h.key) == 0) {
var x = min(h.right);
if (isRed(h))
h = Node.mkRed(x.key, x.value, h.size, h.left, removeMin(h.right));
else
h = Node.mkBlack(x.key, x.value, h.size, h.left, removeMin(h.right));
}
else
h = h.setRight(remove(h.right, key));
}
return h.balance();
}
@Override
public @View int min() {
if (isEmpty()) throw new NoSuchElementException("call to min() with empty symbol table");
return min(root).key;
}
// the smallest key in subtree rooted at x
private static Node min(Node x) {
if (x.left == null) return x;
else return min(x.left);
}
@Override
public @View int max() {
if (isEmpty()) throw new NoSuchElementException("call to max() with empty symbol table");
return max(root).key;
}
// the largest key in the subtree rooted at x
private static Node max(Node x) {
if (x.right == null) return x;
else return max(x.right);
}
@Override
public @View int floorKey(int key) {
if (isEmpty()) throw new NoSuchElementException();
var x = floorKey(root, key);
if (x == null) throw new NoSuchElementException();
else return x.key;
}
// the largest key in the subtree rooted at x less than or equal to the given key
private static Node floorKey(Node x, int key) {
if (x == null) return null;
int cmp = compareTo(key, x.key);
if (cmp == 0) return x;
if (cmp < 0) return floorKey(x.left, key);
var t = floorKey(x.right, key);
if (t != null) return t;
else return x;
}
@Override
public @View int ceilingKey(int key) {
if (isEmpty()) throw new NoSuchElementException();
var x = ceilingKey(root, key);
if (x == null) throw new NoSuchElementException();
else return x.key;
}
// the smallest key in the subtree rooted at x greater than or equal to the given key
private static Node ceilingKey(Node x, int key) {
if (x == null) return null;
int cmp = compareTo(key, x.key);
if (cmp == 0) return x;
if (cmp > 0) return ceilingKey(x.right, key);
var t = ceilingKey(x.left, key);
if (t != null) return t;
else return x;
}
@Override
public @View int select(int k) {
if (k < 0 || k >= size()) throw new IllegalArgumentException("argument to select() is invalid: " + k);
return select(root, k).key;
}
// the key of rank k in the subtree rooted at x
private static Node select(Node x, int k) {
int t = size(x.left);
if (t > k) return select(x.left, k);
else if (t < k) return select(x.right, k-t-1);
else return x;
}
@Override
public @View int rank(int key) {
return rank(key, root);
}
// number of keys less than key in the subtree rooted at x
private static int rank(int key, Node x) {
if (x == null) return 0;
int cmp = compareTo(key, x.key);
if (cmp < 0) return rank(key, x.left);
else if (cmp > 0) return 1 + size(x.left) + rank(key, x.right);
else return size(x.left);
}
@Override
public void update(int key, UnaryOperator how) {
root = update(root, key, how);
mkRootBlack();
}
private static Node update(Node h, int key, UnaryOperator how) {
if (h == null) return Node.mkRed(key, how.apply(null));
int cmp = compareTo(key, h.key);
if (cmp < 0) h = h.setLeft(update(h.left, key, how));
else if (cmp > 0) h = h.setRight(update(h.right, key, how));
else h = h.setValue(how.apply(h.value));
// fix-up any right-leaning links
if (isRed(h.right) && isBlack(h.left)) h = h.rotateLeft();
if (isRed(h.left) && isRed(h.left.left)) h = h.rotateRight();
if (isRed(h.left) && isRed(h.right)) h = h.flipColors();
return h.fixSize();
}
@Override
public void update(int key, V _default, UnaryOperator how) {
root = update(root, key, _default, how);
mkRootBlack();
}
private static Node update(Node h, int key, V _default, UnaryOperator how) {
if (h == null) return Node.mkRed(key, how.apply(_default));
int cmp = compareTo(key, h.key);
if (cmp < 0) h = h.setLeft(update(h.left, key, _default, how));
else if (cmp > 0) h = h.setRight(update(h.right, key, _default, how));
else if (h.value == null)
h = h.setValue(how.apply(_default));
else
h = h.setValue(how.apply(h.value));
// fix-up any right-leaning links
if (isRed(h.right) && isBlack(h.left)) h = h.rotateLeft();
if (isRed(h.left) && isRed(h.left.left)) h = h.rotateRight();
if (isRed(h.left) && isRed(h.right)) h = h.flipColors();
return h.fixSize();
}
@Override
public void update(int key, Supplier _default, UnaryOperator how) {
root = update(root, key, _default, how);
mkRootBlack();
}
private static Node update(Node h, int key, Supplier _default, UnaryOperator how) {
if (h == null) return Node.mkRed(key, how.apply(_default.get()));
int cmp = compareTo(key, h.key);
if (cmp < 0) h = h.setLeft(update(h.left, key, _default, how));
else if (cmp > 0) h = h.setRight(update(h.right, key, _default, how));
else if (h.value == null)
h = h.setValue(how.apply(_default.get()));
else
h = h.setValue(how.apply(h.value));
// fix-up any right-leaning links
if (isRed(h.right) && isBlack(h.left)) h = h.rotateLeft();
if (isRed(h.left) && isRed(h.left.left)) h = h.rotateRight();
if (isRed(h.left) && isRed(h.right)) h = h.flipColors();
return h.fixSize();
}
@Override
public V putIfAbsent(int key, V value) {
class PutIfAbsent {
private V result;
private Node putIfAbsent(Node h) {
// not found: result remains null
if (h == null)
// not found
return Node.mkRed(key, value);
int cmp = compareTo(key, h.key);
if (cmp < 0) h = h.setLeft(putIfAbsent(h.left));
else if (cmp > 0) h = h.setRight(putIfAbsent(h.right));
else if (h.value == null)
// found but was bound to null: result remains null
return h.setValue(value);
else {
// found and was bound to a non-null value
result = h.value;
return h;
}
// fix-up any right-leaning links
if (isRed(h.right) && isBlack(h.left)) h = h.rotateLeft();
if (isRed(h.left) && isRed(h.left.left)) h = h.rotateRight();
if (isRed(h.left) && isRed(h.right)) h = h.flipColors();
return h.fixSize();
}
}
PutIfAbsent pia = new PutIfAbsent();
root = pia.putIfAbsent(root);
mkRootBlack();
return pia.result;
}
@Override
public V computeIfAbsent(int key, Supplier supplier) {
class ComputeIfAbsent {
private V result;
private Node computeIfAbsent(Node h) {
if (h == null)
// not found
return Node.mkRed(key, result = supplier.get());
int cmp = compareTo(key, h.key);
if (cmp < 0) h = h.setLeft(computeIfAbsent(h.left));
else if (cmp > 0) h = h.setRight(computeIfAbsent(h.right));
else if (h.value == null) {
// found but was bound to null
h = h.setValue(supplier.get());
result = h.value;
return h;
}
else {
// found and was bound to a non-null value
result = h.value;
return h;
}
// fix-up any right-leaning links
if (isRed(h.right) && isBlack(h.left)) h = h.rotateLeft();
if (isRed(h.left) && isRed(h.left.left)) h = h.rotateRight();
if (isRed(h.left) && isRed(h.right)) h = h.flipColors();
return h.fixSize();
}
}
ComputeIfAbsent cia = new ComputeIfAbsent();
root = cia.computeIfAbsent(root);
mkRootBlack();
return cia.result;
}
@Override
public V computeIfAbsent(int key, IntFunction supplier) {
class ComputeIfAbsent {
private V result;
private Node computeIfAbsent(Node h) {
if (h == null)
// not found
return Node.mkRed(key, result = supplier.apply(key));
int cmp = compareTo(key, h.key);
if (cmp < 0) h = h.setLeft(computeIfAbsent(h.left));
else if (cmp > 0) h = h.setRight(computeIfAbsent(h.right));
else if (h.value == null) {
// found but was bound to null
h = h.setValue(supplier.apply(key));
result = h.value;
return h;
}
else {
// found and was bound to a non-null value
result = h.value;
return h;
}
// fix-up any right-leaning links
if (isRed(h.right) && isBlack(h.left)) h = h.rotateLeft();
if (isRed(h.left) && isRed(h.left.left)) h = h.rotateRight();
if (isRed(h.left) && isRed(h.right)) h = h.flipColors();
return h.fixSize();
}
}
ComputeIfAbsent cia = new ComputeIfAbsent();
root = cia.computeIfAbsent(root);
mkRootBlack();
return cia.result;
}
@Override
public void clear() {
root = null;
}
@Override
public Iterator> iterator() {
return new StorageMapIterator<>(root);
}
private static class StorageMapIterator implements Iterator> {
// the path under enumeration; it holds that the left children
// have already been enumerated
private final List> stack = new ArrayList<>();
private StorageMapIterator(Node root) {
// initially, the stack contains the leftmost path of the tree
for (var cursor = root; cursor != null; cursor = cursor.left)
stack.add(cursor);
}
@Override
public boolean hasNext() {
return !stack.isEmpty();
}
@Override
public Entry next() {
var topmost = stack.remove(stack.size() - 1);
// we add the leftmost path of the right child of topmost
for (var cursor = topmost.right; cursor != null; cursor = cursor.left)
stack.add(cursor);
return topmost;
}
}
@Override
public Stream> stream() {
return StreamSupport.stream(spliterator(), false);
}
@Override
public List keyList() {
List keys = new ArrayList<>();
if (root != null)
keyList(root, keys);
return keys;
}
private static void keyList(Node x, List keys) {
if (x.left != null)
keyList(x.left, keys);
keys.add(x.key);
if (x.right != null)
keyList(x.right, keys);
}
@Override
public IntStream keys() {
return stream().mapToInt(Entry::getKey);
}
@Override
public Stream values() {
return stream().map(Entry::getValue);
}
@Override
public StorageIntMapView view() {
/**
* A read-only view of a parent storage map. A view contains the same bindings
* as the parent storage map, but does not include modification methods.
* Moreover, a view is exported, so that it can be safely divulged outside
* the store of a node. Calls to the view are simply forwarded to the parent map.
*/
@Exported
class StorageIntMapViewImpl extends Storage implements StorageIntMapView {
@Override
public @View int size() {
return StorageTreeIntMap.this.size();
}
@Override
public @View boolean isEmpty() {
return StorageTreeIntMap.this.isEmpty();
}
@Override
public Iterator> iterator() {
return StorageTreeIntMap.this.iterator();
}
@Override
public V get(int key) {
return StorageTreeIntMap.this.get(key);
}
@Override
public V getOrDefault(int key, V _default) {
return StorageTreeIntMap.this.getOrDefault(key, _default);
}
@Override
public V getOrDefault(int key, Supplier _default) {
return StorageTreeIntMap.this.getOrDefault(key, _default);
}
@Override
public boolean containsKey(int key) {
return StorageTreeIntMap.this.containsKey(key);
}
@Override
public int min() {
return StorageTreeIntMap.this.min();
}
@Override
public int max() {
return StorageTreeIntMap.this.max();
}
@Override
public int floorKey(int key) {
return StorageTreeIntMap.this.floorKey(key);
}
@Override
public int ceilingKey(int key) {
return StorageTreeIntMap.this.ceilingKey(key);
}
@Override
public int select(int k) {
return StorageTreeIntMap.this.select(k);
}
@Override
public int rank(int key) {
return StorageTreeIntMap.this.rank(key);
}
@Override
public String toString() {
return StorageTreeIntMap.this.toString();
}
@Override
public Stream> stream() {
return StorageTreeIntMap.this.stream();
}
@Override
public List keyList() {
return StorageTreeIntMap.this.keyList();
}
@Override
public IntStream keys() {
return StorageTreeIntMap.this.keys();
}
@Override
public StorageIntMapView snapshot() {
return StorageTreeIntMap.this.snapshot();
}
@Override
public Stream values() {
return StorageTreeIntMap.this.values();
}
}
return new StorageIntMapViewImpl();
}
@Override
public StorageIntMapView snapshot() {
return new StorageTreeIntMap<>(this).view();
}
}