erlang各种有用的函数包括一些有用nif封装,还有一些性能测试case。
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

601 lines
14 KiB

5 years ago
  1. // Licensed under the Apache License, Version 2.0 (the "License"); you may not
  2. // use this file except in compliance with the License. You may obtain a copy of
  3. // the License at
  4. //
  5. // http://www.apache.org/licenses/LICENSE-2.0
  6. //
  7. // Unless required by applicable law or agreed to in writing, software
  8. // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  9. // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  10. // License for the specific language governing permissions and limitations under
  11. // the License.
  12. #include <assert.h>
  13. #include <string.h>
  14. #include <stdio.h>
  15. #include "hqueue.h"
  16. typedef struct
  17. {
  18. ERL_NIF_TERM atom_ok;
  19. ERL_NIF_TERM atom_error;
  20. ERL_NIF_TERM atom_value;
  21. ERL_NIF_TERM atom_empty;
  22. ERL_NIF_TERM atom_full;
  23. ERL_NIF_TERM atom_max_elems;
  24. ERL_NIF_TERM atom_heap_size;
  25. ERL_NIF_TERM atom_too_small;
  26. ErlNifResourceType* res_hqueue;
  27. } hqueue_priv;
  28. typedef struct
  29. {
  30. ErlNifEnv* env;
  31. ERL_NIF_TERM value;
  32. } hqnode_nif_t;
  33. typedef struct
  34. {
  35. int version;
  36. uint64_t gen;
  37. hqueue_t* hqueue;
  38. ErlNifPid p;
  39. } hqueue_nif_t;
  40. static const uint32_t default_max_elems = UINT32_MAX-1;
  41. static const uint32_t default_heap_size = 1024;
  42. static inline ERL_NIF_TERM
  43. make_atom(ErlNifEnv* env, const char* name)
  44. {
  45. ERL_NIF_TERM ret;
  46. if(enif_make_existing_atom(env, name, &ret, ERL_NIF_LATIN1)) {
  47. return ret;
  48. }
  49. return enif_make_atom(env, name);
  50. }
  51. static inline ERL_NIF_TERM
  52. make_ok(ErlNifEnv* env, hqueue_priv* priv, ERL_NIF_TERM value)
  53. {
  54. return enif_make_tuple2(env, priv->atom_ok, value);
  55. }
  56. static inline ERL_NIF_TERM
  57. make_error(ErlNifEnv* env, hqueue_priv* priv, ERL_NIF_TERM reason)
  58. {
  59. return enif_make_tuple2(env, priv->atom_error, reason);
  60. }
  61. static inline int
  62. check_pid(ErlNifEnv* env, hqueue_nif_t* hqueue_nif)
  63. {
  64. ErlNifPid pid;
  65. enif_self(env, &pid);
  66. if(enif_compare(pid.pid, hqueue_nif->p.pid) == 0) {
  67. return 1;
  68. }
  69. return 0;
  70. }
  71. void
  72. hqueue_nif_node_free(hqnode_nif_t* hqnode_nif)
  73. {
  74. enif_free_env(hqnode_nif->env);
  75. enif_free(hqnode_nif);
  76. return;
  77. }
  78. void
  79. hqueue_nif_node_free_ext(void* node)
  80. {
  81. hqueue_nif_node_free((hqnode_nif_t*) node);
  82. return;
  83. }
  84. hqnode_nif_t*
  85. hqueue_nif_node_alloc()
  86. {
  87. hqnode_nif_t* node = (hqnode_nif_t*) enif_alloc(sizeof(hqnode_nif_t*));
  88. memset(node, 0, sizeof(hqnode_nif_t));
  89. node->env = enif_alloc_env();
  90. return node;
  91. }
  92. static int
  93. get_uint_param(ErlNifEnv* env, ERL_NIF_TERM value, ERL_NIF_TERM atom, uint32_t* p)
  94. {
  95. const ERL_NIF_TERM* tuple;
  96. int arity;
  97. if(!enif_get_tuple(env, value, &arity, &tuple)) {
  98. return 0;
  99. }
  100. if(arity != 2) {
  101. return 0;
  102. }
  103. if(enif_compare(tuple[0], atom) != 0) {
  104. return 0;
  105. }
  106. if(!enif_get_uint(env, tuple[1], p)) {
  107. return 0;
  108. }
  109. return 1;
  110. }
  111. static inline hqueue_nif_t*
  112. hqueue_nif_create_int(ErlNifEnv* env, hqueue_priv* priv, uint32_t max_elems,
  113. uint32_t heap_size)
  114. {
  115. hqueue_nif_t* hqueue_nif = NULL;
  116. assert(priv != NULL && "missing private data member");
  117. hqueue_nif = (hqueue_nif_t*) enif_alloc_resource(
  118. priv->res_hqueue, sizeof(hqueue_nif_t));
  119. memset(hqueue_nif, 0, sizeof(hqueue_nif_t));
  120. hqueue_nif->version = HQ_VERSION;
  121. hqueue_nif->hqueue = hqueue_new(max_elems, heap_size);
  122. if(hqueue_nif->hqueue == NULL ) {
  123. enif_release_resource(hqueue_nif);
  124. return NULL;
  125. }
  126. enif_self(env, &(hqueue_nif->p));
  127. return hqueue_nif;
  128. }
  129. static ERL_NIF_TERM
  130. hqueue_nif_new(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  131. {
  132. hqueue_priv* priv = enif_priv_data(env);
  133. hqueue_nif_t* hqueue_nif;
  134. ERL_NIF_TERM ret;
  135. ERL_NIF_TERM opts;
  136. ERL_NIF_TERM value;
  137. uint32_t max_elems = default_max_elems;
  138. uint32_t heap_size = default_heap_size;
  139. if(argc != 1) {
  140. return enif_make_badarg(env);
  141. }
  142. opts = argv[0];
  143. if(!enif_is_list(env, opts)) {
  144. return enif_make_badarg(env);
  145. }
  146. while(enif_get_list_cell(env, opts, &value, &opts)) {
  147. if(get_uint_param(env, value, priv->atom_max_elems, &max_elems)) {
  148. continue;
  149. } else if(get_uint_param(env, value, priv->atom_heap_size, &heap_size)) {
  150. continue;
  151. } else {
  152. return enif_make_badarg(env);
  153. }
  154. }
  155. hqueue_nif = hqueue_nif_create_int(env, priv, max_elems, heap_size);
  156. if(hqueue_nif == NULL) {
  157. return enif_make_badarg(env);
  158. }
  159. ret = enif_make_resource(env, hqueue_nif);
  160. enif_release_resource(hqueue_nif);
  161. return make_ok(env, priv, ret);
  162. }
  163. static void
  164. hqueue_nif_free(ErlNifEnv* env, void* obj)
  165. {
  166. hqueue_nif_t* hqueue_nif = (hqueue_nif_t*) obj;
  167. hqueue_free2(hqueue_nif->hqueue, hqueue_nif_node_free_ext);
  168. return;
  169. }
  170. static ERL_NIF_TERM
  171. hqueue_nif_extract_max(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  172. {
  173. hqueue_priv* priv = enif_priv_data(env);
  174. hqueue_nif_t* hqueue_nif;
  175. hqnode_nif_t* hqnode_nif;
  176. double tmp_priority;
  177. ERL_NIF_TERM ret;
  178. ERL_NIF_TERM priority;
  179. ERL_NIF_TERM value;
  180. if(argc != 1) {
  181. return enif_make_badarg(env);
  182. }
  183. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  184. return enif_make_badarg(env);
  185. }
  186. if(!check_pid(env, hqueue_nif)) {
  187. return enif_make_badarg(env);
  188. }
  189. if (!hqueue_extract_max(hqueue_nif->hqueue, &tmp_priority, (void**) &hqnode_nif)) {
  190. return make_error(env, priv, priv->atom_empty);
  191. }
  192. priority = enif_make_double(env, tmp_priority);
  193. value = enif_make_copy(env, hqnode_nif->value);
  194. ret = enif_make_tuple2(env, priority, value);
  195. hqueue_nif_node_free(hqnode_nif);
  196. return ret;
  197. }
  198. static ERL_NIF_TERM
  199. hqueue_nif_insert(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  200. {
  201. hqueue_priv* priv = enif_priv_data(env);
  202. hqueue_nif_t* hqueue_nif;
  203. hqnode_nif_t* hqnode_nif;
  204. ERL_NIF_TERM ret;
  205. double priority;
  206. if(argc != 3) {
  207. return enif_make_badarg(env);
  208. }
  209. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  210. return enif_make_badarg(env);
  211. }
  212. if(!check_pid(env, hqueue_nif)) {
  213. return enif_make_badarg(env);
  214. }
  215. if(!enif_get_double(env, argv[1], &priority)) {
  216. return enif_make_badarg(env);
  217. }
  218. if(priority < 0.0) {
  219. return enif_make_badarg(env);
  220. }
  221. hqnode_nif = hqueue_nif_node_alloc();
  222. hqnode_nif->value = enif_make_copy(hqnode_nif->env, argv[2]);
  223. if (!hqueue_insert(hqueue_nif->hqueue, priority, (void*) hqnode_nif)) {
  224. return make_error(env, priv, priv->atom_full);
  225. }
  226. ret = priv->atom_ok;
  227. return ret;
  228. }
  229. static ERL_NIF_TERM
  230. hqueue_nif_size(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  231. {
  232. hqueue_priv* priv = enif_priv_data(env);
  233. hqueue_nif_t* hqueue_nif;
  234. ERL_NIF_TERM ret;
  235. if(argc != 1) {
  236. return enif_make_badarg(env);
  237. }
  238. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  239. return enif_make_badarg(env);
  240. }
  241. if(!check_pid(env, hqueue_nif)) {
  242. return enif_make_badarg(env);
  243. }
  244. ret = enif_make_uint64(env, hqueue_size(hqueue_nif->hqueue));
  245. return ret;
  246. }
  247. static ERL_NIF_TERM
  248. hqueue_nif_heap_size(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  249. {
  250. hqueue_priv* priv = enif_priv_data(env);
  251. hqueue_nif_t* hqueue_nif;
  252. ERL_NIF_TERM ret;
  253. if(argc != 1) {
  254. return enif_make_badarg(env);
  255. }
  256. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  257. return enif_make_badarg(env);
  258. }
  259. if(!check_pid(env, hqueue_nif)) {
  260. return enif_make_badarg(env);
  261. }
  262. ret = enif_make_uint64(env, hqueue_heap_size(hqueue_nif->hqueue));
  263. return ret;
  264. }
  265. static ERL_NIF_TERM
  266. hqueue_nif_max_elems(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  267. {
  268. hqueue_priv* priv = enif_priv_data(env);
  269. hqueue_nif_t* hqueue_nif;
  270. ERL_NIF_TERM ret;
  271. if(argc != 1) {
  272. return enif_make_badarg(env);
  273. }
  274. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  275. return enif_make_badarg(env);
  276. }
  277. if(!check_pid(env, hqueue_nif)) {
  278. return enif_make_badarg(env);
  279. }
  280. ret = enif_make_uint64(env, hqueue_max_elems(hqueue_nif->hqueue));
  281. return ret;
  282. }
  283. static ERL_NIF_TERM
  284. hqueue_nif_to_list(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  285. {
  286. hqueue_priv* priv = enif_priv_data(env);
  287. hqueue_nif_t* hqueue_nif;
  288. hqueue_t* hqueue;
  289. hqnode_nif_t* hqnode_nif;
  290. double tmp_priority;
  291. ERL_NIF_TERM ret = enif_make_list(env, 0);
  292. ERL_NIF_TERM priority;
  293. ERL_NIF_TERM value;
  294. ERL_NIF_TERM tuple;
  295. uint32_t i;
  296. if(argc != 1) {
  297. return enif_make_badarg(env);
  298. }
  299. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  300. return enif_make_badarg(env);
  301. }
  302. if(!check_pid(env, hqueue_nif)) {
  303. return enif_make_badarg(env);
  304. }
  305. hqueue = hqueue_nif->hqueue;
  306. for (i = 1; i <= hqueue_size(hqueue); i++) {
  307. hqueue_get_elem(hqueue, i, &tmp_priority, (void **) &hqnode_nif);
  308. priority = enif_make_double(env, tmp_priority);
  309. value = enif_make_copy(env, hqnode_nif->value);
  310. tuple = enif_make_tuple2(env, priority, value);
  311. ret = enif_make_list_cell(env, tuple, ret);
  312. }
  313. return ret;
  314. }
  315. static ERL_NIF_TERM
  316. hqueue_nif_scale_by(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  317. {
  318. hqueue_priv* priv = enif_priv_data(env);
  319. hqueue_nif_t* hqueue_nif;
  320. ERL_NIF_TERM ret;
  321. double factor;
  322. if(argc != 2) {
  323. return enif_make_badarg(env);
  324. }
  325. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  326. return enif_make_badarg(env);
  327. }
  328. if(!check_pid(env, hqueue_nif)) {
  329. return enif_make_badarg(env);
  330. }
  331. if(!enif_get_double(env, argv[1], &factor)) {
  332. return enif_make_badarg(env);
  333. }
  334. if(factor < 0.0) {
  335. return enif_make_badarg(env);
  336. }
  337. hqueue_scale_by(hqueue_nif->hqueue, factor);
  338. ret = priv->atom_ok;
  339. return ret;
  340. }
  341. static ERL_NIF_TERM
  342. hqueue_nif_resize_heap(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  343. {
  344. hqueue_priv* priv = enif_priv_data(env);
  345. hqueue_nif_t* hqueue_nif;
  346. ERL_NIF_TERM ret;
  347. uint32_t new_heap_size;
  348. uint32_t old_heap_size;
  349. if(argc != 2) {
  350. return enif_make_badarg(env);
  351. }
  352. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  353. return enif_make_badarg(env);
  354. }
  355. if(!check_pid(env, hqueue_nif)) {
  356. return enif_make_badarg(env);
  357. }
  358. if(!enif_get_uint(env, argv[1], &new_heap_size)) {
  359. return enif_make_badarg(env);
  360. }
  361. if(hqueue_size(hqueue_nif->hqueue) > new_heap_size) {
  362. return make_error(env, priv, priv->atom_too_small);
  363. }
  364. if((old_heap_size = hqueue_resize_heap(hqueue_nif->hqueue, new_heap_size)) == 0) {
  365. return enif_make_badarg(env);
  366. }
  367. ret = enif_make_uint64(env, old_heap_size);
  368. return ret;
  369. }
  370. static ERL_NIF_TERM
  371. hqueue_nif_set_max_elems(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
  372. {
  373. hqueue_priv* priv = enif_priv_data(env);
  374. hqueue_nif_t* hqueue_nif;
  375. ERL_NIF_TERM ret;
  376. uint32_t new_max_elems;
  377. uint32_t old_max_elems;
  378. if(argc != 2) {
  379. return enif_make_badarg(env);
  380. }
  381. if(!enif_get_resource(env, argv[0], priv->res_hqueue, (void**) &hqueue_nif)) {
  382. return enif_make_badarg(env);
  383. }
  384. if(!check_pid(env, hqueue_nif)) {
  385. return enif_make_badarg(env);
  386. }
  387. if(!enif_get_uint(env, argv[1], &new_max_elems)) {
  388. return enif_make_badarg(env);
  389. }
  390. if(hqueue_size(hqueue_nif->hqueue) > new_max_elems) {
  391. return make_error(env, priv, priv->atom_too_small);
  392. }
  393. if ((old_max_elems = hqueue_set_max_elems(hqueue_nif->hqueue, new_max_elems)) == 0) {
  394. return enif_make_badarg(env);
  395. }
  396. ret = enif_make_uint64(env, old_max_elems);
  397. return ret;
  398. }
  399. static int
  400. load(ErlNifEnv* env, void** priv, ERL_NIF_TERM info)
  401. {
  402. int flags = ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER;
  403. ErlNifResourceType* res;
  404. hqueue_priv* new_priv = (hqueue_priv*) enif_alloc(sizeof(hqueue_priv));
  405. if(new_priv == NULL) {
  406. return 1;
  407. }
  408. res = enif_open_resource_type(
  409. env, NULL, "hqueue", hqueue_nif_free, flags, NULL);
  410. if(res == NULL) {
  411. enif_free(new_priv);
  412. return 1;
  413. }
  414. new_priv->res_hqueue = res;
  415. new_priv->atom_ok = make_atom(env, "ok");
  416. new_priv->atom_error = make_atom(env, "error");
  417. new_priv->atom_value = make_atom(env, "value");
  418. new_priv->atom_empty = make_atom(env, "empty");
  419. new_priv->atom_full = make_atom(env, "full");
  420. new_priv->atom_max_elems = make_atom(env, "max_elems");
  421. new_priv->atom_heap_size = make_atom(env, "heap_size");
  422. new_priv->atom_too_small = make_atom(env, "too_small");
  423. *priv = (void*) new_priv;
  424. return 0;
  425. }
  426. static int
  427. upgrade(ErlNifEnv* env, void** priv, void** old_priv, ERL_NIF_TERM info)
  428. {
  429. return load(env, priv, info);
  430. }
  431. static void
  432. unload(ErlNifEnv* env, void* priv)
  433. {
  434. enif_free(priv);
  435. return;
  436. }
  437. static ErlNifFunc funcs[] = {
  438. {"new", 1, hqueue_nif_new},
  439. {"extract_max", 1, hqueue_nif_extract_max},
  440. {"insert", 3, hqueue_nif_insert},
  441. {"size", 1, hqueue_nif_size},
  442. {"heap_size", 1, hqueue_nif_heap_size},
  443. {"max_elems", 1, hqueue_nif_max_elems},
  444. {"set_max_elems", 2, hqueue_nif_set_max_elems},
  445. {"to_list", 1, hqueue_nif_to_list},
  446. {"scale_by", 2, hqueue_nif_scale_by},
  447. {"resize_heap", 2, hqueue_nif_resize_heap}
  448. };
  449. ERL_NIF_INIT(hqueue, funcs, &load, NULL, &upgrade, &unload);