malloc single-call for a 2D array

malloc2d.c

#include <stdlib.h>

// Allocates a 2D array that can be accessed in the form arr[r][c].
// The caller is responsible for calling free() when done.
void** malloc2d(size_t rows, size_t cols, size_t element_size) {
    size_t header = rows * sizeof(void*);
    size_t body = rows * cols * element_size;
    size_t needed = header + body;

    void** mem = malloc(needed);
    if (!mem) {
        return NULL;
    }

    for (size_t i = 0; i < rows; i++) {
        void* col_mem = mem + header + i*rows*cols*element_size;
        mem[i] = col_mem;
    }

    return mem;
}

#include <assert.h>

int main() {
    int rows = 5;
    int cols = 10;

    int** test = (int**) malloc2d(rows, cols, sizeof(int));
    assert(test);

    int i = 1;
    for (int r = 0; r < rows; r++) {
        for (int c = 0; c < cols; c++) {
            test[r][c] = i;
            i++;
        }
    }

    i = 1;
    for (int r = 0; r < rows; r++) {
        for (int c = 0; c < cols; c++) {
            assert(test[r][c] == i);
            i++;
        }
    }

    free(test);
}