Files
ldb/kvs_rbtree_bin.c
2026-01-31 15:38:52 +00:00

585 lines
12 KiB
C

#include "kvstore.h"
#include "kvs_rw_tools.h"
#include "memory/alloc_dispatch.h"
#include "diskuring/diskuring.h"
int kvs_keycmp(const uint8_t *a, uint32_t alen,
const uint8_t *b, uint32_t blen) {
uint32_t min = (alen < blen) ? alen : blen;
int r = 0;
if (min > 0) {
r = memcmp(a, b, min);
if (r != 0) return r;
}
// 前缀相同,则短的更小
if (alen < blen) return -1;
if (alen > blen) return 1;
return 0;
}
rbtree_node *rbtree_mini(rbtree *T, rbtree_node *x) {
while (x->left != T->nil) {
x = x->left;
}
return x;
}
rbtree_node *rbtree_maxi(rbtree *T, rbtree_node *x) {
while (x->right != T->nil) {
x = x->right;
}
return x;
}
rbtree_node *rbtree_successor(rbtree *T, rbtree_node *x) {
rbtree_node *y = x->parent;
if (x->right != T->nil) {
return rbtree_mini(T, x->right);
}
while ((y != T->nil) && (x == y->right)) {
x = y;
y = y->parent;
}
return y;
}
void rbtree_left_rotate(rbtree *T, rbtree_node *x) {
rbtree_node *y = x->right; // x --> y , y --> x, right --> left, left --> right
x->right = y->left; //1 1
if (y->left != T->nil) { //1 2
y->left->parent = x;
}
y->parent = x->parent; //1 3
if (x->parent == T->nil) { //1 4
T->root = y;
} else if (x == x->parent->left) {
x->parent->left = y;
} else {
x->parent->right = y;
}
y->left = x; //1 5
x->parent = y; //1 6
}
void rbtree_right_rotate(rbtree *T, rbtree_node *y) {
rbtree_node *x = y->left;
y->left = x->right;
if (x->right != T->nil) {
x->right->parent = y;
}
x->parent = y->parent;
if (y->parent == T->nil) {
T->root = x;
} else if (y == y->parent->right) {
y->parent->right = x;
} else {
y->parent->left = x;
}
x->right = y;
y->parent = x;
}
void rbtree_insert_fixup(rbtree *T, rbtree_node *z) {
while (z->parent->color == RED) { //z ---> RED
if (z->parent == z->parent->parent->left) {
rbtree_node *y = z->parent->parent->right;
if (y->color == RED) {
z->parent->color = BLACK;
y->color = BLACK;
z->parent->parent->color = RED;
z = z->parent->parent; //z --> RED
} else {
if (z == z->parent->right) {
z = z->parent;
rbtree_left_rotate(T, z);
}
z->parent->color = BLACK;
z->parent->parent->color = RED;
rbtree_right_rotate(T, z->parent->parent);
}
}else {
rbtree_node *y = z->parent->parent->left;
if (y->color == RED) {
z->parent->color = BLACK;
y->color = BLACK;
z->parent->parent->color = RED;
z = z->parent->parent; //z --> RED
} else {
if (z == z->parent->left) {
z = z->parent;
rbtree_right_rotate(T, z);
}
z->parent->color = BLACK;
z->parent->parent->color = RED;
rbtree_left_rotate(T, z->parent->parent);
}
}
}
T->root->color = BLACK;
}
int rbtree_insert(rbtree *T, rbtree_node *z) {
rbtree_node *y = T->nil;
rbtree_node *x = T->root;
while (x != T->nil) {
y = x;
int c = kvs_keycmp(z->key, z->key_len, x->key, x->key_len);
if (c < 0) {
x = x->left;
} else if (c > 0) {
x = x->right;
} else {
return -1;
}
}
z->parent = y;
if (y == T->nil) {
T->root = z;
}else{
int c = kvs_keycmp(z->key, z->key_len, y->key, y->key_len);
if (c < 0) y->left = z;
else y->right = z;
}
z->left = T->nil;
z->right = T->nil;
z->color = RED;
rbtree_insert_fixup(T, z);
return 0;
}
void rbtree_delete_fixup(rbtree *T, rbtree_node *x) {
while ((x != T->root) && (x->color == BLACK)) {
if (x == x->parent->left) {
rbtree_node *w= x->parent->right;
if (w->color == RED) {
w->color = BLACK;
x->parent->color = RED;
rbtree_left_rotate(T, x->parent);
w = x->parent->right;
}
if ((w->left->color == BLACK) && (w->right->color == BLACK)) {
w->color = RED;
x = x->parent;
} else {
if (w->right->color == BLACK) {
w->left->color = BLACK;
w->color = RED;
rbtree_right_rotate(T, w);
w = x->parent->right;
}
w->color = x->parent->color;
x->parent->color = BLACK;
w->right->color = BLACK;
rbtree_left_rotate(T, x->parent);
x = T->root;
}
} else {
rbtree_node *w = x->parent->left;
if (w->color == RED) {
w->color = BLACK;
x->parent->color = RED;
rbtree_right_rotate(T, x->parent);
w = x->parent->left;
}
if ((w->left->color == BLACK) && (w->right->color == BLACK)) {
w->color = RED;
x = x->parent;
} else {
if (w->left->color == BLACK) {
w->right->color = BLACK;
w->color = RED;
rbtree_left_rotate(T, w);
w = x->parent->left;
}
w->color = x->parent->color;
x->parent->color = BLACK;
w->left->color = BLACK;
rbtree_right_rotate(T, x->parent);
x = T->root;
}
}
}
x->color = BLACK;
}
rbtree_node *rbtree_delete(rbtree *T, rbtree_node *z) {
rbtree_node *y = T->nil;
rbtree_node *x = T->nil;
if ((z->left == T->nil) || (z->right == T->nil)) {
y = z;
} else {
y = rbtree_successor(T, z);
}
if (y->left != T->nil) {
x = y->left;
} else if (y->right != T->nil) {
x = y->right;
}
x->parent = y->parent;
if (y->parent == T->nil) {
T->root = x;
} else if (y == y->parent->left) {
y->parent->left = x;
} else {
y->parent->right = x;
}
if (y != z) {
uint8_t *ktmp = z->key; z->key = y->key; y->key = ktmp;
uint32_t ltmp = z->key_len; z->key_len = y->key_len; y->key_len = ltmp;
uint8_t *vtmp = z->value; z->value = y->value; y->value = vtmp;
uint32_t tlen = z->value_len; z->value_len = y->value_len; y->value_len = tlen;
}
if (y->color == BLACK) {
rbtree_delete_fixup(T, x);
}
return y;
}
rbtree_node *rbtree_search(rbtree *T, KEY_TYPE* key, uint32_t keylen) {
rbtree_node *node = T->root;
while (node != T->nil) {
int c = kvs_keycmp(key, keylen, node->key, node->key_len);
if (c < 0) node = node->left;
else if (c > 0) node = node->right;
else return node;
}
return T->nil;
}
void rbtree_traversal(rbtree *T, rbtree_node *node) {
if (node != T->nil) {
rbtree_traversal(T, node->left);
printf("key:%s, color:%d\n", (char*)node->key, node->color);
rbtree_traversal(T, node->right);
}
}
typedef struct _rbtree kvs_rbtree_t;
kvs_rbtree_t global_rbtree;
// 5 + 2
int kvs_rbtree_create(kvs_rbtree_t *inst) {
if (inst == NULL) return 1;
inst->nil = (rbtree_node*)kvs_malloc(sizeof(rbtree_node));
if (!inst->nil) return 2;
inst->nil->color = BLACK;
inst->nil->left = inst->nil->right = inst->nil->parent = inst->nil;
inst->root = inst->nil;
return 0;
}
void kvs_rbtree_destroy(kvs_rbtree_t *inst) {
if (inst == NULL) return ;
rbtree_node *node = NULL;
while (inst->root != inst->nil) {
rbtree_node *mini = rbtree_mini(inst, node);
rbtree_node *cur = rbtree_delete(inst, mini);
if (cur != inst->nil) {
if (cur->key) kvs_free(cur->key);
if (cur->value) kvs_free(cur->value);
kvs_free(cur);
}
}
kvs_free(inst->nil);
inst->nil = NULL;
inst->root = NULL;
return ;
}
/*
* @return: <0 error; 0 success; 1 exist
*/
int kvs_rbtree_set(kvs_rbtree_t *inst, const void *key, uint32_t key_len, const void *value, uint32_t value_len) {
if (!inst || !key || !value) return -1;
rbtree_node *node = (rbtree_node*)kvs_malloc(sizeof(rbtree_node));
if (!node) return -2;
memset(node, 0, sizeof(*node));
node->key = (uint8_t*)kvs_malloc(key_len);
if (!node->key) {
kvs_free(node);return -2;
}
memcpy(node->key, key, key_len);
node->key_len = key_len;
node->value = (uint8_t*)kvs_malloc(value_len);
if (!node->value) { kvs_free(node->key); kvs_free(node); return -2; }
if (value_len) memcpy(node->value, value, value_len);
node->value_len = value_len;
if(rbtree_insert(inst, node) < 0){
kvs_free(node->value);
kvs_free(node->key);
kvs_free(node);
return 1;
}
return 0;
}
/*
* @return: NULL notexist, NOTNULL exist。out_value_len 是长度。
*/
void* kvs_rbtree_get(kvs_rbtree_t *inst, const void *key, uint32_t key_len, uint32_t *out_valuelen) {
if (!inst || !key || key_len == 0 || !out_valuelen) return NULL;
rbtree_node *node = rbtree_search(inst, (uint8_t *)key, key_len);
if (!node) return NULL; // no exist
if (node == inst->nil) return NULL;
*out_valuelen = node->value_len;
return node->value;
}
/*
* @return <0 error; =0 success; >0 no exist
*/
int kvs_rbtree_del(rbtree *inst, const void *key, uint32_t key_len) {
if (!inst || !key || key_len == 0) return -1;
rbtree_node *node = rbtree_search(inst, (uint8_t *)key, key_len);
if (!node) return 1; // no exist
if (node == inst->nil) return 1;
rbtree_node *cur = rbtree_delete(inst, node);
if (cur != inst->nil) {
if (cur->key) kvs_free(cur->key);
if (cur->value) kvs_free(cur->value);
kvs_free(cur);
}
return 0;
}
/*
* @return < 0 error; =0 success; >0 no exist
*/
int kvs_rbtree_mod(kvs_rbtree_t *inst, const void *key, uint32_t key_len, const void *value, uint32_t value_len) {
if (!inst || !key || key_len==0 || !value) return -1;
rbtree_node *node = rbtree_search(inst, (uint8_t *)key, key_len);
if (!node) return 1; // no exist
if (node == inst->nil) return 1;
if (node->value) kvs_free(node->value);
node->value = (uint8_t*)kvs_malloc(value_len);
if (!node->value) { node->value_len = 0; return -2; }
if (value_len) memcpy(node->value, value, value_len);
node->value_len = value_len;
return 0;
}
/*
* @return =0 exist, =1 no exist
*/
int kvs_rbtree_exist(kvs_rbtree_t *inst, const void *key, uint32_t key_len) {
if (!inst || !key || key_len == 0) return -1;
rbtree_node *node = rbtree_search(inst, (uint8_t*)key, key_len);
if (!node) return 1; // no exist
if (node == inst->nil) return 1;
return 0;
}
static int kvs_rbtree_save_node(iouring_ctx_t *uring, int fd, off_t *current_off, kvs_rbtree_t *inst, rbtree_node *node) {
if (!current_off || !inst || !node) return -1;
if (node == inst->nil) return 0;
int rc = 0;
rc = kvs_rbtree_save_node(uring, fd, current_off, inst, node->left);
if (rc < 0) return rc;
uint32_t klen = htonl(node->key_len);
uint32_t vlen = htonl(node->value_len);
void *bufs[4];
size_t lens[4];
int count = 0;
bufs[count] = &klen;
lens[count] = sizeof(klen);
count++;
bufs[count] = &vlen;
lens[count] = sizeof(vlen);
count++;
if (node->key_len > 0) {
bufs[count] = node->key;
lens[count] = node->key_len;
count++;
}
if (node->value_len > 0) {
bufs[count] = node->value;
lens[count] = node->value_len;
count++;
}
size_t total = 0;
for (int i = 0; i < count; i++) total += lens[i];
task_t *t = submit_write(uring, fd, bufs, lens, count, *current_off);
if (!t) { return -4; }
*current_off += (off_t) total;
rc = kvs_rbtree_save_node(uring, fd, current_off, inst, node->right);
if (rc < 0) return rc;
return 0;
}
// 0 success, <0 error
int kvs_rbtree_save(iouring_ctx_t *uring, kvs_rbtree_t *inst, const char* filename){
if (!uring || !inst || !filename) return -1;
int fd = open(filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
if(fd < 0) return -2;
off_t current_off = 0;
int rc = kvs_rbtree_save_node(uring, fd, &current_off, inst, inst->root);
close(fd);
return rc;
}
int kvs_rbtree_load(kvs_rbtree_t *inst, const char* filename){
if (!inst || !filename) return -1;
FILE *fp = fopen(filename, "rb");
if (!fp) return -2;
for (;;) {
uint32_t klen_n = 0, vlen_n = 0;
if (kvs_read_file(fp, &klen_n, 4) < 0) { fclose(fp); return -3; }
if (kvs_read_file(fp, &vlen_n, 4) < 0) { fclose(fp); return -3; }
uint32_t klen = ntohl(klen_n);
uint32_t vlen = ntohl(vlen_n);
if (klen == 0) { fclose(fp); return -3; }
uint8_t *keybuf = (uint8_t*)kvs_malloc((size_t)klen);
if (!keybuf) { fclose(fp); return -4; }
if (kvs_read_file(fp, keybuf, (size_t)klen) < 0) {
kvs_free(keybuf);
fclose(fp);
return -3;
}
uint8_t *valbuf = NULL;
if (vlen > 0) {
valbuf = (uint8_t*)kvs_malloc((size_t)vlen);
if (!valbuf) {
kvs_free(keybuf);
fclose(fp);
return -4;
}
if (kvs_read_file(fp, valbuf, (size_t)vlen) < 0) {
kvs_free(valbuf);
kvs_free(keybuf);
fclose(fp);
return -3;
}
}
int rc = kvs_rbtree_set(inst, keybuf, klen, valbuf, vlen);
if (vlen > 0) kvs_free(valbuf);
if (rc < 0) { // error
fclose(fp);
return -5;
}
}
fclose(fp);
return 0;
}