chromaprint/c_src/chromaprint_nif.cpp

#include <chromaprint.h>
#include <erl_nif.h>

#include <cstdint>
#include <cstring>
#include <vector>

namespace {

ErlNifResourceType *CHROMAPRINT_CTX_RES = nullptr;

struct ChromaprintCtxRes {
  ChromaprintContext *ctx;
};

ERL_NIF_TERM ATOM_OK;
ERL_NIF_TERM ATOM_ERROR;
ERL_NIF_TERM ATOM_TRUE;
ERL_NIF_TERM ATOM_FALSE;
ERL_NIF_TERM ATOM_BADARG;
ERL_NIF_TERM ATOM_ALLOC_FAILED;
ERL_NIF_TERM ATOM_CHROMAPRINT_FAILED;
ERL_NIF_TERM ATOM_NIL;

inline ERL_NIF_TERM mk_atom(ErlNifEnv *env, const char *name) {
  ERL_NIF_TERM atom;
  if (enif_make_existing_atom(env, name, &atom, ERL_NIF_LATIN1)) return atom;
  return enif_make_atom(env, name);
}

inline ERL_NIF_TERM ok_tuple(ErlNifEnv *env, ERL_NIF_TERM value) {
  return enif_make_tuple2(env, ATOM_OK, value);
}

inline ERL_NIF_TERM error_tuple(ErlNifEnv *env, ERL_NIF_TERM reason) {
  return enif_make_tuple2(env, ATOM_ERROR, reason);
}

void destruct_ctx(ErlNifEnv * /*env*/, void *obj) {
  auto *res = static_cast<ChromaprintCtxRes *>(obj);
  if (res->ctx != nullptr) {
    chromaprint_free(res->ctx);
    res->ctx = nullptr;
  }
}

bool get_bool(ErlNifEnv *env, ERL_NIF_TERM term, int *out) {
  if (enif_compare(term, ATOM_TRUE) == 0) {
    *out = 1;
    return true;
  }
  if (enif_compare(term, ATOM_FALSE) == 0) {
    *out = 0;
    return true;
  }
  return false;
}

ERL_NIF_TERM new_context(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  int algorithm;
  if (!enif_get_int(env, argv[0], &algorithm)) {
    return enif_make_badarg(env);
  }
  ChromaprintContext *ctx = chromaprint_new(algorithm);
  if (ctx == nullptr) {
    return error_tuple(env, ATOM_ALLOC_FAILED);
  }
  auto *res = static_cast<ChromaprintCtxRes *>(
      enif_alloc_resource(CHROMAPRINT_CTX_RES, sizeof(ChromaprintCtxRes)));
  res->ctx = ctx;
  ERL_NIF_TERM term = enif_make_resource(env, res);
  enif_release_resource(res);
  return ok_tuple(env, term);
}

ERL_NIF_TERM nif_start(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  ChromaprintCtxRes *res;
  int sample_rate;
  int channels;
  if (!enif_get_resource(env, argv[0], CHROMAPRINT_CTX_RES, (void **)&res) ||
      !enif_get_int(env, argv[1], &sample_rate) ||
      !enif_get_int(env, argv[2], &channels)) {
    return enif_make_badarg(env);
  }
  if (chromaprint_start(res->ctx, sample_rate, channels) != 1) {
    return error_tuple(env, ATOM_CHROMAPRINT_FAILED);
  }
  return ATOM_OK;
}

ERL_NIF_TERM nif_feed(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  ChromaprintCtxRes *res;
  ErlNifBinary bin;
  if (!enif_get_resource(env, argv[0], CHROMAPRINT_CTX_RES, (void **)&res) ||
      !enif_inspect_binary(env, argv[1], &bin)) {
    return enif_make_badarg(env);
  }
  if (bin.size % 2 != 0) {
    return error_tuple(env, ATOM_BADARG);
  }
  const auto *samples = reinterpret_cast<const int16_t *>(bin.data);
  const int sample_count = static_cast<int>(bin.size / 2);
  if (chromaprint_feed(res->ctx, samples, sample_count) != 1) {
    return error_tuple(env, ATOM_CHROMAPRINT_FAILED);
  }
  return ATOM_OK;
}

ERL_NIF_TERM nif_finish(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  ChromaprintCtxRes *res;
  if (!enif_get_resource(env, argv[0], CHROMAPRINT_CTX_RES, (void **)&res)) {
    return enif_make_badarg(env);
  }
  if (chromaprint_finish(res->ctx) != 1) {
    return error_tuple(env, ATOM_CHROMAPRINT_FAILED);
  }
  return ATOM_OK;
}

ERL_NIF_TERM nif_get_fingerprint(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  ChromaprintCtxRes *res;
  if (!enif_get_resource(env, argv[0], CHROMAPRINT_CTX_RES, (void **)&res)) {
    return enif_make_badarg(env);
  }
  char *fp = nullptr;
  if (chromaprint_get_fingerprint(res->ctx, &fp) != 1 || fp == nullptr) {
    if (fp != nullptr) chromaprint_dealloc(fp);
    return error_tuple(env, ATOM_CHROMAPRINT_FAILED);
  }
  const std::size_t len = std::strlen(fp);
  ERL_NIF_TERM bin_term;
  unsigned char *buf = enif_make_new_binary(env, len, &bin_term);
  std::memcpy(buf, fp, len);
  chromaprint_dealloc(fp);
  return ok_tuple(env, bin_term);
}

ERL_NIF_TERM nif_get_raw_fingerprint(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  ChromaprintCtxRes *res;
  if (!enif_get_resource(env, argv[0], CHROMAPRINT_CTX_RES, (void **)&res)) {
    return enif_make_badarg(env);
  }
  uint32_t *fp = nullptr;
  int size = 0;
  if (chromaprint_get_raw_fingerprint(res->ctx, &fp, &size) != 1 || fp == nullptr) {
    if (fp != nullptr) chromaprint_dealloc(fp);
    return error_tuple(env, ATOM_CHROMAPRINT_FAILED);
  }
  std::vector<ERL_NIF_TERM> items;
  items.reserve(size);
  for (int i = 0; i < size; ++i) {
    items.push_back(enif_make_uint(env, fp[i]));
  }
  chromaprint_dealloc(fp);
  ERL_NIF_TERM list =
      items.empty() ? enif_make_list(env, 0)
                    : enif_make_list_from_array(env, items.data(),
                                                static_cast<unsigned>(items.size()));
  return ok_tuple(env, list);
}

ERL_NIF_TERM nif_get_fingerprint_hash(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  ChromaprintCtxRes *res;
  if (!enif_get_resource(env, argv[0], CHROMAPRINT_CTX_RES, (void **)&res)) {
    return enif_make_badarg(env);
  }
  uint32_t hash = 0;
  if (chromaprint_get_fingerprint_hash(res->ctx, &hash) != 1) {
    return error_tuple(env, ATOM_CHROMAPRINT_FAILED);
  }
  return ok_tuple(env, enif_make_uint(env, hash));
}

ERL_NIF_TERM nif_encode_fingerprint(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  unsigned int list_len = 0;
  if (!enif_get_list_length(env, argv[0], &list_len)) {
    return enif_make_badarg(env);
  }
  int algorithm;
  int base64;
  if (!enif_get_int(env, argv[1], &algorithm) || !get_bool(env, argv[2], &base64)) {
    return enif_make_badarg(env);
  }
  std::vector<uint32_t> raw;
  raw.reserve(list_len);
  ERL_NIF_TERM head;
  ERL_NIF_TERM tail = argv[0];
  while (enif_get_list_cell(env, tail, &head, &tail)) {
    unsigned int val;
    if (!enif_get_uint(env, head, &val)) {
      return enif_make_badarg(env);
    }
    raw.push_back(val);
  }
  char *encoded = nullptr;
  int encoded_size = 0;
  if (chromaprint_encode_fingerprint(raw.data(), static_cast<int>(raw.size()), algorithm,
                                     &encoded, &encoded_size, base64) != 1 ||
      encoded == nullptr) {
    if (encoded != nullptr) chromaprint_dealloc(encoded);
    return error_tuple(env, ATOM_CHROMAPRINT_FAILED);
  }
  ERL_NIF_TERM bin_term;
  unsigned char *buf =
      enif_make_new_binary(env, static_cast<std::size_t>(encoded_size), &bin_term);
  std::memcpy(buf, encoded, static_cast<std::size_t>(encoded_size));
  chromaprint_dealloc(encoded);
  return ok_tuple(env, bin_term);
}

ERL_NIF_TERM nif_decode_fingerprint(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM argv[]) {
  ErlNifBinary bin;
  int base64;
  if (!enif_inspect_binary(env, argv[0], &bin) || !get_bool(env, argv[1], &base64)) {
    return enif_make_badarg(env);
  }
  uint32_t *raw = nullptr;
  int size = 0;
  int algorithm = 0;
  if (chromaprint_decode_fingerprint(reinterpret_cast<const char *>(bin.data),
                                     static_cast<int>(bin.size), &raw, &size, &algorithm,
                                     base64) != 1 ||
      raw == nullptr) {
    if (raw != nullptr) chromaprint_dealloc(raw);
    return error_tuple(env, ATOM_CHROMAPRINT_FAILED);
  }
  std::vector<ERL_NIF_TERM> items;
  items.reserve(size);
  for (int i = 0; i < size; ++i) {
    items.push_back(enif_make_uint(env, raw[i]));
  }
  chromaprint_dealloc(raw);
  ERL_NIF_TERM list =
      items.empty() ? enif_make_list(env, 0)
                    : enif_make_list_from_array(env, items.data(),
                                                static_cast<unsigned>(items.size()));
  ERL_NIF_TERM pair = enif_make_tuple2(env, enif_make_int(env, algorithm), list);
  return ok_tuple(env, pair);
}

ERL_NIF_TERM nif_version(ErlNifEnv *env, int /*argc*/, const ERL_NIF_TERM * /*argv*/) {
  const char *v = chromaprint_get_version();
  if (v == nullptr) return ATOM_NIL;
  const std::size_t len = std::strlen(v);
  ERL_NIF_TERM bin_term;
  unsigned char *buf = enif_make_new_binary(env, len, &bin_term);
  std::memcpy(buf, v, len);
  return bin_term;
}

int load(ErlNifEnv *env, void ** /*priv_data*/, ERL_NIF_TERM /*load_info*/) {
  CHROMAPRINT_CTX_RES = enif_open_resource_type(env, nullptr, "ChromaprintContext",
                                                destruct_ctx, ERL_NIF_RT_CREATE, nullptr);
  if (CHROMAPRINT_CTX_RES == nullptr) return -1;

  ATOM_OK = mk_atom(env, "ok");
  ATOM_ERROR = mk_atom(env, "error");
  ATOM_TRUE = mk_atom(env, "true");
  ATOM_FALSE = mk_atom(env, "false");
  ATOM_BADARG = mk_atom(env, "badarg");
  ATOM_ALLOC_FAILED = mk_atom(env, "alloc_failed");
  ATOM_CHROMAPRINT_FAILED = mk_atom(env, "chromaprint_failed");
  ATOM_NIL = mk_atom(env, "nil");
  return 0;
}

ErlNifFunc funcs[] = {
    {"new_context", 1, new_context, 0},
    {"start", 3, nif_start, 0},
    {"feed", 2, nif_feed, ERL_NIF_DIRTY_JOB_CPU_BOUND},
    {"finish", 1, nif_finish, ERL_NIF_DIRTY_JOB_CPU_BOUND},
    {"get_fingerprint", 1, nif_get_fingerprint, 0},
    {"get_raw_fingerprint", 1, nif_get_raw_fingerprint, 0},
    {"get_fingerprint_hash", 1, nif_get_fingerprint_hash, 0},
    {"encode_fingerprint", 3, nif_encode_fingerprint, ERL_NIF_DIRTY_JOB_CPU_BOUND},
    {"decode_fingerprint", 2, nif_decode_fingerprint, ERL_NIF_DIRTY_JOB_CPU_BOUND},
    {"version", 0, nif_version, 0},
};

}  // namespace

ERL_NIF_INIT(Elixir.Chromaprint.NIF, funcs, load, nullptr, nullptr, nullptr)