ucx/mempool.c

/*
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS HEADER.
 *
 * Copyright 2021 Mike Becker, Olaf Wintermann All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   1. Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *
 *   2. Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "cx/mempool.h"

#include <string.h>
#include <errno.h>

static int cx_mempool_ensure_capacity(
        struct cx_mempool_s *pool,
        size_t needed_capacity
) {
    if (needed_capacity <= pool->capacity) return 0;
    size_t newcap = pool->capacity >= 1000 ?
        pool->capacity + 1000 : pool->capacity * 2;
    size_t newmsize;
    // LCOV_EXCL_START
    if (pool->capacity > newcap
            || cx_szmul(newcap, sizeof(void*), &newmsize)) {
        errno = EOVERFLOW;
        return 1;
    } // LCOV_EXCL_STOP
    void **newdata = cxRealloc(pool->base_allocator, pool->data, newmsize);
    if (newdata == NULL) return 1;
    pool->data = newdata;
    pool->capacity = newcap;
    return 0;
}

static int cx_mempool_ensure_registered_capacity(
        struct cx_mempool_s *pool,
        size_t needed_capacity
) {
    if (needed_capacity <= pool->registered_capacity) return 0;
    // we do not expect so many registrations
    size_t newcap = pool->registered_capacity + 8;
    size_t newmsize;
    // LCOV_EXCL_START
    if (pool->registered_capacity > newcap || cx_szmul(newcap,
            sizeof(struct cx_mempool_foreign_memory_s), &newmsize)) {
        errno = EOVERFLOW;
        return 1;
    } // LCOV_EXCL_STOP
    void *newdata = cxRealloc(pool->base_allocator, pool->registered, newmsize);
    if (newdata == NULL) return 1;
    pool->registered = newdata;
    pool->registered_capacity = newcap;
    return 0;
}

static void *cx_mempool_malloc_simple(
        void *p,
        size_t n
) {
    struct cx_mempool_s *pool = p;

    if (cx_mempool_ensure_capacity(pool, pool->size + 1)) {
        return NULL; // LCOV_EXCL_LINE
    }

    struct cx_mempool_memory_s *mem =
        cxMalloc(pool->base_allocator, sizeof(struct cx_mempool_memory_s) + n);
    if (mem == NULL) return NULL;
    mem->destructor = NULL;
    pool->data[pool->size] = mem;
    pool->size++;

    return mem->c;
}

static void *cx_mempool_calloc_simple(
        void *p,
        size_t nelem,
        size_t elsize
) {
    size_t msz;
    if (cx_szmul(nelem, elsize, &msz)) {
        errno = EOVERFLOW;
        return NULL;
    }
    void *ptr = cx_mempool_malloc_simple(p, msz);
    if (ptr == NULL) return NULL;
    memset(ptr, 0, nelem * elsize);
    return ptr;
}

static void cx_mempool_free_simple(
        void *p,
        void *ptr
) {
    if (!ptr) return;
    struct cx_mempool_s *pool = p;

    cx_destructor_func destr = pool->destr;
    cx_destructor_func2 destr2 = pool->destr2;

    struct cx_mempool_memory_s *mem =
        (void*) ((char *) ptr - sizeof(struct cx_mempool_memory_s));

    for (size_t i = 0; i < pool->size; i++) {
        if (mem == pool->data[i]) {
            if (mem->destructor) {
                mem->destructor(mem->c);
            }
            if (destr != NULL) {
                destr(mem->c);
            }
            if (destr2 != NULL) {
                destr2(pool->destr2_data, mem->c);
            }
            cxFree(pool->base_allocator, mem);
            size_t last_index = pool->size - 1;
            if (i != last_index) {
                pool->data[i] = pool->data[last_index];
                pool->data[last_index] = NULL;
            }
            pool->size--;
            return;
        }
    }
    abort(); // LCOV_EXCL_LINE
}

static void *cx_mempool_realloc_simple(
        void *p,
        void *ptr,
        size_t n
) {
    if (ptr == NULL) {
        return cx_mempool_malloc_simple(p, n);
    }
    if (n == 0) {
        cx_mempool_free_simple(p, ptr);
        return NULL;
    }
    struct cx_mempool_s *pool = p;

    const unsigned overhead = sizeof(struct cx_mempool_memory_s);
    struct cx_mempool_memory_s *mem =
        (void *) (((char *) ptr) - overhead);
    struct cx_mempool_memory_s *newm =
        cxRealloc(pool->base_allocator, mem, n + overhead);

    if (newm == NULL) return NULL;
    if (mem != newm) {
        for (size_t i = 0; i < pool->size; i++) {
            if (pool->data[i] == mem) {
                pool->data[i] = newm;
                return ((char*)newm) + overhead;
            }
        }
        abort(); // LCOV_EXCL_LINE
    } else {
        // unfortunately glibc() realloc seems to always move
        return ptr;  // LCOV_EXCL_LINE
    }
}

static void cx_mempool_free_all_simple(const struct cx_mempool_s *pool) {
    cx_destructor_func destr = pool->destr;
    cx_destructor_func2 destr2 = pool->destr2;
    for (size_t i = 0; i < pool->size; i++) {
        struct cx_mempool_memory_s *mem = pool->data[i];
        if (mem->destructor) {
            mem->destructor(mem->c);
        }
        if (destr != NULL) {
            destr(mem->c);
        }
        if (destr2 != NULL) {
            destr2(pool->destr2_data, mem->c);
        }
        cxFree(pool->base_allocator, mem);
    }
}

static cx_allocator_class cx_mempool_simple_allocator_class = {
    cx_mempool_malloc_simple,
    cx_mempool_realloc_simple,
    cx_mempool_calloc_simple,
    cx_mempool_free_simple
};

static void *cx_mempool_malloc_advanced(
        void *p,
        size_t n
) {
    struct cx_mempool_s *pool = p;

    if (cx_mempool_ensure_capacity(pool, pool->size + 1)) {
        return NULL;  // LCOV_EXCL_LINE
    }

    struct cx_mempool_memory2_s *mem =
        cxMalloc(pool->base_allocator, sizeof(struct cx_mempool_memory2_s) + n);
    if (mem == NULL) return NULL;
    mem->destructor = NULL;
    mem->data = NULL;
    pool->data[pool->size] = mem;
    pool->size++;

    return mem->c;
}

static void *cx_mempool_calloc_advanced(
        void *p,
        size_t nelem,
        size_t elsize
) {
    size_t msz;
    if (cx_szmul(nelem, elsize, &msz)) {
        errno = EOVERFLOW;
        return NULL;
    }
    void *ptr = cx_mempool_malloc_advanced(p, msz);
    if (ptr == NULL) return NULL;
    memset(ptr, 0, nelem * elsize);
    return ptr;
}

static void cx_mempool_free_advanced(
        void *p,
        void *ptr
) {
    if (!ptr) return;
    struct cx_mempool_s *pool = p;

    cx_destructor_func destr = pool->destr;
    cx_destructor_func2 destr2 = pool->destr2;

    struct cx_mempool_memory2_s *mem =
        (void*) ((char *) ptr - sizeof(struct cx_mempool_memory2_s));

    for (size_t i = 0; i < pool->size; i++) {
        if (mem == pool->data[i]) {
            if (mem->destructor) {
                mem->destructor(mem->data, mem->c);
            }
            if (destr != NULL) {
                destr(mem->c);
            }
            if (destr2 != NULL) {
                destr2(pool->destr2_data, mem->c);
            }
            cxFree(pool->base_allocator, mem);
            size_t last_index = pool->size - 1;
            if (i != last_index) {
                pool->data[i] = pool->data[last_index];
                pool->data[last_index] = NULL;
            }
            pool->size--;
            return;
        }
    }
    abort(); // LCOV_EXCL_LINE
}

static void *cx_mempool_realloc_advanced(
        void *p,
        void *ptr,
        size_t n
) {
    if (ptr == NULL) {
        return cx_mempool_malloc_advanced(p, n);
    }
    if (n == 0) {
        cx_mempool_free_advanced(p, ptr);
        return NULL;
    }
    struct cx_mempool_s *pool = p;

    const unsigned overhead = sizeof(struct cx_mempool_memory2_s);
    struct cx_mempool_memory2_s *mem =
        (void *) (((char *) ptr) - overhead);
    struct cx_mempool_memory2_s *newm =
        cxRealloc(pool->base_allocator, mem, n + overhead);

    if (newm == NULL) return NULL;
    if (mem != newm) {
        for (size_t i = 0; i < pool->size; i++) {
            if (pool->data[i] == mem) {
                pool->data[i] = newm;
                return ((char*)newm) + overhead;
            }
        }
        abort(); // LCOV_EXCL_LINE
    } else {
        // unfortunately glibc() realloc seems to always move
        return ptr;  // LCOV_EXCL_LINE
    }
}

static void cx_mempool_free_all_advanced(const struct cx_mempool_s *pool) {
    cx_destructor_func destr = pool->destr;
    cx_destructor_func2 destr2 = pool->destr2;
    for (size_t i = 0; i < pool->size; i++) {
        struct cx_mempool_memory2_s *mem = pool->data[i];
        if (mem->destructor) {
            mem->destructor(mem->data, mem->c);
        }
        if (destr != NULL) {
            destr(mem->c);
        }
        if (destr2 != NULL) {
            destr2(pool->destr2_data, mem->c);
        }
        cxFree(pool->base_allocator, mem);
    }
}

static cx_allocator_class cx_mempool_advanced_allocator_class = {
    cx_mempool_malloc_advanced,
    cx_mempool_realloc_advanced,
    cx_mempool_calloc_advanced,
    cx_mempool_free_advanced
};


static void *cx_mempool_malloc_pure(
        void *p,
        size_t n
) {
    struct cx_mempool_s *pool = p;

    if (cx_mempool_ensure_capacity(pool, pool->size + 1)) {
        return NULL; // LCOV_EXCL_LINE
    }

    void *mem = cxMalloc(pool->base_allocator, n);
    if (mem == NULL) return NULL;
    pool->data[pool->size] = mem;
    pool->size++;

    return mem;
}

static void *cx_mempool_calloc_pure(
        void *p,
        size_t nelem,
        size_t elsize
) {
    size_t msz;
    if (cx_szmul(nelem, elsize, &msz)) {
        errno = EOVERFLOW;
        return NULL;
    }
    void *ptr = cx_mempool_malloc_pure(p, msz);
    if (ptr == NULL) return NULL;
    memset(ptr, 0, nelem * elsize);
    return ptr;
}

static void cx_mempool_free_pure(
        void *p,
        void *ptr
) {
    if (!ptr) return;
    struct cx_mempool_s *pool = p;

    cx_destructor_func destr = pool->destr;
    cx_destructor_func2 destr2 = pool->destr2;

    for (size_t i = 0; i < pool->size; i++) {
        if (ptr == pool->data[i]) {
            if (destr != NULL) {
                destr(ptr);
            }
            if (destr2 != NULL) {
                destr2(pool->destr2_data, ptr);
            }
            cxFree(pool->base_allocator, ptr);
            size_t last_index = pool->size - 1;
            if (i != last_index) {
                pool->data[i] = pool->data[last_index];
                pool->data[last_index] = NULL;
            }
            pool->size--;
            return;
        }
    }
    abort(); // LCOV_EXCL_LINE
}

static void *cx_mempool_realloc_pure(
        void *p,
        void *ptr,
        size_t n
) {
    if (ptr == NULL) {
        return cx_mempool_malloc_pure(p, n);
    }
    if (n == 0) {
        cx_mempool_free_pure(p, ptr);
        return NULL;
    }
    struct cx_mempool_s *pool = p;
    void *newm = cxRealloc(pool->base_allocator, ptr, n);
    if (newm == NULL) return NULL;
    if (ptr != newm) {
        for (size_t i = 0; i < pool->size; i++) {
            if (pool->data[i] == ptr) {
                pool->data[i] = newm;
                return newm;
            }
        }
        abort(); // LCOV_EXCL_LINE
    } else {
        // unfortunately glibc() realloc seems to always move
        return ptr;  // LCOV_EXCL_LINE
    }
}

static void cx_mempool_free_all_pure(const struct cx_mempool_s *pool) {
    cx_destructor_func destr = pool->destr;
    cx_destructor_func2 destr2 = pool->destr2;
    for (size_t i = 0; i < pool->size; i++) {
        void *mem = pool->data[i];
        if (destr != NULL) {
            destr(mem);
        }
        if (destr2 != NULL) {
            destr2(pool->destr2_data, mem);
        }
        cxFree(pool->base_allocator, mem);
    }
}

static cx_allocator_class cx_mempool_pure_allocator_class = {
    cx_mempool_malloc_pure,
    cx_mempool_realloc_pure,
    cx_mempool_calloc_pure,
    cx_mempool_free_pure
};

static void cx_mempool_free_foreign(const struct cx_mempool_s *pool) {
    for (size_t i = 0; i < pool->registered_size; i++) {
        struct cx_mempool_foreign_memory_s info = pool->registered[i];
        if (info.destr2_data == NULL) {
            if (info.destr) {
                info.destr(info.mem);
            }
        } else {
            info.destr2(info.destr2_data, info.mem);
        }
    }
}

void cxMempoolFree(CxMempool *pool) {
    if (pool == NULL) return;
    if (pool->allocator->cl == &cx_mempool_simple_allocator_class) {
        cx_mempool_free_all_simple(pool);
    } else if (pool->allocator->cl == &cx_mempool_advanced_allocator_class) {
        cx_mempool_free_all_advanced(pool);
    } else {
        cx_mempool_free_all_pure(pool);
    }
    cx_mempool_free_foreign(pool);
    cxFree(pool->base_allocator, pool->data);
    cxFree(pool->base_allocator, pool->registered);
    cxFree(pool->base_allocator, (void*) pool->allocator);
    cxFree(pool->base_allocator, pool);
}

void cxMempoolSetDestructor(
        void *ptr,
        cx_destructor_func func
) {
    *(cx_destructor_func *) ((char *) ptr - sizeof(cx_destructor_func)) = func;
}

void cxMempoolSetDestructor2(
        void *ptr,
        cx_destructor_func2 func,
        void *data
) {
    struct cx_mempool_memory2_s *info =
        (void*)((char *) ptr - sizeof(struct cx_mempool_memory2_s));
    info->destructor = func;
    info->data = data;
}

void cxMempoolRemoveDestructor(void *ptr) {
    *(cx_destructor_func *) ((char *) ptr - sizeof(cx_destructor_func)) = NULL;
}

void cxMempoolRemoveDestructor2(void *ptr) {
    struct cx_mempool_memory2_s *info =
        (void*)((char *) ptr - sizeof(struct cx_mempool_memory2_s));
    info->destructor = NULL;
    info->data = NULL;
}

int cxMempoolRegister(
        CxMempool *pool,
        void *memory,
        cx_destructor_func destr
) {
    if (cx_mempool_ensure_registered_capacity(pool, pool->registered_size + 1)) {
        return 1; // LCOV_EXCL_LINE
    }

    pool->registered[pool->registered_size++] =
        (struct cx_mempool_foreign_memory_s) {
            .mem = memory,
            .destr = destr,
            .destr2_data = NULL
        };

    return 0;
}

int cxMempoolRegister2(
        CxMempool *pool,
        void *memory,
        cx_destructor_func2 destr,
        void *data
) {
    if (cx_mempool_ensure_registered_capacity(pool, pool->registered_size + 1)) {
        return 1; // LCOV_EXCL_LINE
    }

    pool->registered[pool->registered_size++] =
        (struct cx_mempool_foreign_memory_s) {
            .mem = memory,
            .destr2 = destr,
            .destr2_data = data
        };

    return 0;
}

CxMempool *cxMempoolCreate(
        size_t capacity,
        enum cx_mempool_type type
) {
    if (capacity == 0) capacity = 16;
    size_t poolsize;
    if (cx_szmul(capacity, sizeof(void*), &poolsize)) {
        // LCOV_EXCL_START
        errno = EOVERFLOW;
        return NULL;
    }  // LCOV_EXCL_STOP

    CxAllocator *provided_allocator = cxMallocDefault(sizeof(CxAllocator));
    if (provided_allocator == NULL) { // LCOV_EXCL_START
        return NULL;
    } // LCOV_EXCL_STOP

    CxMempool *pool = cxCallocDefault(1, sizeof(CxMempool));
    if (pool == NULL) {  // LCOV_EXCL_START
        cxFreeDefault(provided_allocator);
        return NULL;
    }  // LCOV_EXCL_STOP

    provided_allocator->data = pool;
    *((const CxAllocator**)&pool->base_allocator) = cxDefaultAllocator;
    pool->allocator = provided_allocator;
    if (type == CX_MEMPOOL_TYPE_SIMPLE) {
        provided_allocator->cl = &cx_mempool_simple_allocator_class;
    } else if (type == CX_MEMPOOL_TYPE_ADVANCED) {
        provided_allocator->cl = &cx_mempool_advanced_allocator_class;
    } else {
        provided_allocator->cl = &cx_mempool_pure_allocator_class;
    }

    pool->data = cxMallocDefault(poolsize);
    if (pool->data == NULL) { // LCOV_EXCL_START
        cxFreeDefault(provided_allocator);
        cxFreeDefault(pool);
        return NULL;
    } // LCOV_EXCL_STOP

    pool->size = 0;
    pool->capacity = capacity;

    return pool;
}

void cxMempoolGlobalDestructor(CxMempool *pool, cx_destructor_func fnc) {
    pool->destr = fnc;
}

void cxMempoolGlobalDestructor2(CxMempool *pool, cx_destructor_func2 fnc, void *data) {
    pool->destr2 = fnc;
    pool->destr2_data = data;
}

static void cx_mempool_free_transferred_allocator(void *base_al, void *al) {
    cxFree(base_al, al);
}

int cxMempoolTransfer(
        CxMempool *source,
        CxMempool *dest
) {
    // safety checks
    if (source == dest) return 1;
    if (source->allocator->cl != dest->allocator->cl) return 1;
    if (source->base_allocator->cl != dest->base_allocator->cl) return 1;

    // ensure enough capacity in the destination pool
    if (cx_mempool_ensure_capacity(dest, dest->size + source->size)) {
        return 1; // LCOV_EXCL_LINE
    }
    if (cx_mempool_ensure_registered_capacity(dest,
            dest->registered_size + source->registered_size)) {
        return 1; // LCOV_EXCL_LINE
    }

    // allocate a replacement allocator for the source pool
    CxAllocator *new_source_allocator =
        cxMalloc(source->base_allocator, sizeof(CxAllocator));
    if (new_source_allocator == NULL) { // LCOV_EXCL_START
        return 1;
    } // LCOV_EXCL_STOP
    new_source_allocator->cl = source->allocator->cl;
    new_source_allocator->data = source;

    // transfer all the data
    if (source->size > 0) {
        memcpy(&dest->data[dest->size], source->data,
            sizeof(void*)*source->size);
        dest->size += source->size;
    }

    // transfer all registered memory
    if (source->registered_size > 0) {
        memcpy(&dest->registered[dest->registered_size], source->registered,
               sizeof(struct cx_mempool_foreign_memory_s)
               * source->registered_size);
        dest->registered_size += source->registered_size;
    }

    // register the old allocator with the new pool
    // we have to remove const-ness for this, but that's okay here
    // also register the base allocator, s.t. the pool knows how to free it
    CxAllocator *transferred_allocator = (CxAllocator*) source->allocator;
    transferred_allocator->data = dest;
    cxMempoolRegister2(dest, transferred_allocator,
        cx_mempool_free_transferred_allocator, (void*)source->base_allocator);

    // prepare the source pool for re-use
    source->allocator = new_source_allocator;
    memset(source->data, 0, source->size * sizeof(void*));
    memset(source->registered, 0,
        source->registered_size * sizeof(struct cx_mempool_foreign_memory_s));
    source->size = 0;
    source->registered_size = 0;

    return 0;
}

int cxMempoolTransferObject(
        CxMempool *source,
        CxMempool *dest,
        const void *obj
) {
    // safety checks
    if (source == dest) return 1;
    if (source->allocator->cl != dest->allocator->cl) return 1;
    if (source->base_allocator->cl != dest->base_allocator->cl) return 1;

    // search for the object
    for (size_t i = 0; i < source->size; i++) {
        struct cx_mempool_memory_s *mem = source->data[i];
        if (mem->c == obj) {
            // first, make sure that the dest pool can take the object
            if (cx_mempool_ensure_capacity(dest, dest->size + 1)) {
                return 1; // LCOV_EXCL_LINE
            }
            // remove from the source pool
            size_t last_index = source->size - 1;
            if (i != last_index) {
                source->data[i] = source->data[last_index];
                source->data[last_index] = NULL;
            }
            source->size--;
            // add to the target pool
            dest->data[dest->size++] = mem;
            return 0;
        }
    }
    // search in the registered objects
    for (size_t i = 0; i < source->registered_size; i++) {
        struct cx_mempool_foreign_memory_s *mem = &source->registered[i];
        if (mem->mem == obj) {
            // first, make sure that the dest pool can take the object
            if (cx_mempool_ensure_registered_capacity(dest,
                    dest->registered_size + 1)) {
                return 1; // LCOV_EXCL_LINE
            }
            dest->registered[dest->registered_size++] = *mem;
            // remove from the source pool
            size_t last_index = source->registered_size - 1;
            if (i != last_index) {
                source->registered[i] = source->registered[last_index];
                memset(&source->registered[last_index], 0,
                    sizeof(struct cx_mempool_foreign_memory_s));
            }
            source->registered_size--;
            return 0;
        }
    }
    // not found
    return 1;
}