//===========================================================================
//  $Name: arts++-1-1-a12 $
//  $Id: Ipv4PrefixPatricia.hh.in,v 1.5 2004/12/14 02:23:08 kkeys Exp $
//
//  Ken Keys
//  CAIDA
//  December 2002
//===========================================================================

extern "C" {
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdlib.h>

#include "caida_t.h"
}

#include <utility>

#include "Ipv4Network.hh"

#undef WORDS_BIGENDIAN

//----------------------------------------------------------------------------
//                          template <class Type>
//                          class Ipv4PrefixPatricia
//----------------------------------------------------------------------------
//  This template class implements a Patricia tree keyed by IPv4
//  prefix. It is particularly useful for storing any type of data keyed
//  by IP address.  Because it's Patricia, exact match searches are
//  relatively speedy.  So are longest-match searches.
//
//  The interfaces here are very STL-like.  For the most part, missing
//  features are missing only because they weren't needed in skitter.  Unlike
//  dwm's implementation, no features are left out for performance reasons; in
//  particular, all iterators _can_ be incremented and decremented (and
//  iterating from begin() to end() has performance equal to that of visit()).
//
//  Iterator traversal is preorder, not inorder; so it would give, e.g.,
//  42.0.0.0/8 42.0.0.0/16 42.128.0.0/16
//  whereas dwm's implementation would have given the less useful
//  42.0.0.0/16 42.0.0.0/8 42.128.0.0/16.
//
//  operator[] works just like it does for STL's <map>: it will insert if the
//  key was not found.  This permits using it on the left hand side of a
//  statement for assignment of a value in the tree.
//
//  erase() members work as expected, removing a key or value from the
//  tree.  They do not invalidate any iterators other than those pointing to
//  the values being erased.
//
//  TODO:  aggregate() member that eliminates prefixes inside other prefixes
//  and joins neighboring prefixes into a single prefix to produce the
//  smallest set of prefixes that covers the same space as the original set.
//  (But what do we do with the data when we merge/delete prefixes?)
//
//  Ken Keys
//  CAIDA
//  December 2002
//
//----------------------------------------------------------------------------

#define DEBUGGING 0
#if !DEBUGGING
#define IPPDATATYPE Type
#define Ipv4PrefixPatriciaClass Ipv4PrefixPatricia<Type>
template <class Type>
#else
class PrefixData; // forward decl
#define IPPDATATYPE PrefixData*
#define Ipv4PrefixPatriciaClass Ipv4PrefixPatricia
#endif
class Ipv4PrefixPatricia
{
public:
    typedef Ipv4Network				key_type;
    typedef IPPDATATYPE				data_type;
    typedef std::pair<const key_type, data_type>	value_type;
    typedef uint32_t				size_type;

private:
    enum { useNode = 0x1, useLeft = 0x2, useRight = 0x4 };

    class Node {
    public:
	typename Ipv4PrefixPatricia::value_type	val;
	bool		has_data;   // val.second (i.e., data) is valid
	Node		*left;	    // left child node
	Node		*right;	    // right child node
	Node		*parent;    // parent node

	Node() : val(), left(), right(), parent() {};

	Node(const Ipv4PrefixPatricia::value_type &value, Node *p = 0) :
	    val(value), has_data(true), left(), right(), parent(p) { };

	Node(const Ipv4PrefixPatricia::key_type &key, Node *p = 0) :
	    val(value_type(key, data_type())), has_data(false),
	    left(), right(), parent(p) { };

	// deep-copy constructor from the tree rooted at src, with newparent as
	// the new tree's parent node.
	Node(const Node *src, Node *newparent) :
	    val(src->val), has_data(src->has_data), parent(newparent)
	{
	    this->left = src->left ? new Node(src->left, this) : NULL;
	    this->right = src->right ? new Node(src->right, this) : NULL;
	}

	// like Ipv4Network::Matches(), but also handles maskLen 0
	bool Matches(ipv4addr_t addr) const {
	    return (val.first.maskLen == 0) || val.first.Matches(addr);
	}

