/* -*- Mode: C; c-basic-offset:4 ; -*- */
/*
 * Copyright (c) 2004-2006 The Trustees of Indiana University and Indiana
 *                         University Research and Technology
 *                         Corporation.  All rights reserved.
 * Copyright (c) 2004-2019 The University of Tennessee and The University
 *                         of Tennessee Research Foundation.  All rights
 *                         reserved.
 * Copyright (c) 2004-2006 High Performance Computing Center Stuttgart,
 *                         University of Stuttgart.  All rights reserved.
 * Copyright (c) 2004-2006 The Regents of the University of California.
 *                         All rights reserved.
 * Copyright (c) 2009      Oak Ridge National Labs.  All rights reserved.
 * Copyright (c) 2014      Cisco Systems, Inc.  All rights reserved.
 * Copyright (c) 2015-2018 Research Organization for Information Science
 *                         and Technology (RIST).  All rights reserved.
 * $COPYRIGHT$
 *
 * Additional copyrights may follow
 *
 * $HEADER$
 */

#include "opal_config.h"

#include <stddef.h>
#include <stdlib.h>

#include "opal/datatype/opal_datatype.h"
#include "opal/datatype/opal_convertor.h"
#include "opal/datatype/opal_datatype_internal.h"

static int32_t
opal_datatype_optimize_short( opal_datatype_t* pData,
                              size_t count,
                              dt_type_desc_t* pTypeDesc )
{
    dt_elem_desc_t* pElemDesc;
    dt_stack_t *pOrigStack, *pStack; /* pointer to the position on the stack */
    int32_t pos_desc = 0;            /* actual position in the description of the derived datatype */
    int32_t stack_pos = 0;
    int32_t nbElems = 0;
    ptrdiff_t total_disp = 0;
    ddt_elem_desc_t last = {.common.flags = 0xFFFF /* all on */, .count = 0, .disp = 0}, compress;
    ddt_elem_desc_t* current;

    pOrigStack = pStack = (dt_stack_t*)malloc( sizeof(dt_stack_t) * (pData->loops+2) );
    SAVE_STACK( pStack, -1, 0, count, 0 );

    pTypeDesc->length = 2 * pData->desc.used + 1 /* for the fake OPAL_DATATYPE_END_LOOP at the end */;
    pTypeDesc->desc = pElemDesc = (dt_elem_desc_t*)malloc( sizeof(dt_elem_desc_t) * pTypeDesc->length );
    pTypeDesc->used = 0;

    assert( OPAL_DATATYPE_END_LOOP == pData->desc.desc[pData->desc.used].elem.common.type );

    while( stack_pos >= 0 ) {
        if( OPAL_DATATYPE_END_LOOP == pData->desc.desc[pos_desc].elem.common.type ) { /* end of the current loop */
            ddt_endloop_desc_t* end_loop = &(pData->desc.desc[pos_desc].end_loop);
            if( 0 != last.count ) {
                CREATE_ELEM( pElemDesc, last.common.type, OPAL_DATATYPE_FLAG_BASIC,
                             last.blocklen, last.count, last.disp, last.extent );
                pElemDesc++; nbElems++;
                last.count= 0;
            }
            CREATE_LOOP_END( pElemDesc, nbElems - pStack->index + 1,  /* # of elems in this loop */
                             end_loop->first_elem_disp, end_loop->size, end_loop->common.flags );
            if( --stack_pos >= 0 ) {  /* still something to do ? */
                ddt_loop_desc_t* pStartLoop = &(pTypeDesc->desc[pStack->index - 1].loop);
                pStartLoop->items = pElemDesc->end_loop.items;
                total_disp = pStack->disp;  /* update the displacement position */
            }
            pElemDesc++; nbElems++;
            pStack--;  /* go down one position on the stack */
            pos_desc++;
            continue;
        }
        if( OPAL_DATATYPE_LOOP == pData->desc.desc[pos_desc].elem.common.type ) {
            ddt_loop_desc_t* loop = (ddt_loop_desc_t*)&(pData->desc.desc[pos_desc]);
            int index = GET_FIRST_NON_LOOP( &(pData->desc.desc[pos_desc]) );

            if( loop->common.flags & OPAL_DATATYPE_FLAG_CONTIGUOUS ) {
                ddt_endloop_desc_t* end_loop = (ddt_endloop_desc_t*)&(pData->desc.desc[pos_desc + loop->items]);

                assert(pData->desc.desc[pos_desc + index].elem.disp == end_loop->first_elem_disp);
                compress.common.flags = loop->common.flags;
                compress.common.type =  pData->desc.desc[pos_desc + index].elem.common.type;
                compress.blocklen = pData->desc.desc[pos_desc + index].elem.blocklen;
                for( uint32_t i = index+1; i < loop->items; i++ ) {
                    current = &pData->desc.desc[pos_desc + i].elem;
                    assert(1 == current->count);
                    if ((current->common.type == OPAL_DATATYPE_LOOP)
                        || compress.common.type != current->common.type) {
                        compress.common.type   = OPAL_DATATYPE_UINT1;
                        compress.common.flags |= OPAL_DATATYPE_OPTIMIZED_RESTRICTED;
                        pData->flags          |= OPAL_DATATYPE_OPTIMIZED_RESTRICTED;
                        compress.blocklen = end_loop->size;
                        break;
                    }
                    compress.blocklen += current->blocklen;
                }
                compress.count = loop->loops;
                compress.extent = loop->extent;
                compress.disp = end_loop->first_elem_disp;
                if( compress.extent == (ptrdiff_t)(compress.blocklen * opal_datatype_basicDatatypes[compress.common.type]->size) ) {
                    /* The compressed element is contiguous: collapse it into a single large blocklen */
                    compress.blocklen *= compress.count;
                    compress.extent   *= compress.count;
                    compress.count     = 1;
                }
                /**
                 * The current loop has been compressed and can now be treated as if it
                 * was a data element. We can now look if it can be fused with last,
                 * as done in the fusion of 2 elements below. Let's use the same code.
                 */
                pos_desc += loop->items + 1;
                current = &compress;
                goto fuse_loops;
            }

            /**
             * If the content of the loop is not contiguous there is little we can do
             * that would not incur significant optimization cost and still be beneficial
             * in reducing the number of memcpy during pack/unpack.
             */

            if( 0 != last.count ) {  /* Generate the pending element */
                CREATE_ELEM( pElemDesc, last.common.type, OPAL_DATATYPE_FLAG_BASIC,
                             last.blocklen, last.count, last.disp, last.extent );
                pElemDesc++; nbElems++;
                last.count       = 0;
                last.common.type = OPAL_DATATYPE_LOOP;
            }

            /* Can we unroll the loop? */
            if( (loop->items <= 3) && (loop->loops <= 2) ) {
                ptrdiff_t elem_displ = 0;
                for( uint32_t i = 0; i < loop->loops; i++ ) {
                    for( uint32_t j = 0; j < (loop->items - 1); j++ ) {
                        current = &pData->desc.desc[pos_desc + index + j].elem;
                        CREATE_ELEM( pElemDesc, current->common.type, current->common.flags,
                                     current->blocklen, current->count, current->disp + elem_displ, current->extent );
                        pElemDesc++; nbElems++;
                    }
                    elem_displ += loop->extent;
                }
                pos_desc += loop->items + 1;
                goto complete_loop;
            }

            CREATE_LOOP_START( pElemDesc, loop->loops, loop->items, loop->extent, loop->common.flags );
            pElemDesc++; nbElems++;
            PUSH_STACK( pStack, stack_pos, nbElems, OPAL_DATATYPE_LOOP, loop->loops, total_disp );
            pos_desc++;
            DDT_DUMP_STACK( pStack, stack_pos, pData->desc.desc, "advance loops" );

        complete_loop:
            total_disp = pStack->disp;  /* update the displacement */
            continue;
        }
        while( pData->desc.desc[pos_desc].elem.common.flags & OPAL_DATATYPE_FLAG_DATA ) {  /* go over all basic datatype elements */
            current = &pData->desc.desc[pos_desc].elem;
            pos_desc++;  /* point to the next element as current points to the current one */

          fuse_loops:
            if( 0 == last.count ) {  /* first data of the datatype */
                last = *current;
                continue;  /* next data */
            } else {  /* can we merge it in order to decrease count */
                if( (ptrdiff_t)last.blocklen * (ptrdiff_t)opal_datatype_basicDatatypes[last.common.type]->size == last.extent ) {
                    last.extent *= last.count;
                    last.blocklen *= last.count;
                    last.count = 1;
                }
            }

            /* are the two elements compatible: aka they have very similar values and they
             * can be merged together by increasing the count, and/or changing the extent.
             */
            if ((last.blocklen * opal_datatype_basicDatatypes[last.common.type]->size)
                == (current->blocklen * opal_datatype_basicDatatypes[current->common.type]->size)) {
                ddt_elem_desc_t save = last; /* safekeep the type and blocklen */
                if (last.common.type != current->common.type) {
                    last.blocklen *= opal_datatype_basicDatatypes[last.common.type]->size;
                    last.common.type   = OPAL_DATATYPE_UINT1;
                    last.common.flags |= OPAL_DATATYPE_OPTIMIZED_RESTRICTED;
                    pData->flags      |= OPAL_DATATYPE_OPTIMIZED_RESTRICTED;
                }

                if( (last.extent * (ptrdiff_t)last.count + last.disp) == current->disp ) {
                    if( 1 == current->count ) {
                        last.count++;
                        continue;
                    }
                    if( last.extent == current->extent ) {
                        last.count += current->count;
                        continue;
                    }
                }
                if( 1 == last.count ) {
                    /* we can ignore the extent of the element with count == 1 and merge them together if their displacements match */
                    if( 1 == current->count ) {
                        last.extent = current->disp - last.disp;
                        last.count++;
                        continue;
                    }
                    /* can we compute a matching displacement ? */
                    if( (last.disp + current->extent) == current->disp ) {
                        last.extent = current->extent;
                        last.count = current->count + last.count;
                        continue;
                    }
                }
                last.blocklen = save.blocklen;
                last.common.type = save.common.type;
                /* try other optimizations */
            }
            /* are the elements fusionable such that we can fusion the last blocklen of one with the first
             * blocklen of the other.
             */
            if( (ptrdiff_t)(last.disp + (last.count - 1) * last.extent + last.blocklen * opal_datatype_basicDatatypes[last.common.type]->size) ==
                current->disp ) {
                if( last.count != 1 ) {
                    CREATE_ELEM( pElemDesc, last.common.type, OPAL_DATATYPE_FLAG_BASIC,
                                 last.blocklen, last.count - 1, last.disp, last.extent );
                    pElemDesc++; nbElems++;
                    last.disp += (last.count - 1) * last.extent;
                    last.count = 1;
                }
                if( last.common.type == current->common.type ) {
                    last.blocklen += current->blocklen;
                } else {
                    last.blocklen = ((last.blocklen
                                      * opal_datatype_basicDatatypes[last.common.type]->size)
                                     + (current->blocklen
                                        * opal_datatype_basicDatatypes[current->common.type]
                                              ->size));
                    last.common.type   = OPAL_DATATYPE_UINT1;
                    last.common.flags |= OPAL_DATATYPE_OPTIMIZED_RESTRICTED;
                    pData->flags      |= OPAL_DATATYPE_OPTIMIZED_RESTRICTED;
                }
                last.extent += current->extent;
                if( current->count != 1 ) {
                    CREATE_ELEM( pElemDesc, last.common.type, OPAL_DATATYPE_FLAG_BASIC,
                                 last.blocklen, last.count, last.disp, last.extent );
                    pElemDesc++; nbElems++;
                    last = *current;
                    last.count -= 1;
                    last.disp += last.extent;
                }
                continue;
            }
            CREATE_ELEM( pElemDesc, last.common.type, OPAL_DATATYPE_FLAG_BASIC,
                         last.blocklen, last.count, last.disp, last.extent );
            pElemDesc++; nbElems++;
            last = *current;
        }
    }

    if( 0 != last.count ) {
        CREATE_ELEM( pElemDesc, last.common.type, OPAL_DATATYPE_FLAG_BASIC,
                     last.blocklen, last.count, last.disp, last.extent );
        pElemDesc++; nbElems++;
    }
    /* cleanup the stack */
    pTypeDesc->used = nbElems - 1;  /* except the last fake END_LOOP */
    free(pOrigStack);
    return OPAL_SUCCESS;
}

