src/main/java/net/onrc/onos/apps/bgproute/PatriciaTrie.java

package net.onrc.onos.apps.bgproute;

import java.util.Iterator;
import java.util.NoSuchElementException;

public class PatriciaTrie<V> implements IPatriciaTrie<V> {
    private final byte[] maskBits = {(byte) 0x00, (byte) 0x80, (byte) 0xc0, (byte) 0xe0, (byte) 0xf0,
            (byte) 0xf8, (byte) 0xfc, (byte) 0xfe, (byte) 0xff};

    private int maxPrefixLength;

    private Node top;

    public PatriciaTrie(int maxPrefixLength) {
        this.maxPrefixLength = maxPrefixLength;
    }

    @Override
    public synchronized V put(Prefix prefix, V value) {
        if (prefix == null || value == null) {
            throw new NullPointerException();
        }

        if (prefix.getPrefixLength() > maxPrefixLength) {
            throw new IllegalArgumentException(String.format(
                    "Prefix length %d is greater than max prefix length %d",
                    prefix.getPrefixLength(), maxPrefixLength));
        }

        Node node = top;
        Node match = null;

        while (node != null
                && node.prefix.getPrefixLength() <= prefix.getPrefixLength()
                && key_match(node.prefix.getAddress(), node.prefix.getPrefixLength(), prefix.getAddress(), prefix.getPrefixLength()) == true) {
            if (node.prefix.getPrefixLength() == prefix.getPrefixLength()) {
                /*
                 * Prefix is already in tree. This may be an aggregate node, in which case
                 * we are inserting a new prefix, or it could be an actual node, in which
                 * case we are inserting a new nexthop for the prefix and should return
                 * the old nexthop.
                 */
                V oldValue = node.value;
                node.value = value;
                return oldValue;
            }

            match = node;

            if (bit_check(prefix.getAddress(), node.prefix.getPrefixLength()) == true) {
                node = node.right;
            } else {
                node = node.left;
            }
        }

        Node add = null;

        if (node == null) {
            //add = new Node(p, r);
            add = new Node(prefix);
            add.value = value;

            if (match != null) {
                node_link(match, add);
            } else {
                top = add;
            }
        } else {
            add = node_common(node, prefix.getAddress(), prefix.getPrefixLength());
            if (add == null) {
                //I think this is -ENOMEM?
                //return null;
            }

            if (match != null) {
                node_link(match, add);
            } else {
                top = add;
            }
            node_link(add, node);

            if (add.prefix.getPrefixLength() != prefix.getPrefixLength()) {
                match = add;

                //add = new Node(p, r);
                add = new Node(prefix);
                add.value = value;
                node_link(match, add);
            } else {
                add.value = value;
            }
        }

        //If we added a new Node, there was no previous mapping
        return null;
        //return addReference(add);
    }

    /*exact match*/
    @Override
    public synchronized V lookup(Prefix prefix) {
        if (prefix.getPrefixLength() > maxPrefixLength) {
            return null;
        }

        /*
        Node node = top;

        while (node != null
                && node.prefix.getPrefixLength() <= p.getPrefixLength()
                && key_match(node.prefix.getAddress(), node.prefix.getPrefixLength(), p.getAddress(), p.getPrefixLength()) == true) {
            if (node.prefix.getPrefixLength() == p.getPrefixLength()) {
                //return addReference(node);
                return node.rib;
            }

            if (bit_check(p.getAddress(), node.prefix.getPrefixLength()) == true) {
                node = node.right;
            } else {
                node = node.left;
            }
        }
        */

        Node node = findNode(prefix);

        return node == null ? null : node.value;
    }

    /*closest containing prefix*/
    @Override
    public synchronized V match(Prefix prefix) {
        //TODO
        if (prefix.getPrefixLength() > maxPrefixLength) {
            return null;
        }

        Node closestNode = findClosestNode(prefix);

        return closestNode == null ? null : closestNode.value;
    }

