mecab_drv.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mecab.h>
#include "erl_driver.h"
#include "ei.h"
#include "erl_interface.h"

#define OP_BEGIN_PARSING 1
#define OP_GET_VERSION   2
#define OP_GET_DIC_INFO  3

#define ERROR(reason) \
  ei_x_new_with_version(x); \
  ei_x_encode_tuple_header(x, 2); \
  ei_x_encode_atom(x, "error"); \
  reason; \
  return

#define NOT_ALLOC ERROR(ei_x_encode_atom(x, "not_alloc"))
#define BADARG ERROR(ei_x_encode_atom(x, "badarg"))

#define REPLY(a, message) \
  ei_x_new_with_version(&x); \
  ei_x_encode_tuple_header(&x, 2); \
  ei_x_encode_atom(&x, "mecab"); \
  ei_x_encode_tuple_header(&x, 2); \
  ei_x_encode_ref(&x, a->id); \
  message; \
  driver_output(a->port, x.buff, x.index);

typedef struct {
  ErlDrvPort port;
} mecab_data;

typedef struct {
  ErlDrvPort port;
  mecab_t* mecab;
  erlang_ref* id;
  char* sentence;
} mecab_async_data;

static void begin_parsing(ei_x_buff*, const char*, const ErlDrvPort);
static void get_version(ei_x_buff*);
static void get_dic_info(ei_x_buff*, const char*);

static int get_dic_info_args(const char*, int*, char***);
static int get_dicdir_args(char*, int*, char***);
static void encode_ok(ei_x_buff*);
static void encode_error(ei_x_buff*);

static ErlDrvData start(ErlDrvPort port, char* buff) {
  mecab_data* d = (mecab_data*)driver_alloc(sizeof(mecab_data));
  d->port = port;
  set_port_control_flags(port, PORT_CONTROL_FLAG_BINARY);
  return (ErlDrvData)d;
}

static void stop(ErlDrvData handle) {
  driver_free((char*)handle);
}

static int control(ErlDrvData handle, unsigned int command,
                   char* buff, int len, char** rbuff, int rlen) {
  mecab_data* d = (mecab_data*)handle;
  ErlDrvBinary* bin;
  ei_x_buff x;

  switch (command) {
  case OP_BEGIN_PARSING:
    begin_parsing(&x, buff, d->port);
    break;
  case OP_GET_VERSION:
    get_version(&x);
    break;
  case OP_GET_DIC_INFO:
    get_dic_info(&x, buff);
    break;
  default:
    break;
  }

  bin = driver_alloc_binary(x.index);

  if (bin != NULL) {
    memcpy(bin->orig_bytes, x.buff, x.index);
    *rbuff = (char*)bin;
    rlen = bin->orig_size;
  }
  else {
    *rbuff = NULL;
  }

  ei_x_free(&x);
  return rlen;
}

static void begin_parsing(ei_x_buff* x, const char* buff, const ErlDrvPort port) {
  // Not implement yet
}

static void get_version(ei_x_buff* x) {
  ei_x_new_with_version(x);
  ei_x_encode_string(x, mecab_version());
}

static void get_dic_info(ei_x_buff* x, const char* buff) {
  int argc;
  char** argv;
  mecab_t* mecab;
  const mecab_dictionary_info_t* info;

  switch (get_dic_info_args(buff, &argc, &argv)) {
  case 1:
    NOT_ALLOC;
  case 2:
    BADARG;
  default:
    break;
  }

  mecab = mecab_new(argc, argv);

  if (mecab != NULL) {
    info = mecab_dictionary_info(mecab);
    encode_ok(x);
    ei_x_encode_tuple_header(x, 7);
    ei_x_encode_string(x, info->filename);
    ei_x_encode_string(x, info->charset);
    ei_x_encode_ulong(x, info->size);

    switch (info->type) {
    case MECAB_USR_DIC:
      ei_x_encode_atom(x, "usr");
      break;
    case MECAB_SYS_DIC:
      ei_x_encode_atom(x, "sys");
      break;
    case MECAB_UNK_DIC:
      ei_x_encode_atom(x, "unk");
      break;
    default:
      ei_x_encode_atom(x, "unknown");
      break;
    }

    ei_x_encode_ulong(x, info->lsize);
    ei_x_encode_ulong(x, info->rsize);
    ei_x_encode_ulong(x, info->version);
    mecab_destroy(mecab);
  }
  else {
    encode_error(x);
    ei_x_encode_tuple_header(x, 2);
    ei_x_encode_atom(x, "mecab_error");
    ei_x_encode_string(x, mecab_strerror(NULL));
  }

  if (argc == 3) free(*(argv + 2));
  free(argv);
}

static int get_dic_info_args(const char* buff, int* argc, char*** argv) {
  int index = 1;
  int type;
  int size;
  char* data;

  ei_get_type(buff, &index, &type, &size);

  switch (type) {
  case 100:
    data = (char*)malloc(sizeof(char) * (size + 1));
    if (data == NULL) return 1;

    if (ei_decode_atom(buff, &index, data) || strcmp(data, "default")) {
      free(data);
      return 2;
    }

    free(data);
    *argv = (char**)malloc(sizeof(char*) * (*argc = 1));
    **argv = "";
    return 0;
  case 106:
    data = (char*)malloc(sizeof(char) * 2);
    if (data == NULL) return 1;
    strcpy(data, ".");
    return get_dicdir_args(data, argc, argv);
  case 107:
    data = (char*)malloc(sizeof(char) * (size + 1));
    if (data == NULL) return 1;

    if (ei_decode_string(buff, &index, data)) {
      free(data);
      return 2;
    }

    return get_dicdir_args(data, argc, argv);
  default:
    *argc = 0;
    return 2;
  }
}

static int get_dicdir_args(char* dicdir, int* argc, char*** argv) {
  *argv = (char**)malloc(sizeof(char*) * (*argc = 3));
  (*argv)[0] = "";
  (*argv)[1] = "-d";
  (*argv)[2] = dicdir;
  return 0;
}

ErlDrvEntry mecab_driver_entry = {
  .init            = NULL,
  .start           = start,
  .stop            = stop,
  .output          = NULL,
  .ready_input     = NULL,
  .ready_output    = NULL,
  .driver_name     = "mecab_drv",
  .finish          = NULL,
  .handle          = NULL,
  .control         = control,
  .timeout         = NULL,
  .outputv         = NULL,
  .ready_async     = NULL,
  .flush           = NULL,
  .call            = NULL,
  .event           = NULL,
  .extended_marker = ERL_DRV_EXTENDED_MARKER,
  .major_version   = ERL_DRV_EXTENDED_MAJOR_VERSION,
  .minor_version   = ERL_DRV_EXTENDED_MINOR_VERSION,
  .driver_flags    = 0,
  .handle2         = NULL,
  .process_exit    = NULL,
  .stop_select     = NULL
};

DRIVER_INIT(mecab_drv) {
  return &mecab_driver_entry;
}

static void encode_ok(ei_x_buff* x) {
  ei_x_new_with_version(x);
  ei_x_encode_tuple_header(x, 2);
  ei_x_encode_atom(x, "ok");
}

static void encode_error(ei_x_buff* x) {
  ei_x_new_with_version(x);
  ei_x_encode_tuple_header(x, 2);
  ei_x_encode_atom(x, "error");
}