#include <Rinternals.h>
#include <string.h>             /* for memcpy */

static void _add_names(SEXP ans, SEXP v)
{
    SEXP ans_nms, v_nms = Rf_getAttrib(v, R_NamesSymbol);
    int i, len = LENGTH(ans), *idx;
    if (v_nms != R_NilValue) {
        PROTECT(ans_nms = Rf_allocVector(STRSXP, len));
        idx = INTEGER(ans);
        for (i = 0; i < len; i++) {
            SET_STRING_ELT(ans_nms, i, STRING_ELT(v_nms, idx[i] - 1));
        }
        Rf_setAttrib(ans, R_NamesSymbol, ans_nms);
        UNPROTECT(1);
    }
}

SEXP nid_which(SEXP v)
{
    SEXP ans;
    int i, j = 0, len, *buf, *tf;

    if (!Rf_isLogical(v))
        Rf_error("argument to 'nid_which' is not logical");
    len = Rf_length(v);
    buf = (int *) R_alloc(len, sizeof(int));
    tf = LOGICAL(v);

    for (i = 0; i < len; i++) {
        if (tf[i] == TRUE) {
            buf[j] = i + 1;
            j++;
        }
    }

    len = j;
    PROTECT(ans = Rf_allocVector(INTSXP, len));
    memcpy(INTEGER(ans), buf, sizeof(int) * len);

    _add_names(ans, v);

    UNPROTECT(1);
    return ans;
}
