Add get_mut_ref

tidwall · Nov 20, 2024 · c4c6fd0 · c4c6fd0
1 parent 53c43fb
commit c4c6fd0
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 7 deletions.
diff --git a/bgen.h b/bgen.h
@@ -343,6 +343,11 @@ BGEN_EXTERN int BGEN_API(seek_at_mut)(BGEN_NODE **root, size_t index,
 BGEN_EXTERN int BGEN_API(seek_at_desc_mut)(BGEN_NODE **root, size_t index,
     bool(*iter)(BGEN_ITEM item, void *udata), void *udata);
 
+// Access direct mutable references... use with care.
+BGEN_EXTERN int BGEN_API(get_mut_ref)(BGEN_NODE **root, BGEN_ITEM key,
+    BGEN_ITEM **item, void *udata);
+
+
 #endif // !BGEN_SOURCE
 
 #ifndef BGEN_HEADER
@@ -1928,13 +1933,11 @@ static int BGEN_SYM(get)(BGEN_NODE **root, BGEN_ITEM key, BGEN_ITEM *item_out,
 #endif
 }
 
-// Work like (get) but, if needed, performs copy-on-write operations.
-// returns FOUND or NOTFOUND or NOMEM
-static int BGEN_SYM(get_mut)(BGEN_NODE **root, BGEN_ITEM key, 
-    BGEN_ITEM *item_out, void *udata)
+static int BGEN_SYM(get_mut_ref)(BGEN_NODE **root, BGEN_ITEM key,
+    BGEN_ITEM **item, void *udata)
 {
 #ifdef BGEN_NOORDER
-    (void)root, (void)key, (void)item_out, (void)udata;
+    (void)root, (void)key, (void)item, (void)udata;
     return BGEN_UNSUPPORTED;
 #else
     if (!*root) {
@@ -1950,8 +1953,8 @@ static int BGEN_SYM(get_mut)(BGEN_NODE **root, BGEN_ITEM key,
         int i, found;
         i = BGEN_SYM(search)(node, key, udata, &found, depth);
         if (found) {
-            if (item_out) {
-                *item_out = node->items[i];
+            if (item) {
+                *item = &node->items[i];
             }
             return BGEN_FOUND;
         } else if (node->isleaf) {
@@ -1966,6 +1969,20 @@ static int BGEN_SYM(get_mut)(BGEN_NODE **root, BGEN_ITEM key,
 #endif
 }
 
+
+// Work like (get) but, if needed, performs copy-on-write operations.
+// returns FOUND or NOTFOUND or NOMEM
+static int BGEN_SYM(get_mut)(BGEN_NODE **root, BGEN_ITEM key, 
+    BGEN_ITEM *item_out, void *udata)
+{
+    BGEN_ITEM *item;
+    int ret = BGEN_SYM(get_mut_ref)(root, key, &item, udata);
+    if (ret == BGEN_FOUND && item_out) {
+        *item_out = *item;
+    }
+    return ret;
+}
+
 // returns true if key is found
 static bool BGEN_SYM(contains)(BGEN_NODE **root, BGEN_ITEM key, void *udata) {
     return BGEN_SYM(get)(root, key, 0, udata) == BGEN_FOUND;
@@ -4600,6 +4617,7 @@ static inline void BGEN_SYM(all_sym_calls)(void) {
     (void)BGEN_SYM(feat_dims);
     (void)BGEN_SYM(print);
     (void)BGEN_SYM(get_mut);
+    (void)BGEN_SYM(get_mut_ref);
     (void)BGEN_SYM(get_at_mut);
     (void)BGEN_SYM(insert);
     (void)BGEN_SYM(get);
@@ -4678,6 +4696,7 @@ static inline void BGEN_SYM(all_api_calls)(void) {
     (void)BGEN_API(feat_atomics);
     (void)BGEN_API(feat_dims);
     (void)BGEN_API(get_mut);
+    (void)BGEN_API(get_mut_ref);
     (void)BGEN_API(get_at_mut);
     (void)BGEN_API(insert);
     (void)BGEN_API(get);
@@ -4823,6 +4842,12 @@ int BGEN_API(get_mut)(BGEN_NODE **root, BGEN_ITEM key, BGEN_ITEM *item_out,
     return BGEN_SYM(get_mut)(root, key, item_out, udata);
 }
 
+int BGEN_API(get_mut_ref)(BGEN_NODE **root, BGEN_ITEM key, BGEN_ITEM **item,
+    void *udata)
+{
+    return BGEN_SYM(get_mut_ref)(root, key, item, udata);
+}
+
 bool BGEN_API(contains)(BGEN_NODE **root, BGEN_ITEM key, void *udata) {
     return BGEN_SYM(contains)(root, key, udata);
 }

diff --git a/tests/test_base.h b/tests/test_base.h
@@ -351,6 +351,8 @@ void test_various(void) {
     assert(val == -1);
     assert(kv_get_mut(&tree, 0, &val, 0) == kv_NOTFOUND);
     assert(val == -1);
+    assert(kv_get_mut_ref(&tree, 0, 0, 0) == kv_NOTFOUND);
+    assert(val == -1);
     assert(kv_front(&tree, &val, 0) == kv_NOTFOUND);
     assert(val == -1);
     assert(kv_front_mut(&tree, &val, 0) == kv_NOTFOUND);