Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add state to the iterator to avoid re-scans for children #44

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Add state to the iterator to avoid re-scans for children #44

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3d171af
Add state to the iterator to avoid re-scans for children
michaelplaing May 1, 2022
1d7da93
reset child offset stack in raxSeek()
michaelplaing May 2, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 110 additions & 51 deletions rax.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
#include <stdio.h>
#include <errno.h>
#include <math.h>
#include <ctype.h>
#include "rax.h"

#ifndef RAX_MALLOC_INCLUDE
Expand Down Expand Up @@ -499,6 +500,88 @@ static inline size_t raxLowWalk(rax *rax, unsigned char *s, size_t len, raxNode
return i;
}

int raxIteratorPushChildOffset(raxIterator *it, uint8_t child_offset) {
if (it->cpos == it->cpos_max) {
uint8_t *old = (
(it->child_offset_stack == it->child_offset_stack_static) ?
NULL :
it->child_offset_stack
);

size_t new_max = it->cpos_max * 2;
it->child_offset_stack = rax_realloc(old, new_max);

if (it->child_offset_stack == NULL) {
it->child_offset_stack = (!old) ? it->child_offset_stack_static : old;
errno = ENOMEM;
return 0;
}

if (old == NULL) memcpy(it->child_offset_stack, it->child_offset_stack_static, it->cpos_max);
it->cpos_max = new_max;
}

it->child_offset_stack[it->cpos] = child_offset;
it->cpos++;
return 1;
}

uint8_t raxIteratorPopChildOffset(raxIterator *it) {
if (it->cpos == 0) return 0;
it->cpos--;
return it->child_offset_stack[it->cpos];
}

/* See cmts in raxLowWalk() above */
static inline size_t raxLowWalkSeek(unsigned char *s, size_t len, int *splitpos, raxIterator* it) {
raxNode *h = it->rt->head;

size_t i = 0;
size_t j = 0;
size_t k = 0; /* actual child offset or mismatch offset */
while(h->size && i < len) {
unsigned char *v = h->data;

if (h->iscompr) {
for (j = 0; j < h->size && i < len; j++, i++) {
if (v[j] != s[i]) break;
}

if (j != h->size) break;
}
else {
for (j = 0; j < h->size; j++) {
if (v[j] >= s[i]) break;
}

if (v[j] > s[i]) {
k = j;
j = h->size;
break;
}

if (j == h->size) {
k = j;
break;
}

i++;
}

if (h->iscompr) j = 0;
raxIteratorPushChildOffset(it, j);
raxStackPush(&it->stack,h);
raxNode **children = raxNodeFirstChildPtr(h);
memcpy(&h,children+j,sizeof(h));
j = 0;
}

it->node = h;
if (h->iscompr) *splitpos = j;
it->child_offset = k;
return i;
}

