/*   Copyright (c) 2008 Sun Microsystems, Inc.
 *
 *   This file is part of Lustre, http://www.lustre.org
 *
 *   Lustre is free software; you can redistribute it and/or
 *   modify it under the terms of version 2 of the GNU General Public
 *   License as published by the Free Software Foundation.
 *
 *   Lustre is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with Lustre; if not, write to the Free Software
 *   Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/time.h>

#define LIBCFS_ALLOC(v, nob) do {               \
        (v) = (void *)calloc(nob, 1);           \
} while (0)

#define LIBCFS_FREE(v, nob) do {                \
        free((void *)v);                        \
} while (0)

#if 1
#define LASSERT(e)
#else
#define LASSERT(e) do {                                         \
        if (!(e)) {                                             \
                printf("Assertion %s failed at %s:%d\n",        \
                       #e, __FILE__, __LINE__);                 \
                abort();                                        \
        }                                                       \
} while (0)
#endif

#define CAA_SHIFT  10
#define CAA_SIZE   (1 << CAA_SHIFT)             /* # ptrs per level */
#define CAA_MASK   (CAA_SIZE - 1)
#define CAA_NOB    (CAA_SIZE * sizeof(void *))

typedef struct
{
        void         ****aa_3d;                 /* Triple indirect */
        void          ***aa_2d;                 /* double indirect */
        void           **aa_1d;                 /* single indirect */
} cfs_autoarray_t;

void   cfs_aa_init(cfs_autoarray_t *aa);
void   cfs_aa_fini(cfs_autoarray_t *aa);
void **cfs_aa_index(cfs_autoarray_t *aa, unsigned int idx);
void **cfs_aa_lookup(cfs_autoarray_t *aa, unsigned int idx, int grow);


void
cfs_aa_init(cfs_autoarray_t *aa)
{
        aa->aa_1d = NULL;
        aa->aa_2d = NULL;
        aa->aa_3d = NULL;
}

void
cfs_aa_fini(cfs_autoarray_t *aa)
{
        int idx0;
        int idx1;

        if (aa->aa_1d != NULL) {
                LIBCFS_FREE(aa->aa_1d, CAA_NOB);
                aa->aa_1d = NULL;
        }

        if (aa->aa_2d != NULL) {
                for (idx0 = 0; idx0 < CAA_SIZE; idx0++) {
                        if (aa->aa_2d[idx0] == NULL)
                                continue;
                        
                        LIBCFS_FREE(aa->aa_2d[idx0], CAA_NOB);
                }

                LIBCFS_FREE(aa->aa_2d, CAA_NOB);
                aa->aa_2d = NULL;
        }
        
        if (aa->aa_3d != NULL) {
                for (idx0 = 0; idx0 < CAA_SIZE; idx0++) {

                        if (aa->aa_3d[idx0] == NULL)
                                continue;
                        
                        for (idx1 = 0; idx1 < CAA_SIZE; idx1++) {
                                if (aa->aa_3d[idx0][idx1] == NULL)
                                        continue;
                                
                                LIBCFS_FREE(aa->aa_3d[idx0][idx1], CAA_NOB);
                        }

                        LIBCFS_FREE(aa->aa_3d[idx0], CAA_NOB);
                }

                LIBCFS_FREE(aa->aa_3d, CAA_NOB);
                aa->aa_3d = NULL;
        }
}

void **
cfs_aa_index(cfs_autoarray_t *a, unsigned int idx)
{
        unsigned int idx0;
        unsigned int idx1;
        
        if (idx < CAA_SIZE)
                return &(a->aa_1d[idx]);
        
        idx -= CAA_SIZE;
        idx0 = idx >> CAA_SHIFT;
        if (idx0 < CAA_SIZE)
                return &(a->aa_2d[idx0][idx & CAA_MASK]);

        idx -= CAA_SIZE * CAA_SIZE;
        idx0 = idx >> CAA_SHIFT;
        idx1 = idx0 >> CAA_SHIFT;

        return &(a->aa_3d[idx1][idx0 & CAA_MASK][idx & CAA_MASK]);
}

