diff options
Diffstat (limited to 'erts/example/matrix_nif.c')
| -rw-r--r-- | erts/example/matrix_nif.c | 210 |
1 files changed, 210 insertions, 0 deletions
diff --git a/erts/example/matrix_nif.c b/erts/example/matrix_nif.c new file mode 100644 index 0000000000..c5e01dade5 --- /dev/null +++ b/erts/example/matrix_nif.c @@ -0,0 +1,210 @@ +/* + * %CopyrightBegin% + * + * Copyright Ericsson AB 2010. All Rights Reserved. + * + * The contents of this file are subject to the Erlang Public License, + * Version 1.1, (the "License"); you may not use this file except in + * compliance with the License. You should have received a copy of the + * Erlang Public License along with this software. If not, it can be + * retrieved online at http://www.erlang.org/. + * + * Software distributed under the License is distributed on an "AS IS" + * basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See + * the License for the specific language governing rights and limitations + * under the License. + * + * %CopyrightEnd% + */ +/* + * Purpose: Simple example of NIFs using resource objects to implement functions + * for matrix calculations. + */ + +#include "erl_nif.h" + +#include <stddef.h> +#include <assert.h> + +typedef struct +{ + unsigned nrows; + unsigned ncols; + double* data; +}Matrix; + +#define POS(MX, ROW, COL) ((MX)->data[(ROW)* (MX)->ncols + (COL)]) + +static int get_number(ErlNifEnv* env, ERL_NIF_TERM term, double* dp); +static Matrix* alloc_matrix(ErlNifEnv* env, unsigned nrows, unsigned ncols); +static void matrix_dtor(ErlNifEnv* env, void* obj); + + +static ErlNifResourceType* resource_type = NULL; + +static int load(ErlNifEnv* env, void** priv_data, ERL_NIF_TERM load_info) +{ + ErlNifResourceType* rt = enif_open_resource_type(env, "matrix_nif_example", + matrix_dtor, + ERL_NIF_RT_CREATE, NULL); + if (rt == NULL) { + return -1; + } + assert(resource_type == NULL); + resource_type = rt; + return 0; +} + +static ERL_NIF_TERM create(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) +{ + /* create(Nrows, Ncolumns, [[first row],[second row],...,[last row]]) -> Matrix */ + unsigned nrows, ncols; + unsigned i, j; + ERL_NIF_TERM list, row, ret; + Matrix* mx = NULL; + + if (!enif_get_uint(env, argv[0], &nrows) || nrows < 1 || + !enif_get_uint(env, argv[1], &ncols) || ncols < 1) { + + goto badarg; + } + mx = alloc_matrix(env, nrows, ncols); + list = argv[2]; + for (i = 0; i<nrows; i++) { + if (!enif_get_list_cell(env, list, &row, &list)) { + goto badarg; + } + for (j = 0; j<ncols; j++) { + ERL_NIF_TERM v; + if (!enif_get_list_cell(env, row, &v, &row) || + !get_number(env, v, &POS(mx,i,j))) { + goto badarg; + } + } + if (!enif_is_empty_list(env, row)) { + goto badarg; + } + } + if (!enif_is_empty_list(env, list)) { + goto badarg; + } + + ret = enif_make_resource(env, mx); + enif_release_resource(env, mx); + return ret; + +badarg: + if (mx != NULL) { + enif_release_resource(env,mx); + } + return enif_make_badarg(env); +} + + +static ERL_NIF_TERM pos(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) +{ + /* pos(Matrix, Row, Column) -> float() */ + Matrix* mx; + unsigned i, j; + if (!enif_get_resource(env, argv[0], resource_type, (void**)&mx) || + !enif_get_uint(env, argv[1], &i) || (--i >= mx->nrows) || + !enif_get_uint(env, argv[2], &j) || (--j >= mx->ncols)) { + return enif_make_badarg(env); + } + return enif_make_double(env, POS(mx, i,j)); +} + +static ERL_NIF_TERM add(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) +{ + /* add(Matrix_A, Matrix_B) -> Matrix_Sum */ + unsigned i, j; + ERL_NIF_TERM ret; + Matrix* mxA = NULL; + Matrix* mxB = NULL; + Matrix* mxS = NULL; + + if (!enif_get_resource(env, argv[0], resource_type, (void**)&mxA) || + !enif_get_resource(env, argv[1], resource_type, (void**)&mxB) || + mxA->nrows != mxB->nrows || + mxB->ncols != mxB->ncols) { + + return enif_make_badarg(env); + } + mxS = alloc_matrix(env, mxA->nrows, mxA->ncols); + for (i = 0; i < mxA->nrows; i++) { + for (j = 0; j < mxA->ncols; j++) { + POS(mxS, i, j) = POS(mxA, i, j) + POS(mxB, i, j); + } + } + ret = enif_make_resource(env, mxS); + enif_release_resource(env, mxS); + return ret; +} + +static ERL_NIF_TERM size_of(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) +{ + /* size(Matrix) -> {Nrows, Ncols} */ + Matrix* mx; + if (!enif_get_resource(env, argv[0], resource_type, (void**)&mx)) { + return enif_make_badarg(env); + } + return enif_make_tuple2(env, enif_make_uint(env, mx->nrows), + enif_make_uint(env, mx->ncols)); +} + +static ERL_NIF_TERM to_term(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) +{ + /* to_term(Matrix) -> [[first row], [second row], ...,[last row]] */ + unsigned i, j; + ERL_NIF_TERM res; + Matrix* mx = NULL; + + if (!enif_get_resource(env, argv[0], resource_type, (void**)&mx)) { + return enif_make_badarg(env); + } + res = enif_make_list(env, 0); + for (i = mx->nrows; i-- > 0; ) { + ERL_NIF_TERM row = enif_make_list(env, 0); + for (j = mx->ncols; j-- > 0; ) { + row = enif_make_list_cell(env, enif_make_double(env, POS(mx,i,j)), + row); + } + res = enif_make_list_cell(env, row, res); + } + return res; +} + +static int get_number(ErlNifEnv* env, ERL_NIF_TERM term, double* dp) +{ + long i; + return enif_get_double(env, term, dp) || + (enif_get_long(env, term, &i) && (*dp=(double)i, 1)); +} + +static Matrix* alloc_matrix(ErlNifEnv* env, unsigned nrows, unsigned ncols) +{ + Matrix* mx = enif_alloc_resource(env, resource_type, sizeof(Matrix)); + mx->nrows = nrows; + mx->ncols = ncols; + mx->data = enif_alloc(env, nrows*ncols*sizeof(double)); + return mx; +} + +static void matrix_dtor(ErlNifEnv* env, void* obj) +{ + Matrix* mx = (Matrix*) obj; + enif_free(env, mx->data); + mx->data = NULL; +} + +static ErlNifFunc nif_funcs[] = +{ + {"create", 3, create}, + {"pos", 3, pos}, + {"add", 2, add}, + {"size_of", 1, size_of}, + {"to_term", 1, to_term} +}; + +ERL_NIF_INIT(matrix_nif,nif_funcs,load,NULL,NULL,NULL); + |