[PATCH] STree ForEach macro Implement a ForEach style macro for iterating over string trees. Comments on the usage and implementation in the code. I hadn't seen a construct like this before - it was an interesting challenge... The stack required for keeping state information is allocated dynamically, so a special "break" command is also provided to make sure the memory is freed properly. The _After version allows you to start iterating from a given entry in the tree. This was implemented specially for one case where it was required in NQ for cvars. Not sure if it will get a lot of use elsewhere... diff -urN a/common/shell.c head/common/shell.c --- a/common/shell.c 2006-02-21 20:11:30.000000000 +1030 +++ head/common/shell.c 2006-02-21 17:50:57.000000000 +1030 @@ -250,12 +250,162 @@ } /* An R-B Tree with n entries has a maximum height of 2log(n +1) */ -int +static int STree_MaxDepth(struct stree_root *root) { return 2 * Q_log2(root->entries + 1); } +void +STree_ForEach_Init__(struct stree_root *root, struct stree_node **n) +{ + /* Allocate the stack */ + root->stack = Z_Malloc(sizeof(struct stree_stack)); + if (root->stack) { + struct stree_stack *s = root->stack; + s->depth = 0; + s->max_depth = STree_MaxDepth(root); + s->stack = Z_Malloc(s->max_depth * sizeof(struct rb_node *)); + if (!s->stack) { + Z_Free(s); + root->stack = NULL; + } + } + /* Possibly this harsh failure is not suitable in all cases? */ + if (!root->stack) + Sys_Error("%s: not enough mem for stack! (%i)", __func__, + STree_MaxDepth(root)); + + /* Point to the first node */ + if (root->root.rb_node) + *n = stree_entry(root->root.rb_node); + else + *n = NULL; +} + +void +STree_ForEach_Cleanup__(struct stree_root *root) +{ + if (root->stack) { + Z_Free(root->stack->stack); + Z_Free(root->stack); + root->stack = NULL; + } +} + +static void +STree_StackPush(struct stree_root *root, struct rb_node *n) +{ + struct stree_stack *s = root->stack; + assert(s->depth < s->max_depth); + s->stack[s->depth++] = n; +} + +static struct rb_node * +STree_StackPop(struct stree_root *root) +{ + struct stree_stack *s = root->stack; + assert(s->depth > 0); + return s->stack[--s->depth]; +} + +/* + * STree_WalkLeft - Helper for STree_ForEach + * + * Explanation of implied semantics: + * 1. If *n is not null, we haven't looked at it at all yet + * - Check the left child. + * - If child is non-null, push *n onto the stack, *n = (*n)->left; Goto 1. + * - If left child is null, keep this *n and exit (true) + * 2. If *n is null, we need to grab the node from the top of the stack + * - If the stack is empty, we're finished + * - Free the stack and exit (false) + * - Otherwise, *n = and exit (true) + */ +qboolean +STree_WalkLeft__(struct stree_root *root, struct stree_node **n) +{ + struct rb_node *rb; + + if (*n) { + rb = &(*n)->node; + while (rb->rb_left) { + STree_StackPush(root, rb); + rb = rb->rb_left; + } + *n = stree_entry(rb); + } else { + /* Null signifies that we need to pop from the stack */ + if (root->stack->depth > 0) { + rb = STree_StackPop(root); + *n = stree_entry(rb); + } else + STree_ForEach_Cleanup__(root); + } + + return *n != NULL; +} + +/* + * STree_WalkRight__ - Helper for STree_ForEach + * + * Called at the end of a loop iteration. So, *n has been processed. + * - If *n has a right child, *n = right child. Exit. + * - Otherwise, *n = NULL (tells WalkLeft to grab parent next). Exit. + */ +void +STree_WalkRight__(struct stree_node **n) +{ + struct rb_node *rb; + + rb = &(*n)->node; + if (rb->rb_right) + *n = stree_entry(rb->rb_right); + else + *n = NULL; +} + +/* + * Skip through the tree to a specified entry, without iterating + * through everything. This is basically STree_Find, but with stack + * pushes so the iteration can be continued. We then move n to the + * next node. + */ +void +STree_ForEach_After__(struct stree_root *root, struct stree_node **n, + const char *s) +{ + struct rb_node *p; + struct stree_node *node; + int cmp; + + *n = NULL; + p = root->root.rb_node; + while (p) { + node = stree_entry(p); + cmp = strcasecmp(s, node->string); + if (cmp < 0) { + STree_StackPush(root, p); + p = p->rb_left; + } else if (cmp > 0) { + p = p->rb_right; + } else { + *n = node; + break; + } + } + + if (*n) { + /* found the exact node; skip on to the next one */ + if (p->rb_right) + *n = stree_entry(p->rb_right); + else + *n = NULL; + } else { + /* Didn't find str. Don't walk the tree at all! */ + root->stack->depth = 0; + } +} const char **completions_list = NULL; static int num_completions = 0; diff -urN a/include/shell.h head/include/shell.h --- a/include/shell.h 2006-02-19 19:36:51.000000000 +1030 +++ head/include/shell.h 2006-02-21 17:37:12.000000000 +1030 @@ -25,6 +25,25 @@ #include "rb_tree.h" /* + * String node is simply an rb_tree node using the string as the index (and + * the only data as well). + */ +struct stree_node { + const char *string; + struct rb_node node; +}; + +/* stree_entry - Gets the stree_node ptr from the internal rb_node ptr */ +#define stree_entry(ptr) container_of(ptr, struct stree_node, node) + +/* For iterative tree walking */ +struct stree_stack { + struct rb_node **stack; + int depth; + int max_depth; +}; + +/* * We keep track of the number of entries in the tree as well as the longest * entry in the tree. This info is used for displaying lists on the console. */ @@ -33,29 +52,19 @@ unsigned int maxlen; unsigned int minlen; struct rb_root root; + struct stree_stack *stack; /* used in STree_ForEach() */ }; -#define STREE_ROOT (struct stree_root) { 0, 0, -1, RB_ROOT } +#define STREE_ROOT (struct stree_root) { 0, 0, -1, RB_ROOT, NULL } #define DECLARE_STREE_ROOT(_x) \ struct stree_root _x = { \ .entries = 0, \ .maxlen = 0, \ .minlen = -1, \ - .root = RB_ROOT \ + .root = RB_ROOT, \ + .stack = NULL \ } -/* - * String node is simply an rb_tree node using the string as the index (and - * the only data as well). - */ -struct stree_node { - const char *string; - struct rb_node node; -}; - -/* stree_entry - Gets the stree_node ptr from the internal rb_node ptr */ -#define stree_entry(ptr) container_of(ptr, struct stree_node, node) - void STree_AllocInit(void); qboolean STree_Insert(struct stree_root *root, struct stree_node *node); qboolean STree_InsertAlloc(struct stree_root *root, const char *s, @@ -63,7 +72,39 @@ void STree_Remove(struct stree_root *root, struct stree_node *node); char *STree_MaxMatch(struct stree_root *root, const char *pfx); struct stree_node *STree_Find(struct stree_root *root, const char *s); -int STree_MaxDepth(struct stree_root *root); + +/* Private helper functions for the STree_ForEach macro */ +void STree_ForEach_Init__(struct stree_root *root, struct stree_node **n); +void STree_ForEach_After__(struct stree_root *root, struct stree_node **n, + const char *s); +qboolean STree_WalkLeft__(struct stree_root *root, struct stree_node **n); +void STree_WalkRight__(struct stree_node **n); +void STree_ForEach_Cleanup__(struct stree_root *root); + +/* + * Walk an STree root__ in order (LNR), using node__ as the iterator + * - struct stree_root *root__ + * - struct stree_node *node__ + * + * Note that we need to allocate a stack to walk the tree (on the zone). If + * this fails, the program will exit with an error. This is only of size + * 2*log(n+1), so if that allocation fails, we're in a bad situation anyway... + * + * STree_ForEach_After will start walking the tree from the entry _after_ the + * one matching str__. + * + * NOTE - don't use a regular 'break' inside this loop, or you'll leak the + * allocated stack memory. Use STree_ForEach_break(root). + */ +#define STree_ForEach_After(root__, node__, str__) for ( \ + STree_ForEach_Init__((root__), &(node__)), \ + ({ if (str__) STree_ForEach_After__((root__), &(node__), (str__));}); \ + STree_WalkLeft__((root__), &(node__)) ; \ + STree_WalkRight__(&(node__)) ) +#define STree_ForEach(root__, node__) \ + STree_ForEach_After(root__, node__, NULL) +#define STree_ForEach_break(root__) \ + ({ STree_ForEach_Cleanup__(root__); break; }) /* * Set up some basic completion helpers