	void clear() {
	    if (left) { left->clear(); delete left; }
	    if (right) { right->clear(); delete right; }
	}
    };

    size_type	_size;
    Node	*root;

public:
    class iterator
    {
	Node *node;
	friend class Ipv4PrefixPatriciaClass;

    public:
	// constructor
	iterator(Node *n = 0) : node(n) {};

	// destructor
	virtual ~iterator() { node = (Node *)0; }

	virtual iterator & operator= (const iterator & src)
	{
	    node = src.node;
	    return *this;
	}

	iterator & operator++ ()  // prefix ++
	{
	    if (!node) return *this;
	    unsigned int state = useLeft | useRight;

	    while (1) {
		if ((state & useNode) && node->has_data) {
		    return *this;
		} else if ((state & useLeft) && node->left) {
		    state = useNode | useLeft | useRight;
		    node = node->left;
		} else if ((state & useRight) && node->right) {
		    state = useNode | useLeft | useRight;
		    node = node->right;
		} else if (node->parent) {
		    state = 0;
		    if (node == node->parent->left)
			state |= useRight;
		    node = node->parent;
		} else {
		    node = NULL;
		    return *this;
		}
	    }
	}

	iterator & operator-- ()  // prefix --
	{
	    if (!node) return *this;
	    unsigned int state = 0;

	    while (1) {
		if ((state & useRight) && node->right) {
		    state = useNode | useLeft | useRight;
		    node = node->right;
		} else if ((state & useLeft) && node->left) {
		    state = useNode | useLeft | useRight;
		    node = node->left;
		} else if ((state & useNode) && node->has_data) {
		    return *this;
		} else if (node->parent) {
		    state = useNode;
		    if (node == node->parent->right)
			state |= useLeft;
		    node = node->parent;
		} else {
		    node = NULL;
		    return *this;
		}
	    }
	}

	void operator++ (int)  // postfix ++
	{
	    ++*this;
	}

	void operator-- (int)  // postfix --
	{
	    --*this;
	}

	virtual value_type & operator* ()
	{
	    return node->val;
	}

	virtual const value_type & operator* () const
	{
	    return node->val;
	}

	bool operator!= (const iterator & rhs) const
	{
	    return this->node != rhs.node;
	}

	bool operator== (const iterator & rhs) const
	{
	    return this->node == rhs.node;
	}
    };

    class reverse_iterator
    {
	iterator fwd;
    public:
	reverse_iterator() : fwd() {};
	reverse_iterator(Node *node) : fwd(node) {};
	virtual ~reverse_iterator() {};
	reverse_iterator & operator++ () { --fwd; return *this; }
	void operator++ (int) { ++*this; }
	reverse_iterator & operator-- () { ++fwd; return *this; }
	void operator-- (int) { --*this; }
	const value_type & operator* () const { return *fwd; }
	bool operator!= (const reverse_iterator & rit) const
	    { return fwd != rit.fwd; }
	bool operator== (const reverse_iterator & rit) const
	    { return fwd == rit.fwd; }
    };

    class const_iterator
    {
	iterator it;
    public:
	const_iterator() : it() {};
	const_iterator(Node *node) : it(node) {};
	~const_iterator() {};
	const_iterator & operator++ () { ++it; return *this; }
	const_iterator & operator-- () { --it; return *this; }
	const_iterator & operator++ (int) { ++it; }
	const_iterator & operator-- (int) { --it; }
	const value_type & operator* () const { return *it; }
	bool operator!= (const iterator & rhs) const
	    { return it != rhs; }
	bool operator!= (const const_iterator & rhs) const
	    { return it != rhs.it; }
	bool operator== (const iterator & rhs) const
	    { return it == it; }
	bool operator== (const const_iterator & rhs) const
	    { return it == rhs.it; }
    };

