cursor set key and set range

git-svn-id: file:///svn/tokudb@284 c7de825b-a66e-492c-adef-691d508d4ae1

newbrt/pma-internal.h

@@ -60,6 +60,12 @@ void __pma_update_my_cursors(PMA pma, struct kv_pair_tag *tpairs, int n);
  */
 void __pma_delete_at(PMA pma, int here);
 
+/*
+ * if the pma entry at here is deleted and there are no more references to it
+ * then finish the deletion
+ */
+void __pma_delete_resume(PMA pma, int here);
+
 /*
  * finish a deletion from the pma. called when there are no cursor references
  * to the kv pair.

newbrt/pma-test.c

@@ -1465,6 +1465,118 @@ void test_pma_already_there() {
     assert(error == 0);
 }
 
+void test_pma_cursor_set_key() {
+    printf("test_pma_cursor_set_key\n");
+
+    int error;
+    PMA pma;
+
+    error = pma_create(&pma, default_compare_fun);
+    assert(error == 0);
+
+    DBT key, val;
+    int k, v;
+
+    const int n = 100;
+    int i;
+    for (i=0; i<n; i += 10) {
+        k = htonl(i);
+        v = i;
+        fill_dbt(&key, &k, sizeof k);
+        fill_dbt(&val, &v, sizeof v);
+        error = pma_insert(pma, &key, &val, 0);
+        assert(error == 0);
+    }
+
+    PMA_CURSOR cursor;
+    error = pma_cursor(pma, &cursor);
+    assert(error == 0);
+
+    for (i=0; i<n; i += 1) {
+        k = htonl(i);
+        fill_dbt(&key, &k, sizeof k);
+        error = pma_cursor_set_key(cursor, &key, 0);
+        if (i % 10 == 0) {
+            assert(error == 0);
+            init_dbt(&val); val.flags = DB_DBT_MALLOC;
+            error = pma_cget_current(cursor, &key, &val);
+            assert(error == 0);
+            int vv;
+            assert(val.size == sizeof vv);
+            memcpy(&vv, val.data, val.size);
+            assert(vv == i);
+            toku_free(val.data);
+        } else 
+            assert(error == DB_NOTFOUND);
+    }
+
+    error = pma_cursor_free(&cursor);
+    assert(error == 0);
+
+    error = pma_free(&pma);
+    assert(error == 0);
+}
+
+/*
+ * verify that set range works with a pma with keys 10, 20, 30 ... 90
+ */
+void test_pma_cursor_set_range() {
+    printf("test_pma_cursor_set_range\n");
+
+    int error;
+    PMA pma;
+
+    error = pma_create(&pma, default_compare_fun);
+    assert(error == 0);
+
+    DBT key, val;
+    int k, v;
+
+    const int smallest_key = 10;
+    const int largest_key = 90;
+    int i;
+    for (i=smallest_key; i<=largest_key; i += 10) {
+        k = htonl(i);
+        v = i;
+        fill_dbt(&key, &k, sizeof k);
+        fill_dbt(&val, &v, sizeof v);
+        error = pma_insert(pma, &key, &val, 0);
+        assert(error == 0);
+    }
+
+    PMA_CURSOR cursor;
+    error = pma_cursor(pma, &cursor);
+    assert(error == 0);
+
+    for (i=0; i<100; i += 1) {
+        k = htonl(i);
+        fill_dbt(&key, &k, sizeof k);
+        error = pma_cursor_set_range(cursor, &key, 0);
+        if (error != 0) {
+            assert(error == DB_NOTFOUND);
+            assert(i > largest_key);
+        } else {
+            init_dbt(&val); val.flags = DB_DBT_MALLOC;
+            error = pma_cget_current(cursor, &key, &val);
+            assert(error == 0);
+            int vv;
+            assert(val.size == sizeof vv);
+            memcpy(&vv, val.data, val.size);
+            if (i <= smallest_key)
+                assert(vv == smallest_key);
+            else
+                assert(vv == (((i+9)/10)*10));
+            toku_free(val.data);
+        }
+    }
+
+    error = pma_cursor_free(&cursor);
+    assert(error == 0);
+
+    error = pma_free(&pma);
+    assert(error == 0);
+}
+
 void pma_tests (void) {
     memory_check=1;
     test_keycompare();            memory_check_all_free();
@@ -1485,6 +1597,8 @@ void pma_tests (void) {
     test_pma_insert_or_replace(); memory_check_all_free();
     test_pma_delete();
     test_pma_already_there();     memory_check_all_free();
+    test_pma_cursor_set_key();    memory_check_all_free();
+    test_pma_cursor_set_range();  memory_check_all_free();
 }
 
 int main (int argc __attribute__((__unused__)), char *argv[] __attribute__((__unused__))) {

newbrt/pma.c

@@ -390,8 +390,7 @@ int pma_cursor_set_position_prev (PMA_CURSOR c) {
     c->position--;
     while (c->position >= 0) {
         if (kv_pair_valid(pma->pairs[c->position])) {
-            if (old_position >= 0 && kv_pair_deleted(pma->pairs[old_position]) &&__pma_count_cursor_refs(pma, old_position) == 0)
-                __pma_delete_finish(pma, old_position);
+            __pma_delete_resume(pma, old_position);
             return 0;
         }
         c->position--;
@@ -416,8 +415,7 @@ int pma_cursor_set_position_next (PMA_CURSOR c) {
     c->position++;
     while (c->position<pma->N) {
 	if (kv_pair_valid(c->pma->pairs[c->position])) {
-            if (old_position >= 0 && kv_pair_deleted(pma->pairs[old_position]) && __pma_count_cursor_refs(pma, old_position) == 0)
-                __pma_delete_finish(pma, old_position);
+            __pma_delete_resume(pma, old_position);
             return 0;
 	}
         c->position++;
@@ -438,6 +436,40 @@ int pma_cget_current (PMA_CURSOR c, DBT *key, DBT *val) {
     return 0;
 }
 
+int pma_cursor_set_key(PMA_CURSOR c, DBT *key, DB *db) {
+    PMA pma = c->pma;
+    int here = pmainternal_find(pma, key, db);
+    assert(0<=here ); assert(here<=pma_index_limit(pma));
+    if (here==pma_index_limit(pma)) 
+        return DB_NOTFOUND;
+    DBT k2;
+    struct kv_pair *pair = pma->pairs[here];
+    if (kv_pair_valid(pair) && pma->compare_fun(db, key, fill_dbt(&k2, kv_pair_key(pair), kv_pair_keylen(pair)))==0) {
+        __pma_delete_resume(c->pma, c->position);
+        c->position = here;
+        return 0;
+    } else
+        return DB_NOTFOUND;
+}
+
+int pma_cursor_set_range(PMA_CURSOR c, DBT *key, DB *db) {
+    PMA pma = c->pma;
+    int here = pmainternal_find(pma, key, db);
+    assert(0<=here ); assert(here<=pma_index_limit(pma));
+
+    /* find the first valid pair where key[here] >= key */
+    while (here < pma->N) {
+        struct kv_pair *pair = pma->pairs[here];
+        if (kv_pair_valid(pair)) {
+            __pma_delete_resume(c->pma, c->position);
+            c->position = here;
+            return 0;
+        }
+        here += 1;
+    }
+
+    return DB_NOTFOUND;
+}
 
 #if 0
 int pma_cget_first (PMA_CURSOR c, YBT *key, YBT *val) {
@@ -613,6 +645,11 @@ int pma_delete (PMA pma, DBT *k, DB *db) {
     return BRT_OK;
 }
 
+void __pma_delete_resume(PMA pma, int here) {
+    if (here >= 0 && kv_pair_deleted(pma->pairs[here]) &&__pma_count_cursor_refs(pma, here) == 0)
+        __pma_delete_finish(pma, here);
+}
+
 void __pma_delete_finish(PMA pma, int here) {
     struct kv_pair *kv = pma->pairs[here];
     if (!kv_pair_inuse(kv))

newbrt/pma.h

@@ -75,7 +75,7 @@ int pma_cursor_free (PMA_CURSOR*);
  * get the pma that a pma cursor is bound to
  *
  * c - the pma cursor
- * pma - the location that the bound pma is returned
+ * pma - the pma that the cursor is bound to
  */
 int pma_cursor_get_pma(PMA_CURSOR c, PMA *pma);
 int pma_cursor_set_position_last (PMA_CURSOR c);
@@ -84,6 +84,12 @@ int pma_cursor_set_position_next (PMA_CURSOR c); /* Requires the cursor is init'
 int pma_cursor_set_position_prev (PMA_CURSOR c);
 int pma_cget_current (PMA_CURSOR c, DBT *key, DBT *val);
 
+/* set the cursor by key */
+int pma_cursor_set_key(PMA_CURSOR c, DBT *key, DB *db);
+
+/* set the cursor to the smallest key >= requested key */
+int pma_cursor_set_range(PMA_CURSOR c, DBT *key, DB *db);
+
 /*
  * Get the last key and value in the pma
  */