//========= Copyright Valve Corporation, All rights reserved. ============//
//
// Purpose: Bi-directional set. A Bucket knows about the elements that lie
// in it, and the elements know about the buckets they lie in.
//
// $Revision: $
// $NoKeywords: $
//=============================================================================//

#ifndef UTLBIDIRECTIONALSET_H
#define UTLBIDIRECTIONALSET_H

#ifdef _WIN32
#pragma once
#endif

#include "tier0/dbg.h"
#include "utllinkedlist.h"

//-----------------------------------------------------------------------------
// Templatized helper class to deal with the kinds of things that spatial
// partition code always seems to have; buckets with lists of lots of elements
// and elements that live in lots of buckets. This makes it really quick to
// add and remove elements, and to iterate over all elements in a bucket.
//
// For this to work, you must initialize the set with two functions one that
// maps from bucket to the index of the first element in that bucket, and one
// that maps from element to the index of the first bucket that element lies in.
// The set will completely manage the index, it's just expected that those
// indices will be stored outside the set.
//
// S is the storage type of the index; it is the type that you may use to
// save indices into memory. I is the local iterator type, which you should
// use in any local scope (eg, inside a for() loop.) The reason for this is
// that you may wish to use unsigned shorts inside the structs you are
// saving with a CBidirectionalSet; but 16-bit arithmetic is catastrophically
// slow on a PowerPC -- during testing we saw CBidirectionalSet:: operations
// consume as much as 8% of the frame.
//
// For this reason, on the 360, the handles have been typedef'd to native
// register types (U32) which are accepted as parameters by the functions.
// The implicit assumption is that CBucketHandle and CElementHandle can
// be safely cast to ints! You can increase to U64 without performance
// penalty if necessary; the PowerPC is a 64-bit processor.
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I = S>
class CBidirectionalSet {
   public:
    // Install methods to get at the first bucket given a element
    // and vice versa...
    typedef S& (*FirstElementFunc_t)(CBucketHandle);
    typedef S& (*FirstBucketFunc_t)(CElementHandle);

#ifdef _X360
    typedef uint32 CBucketHandlePram;
    typedef uint32 CElementHandlePram;
#else
    typedef CBucketHandle CBucketHandlePram;
    typedef CElementHandle CElementHandlePram;
#endif

    // Constructor
    CBidirectionalSet();

    // Call this before using the set
    void Init(FirstElementFunc_t elemFunc, FirstBucketFunc_t bucketFunc);

    // Add an element to a particular bucket
    void AddElementToBucket(CBucketHandlePram bucket,
                            CElementHandlePram element);

    // Prevalidate an add to a particular bucket
    // NOTE: EXPENSIVE!!!
    void ValidateAddElementToBucket(CBucketHandlePram bucket,
                                    CElementHandlePram element);

    // Test if an element is in a particular bucket.
    // NOTE: EXPENSIVE!!!
    bool IsElementInBucket(CBucketHandlePram bucket,
                           CElementHandlePram element);

    // Remove an element from a particular bucket
    void RemoveElementFromBucket(CBucketHandlePram bucket,
                                 CElementHandlePram element);

    // Remove an element from all buckets
    void RemoveElement(CElementHandlePram element);
    void RemoveBucket(CBucketHandlePram element);

    // Used to iterate elements in a bucket; I is the iterator
    I FirstElement(CBucketHandlePram bucket) const;
    I NextElement(I idx) const;
    CElementHandle Element(I idx) const;

    // Used to iterate buckets associated with an element; I is the iterator
    I FirstBucket(CElementHandlePram bucket) const;
    I NextBucket(I idx) const;
    CBucketHandle Bucket(I idx) const;

    static S InvalidIndex();

    // Ensure capacity
    void EnsureCapacity(int count);

    // Deallocate....
    void Purge();

    int NumAllocated(void) const;

   private:
    struct BucketListInfo_t {
        CElementHandle m_Element;
        S m_BucketListIndex;  // what's the m_BucketsUsedByElement index of the
                              // entry?
    };

    struct ElementListInfo_t {
        CBucketHandle m_Bucket;
        S m_ElementListIndex;  // what's the m_ElementsInBucket index of the
                               // entry?
    };

    // Maintains a list of all elements in a particular bucket
    CUtlLinkedList<BucketListInfo_t, S, true, I> m_ElementsInBucket;

    // Maintains a list of all buckets a particular element lives in
    CUtlLinkedList<ElementListInfo_t, S, true, I> m_BucketsUsedByElement;

