Add support for rank

Author: silviucpp

c_src/epqueue.cc

@@ -43,6 +43,17 @@ bool internal_remove(epqueue* q, queue_item* item)
     return q->queue->remove(item, item->heap_index);
 }
 
+bool internal_rank(epqueue* q, queue_item* item, int32_t* rank)
+{
+    CritScope ss(q->crit);
+
+    if(item->heap_index == -1)
+        return false;
+
+    *rank = q->queue->rank(item->heap_index);
+    return true;
+}
+
 queue_item* internal_pop(epqueue* q)
 {
     CritScope ss(q->crit);
@@ -255,3 +266,28 @@ ERL_NIF_TERM nif_epqueue_peek(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[
 
     return enif_make_tuple3(env, ATOMS.atomOk, bin_term, enif_make_int64(env, item->priority));
 }
+
+ERL_NIF_TERM nif_epqueue_rank(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
+{
+    UNUSED(argc);
+
+    epqueue_data* data = static_cast<epqueue_data*>(enif_priv_data(env));
+    epqueue* inst = NULL;
+    queue_item* item = NULL;
+
+    if(!enif_get_resource(env, argv[0], data->resPQueueInstance, reinterpret_cast<void**>(&inst)))
+        return enif_make_badarg(env);
+
+    if(!is_owner(env, inst))
+        return enif_make_badarg(env);
+
+    if(!enif_get_resource(env, argv[1], data->resPQueueItem, reinterpret_cast<void**>(&item)))
+        return enif_make_badarg(env);
+
+    int32_t rank;
+
+    if(!internal_rank(inst, item, &rank))
+        return make_error(env, ATOMS.atomNotFound);
+
+    return enif_make_tuple2(env, ATOMS.atomOk, enif_make_int(env, rank));
+}

c_src/epqueue.h

@@ -10,5 +10,6 @@ ERL_NIF_TERM nif_epqueue_insert(ErlNifEnv* env, int argc, const ERL_NIF_TERM arg
 ERL_NIF_TERM nif_epqueue_remove(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 ERL_NIF_TERM nif_epqueue_pop(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 ERL_NIF_TERM nif_epqueue_peek(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
+ERL_NIF_TERM nif_epqueue_rank(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]);
 
 #endif  // C_SRC_EPQUEUE_H_

c_src/epqueue_nif.cc

@@ -10,6 +10,7 @@ const char kAtomTrue[] = "true";
 const char kAtomFalse[] = "false";
 const char kAtomUndefined[] = "undefined";
 const char kAtomGlobalLock[] = "global_lock";
+const char kAtomNotFound[] = "not_found";
 
 atoms ATOMS;
 
@@ -29,6 +30,7 @@ int on_nif_load(ErlNifEnv* env, void** priv_data, ERL_NIF_TERM load_info)
     ATOMS.atomTrue = make_atom(env, kAtomTrue);
     ATOMS.atomFalse = make_atom(env, kAtomFalse);
     ATOMS.atomUndefined = make_atom(env, kAtomUndefined);
+    ATOMS.atomNotFound = make_atom(env, kAtomNotFound);
 
     ATOMS.atomGlobalLock = make_atom(env, kAtomGlobalLock);
 
@@ -66,7 +68,8 @@ static ErlNifFunc nif_funcs[] =
     {"remove", 2, nif_epqueue_remove},
     {"pop", 1, nif_epqueue_pop},
     {"peek", 1, nif_epqueue_peek},
-    {"size", 1, nif_epqueue_size}
+    {"size", 1, nif_epqueue_size},
+    {"rank", 2, nif_epqueue_rank},
 };
 
 ERL_NIF_INIT(epqueue_nif, nif_funcs, on_nif_load, NULL, on_nif_upgrade, on_nif_unload)

c_src/epqueue_nif.h

@@ -10,6 +10,7 @@ struct atoms
     ERL_NIF_TERM atomTrue;
     ERL_NIF_TERM atomFalse;
     ERL_NIF_TERM atomUndefined;
+    ERL_NIF_TERM atomNotFound;
 
     ERL_NIF_TERM atomGlobalLock;
 };

c_src/nif_utils.cc

@@ -25,3 +25,8 @@ ERL_NIF_TERM make_error(ErlNifEnv* env, const char* error)
 {
     return enif_make_tuple2(env, ATOMS.atomError, make_binary(env, error, strlen(error)));
 }
+
+ERL_NIF_TERM make_error(ErlNifEnv* env, ERL_NIF_TERM term)
+{
+    return enif_make_tuple2(env, ATOMS.atomError, term);
+}

c_src/nif_utils.h

@@ -5,6 +5,7 @@
 
 ERL_NIF_TERM make_atom(ErlNifEnv* env, const char* name);
 ERL_NIF_TERM make_error(ErlNifEnv* env, const char* error);
+ERL_NIF_TERM make_error(ErlNifEnv* env, ERL_NIF_TERM term);
 ERL_NIF_TERM make_binary(ErlNifEnv* env, const char* buff, size_t length);
 
 #endif  // C_SRC_NIF_UTILS_H_

c_src/priority_queue.cc

@@ -125,6 +125,22 @@ void PriorityQueue::bubble_down(int pos)
     }
 }
 
+int PriorityQueue::rank_r(int target_index, int cur_node_index)
+{
+    if (cur_node_index >= length_)
+        return 0;
+
+    // if the current node priority is <= target_node, it is part of the rank count
+    if (less(cur_node_index, target_index))
+    {
+        int left_count = rank_r(target_index, left(cur_node_index));
+        int right_count = rank_r(target_index, right(cur_node_index));
+        return 1 + left_count + right_count;
+    }
+
+    return 0;
+}
+
 void PriorityQueue::bubble_up(int pos)
 {
     while (true)

c_src/priority_queue.h

@@ -20,13 +20,15 @@ class PriorityQueue
     void* pop() {return remove(0);}
     void* peek();
     int size() const { return length_;}
+    int rank(int pos) {return rank_r(pos, 0)+1;}
 
 private:
 
     inline void set(int pos, void* item);
     inline void pos_swap(int pos1, int pos2);
     void bubble_down(int pos);
     void bubble_up(int pos);
+    int rank_r(int x, int cur_node_index);
 
     int capacity_;
     int length_;

src/epqueue.erl

@@ -7,7 +7,8 @@
     insert/3,
     remove/2,
     pop/1,
-    peek/1
+    peek/1,
+    rank/2
 ]).
 
 -type error() :: badarg | {error, binary()}.
@@ -58,3 +59,9 @@ pop(QueueRef) ->
 
 peek(QueueRef) ->
     epqueue_nif:peek(QueueRef).
+
+-spec rank(queue_ref(), data_ref()) ->
+    {ok, non_neg_integer()} | error().
+
+rank(QueueRef, Ref) ->
+    epqueue_nif:rank(QueueRef, Ref).

src/epqueue_nif.erl

@@ -10,7 +10,8 @@
     insert/3,
     remove/2,
     pop/1,
-    peek/1
+    peek/1,
+    rank/2
 ]).
 
 load_nif() ->