    @Override
    public synchronized boolean remove(Prefix prefix, V value) {
        Node child;
        Node parent;

        if (prefix == null || value == null) {
            return false;
        }

        Node node = findNode(prefix);

        if (node == null || node.isAggregate() || !node.value.equals(value)) {
            //Given <prefix, nexthop> mapping is not in the tree
            return false;
        }

        if (node.left != null && node.right != null) {
            //Remove the RibEntry entry and leave this node as an aggregate node
            //In the future, maybe we should re-evaluate what the aggregate prefix should be?
            //It shouldn't necessarily stay the same.
            //More complicated if the above prefix is also aggregate.
            node.value = null;
            return true;
        }

        if (node.left != null) {
            child = node.left;
        } else {
            child = node.right;
        }

        parent = node.parent;

        if (child != null) {
            child.parent = parent;
        }

        if (parent != null) {
            if (parent.left == node) {
                parent.left = child;
            } else {
                parent.right = child;
            }
        } else {
            top = child;
        }

        /*
         * TODO not sure what to do here. I think this is lazily deleting aggregate nodes,
         * notice that it used to do nothing if it detected both children were not null earlier.
         * But here, what we really should do is reevaluate the aggregate prefix of the parent
         * node (if it is indeed an aggregate). Because at the moment, no aggregate node will ever
         * be removed. BUT, I don't actually think this presents a correctness problem, at
         * least from an external point of view.
         */
        //if (parent != null && parent.refCount == 0) {
        //node_remove(parent);
        //}

        return true;
    }

    @Override
    public Iterator<Entry<V>> iterator() {
        return new PatriciaTrieIterator(top);
    }

    private Node findNode(Prefix prefix) {
        Node node = top;

        while (node != null
                && node.prefix.getPrefixLength() <= prefix.getPrefixLength()
                && key_match(node.prefix.getAddress(), node.prefix.getPrefixLength(), prefix.getAddress(), prefix.getPrefixLength()) == true) {
            if (node.prefix.getPrefixLength() == prefix.getPrefixLength()) {
                //return addReference(node);
                return node;
            }

            if (bit_check(prefix.getAddress(), node.prefix.getPrefixLength()) == true) {
                node = node.right;
            } else {
                node = node.left;
            }
        }

        return null;
    }

    private Node findClosestNode(Prefix prefix) {
        Node node = top;
        Node match = null;

        while (node != null
                && node.prefix.getPrefixLength() <= prefix.getPrefixLength()
                && key_match(node.prefix.getAddress(), node.prefix.getPrefixLength(), prefix.getAddress(), prefix.getPrefixLength()) == true) {
            if (!node.isAggregate()) {
                match = node;
            }

            if (bit_check(prefix.getAddress(), node.prefix.getPrefixLength()) == true) {
                node = node.right;
            } else {
                node = node.left;
            }
        }

        return match;
    }

    /*
     * Receives a 1-based bit index
     * Returns a 1-based byte index
     * eg. (0 => 1), 1 => 1, 8 => 1, 9 => 2, 17 => 3
     */
    private int getByteContainingBit(int bitNumber) {
        return Math.max((bitNumber + 7) / 8, 1);
    }

    private boolean key_match(byte[] key1, int key1_len, byte[] key2, int key2_len) {
        //int offset;
        //int shift;

        if (key1_len > key2_len) {
            return false;
        }

        int offset = (Math.min(key1_len, key2_len)) / 8;
        int shift = (Math.min(key1_len, key2_len)) % 8;

        if (shift != 0) {
            if ((maskBits[shift] & (key1[offset] ^ key2[offset])) != 0) {
                return false;
            }
        }

        while (offset != 0) {
            offset--;
            if (key1[offset] != key2[offset]) {
                return false;
            }
        }
        return true;
    }

    private boolean bit_check(byte[] key, int key_bits) {
        int offset = key_bits / 8;
        int shift = 7 - (key_bits % 8);
        int bit = key[offset] & 0xff;

        bit >>= shift;

        if ((bit & 1) == 1) {
            return true;
        } else {
            return false;
        }
    }

    private void node_link(Node node, Node add) {
        boolean bit = bit_check(add.prefix.getAddress(), node.prefix.getPrefixLength());

        if (bit == true) {
            node.right = add;
        } else {
            node.left = add;
        }
        add.parent = node;
    }