    class const_reverse_iterator
    {
	reverse_iterator rit;
    public:
	const_reverse_iterator() : rit() {};
	const_reverse_iterator(Node *node) : rit(node) {};
	~const_reverse_iterator() {};
	const_reverse_iterator & operator++ () { ++rit; return *this; }
	const_reverse_iterator & operator-- () { --rit; return *this; }
	const_reverse_iterator & operator++ (int) { ++rit; }
	const_reverse_iterator & operator-- (int) { --rit; }
	const value_type & operator* () const { return *rit; }
	bool operator!= (const reverse_iterator & rhs) const
	    { return rit != rhs; }
	bool operator!= (const const_reverse_iterator & rhs) const
	    { return rit != rhs.rit; }
	bool operator== (const reverse_iterator & rhs) const
	    { return rit == rit; }
	bool operator== (const const_reverse_iterator & rhs) const
	    { return rit == rhs.rit; }
    };

    // constructor
    Ipv4PrefixPatricia() : _size(), root() {};

    // destructor
    ~Ipv4PrefixPatriciaClass() { clear(); }

    void clear() { if (root) { root->clear(); delete root; } }

    // operator= to deep-copy an Ipv4PrefixPatricia<Type>.
    Ipv4PrefixPatricia & operator= (Ipv4PrefixPatriciaClass & src)
    {
	root = new Node(src.root, NULL);
	_size = src._size;
	return *this;
    }

    inline iterator begin()
    {
	iterator it(root);
	return (root->has_data) ? it : ++it;
    }

    inline iterator end()
    {
	iterator it(NULL);
	return it;
    }

    inline const_iterator begin() const
    {
	const_iterator it(root);
	return (root->has_data) ? it : ++it;
    }

    inline const_iterator end() const
    {
	const_iterator it(NULL);
	return it;
    }

    inline reverse_iterator rbegin()
    {
	Node *node = root;
	if (node) {
	    while (node->right || node->left)
		node = node->right ? node->right : node->left;
	    assert(node->has_data);
	}
	return reverse_iterator(node);
    }

    inline reverse_iterator rend()
    {
	reverse_iterator it(NULL);
	return it;
    }

    iterator find(const key_type & key)
    {
	Node *node = root;

	while (node && node->Matches(key.net)) {
	    if (node->has_data && node->val.first.maskLen == key.maskLen)
		return iterator(node);
	    if (node->val.first.maskLen >= key.maskLen)
		break;
	    node = (bitset(key.net, 31 - node->val.first.maskLen)) ?
		node->right : node->left;
	}

	return iterator(NULL);
    }

    size_type count(const key_type & key) { return find(key) == end() ? 0 : 1; }

    //------------------------------------------------------------------------
    //         inline iterator LongestMatch(const ipv4addr_t & ipv4addr)
    //........................................................................
    //!  This member function permits searching for a longest-match IP
    //!  address.  It returns an iterator pointing to the matching entry if
    //!  a match is found, else returns end().
    //------------------------------------------------------------------------
    inline iterator LongestMatch(const ipv4addr_t & ipv4addr, int len = 32)
    {
        Node *node, *match;

        node = root;
	match = NULL;

        while (node && node->val.first.maskLen <= len &&
	    node->Matches(ipv4addr))
	{
	    if (node->has_data)
		match = node;
	    if (node->val.first.maskLen == len)
		return iterator(match);
	    node = (bitset(ipv4addr, 31 - node->val.first.maskLen)) ?
		node->right : node->left;
        }

        return iterator(match);
    }


    bool empty() const { return root == 0; }

    const uint32_t size() const { return _size; }

    IPPDATATYPE & operator[] (const key_type & key)
    {
        return (*((insert(value_type(key, data_type()))).first)).second;
    }

    // kept for backward compatability
    inline void visit(void (*visitFunc)(iterator &, void *),
        void *other)
    {
	for (iterator it = begin(); it != end(); it++)
	    visitFunc(it, other);
    }

    size_type erase(const key_type & key)
    {
	iterator it = find(key);

	if (it == end()) return 0;
	erase(it);
	return 1;
    }