/* Insert the element 's' of size 'len', setting as auxiliary data
* the pointer 'data'. If the element is already present, the associated
* data is updated (only if 'overwrite' is set to 1), and 0 is returned,
Expand Down Expand Up @@ -1263,6 +1346,10 @@ void raxStart(raxIterator *it, rax *rt) {
it->key_max = RAX_ITER_STATIC_LEN;
it->data = NULL;
it->node_cb = NULL;
it->child_offset = 0;
it->cpos = 0;
it->cpos_max = RAX_ITER_CHILD_STATIC_LEN;
it->child_offset_stack = it->child_offset_stack_static;
raxStackInit(&it->stack);
}

Expand Down Expand Up @@ -1331,6 +1418,7 @@ int raxIteratorNextStep(raxIterator *it, int noup) {
/* Seek the lexicographically smaller key in this subtree, which
* is the first one found always going torwards the first child
* of every successive node. */
if (!raxIteratorPushChildOffset(it, 0)) return 0; // push 0 to offsets
if (!raxStackPush(&it->stack,it->node)) return 0;
raxNode **cp = raxNodeFirstChildPtr(it->node);
if (!raxIteratorAddChars(it,it->node->data,
Expand Down Expand Up @@ -1365,9 +1453,9 @@ int raxIteratorNextStep(raxIterator *it, int noup) {
}
/* If there are no children at the current node, try parent's
* next child. */
unsigned char prevchild = it->key[it->key_len-1];
if (!noup) {
it->node = raxStackPop(&it->stack);
it->child_offset = raxIteratorPopChildOffset(it);
} else {
noup = 0;
}
Expand All @@ -1379,17 +1467,14 @@ int raxIteratorNextStep(raxIterator *it, int noup) {
/* Try visiting the next child if there was at least one
* additional child. */
if (!it->node->iscompr && it->node->size > (old_noup ? 0 : 1)) {
raxNode **cp = raxNodeFirstChildPtr(it->node);
int i = 0;
while (i < it->node->size) {
debugf("SCAN NEXT %c\n", it->node->data[i]);
if (it->node->data[i] > prevchild) break;
i++;
cp++;
}
if (i != it->node->size) {
debugf("it->node->size: %d\n", it->node->size);
if (it->child_offset + (old_noup ? 0 : 1) != it->node->size) {
it->child_offset += (old_noup ? 0 : 1); // set parent offset to current child
raxNode **cp = raxNodeFirstChildPtr(it->node) + it->child_offset;
debugf("SCAN found a new node\n");
raxIteratorAddChars(it,it->node->data+i,1);
raxIteratorAddChars(it,it->node->data + it->child_offset, 1);
if (!raxIteratorPushChildOffset(it, it->child_offset)) return 0;
it->child_offset = 0; // set current_offset to 0
if (!raxStackPush(&it->stack,it->node)) return 0;
memcpy(&it->node,cp,sizeof(it->node));
/* Call the node callback if any, and replace the node
Expand All @@ -1416,9 +1501,11 @@ int raxSeekGreatest(raxIterator *it) {
if (it->node->iscompr) {
if (!raxIteratorAddChars(it,it->node->data,
it->node->size)) return 0;
if (!raxIteratorPushChildOffset(it, 0)) return 0;
} else {
if (!raxIteratorAddChars(it,it->node->data+it->node->size-1,1))
return 0;
if (!raxIteratorPushChildOffset(it, it->node->size - 1)) return 0;
}
raxNode **cp = raxNodeLastChildPtr(it->node);
if (!raxStackPush(&it->stack,it->node)) return 0;
Expand Down Expand Up @@ -1456,9 +1543,9 @@ int raxIteratorPrevStep(raxIterator *it, int noup) {
return 1;
}

unsigned char prevchild = it->key[it->key_len-1];
if (!noup) {
it->node = raxStackPop(&it->stack);
it->child_offset = raxIteratorPopChildOffset(it);
} else {
noup = 0;
}
Expand All @@ -1471,21 +1558,13 @@ int raxIteratorPrevStep(raxIterator *it, int noup) {
/* Try visiting the prev child if there is at least one
* child. */
if (!it->node->iscompr && it->node->size > (old_noup ? 0 : 1)) {
raxNode **cp = raxNodeLastChildPtr(it->node);
int i = it->node->size-1;
while (i >= 0) {
debugf("SCAN PREV %c\n", it->node->data[i]);
if (it->node->data[i] < prevchild) break;
i--;
cp--;
}
/* If we found a new subtree to explore in this node,
* go deeper following all the last children in order to
* find the key lexicographically greater. */
if (i != -1) {
if (it->child_offset != 0) {
it->child_offset -= 1;
if (!raxIteratorPushChildOffset(it, it->child_offset)) return 0;
raxNode **cp = raxNodeFirstChildPtr(it->node) + it->child_offset;
debugf("SCAN found a new node\n");
/* Enter the node we just found. */
if (!raxIteratorAddChars(it,it->node->data+i,1)) return 0;
if (!raxIteratorAddChars(it,it->node->data + it->child_offset, 1)) return 0;
if (!raxStackPush(&it->stack,it->node)) return 0;
memcpy(&it->node,cp,sizeof(it->node));
/* Seek sub-tree max. */
Expand All @@ -1511,6 +1590,8 @@ int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len) {
int eq = 0, lt = 0, gt = 0, first = 0, last = 0;

it->stack.items = 0; /* Just resetting. Intialized by raxStart(). */
it->child_offset = 0;
it->cpos = 0;
it->flags |= RAX_ITER_JUST_SEEKED;
it->flags &= ~RAX_ITER_EOF;
it->key_len = 0;
Expand Down Expand Up @@ -1561,7 +1642,7 @@ int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len) {
* perform a lookup, and later invoke the prev/next key code that
* we already use for iteration. */
int splitpos = 0;
size_t i = raxLowWalk(it->rt,ele,len,&it->node,NULL,&splitpos,&it->stack);
size_t i = raxLowWalkSeek(ele,len,&splitpos, it);

/* Return OOM on incomplete stack info. */
if (it->stack.oom) return 0;
Expand All @@ -1573,33 +1654,11 @@ int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len) {
* "equal" condition. */
if (!raxIteratorAddChars(it,ele,len)) return 0; /* OOM. */
it->data = raxGetData(it->node);
} else if (lt || gt) {
} else if (lt || gt) { /* matches Redis 7.0 */
/* Exact key not found or eq flag not set. We have to set as current
* key the one represented by the node we stopped at, and perform
* a next/prev operation to seek. To reconstruct the key at this node
* we start from the parent and go to the current node, accumulating
* the characters found along the way. */
if (!raxStackPush(&it->stack,it->node)) return 0;
for (size_t j = 1; j < it->stack.items; j++) {
raxNode *parent = it->stack.stack[j-1];
raxNode *child = it->stack.stack[j];
if (parent->iscompr) {
if (!raxIteratorAddChars(it,parent->data,parent->size))
return 0;
} else {
raxNode **cp = raxNodeFirstChildPtr(parent);
unsigned char *p = parent->data;
while(1) {
raxNode *aux;
memcpy(&aux,cp,sizeof(aux));
if (aux == child) break;
cp++;
p++;
}
if (!raxIteratorAddChars(it,p,1)) return 0;
}
}
raxStackPop(&it->stack);
* a next/prev operation to seek. */
raxIteratorAddChars(it, ele, i-splitpos);

/* We need to set the iterator in the correct state to call next/prev
* step in order to seek the desired element. */
Expand Down
6 changes: 6 additions & 0 deletions rax.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,7 @@ typedef int (*raxNodeCallback)(raxNode **noderef);

/* Radix tree iterator state is encapsulated into this data structure. */
#define RAX_ITER_STATIC_LEN 128
#define RAX_ITER_CHILD_STATIC_LEN 32
#define RAX_ITER_JUST_SEEKED (1<<0) /* Iterator was just seeked. Return current
element for the first iteration and
clear the flag. */
Expand All @@ -182,6 +183,11 @@ typedef struct raxIterator {
unsigned char key_static_string[RAX_ITER_STATIC_LEN];
raxNode *node; /* Current node. Only for unsafe iteration. */
raxStack stack; /* Stack used for unsafe iteration. */
size_t child_offset; /* current child offset */
size_t cpos_max; /* max number of child offsets the current stack can hold */
size_t cpos; /* current position in the stack of child offsets */
uint8_t *child_offset_stack; /* stack (vector) of child offsets */
uint8_t child_offset_stack_static[RAX_ITER_CHILD_STATIC_LEN];
raxNodeCallback node_cb; /* Optional node callback. Normally set to NULL. */
} raxIterator;

Expand Down