    FirstBucketFunc_t m_FirstBucket;
    FirstElementFunc_t m_FirstElement;
};

//-----------------------------------------------------------------------------
// Constructor
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::CBidirectionalSet() {
    m_FirstBucket = NULL;
    m_FirstElement = NULL;
}

//-----------------------------------------------------------------------------
// Call this before using the set
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
void CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::Init(
    FirstElementFunc_t elemFunc, FirstBucketFunc_t bucketFunc) {
    m_FirstBucket = bucketFunc;
    m_FirstElement = elemFunc;
}

//-----------------------------------------------------------------------------
// Adds an element to the bucket
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
void CBidirectionalSet<CBucketHandle, CElementHandle, S,
                       I>::ValidateAddElementToBucket(CBucketHandlePram bucket,
                                                      CElementHandlePram
                                                          element) {
#ifdef _DEBUG
    // Make sure that this element doesn't already exist in the list of elements
    // in the bucket
    I elementInBucket = m_FirstElement(bucket);
    while (elementInBucket != m_ElementsInBucket.InvalidIndex()) {
        // If you hit an Assert here, fix the calling code.  It's too expensive
        // to ensure that each item only shows up once here.  Hopefully you can
        // do something better outside of here.
        Assert(m_ElementsInBucket[elementInBucket].m_Element != element);
        elementInBucket = m_ElementsInBucket.Next(elementInBucket);
    }
    // Make sure that this bucket doesn't already exist in the element's list of
    // buckets.
    I bucketInElement = m_FirstBucket(element);
    while (bucketInElement != m_BucketsUsedByElement.InvalidIndex()) {
        // If you hit an Assert here, fix the calling code.  It's too expensive
        // to ensure that each item only shows up once here.  Hopefully you can
        // do something better outside of here.
        Assert(m_BucketsUsedByElement[bucketInElement].m_Bucket != bucket);
        bucketInElement = m_BucketsUsedByElement.Next(bucketInElement);
    }
#endif
}

//-----------------------------------------------------------------------------
// Adds an element to the bucket
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
void CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::AddElementToBucket(
    CBucketHandlePram bucket, CElementHandlePram element) {
    Assert(m_FirstBucket && m_FirstElement);

    // Allocate new element + bucket entries
    I idx = m_ElementsInBucket.Alloc(true);
    I list = m_BucketsUsedByElement.Alloc(true);

    // Store off the element data
    m_ElementsInBucket[idx].m_Element = element;
    m_ElementsInBucket[idx].m_BucketListIndex = list;

    // Here's the bucket data
    m_BucketsUsedByElement[list].m_Bucket = bucket;
    m_BucketsUsedByElement[list].m_ElementListIndex = idx;

    // Insert the element into the list of elements in the bucket
    S& firstElementInBucket = m_FirstElement(bucket);
    if (firstElementInBucket != m_ElementsInBucket.InvalidIndex())
        m_ElementsInBucket.LinkBefore(firstElementInBucket, idx);
    firstElementInBucket = idx;

    // Insert the bucket into the element's list of buckets
    S& firstBucketInElement = m_FirstBucket(element);
    if (firstBucketInElement != m_BucketsUsedByElement.InvalidIndex())
        m_BucketsUsedByElement.LinkBefore(firstBucketInElement, list);
    firstBucketInElement = list;
}

//-----------------------------------------------------------------------------
// Test if an element is in a particular bucket.
// NOTE: EXPENSIVE!!!
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
bool CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::IsElementInBucket(
    CBucketHandlePram bucket, CElementHandlePram element) {
    // Search through all elements in this bucket to see if element is in there.
    I elementInBucket = m_FirstElement(bucket);
    while (elementInBucket != m_ElementsInBucket.InvalidIndex()) {
        if (m_ElementsInBucket[elementInBucket].m_Element == element) {
            return true;
        }
        elementInBucket = m_ElementsInBucket.Next(elementInBucket);
    }
    return false;
}

//-----------------------------------------------------------------------------
// Remove an element from a particular bucket
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
void CBidirectionalSet<CBucketHandle, CElementHandle, S,
                       I>::RemoveElementFromBucket(CBucketHandlePram bucket,
                                                   CElementHandlePram element) {
    // FIXME: Implement me!
    Assert(0);
}