    // does not invalidate iterators other than itself
    void erase(iterator it)
    {
	Node *node = it.node;
	Node **inheritor;

	assert(node->has_data);
	node->has_data = false;
	_size--;
	if (node->left && node->right)
	    return; // must keep node with 2 children
	// delete node with 0 or 1 children
	inheritor = (node->parent == NULL) ? &root :
	    (node == node->parent->left) ? &node->parent->left :
	    &node->parent->right;
	*inheritor = node->left ? node->left : node->right;
	delete node;
    }


    void erase(const iterator & first, const iterator & last)
    {
	// depends on erase(iterator) not invalidating other iterators
	for (iterator it = first; it != last; it++)
	    erase(it);
    }

    //------------------------------------------------------------------------
    //                bool bitset(ipv4addr_t x, unsigned int k)
    //------------------------------------------------------------------------
    //!  Returns true if the 2^k bit of an IP address is set.  NOTE WE
    //!  HANDLE OUR ENDIAN ISSUES HERE!  We do this because we want support
    //!  for traversal of the trie by key in host byte order (IP address),
    //!  but we want the key values to be in network byte order.
    //------------------------------------------------------------------------
    bool bitset(ipv4addr_t x, unsigned int k) const
    {
	// TODO: optimize for little endian hosts by using precomputed table
	// of ntohl(0x01<<k).
	return (k >= 32) ? false : (x & ntohl(0x01 << k));
    }

    std::pair<iterator,bool> insert(const value_type & value)
    {
	Node **nodep;
	Node *parent, *datanode, *newnode;

	// Search for the child pointer that we'd want to follow but can't
	// because there is no child there, the child's length is too long,
	// or the child's prefix doesn't match.
	parent = NULL;
	nodep = &root;
        while (*nodep && (*nodep)->val.first.maskLen < value.first.maskLen &&
	    (*nodep)->Matches(value.first.net))
	{
	    parent = *nodep;
	    nodep = bitset(value.first.net, 31 - (*nodep)->val.first.maskLen) ?
		&(*nodep)->right : &(*nodep)->left;
        }

	if (!*nodep) {
	    // The branch where value's key belongs does not exist
	    *nodep = new Node(value, parent);
	    _size++;
	    return make_pair(iterator(*nodep), true);
	}

        if (value.first.Matches((*nodep)->val.first.net)) {
	    if ((*nodep)->val.first.maskLen == value.first.maskLen) {
		// The correct node already exists at (*nodep)
		bool is_new = !((*nodep)->has_data);
		if (is_new)
		    _size++;
		(*nodep)->val.second = value.second;
		(*nodep)->has_data = true;
		return make_pair(iterator(*nodep), is_new);
	    } else {
		// value's key is a prefix of (*nodep)'s key; insert value
		// in new node above (*nodep).
		datanode = new Node(value, (*nodep)->parent);
		if (bitset((*nodep)->val.first.net, 31 - value.first.maskLen))
		    datanode->right = *nodep;
		else
		    datanode->left = *nodep;
		(*nodep)->parent = datanode;
		*nodep = datanode;
		_size++;
		return make_pair(iterator(datanode), true);
	    }
	}

	// Value and (*nodep) may share a prefix, but neither is a prefix of
	// the other.  Insert a new node corresponding to the shared prefix
	// and make value and (*nodep) the children of the new node.
	int i;
	for (i = (*nodep)->parent ? (*nodep)->parent->val.first.maskLen : 0;
	    i < (*nodep)->val.first.maskLen; i++)
	{
	    uint32_t mask = htonl(0x80000000 >> i);
	    if (((*nodep)->val.first.net & mask) != (value.first.net & mask)) {
		newnode = new Node(key_type(value.first.net, i),
		    (*nodep)->parent);
		(*nodep)->parent = newnode;
		datanode = new Node(value, newnode);
		if (((*nodep)->val.first.net & mask) == mask) {
		    newnode->right = *nodep;
		    newnode->left = datanode;
		} else {
		    newnode->left = *nodep;
		    newnode->right = datanode;
		}
		*nodep = newnode;
		_size++;
		return make_pair(iterator(datanode), true);
	    }
	}

	assert(0);
	return make_pair(iterator(), false); // shouldn't happen
    }
};

