/*
 * Copyright (C) by Argonne National Laboratory
 *     See COPYRIGHT in top-level directory
 */

#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include "mpitest.h"

static MPI_Datatype make_largexfer_type_struct(MPI_Offset nbytes)
{
    int typechunk_size = 1024 * 1024;   /* in bytes: TODO: figure out how big a
                                         * chunk is really needed */
    int chunk_count;
    int remainder = 0;
    MPI_Datatype memtype, chunktype;

    /* need to cook up a new datatype to accommodate large datatypes */
    /* first pass: chunks of 1 MiB plus an additional remainder.  Does require
     * 8 byte MPI_Aint, which should have been checked for earlier */

    chunk_count = nbytes / typechunk_size;
    remainder = nbytes % typechunk_size;
    MPI_Type_contiguous(typechunk_size, MPI_BYTE, &chunktype);
    MPI_Type_commit(&chunktype);

    /* a zero remainder means we can just count contigs */
    if (remainder == 0) {
        MPI_Type_contiguous(chunk_count, chunktype, &memtype);
        MPI_Type_free(&chunktype);
    } else {
        if (sizeof(MPI_Aint) <= sizeof(int)) {
            return MPI_DATATYPE_NULL;
        }
        /* struct type: some number of chunks plus remaining bytes tacked
         * on at end */
        int lens[] = { chunk_count, remainder };
        MPI_Aint disp[] = { 0, (MPI_Aint) typechunk_size * (MPI_Aint) chunk_count };
        MPI_Datatype types[] = { chunktype, MPI_BYTE };

        MPI_Type_create_struct(2, lens, disp, types, &memtype);
        MPI_Type_free(&chunktype);
    }
    MPI_Type_commit(&memtype);
    return memtype;
}

static MPI_Datatype make_largexfer_type_hindexed(MPI_Offset nbytes)
{
    int i, count;
    int chunk_size = 1024 * 1024;
    int *blocklens;
    MPI_Aint *disp;
    MPI_Datatype memtype;

    /* need to cook up a new datatype to accommodate large datatypes */
    /* Does require 8 byte MPI_Aint, which should have been checked for earlier
     */

    if (sizeof(MPI_Aint) <= sizeof(int)) {
        return MPI_DATATYPE_NULL;
    }

    /* ceiling division */
    count = 1 + ((nbytes - 1) / chunk_size);

    blocklens = calloc(count, sizeof(int));
    disp = calloc(count, sizeof(MPI_Aint));


    for (i = 0; i < (count - 1); i++) {
        blocklens[i] = chunk_size;
        disp[i] = (MPI_Aint) chunk_size *i;
    }
    blocklens[count - 1] = nbytes - ((MPI_Aint) chunk_size * i);
    disp[count - 1] = (MPI_Aint) chunk_size *(count - 1);

    MPI_Type_create_hindexed(count, blocklens, disp, MPI_BYTE, &memtype);
    MPI_Type_commit(&memtype);

    free(blocklens);
    free(disp);
    return memtype;
}


static int testtype(MPI_Datatype type, MPI_Offset expected)
{
    MPI_Count size, lb, extent;
    int errs = 0;
    MPI_Type_size_x(type, &size);

    if (size < 0) {
        printf("ERROR: type size apparently overflowed integer\n");
        errs++;
    }

    if (size != expected) {
        printf("reported type size %lld does not match expected %lld\n", (long long) size,
               (long long) expected);
        errs++;
    }

    MPI_Type_get_true_extent_x(type, &lb, &extent);
    if (lb != 0) {
        printf("ERROR: type should have lb of 0, reported %lld\n", (long long) lb);
        errs++;
    }

    if (extent != size) {
        printf("ERROR: extent should match size, not %lld\n", (long long) extent);
        errs++;
    }
    return errs;
}


int main(int argc, char **argv)
{

    int errs = 0, i;
    int rank, size;
    MTest_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

#define NR_TYPES 3
    MPI_Offset expected_sizes[NR_TYPES] = { 1024UL * 1024UL * 2400UL,
        2346319872,
        2346319872
    };
    MPI_Datatype types[NR_TYPES];

    /* MPICH does not support datatype that can not fit in the address space */
    if (sizeof(MPI_Aint) == sizeof(int))
        goto epilogue;

    /* a contig type, itself large, but does not need 8 byte aints */
    types[0] = make_largexfer_type_struct(expected_sizes[0]);
    /* struct with addresses out past 2 GiB */
    types[1] = make_largexfer_type_struct(expected_sizes[1]);
    /* similar, but with hindexed type */
    types[2] = make_largexfer_type_hindexed(expected_sizes[2]);

    for (i = 0; i < NR_TYPES; i++) {
        if (types[i] != MPI_DATATYPE_NULL) {
            errs += testtype(types[i], expected_sizes[i]);
            MPI_Type_free(&(types[i]));
        }
    }

  epilogue:
    MTest_Finalize(errs);

    return MTestReturnValue(errs);
}