//-----------------------------------------------------------------------------
// Removes an element from all buckets
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
void CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::RemoveElement(
    CElementHandlePram element) {
    Assert(m_FirstBucket && m_FirstElement);

    // Iterate over the list of all buckets the element is in
    I i = m_FirstBucket(element);
    while (i != m_BucketsUsedByElement.InvalidIndex()) {
        CBucketHandlePram bucket = m_BucketsUsedByElement[i].m_Bucket;
        I elementListIndex = m_BucketsUsedByElement[i].m_ElementListIndex;

        // Unhook the element from the bucket's list of elements
        if (elementListIndex == m_FirstElement(bucket))
            m_FirstElement(bucket) = m_ElementsInBucket.Next(elementListIndex);
        m_ElementsInBucket.Free(elementListIndex);

        I prevNode = i;
        i = m_BucketsUsedByElement.Next(i);
        m_BucketsUsedByElement.Free(prevNode);
    }

    // Mark the list as empty
    m_FirstBucket(element) = m_BucketsUsedByElement.InvalidIndex();
}

//-----------------------------------------------------------------------------
// Removes a bucket from all elements
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
void CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::RemoveBucket(
    CBucketHandlePram bucket) {
    // Iterate over the list of all elements in the bucket
    I i = m_FirstElement(bucket);
    while (i != m_ElementsInBucket.InvalidIndex()) {
        CElementHandlePram element = m_ElementsInBucket[i].m_Element;
        I bucketListIndex = m_ElementsInBucket[i].m_BucketListIndex;

        // Unhook the bucket from the element's list of buckets
        if (bucketListIndex == m_FirstBucket(element))
            m_FirstBucket(element) =
                m_BucketsUsedByElement.Next(bucketListIndex);
        m_BucketsUsedByElement.Free(bucketListIndex);

        // Remove the list element
        I prevNode = i;
        i = m_ElementsInBucket.Next(i);
        m_ElementsInBucket.Free(prevNode);
    }

    // Mark the bucket list as empty
    m_FirstElement(bucket) = m_ElementsInBucket.InvalidIndex();
}

//-----------------------------------------------------------------------------
// Ensure capacity
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
void CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::EnsureCapacity(
    int count) {
    m_ElementsInBucket.EnsureCapacity(count);
    m_BucketsUsedByElement.EnsureCapacity(count);
}

//-----------------------------------------------------------------------------
// Deallocate....
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
void CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::Purge() {
    m_ElementsInBucket.Purge();
    m_BucketsUsedByElement.Purge();
}

//-----------------------------------------------------------------------------
// Number of elements allocated in each linked list (should be the same)
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
int CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::NumAllocated(
    void) const {
    Assert(m_ElementsInBucket.NumAllocated() ==
           m_BucketsUsedByElement.NumAllocated());
    return m_ElementsInBucket.NumAllocated();
}

//-----------------------------------------------------------------------------
// Invalid index for iteration..
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
inline S
CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::InvalidIndex() {
    return CUtlLinkedList<CElementHandle, I>::InvalidIndex();
}

//-----------------------------------------------------------------------------
// Used to iterate elements in a bucket; I is the iterator
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
inline I CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::FirstElement(
    CBucketHandlePram bucket) const {
    Assert(m_FirstElement);
    return m_FirstElement(bucket);
}

template <class CBucketHandle, class CElementHandle, class S, class I>
inline I CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::NextElement(
    I idx) const {
    return m_ElementsInBucket.Next(idx);
}

template <class CBucketHandle, class CElementHandle, class S, class I>
inline CElementHandle
CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::Element(I idx) const {
    return m_ElementsInBucket[idx].m_Element;
}

//-----------------------------------------------------------------------------
// Used to iterate buckets an element lies in; I is the iterator
//-----------------------------------------------------------------------------
template <class CBucketHandle, class CElementHandle, class S, class I>
inline I CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::FirstBucket(
    CElementHandlePram element) const {
    Assert(m_FirstBucket);
    return m_FirstBucket(element);
}

template <class CBucketHandle, class CElementHandle, class S, class I>
inline I CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::NextBucket(
    I idx) const {
    return m_BucketsUsedByElement.Next(idx);
}

template <class CBucketHandle, class CElementHandle, class S, class I>
inline CBucketHandle
CBidirectionalSet<CBucketHandle, CElementHandle, S, I>::Bucket(I idx) const {
    return m_BucketsUsedByElement[idx].m_Bucket;
}

#endif  // UTLBIDIRECTIONALSET_H