    private Node node_common(Node node, byte[] key, int key_bits) {
        int i;
        int limit = Math.min(node.prefix.getPrefixLength(), key_bits) / 8;

        for (i = 0; i < limit; i++) {
            if (node.prefix.getAddress()[i] != key[i]) {
                break;
            }
        }

        int common_len = i * 8;
        int boundary = 0;

        if (common_len != key_bits) {
            byte diff = (byte) (node.prefix.getAddress()[i] ^ key[i]);
            byte mask = (byte) 0x80;
            int shift_mask = 0;

            while (common_len < key_bits && ((mask & diff) == 0)) {
                boundary = 1;

                shift_mask = (mask & 0xff);
                shift_mask >>= 1;
                mask = (byte) shift_mask;

                common_len++;
            }
        }

        //Node add = new Node(null, common_len, maxKeyOctets);
        //if (add == null)
        //Another -ENOMEM;
        //return null;

        //Creating a new Prefix with a prefix length of common_len
        //Bits are copied from node's up until the common_len'th bit
        //RibEntry is null, because this is an aggregate prefix - it's not
        //actually been added to the trie.

        byte[] newPrefix = new byte[getByteContainingBit(maxPrefixLength)];

        int j;
        for (j = 0; j < i; j++)
            newPrefix[j] = node.prefix.getAddress()[j];

        if (boundary != 0)
            newPrefix[j] = (byte) (node.prefix.getAddress()[j] & maskBits[common_len % 8]);

        //return new Node(new Prefix(newPrefix, common_len), null);
        return new Node(new Prefix(newPrefix, common_len));
        //return add;
    }

    private class Node {
        public Node parent = null;
        public Node left = null;
        public Node right = null;

        public final Prefix prefix;
        public V value;

        //public Node(Prefix p, RibEntry r) {
        //      this.prefix = p;
        //      this.rib = r;
        //}
        public Node(Prefix p) {
            this.prefix = p;
        }

        public boolean isAggregate() {
            return value == null;
        }

        public Entry<V> getEntry() {
            return new PatriciaTrieEntry(prefix, value);
        }
    }

    private class PatriciaTrieEntry implements Entry<V> {
        private Prefix prefix;
        private V value;

        public PatriciaTrieEntry(Prefix prefix, V value) {
            this.prefix = prefix;
            this.value = value;
        }

        @Override
        public Prefix getPrefix() {
            return prefix;
        }

        @Override
        public V getValue() {
            return value;
        }
    }

    private class PatriciaTrieIterator implements Iterator<Entry<V>> {

        private Node current;
        private boolean started = false;

        public PatriciaTrieIterator(Node start) {
            current = start;

            //If the start is an aggregate node fast forward to find the next valid node
            if (current != null && current.isAggregate()) {
                current = findNext(current);
            }
        }

        @Override
        public boolean hasNext() {
            if (current == null) {
                return false;
            }

            if (!started) {
                return true;
            }

            return findNext(current) != null;
        }

        @Override
        public Entry<V> next() {
            if (current == null) {
                throw new NoSuchElementException();
            }

            if (!started) {
                started = true;
                return current.getEntry();
            }

            current = findNext(current);
            if (current == null) {
                throw new NoSuchElementException();
            }

            return current.getEntry();
        }

        @Override
        public void remove() {
            // TODO This could be implemented, if it were needed
            throw new NoSuchElementException();
        }

        private Node findNext(Node node) {
            Node next = null;

            if (node.left != null) {
                next = node.left;
                //addReference(next);
                //delReference(node);
                //return next;
            } else if (node.right != null) {
                next = node.right;
                //addReference(next);
                //delReference(node);
                //return next;
            } else {
                //Node start = node;
                while (node.parent != null) {
                    if (node.parent.left == node && node.parent.right != null) {
                        next = node.parent.right;
                        //addReference(next);
                        //delReference(start);
                        //return next;
                        break;
                    }
                    node = node.parent;
                }
            }

            if (next == null) {
                return null;
            }

            //If the node doesn't have a value, it's not an actual node, it's an artifically
            //inserted aggregate node. We don't want to return these to the user.
            if (next.isAggregate()) {
                return findNext(next);
            }

            return next;
        }
    }
}