static int
cfs_aa_lookup1(void **array, int grow)
{
        void *a;
        
        if (*array != NULL)
                return 1;
        
        if (!grow)
                return 0;
        
        LIBCFS_ALLOC(a, CAA_NOB);
        if (a == NULL)
                return 0;
        
        *array = a;
        return 1;
}

void **
cfs_aa_lookup(cfs_autoarray_t *a, unsigned int idx, int grow)
{
        unsigned int idx0;
        unsigned int idx1;
        
        if (idx < CAA_SIZE) {
                if (!cfs_aa_lookup1((void **)&a->aa_1d, grow))
                        return NULL;
                return &(a->aa_1d[idx]);
        }
        
        idx -= CAA_SIZE;
        idx0 = idx >> CAA_SHIFT;
        if (idx0 < CAA_SIZE) {
                if (!cfs_aa_lookup1((void **)&a->aa_2d, grow) ||
                    !cfs_aa_lookup1((void **)&(a->aa_2d[idx0]), grow))
                        return NULL;
                
                return &(a->aa_2d[idx0][idx & CAA_MASK]);
        }
        
        idx -= CAA_SIZE * CAA_SIZE;
        idx0 = idx >> CAA_SHIFT;
        idx1 = idx0 >> CAA_SHIFT;
        idx0 &= CAA_MASK;

        if (!cfs_aa_lookup1((void **)&a->aa_3d, grow) ||
            !cfs_aa_lookup1((void **)&(a->aa_3d[idx1]), grow) ||
            !cfs_aa_lookup1((void **)&(a->aa_3d[idx1][idx0]), grow))
                return NULL;
        
        return &(a->aa_3d[idx1][idx0 & CAA_MASK][idx & CAA_MASK]);
}

typedef struct
{
        cfs_autoarray_t  cbh_members;
        unsigned int     cbh_size;
        unsigned int     cbh_maxsize;
        int            (*cbh_compare)(void *a, void *b);
} cfs_bheap_t;

void
cfs_bheap_init(cfs_bheap_t *h, int (*compare)(void *a, void *b)) 
{
        cfs_aa_init(&h->cbh_members);
        
        h->cbh_compare = compare;
        h->cbh_maxsize = h->cbh_size = 0;
}

void
cfs_bheap_fini(cfs_bheap_t *h)
{
        cfs_aa_fini(&h->cbh_members);

        h->cbh_compare = NULL;
        h->cbh_maxsize = h->cbh_size = 0;
}

int
cfs_bheap_size(cfs_bheap_t *h) 
{
        return h->cbh_size;
}

int
cfs_bheap_insert(cfs_bheap_t *h, void *e)
{
        int            (*compare)(void *a, void *b) = h->cbh_compare;
        unsigned int     child_idx = h->cbh_size;
        void           **child_ptr;
        unsigned int     parent_idx;
        void           **parent_ptr;
        int              rc;

        LASSERT (child_idx <= h->cbh_maxsize);

        /* add new child at the very bottom of the tree == last slot in the members array */
        h->cbh_size = child_idx + 1;

        if (child_idx < h->cbh_maxsize) {
                /* I know there is a slot in the members array */
                child_ptr = cfs_aa_index(&h->cbh_members, child_idx);
        } else {
                /* Ensure auto-array grows to accomodate new member */
                child_ptr = cfs_aa_lookup(&h->cbh_members, child_idx, 1);
                if (child_ptr == NULL)
                        return -ENOMEM;
                h->cbh_maxsize = h->cbh_size;
        }
        
        while (child_idx > 0) {
                parent_idx = (child_idx - 1) >> 1;
                
                parent_ptr = cfs_aa_index(&h->cbh_members, parent_idx);
                
                if (compare(*parent_ptr, e))
                        break;
                
                *child_ptr = *parent_ptr;
                child_ptr = parent_ptr;
                child_idx = parent_idx;
        }
        
        *child_ptr = e;
        return 0;
}