int32_t opal_datatype_commit( opal_datatype_t * pData )
{
    ddt_endloop_desc_t* pLast = &(pData->desc.desc[pData->desc.used].end_loop);
    ptrdiff_t first_elem_disp = 0;

    if( pData->flags & OPAL_DATATYPE_FLAG_COMMITTED ) return OPAL_SUCCESS;
    pData->flags |= OPAL_DATATYPE_FLAG_COMMITTED;

    /* We have to compute the displacement of the first non loop item in the
     * description.
     */
    if( 0 != pData->size ) {
        int index;
        dt_elem_desc_t* pElem = pData->desc.desc;

        index = GET_FIRST_NON_LOOP( pElem );
        assert( pElem[index].elem.common.flags & OPAL_DATATYPE_FLAG_DATA );
        first_elem_disp = pElem[index].elem.disp;
    }

    /* let's add a fake element at the end just to avoid useless comparaisons
     * in pack/unpack functions.
     */
    pLast->common.type     = OPAL_DATATYPE_END_LOOP;
    pLast->common.flags    = 0;
    pLast->items           = pData->desc.used;
    pLast->first_elem_disp = first_elem_disp;
    pLast->size            = pData->size;

    /* If there is no datatype description how can we have an optimized description ? */
    if( 0 == pData->desc.used ) {
        pData->opt_desc.length = 0;
        pData->opt_desc.desc   = NULL;
        pData->opt_desc.used   = 0;
        return OPAL_SUCCESS;
    }

    /* If the data is contiguous is useless to generate an optimized version. */
    /*if( pData->size == (pData->true_ub - pData->true_lb) ) return OPAL_SUCCESS; */

    (void)opal_datatype_optimize_short( pData, 1, &(pData->opt_desc) );
    if( 0 != pData->opt_desc.used ) {
        /* let's add a fake element at the end just to avoid useless comparaisons
         * in pack/unpack functions.
         */
        pLast = &(pData->opt_desc.desc[pData->opt_desc.used].end_loop);
        pLast->common.type     = OPAL_DATATYPE_END_LOOP;
        pLast->common.flags    = 0;
        pLast->items           = pData->opt_desc.used;
        pLast->first_elem_disp = first_elem_disp;
        pLast->size            = pData->size;
    }
    return OPAL_SUCCESS;
}