@@ -36,6 +37,9 @@ pop(_QueueRef) ->
 peek(_QueueRef) ->
     ?NOT_LOADED.
 
+rank(_QueueRef, _ItemRef) ->
+    ?NOT_LOADED.
+
 % internals
 
 get_priv_path(File) ->

test/integrity_test.erl

@@ -58,3 +58,26 @@ pop_test() ->
     ?assertEqual({ok, 3, 3}, epqueue:pop(Q1)),
     ?assertEqual({ok, 5, 5}, epqueue:pop(Q1)),
     ok.
+
+rank_test() ->
+    {ok, Q1} = epqueue:new([]),
+    {ok, R1} = epqueue:insert(Q1, 3, 1),
+    {ok, R2} = epqueue:insert(Q1, 5, 3),
+    {ok, R3} = epqueue:insert(Q1, 1, 2),
+    {ok, R4} = epqueue:insert(Q1, 1, 1),
+
+    ?assertEqual({ok, 1}, epqueue:rank(Q1, R1)),
+    ?assertEqual({ok, 4}, epqueue:rank(Q1, R2)),
+    ?assertEqual({ok, 3}, epqueue:rank(Q1, R3)),
+    ?assertEqual({ok, 2}, epqueue:rank(Q1, R4)),
+
+    ?assertEqual(true, epqueue:remove(Q1, R1)),
+    ?assertEqual({ok, 1}, epqueue:rank(Q1, R4)),
+
+    ?assertEqual(true, epqueue:remove(Q1, R2)),
+    ?assertEqual(true, epqueue:remove(Q1, R3)),
+
+    ?assertEqual({ok, 1}, epqueue:rank(Q1, R4)),
+    ?assertEqual(true, epqueue:remove(Q1, R4)),
+    ?assertEqual({error, not_found}, epqueue:rank(Q1, R4)),
+    ok.