void *
cfs_bheap_remove_root(cfs_bheap_t *h)
{
        int         (*compare)(void *a, void *b) = h->cbh_compare;
        unsigned int  n = h->cbh_size;
        void         *r;
        void         *e;
        unsigned int  child_idx;
        void        **child_ptr;
        void         *child;
        unsigned int  child2_idx;
        void        **child2_ptr;
        void         *child2;
        unsigned int  parent_idx;
        void        **parent_ptr;
        
        if (n == 0)
                return NULL;

        parent_idx = 0;
        parent_ptr = cfs_aa_index(&h->cbh_members, parent_idx);
        r = *parent_ptr;

        n--;
        e = *cfs_aa_index(&h->cbh_members, n);
        h->cbh_size = n;
        
        while (parent_idx < n) {
                child_idx = (parent_idx << 1) + 1;
                if (child_idx >= n)
                        break;

                child_ptr = cfs_aa_index(&h->cbh_members, child_idx);
                child = *child_ptr;

                child2_idx = child_idx + 1;
                if (child2_idx < n) {
                        child2_ptr = cfs_aa_index(&h->cbh_members, child2_idx);
                        child2 = *child2_ptr;
                        
                        if (compare(child2, child)) {
                                child_idx = child2_idx;
                                child_ptr = child2_ptr;
                                child = child2;
                        }
                }
                
                if (compare(e, child))
                        break;
                
                *parent_ptr = child;
                parent_ptr = child_ptr;
                parent_idx = child_idx;
        }
        
        *parent_ptr = e;
        return r;
}

typedef struct 
{
        int        val;
        
} thing_t;

thing_t *
create_thing(int n)
{
        thing_t *t = malloc(sizeof(thing_t));

        t->val = n;
        return t;
}

void
destroy_thing(thing_t *t)
{
        free(t);
}

int
compare_things(void *a, void *b) 
{
        return (((thing_t *)a)->val < ((thing_t *)b)->val);
}

int
do_print_thing_heap(cfs_bheap_t *h, int idx) 
{
        int      indent;
        int      i;
        int      next_idx = (idx << 1) + 1;
        thing_t *t;
        
        if (next_idx < h->cbh_size)
                indent = do_print_thing_heap(h, next_idx);
        else
                indent = 1;
        
        while (idx < next_idx && idx < h->cbh_size) {
                t = (thing_t *)*cfs_aa_index(&h->cbh_members, idx);
                idx++;
                
                printf("%3d ", t->val);
                for (i = 0; i < (indent - 1) * 4; i++)
                        printf(" ");
        }
        printf("\n");

        return indent * 2;
}

void
print_thing_heap(cfs_bheap_t *h) 
{
        printf ("Heap size %d\n", cfs_bheap_size(h));
        do_print_thing_heap(h, 0);
}

double
tnow() 
{
        struct timeval tv;

        gettimeofday(&tv, NULL);
        return tv.tv_sec + tv.tv_usec/1000000.0;
}

int
main(int argc, char **argv) 
{
        int             size = atoi(argv[1]);
        int             loops = atoi(argv[2]);
        cfs_bheap_t   h;
        thing_t        *t;
        int             n = size;
        int             i;
        double          t0;
        double          secs_overhead;
        double          secs;

        cfs_bheap_init(&h, compare_things);

        for (i = 0; i < size; i++) {
                t = create_thing((int)(drand48() * size));
                
                if (cfs_bheap_insert(&h, t) != 0) {
                        printf ("FULL with %d to go\n", size - i);
                        destroy_thing(t);
                        break;
                }

                //print_thing_heap(h);
        }

        t0 = tnow();
        
        for (i = 0; i < loops; i++) {
                n = i + drand48() * size;
        }
        
        secs_overhead = tnow() - t0;

        t0 = tnow();
        
        for (i = 0; i < loops; i++) {
                t = cfs_bheap_remove_root(&h);
                
                t->val = i + drand48() * 1000;
                cfs_bheap_insert(&h, t);
                //print_thing_heap(h);
        }
        
        secs = tnow() - t0;
        
        printf ("Size %d: %0.3fuS\n", size, ((secs-secs_overhead)/loops) * 1000000);

        for (;;) {
                t = cfs_bheap_remove_root(&h);
                //print_thing_heap(h);
                if (t == NULL)
                        break;
                
                destroy_thing(t);
        }

        cfs_bheap_fini(&h);
        return 0;
}