Merge pull request #3 from recp/red-black-tree-cmp

Red Black Tree: Provide extra function to override compare(cmp) function
recp · Dec 8, 2020 · ad33882 · ad33882
2 parents f1f84ef + 1c3cc76
commit ad33882
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 12 deletions.
diff --git a/include/ds/common.h b/include/ds/common.h
@@ -25,4 +25,8 @@
 #define DS_ITOP(I)    ((void *)(uintptr_t)I)
 #define DS_PTOI(I, T) ((T)(uintptr_t)I)
 
+/* pack key (e.g non-null terminated string) and its length */
+#define DS_KEYnLEN(KEY, LEN)                                                  \
+   ((void *[]){(void *)KEY, (void *)(uintptr_t)LEN})
+
 #endif /* common_h */
diff --git a/include/ds/rb.h b/include/ds/rb.h
@@ -139,8 +139,7 @@ rb_remove(RBTree *tree,
  */
 DS_EXPORT
 void*
-rb_find(RBTree *tree,
-        void   *key);
+rb_find(RBTree * __restrict tree, void * __restrict key);
 
 /*!
  * @brief find node by key. It returns RBnode you can get value from node
@@ -152,8 +151,40 @@ rb_find(RBTree *tree,
  */
 DS_EXPORT
 RBNode*
-rb_find_node(RBTree *tree,
-             void   *key);
+rb_find_node(RBTree * __restrict tree, void * __restrict key);
+
+/*!
+ * @brief find value by key with cmp
+ *
+ * for instance you can compare null-terminated string key with non-null
+ * terminated string by providing cmp func for find
+ *
+ * @param[in] tree rbtree
+ * @param[in] key  key
+ *
+ * @return found value or NULL
+ */
+DS_EXPORT
+void*
+rb_find_by(RBTree * __restrict tree, void * __restrict key, const DsCmpFn cmp);
+
+/*!
+ * @brief find node by key with cmp.
+ *        It returns RBnode you can get value from node
+ *
+ * for instance you can compare null-terminated string key with non-null
+ * terminated string by providing cmp func for find
+ *
+ * @param[in] tree rbtree
+ * @param[in] key  key
+ *
+ * @return found RBNode or NULL
+ */
+DS_EXPORT
+RBNode*
+rb_find_node_by(RBTree * __restrict tree,
+                void   * __restrict key,
+                const DsCmpFn       cmp);
 
 /*!
  * @brief get parent node of found node (of key). This will lookup for node

diff --git a/src/rb.c b/src/rb.c
@@ -541,9 +541,9 @@ rb_remove(RBTree *tree, void *key) {
 
 DS_EXPORT
 void*
-rb_find(RBTree *tree, void *key) {
+rb_find(RBTree * __restrict tree, void * __restrict key) {
   RBNode *found;
-  found = rb_find_node(tree, key);
+  found = rb_find_node_by(tree, key, tree->cmp);
   if (found == NULL)
     return NULL;
 
@@ -552,23 +552,45 @@ rb_find(RBTree *tree, void *key) {
 
 DS_EXPORT
 RBNode*
-rb_find_node(RBTree *tree, void *key) {
-  RBNode *iter;
+rb_find_node(RBTree * __restrict tree, void * __restrict key) {
+  return rb_find_node_by(tree, key, tree->cmp);
+}
 
-  iter = tree->root->chld[RB_RIGHT];
+DS_EXPORT
+void*
+rb_find_by(RBTree * __restrict tree,
+           void   * __restrict key,
+           const DsCmpFn       cmp) {
+  RBNode *found;
+  found = rb_find_node_by(tree, key, cmp);
+  if (found == NULL)
+    return NULL;
 
-  while (iter != tree->nullNode) {
+  return found->val;
+}
+
+DS_EXPORT
+RBNode*
+rb_find_node_by(RBTree * __restrict tree,
+                void   * __restrict key,
+                const DsCmpFn       cmp) {
+  RBNode *iter, *nullNode;
+
+  iter     = tree->root->chld[RB_RIGHT];
+  nullNode = tree->nullNode;
+
+  while (iter != nullNode) {
     int cmpRet;
 
-    cmpRet = tree->cmp(iter->key, key);
+    cmpRet = cmp(iter->key, key);
 
     if (cmpRet == 0)
       break;
 
     iter = iter->chld[cmpRet < 0];
   }
 
-  if (!iter || iter == tree->nullNode)
+  if (!iter || iter == nullNode)
     return NULL;
 
   return iter;