foreign.c

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <stdlib.h>

static double** python_matrix_to_c(PyObject* py_matrix, Py_ssize_t* N, Py_ssize_t* M) {
    *N = PyList_Size(py_matrix);
    *M = PyList_Size(PyList_GetItem(py_matrix, 0));
    double** matrix = (double**) malloc(*N * sizeof(double*));
    if (!matrix) {
        PyErr_NoMemory();
        return NULL;
    }
    for (Py_ssize_t i = 0; i < *N; i++) {
        matrix[i] = (double*) malloc(*M * sizeof(double));
        if (!matrix[i]) {
            for (Py_ssize_t k = 0; k < i; k++) {
                free(matrix[k]);
            }
            free(matrix);
            PyErr_NoMemory();
            return NULL;
        }
        PyObject* row = PyList_GetItem(py_matrix, i);
        for (Py_ssize_t j = 0; j < *M; j++) {
            PyObject* elem = PyList_GetItem(row, j);
            matrix[i][j] = PyFloat_AsDouble(elem);
        }
    }
    return matrix;
}

static PyObject* c_matrix_to_py(double** c_matrix, Py_ssize_t N, Py_ssize_t M) {
    PyObject* matrix = PyList_New(N);
    if (!matrix) {
        return NULL;
    }
    for (Py_ssize_t i = 0; i < N; i++) {
        PyObject* row = PyList_New(M);
        if (!row) {
            Py_DECREF(matrix);
            return NULL;
        }

        for (Py_ssize_t j = 0; j < M; j++) {
            PyObject* elem = PyFloat_FromDouble(c_matrix[i][j]);
            if (!elem) {
                Py_DECREF(row);
                Py_DECREF(matrix);
                return NULL;
            }
            PyList_SetItem(row, j, elem);
        }
        PyList_SetItem(matrix, i, row);
    }
    return matrix;
}

static double** matrix_multiply(double** first_matrix, double** second_matrix, 
    Py_ssize_t first_rows, Py_ssize_t first_columns, Py_ssize_t second_rows, Py_ssize_t second_columns) {
        if (first_columns != second_rows) {
            return NULL;
        }
        double** res_matrix = (double**) malloc(first_rows * sizeof(double*));
        if (!res_matrix) {
            return NULL;
        }
        for (Py_ssize_t row = 0; row < first_rows; row++) {
            double* new_row = (double*) malloc(second_columns * sizeof(double));
            for (Py_ssize_t j = 0; j < second_columns; j++) {
                double new_elem = 0;
                for (Py_ssize_t k = 0; k < second_rows; k++) {
                    new_elem = new_elem + (first_matrix[row][k] * second_matrix[k][j]);
                }
                new_row[j] = new_elem;
                new_elem = 0; 
            }
            res_matrix[row] = new_row;
        }
        return res_matrix;
    }


static double** identity_matrix(Py_ssize_t N, Py_ssize_t M) {
    double** res_matrix = (double**) malloc(N * sizeof(double*));
    if (!res_matrix) {
        return NULL;
    }
    for (Py_ssize_t i = 0; i < N; i++) {
        res_matrix[i] = (double*) calloc(M, sizeof(double));
        if (!res_matrix[i]) {
            for (Py_ssize_t k = 0; k < i; k++) {
                free(res_matrix[k]);
            }
            free(res_matrix);
            return NULL;
        }
        res_matrix[i][i] = 1.0;
    }
    return res_matrix;
}

static void free_c_matrix(double** matrix, Py_ssize_t N, Py_ssize_t M) {
    for (Py_ssize_t i = 0; i < N; i++) {
        free(matrix[i]);
    }
    free(matrix);
}

static PyObject* foreign_matrix_power(PyObject* self, PyObject* args) {
    PyObject* matrix;
    size_t power;

    if (!PyArg_ParseTuple(args, "Ol", &matrix, &power)) {
        return NULL;
    }

    Py_ssize_t N = PyList_Size(matrix);
    if (!N) {
        return PyList_New(0);
    }
    Py_ssize_t M = PyList_Size(PyList_GetItem(matrix, 0));
    if (N != M) {
        return NULL;
    }
    double** c_matrix = python_matrix_to_c(matrix, &N, &M);
    double** res_matrix = identity_matrix(N, M);
    if (!res_matrix) {
        free_c_matrix(c_matrix, N, M);
        PyErr_NoMemory();
        return NULL;
    }
    for (size_t i = 0; i < power; i++) {
        double** temp = matrix_multiply(res_matrix, c_matrix, N, M, N, M);
        if (!temp) {
            free_c_matrix(c_matrix, N, M);
            free_c_matrix(res_matrix, N, M);
            PyErr_NoMemory();
            return NULL;
        }
        free_c_matrix(res_matrix, N, M);
        res_matrix = temp;
    }
    PyObject* result = c_matrix_to_py(res_matrix, N, N);
    free_c_matrix(res_matrix, N, M);
    free_c_matrix(c_matrix, N, M);
    return result;
}

static PyObject* foreign_matrix_multiply(PyObject* self, PyObject* args) {
    PyObject* first_matrix;
    PyObject* second_matrix;

    if (!PyArg_ParseTuple(args, "OO", &first_matrix, &second_matrix)) {
        return NULL;
    }

    Py_ssize_t first_n = PyList_Size(first_matrix);
    Py_ssize_t first_m = PyList_Size(PyList_GetItem(first_matrix, 0));
    Py_ssize_t second_n = PyList_Size(second_matrix);
    Py_ssize_t second_m = PyList_Size(PyList_GetItem(second_matrix, 0));

    double** first_c_matrix = python_matrix_to_c(first_matrix, &first_n, &first_m);
    double** second_c_matrix = python_matrix_to_c(second_matrix, &second_n, &second_m);
    double** res_matrix = matrix_multiply(first_c_matrix, second_c_matrix, first_n, first_m, second_n, second_m);
    PyObject* result = c_matrix_to_py(res_matrix, first_n, second_m);
    free_c_matrix(first_c_matrix, first_n, first_m);
    free_c_matrix(second_c_matrix, second_n, second_m);
    free_c_matrix(res_matrix, first_n, second_m);
    return result;
}


static PyMethodDef ForeignMethods[] = {
    {"foreign_matrix_power", 
    foreign_matrix_power, METH_VARARGS,
    "Raises matrix to a given power"}, 
    {"foreign_matrix_multiply", 
    foreign_matrix_multiply, METH_VARARGS, 
    "Multiplies two matrix"},

    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef foreignmodule = {
    PyModuleDef_HEAD_INIT, 
    "extension", 
    NULL, 
    -1, 

    ForeignMethods
};

PyMODINIT_FUNC PyInit_foreign(void) {
    return PyModule_Create(&foreignmodule);
}