Unbalanced binary search tree

Question

I wrote this unbalanced binary tree and would like to know how to improve the code and performance. If you can point out any situations that are not being handled appropriately, that would be great too.

It supports element insertion, removal, search, iteration, tree balancing and encoding/decoding.

bst.h

#ifndef BST_H
#define BST_H

#include <stdint.h>

//Encoding
#define ENCODING_SIZE_T uint32_t

//Maximum size supported by encoding/decoding
#define MAX_CONTENT_LENGTH ((ENCODING_SIZE_T) -1)

//Return codes
#define BST_SUCCESS 0
#define BST_ERROR 1
#define BST_DUPLICATE 2
#define BST_NO_MEMORY 4

typedef struct BST_Node BST_Node;
struct BST_Node {
    void *content;
    BST_Node *smaller;
    BST_Node *greater;
};

typedef struct {
    BST_Node *root;
    size_t node_count;
    int (*compare)(void *, void *); 
} BST;


static inline void bst_init(BST *bst, int (*compare)(void *, void *))
{
    bst->root = NULL;    
    bst->node_count = 0;
    bst->compare = compare;
}

int bst_insert(BST *bst, void *content);
void bst_remove(BST *bst, void *content);
void *bst_find(BST *bst, void *content);
void bst_free(BST *bst);
void bst_iterate(BST *bst, void (*callback)(void *));
void bst_iterate_reverse(BST *bst, void (*callback)(void *));
void bst_balance(BST *bst);



/* Return total size of the encoded tree. container will point to the allocated
memory on success. size_after_encoding must return how much space the element
will require after encoded. encoder will be called like encoder(destination,
element, element_size) and should write the encoded element at destination and
return one past it, if it returns NULL, the encoding will be aborted. */
size_t bst_encode
(
    BST *bst,
    void **container,
    ENCODING_SIZE_T (*size_after_encoding)(void *), //Must not return 0
    void *(*encoder)(void *, void *, ENCODING_SIZE_T) //Must return NULL if it fails
);

/* Decode tree. Both constructor and destructor must be set. Constructor is
expected to return a pointer to the decoded element. Destructor should free
the element created by constructor, it will be called in case of errors. */
int bst_decode
(
    BST *bst,
    void *encoded,
    void *(*constructor)(void *, ENCODING_SIZE_T), //Return NULL if any error occurs
    void (destructor)(void *)
);

#endif

bst_internal.h

#include <stdlib.h>
#include "bst.h"


////////////////////////////
/////////   Declarations
///////////////////////////

static inline int choose_smaller(BST *bst);
static BST_Node **a_z(BST *bst);
static void recursive_free(BST_Node *node);
static void node_traverse(BST_Node *node, void (*callback)(BST_Node *));
static void content_traverse(BST_Node *node, void (*callback)(void *));             
static void reverse_content_traverse(BST_Node *node, void (*callback)(void *));
static inline void free_all_nodes_from_list(BST_Node **nodes, size_t node_count);
static int insert(BST *bst, BST_Node **where, void *what);
static BST_Node **find_node_and_its_parent( BST *bst, void *content);
static BST_Node **find_closest_smaller_and_its_parent(BST_Node *node);
static BST_Node **find_closest_greater_and_its_parent(BST_Node *node);
static inline void free_and_unlink(BST_Node **node);
static inline void free_and_link_smaller(BST_Node **node);
static void free_and_link_greater(BST_Node **node);
static void remove_node_with_children(BST *bst, BST_Node *node);
static inline void remove(BST *bst, BST_Node **node);
static BST_Node *link_middle(BST_Node **start, BST_Node **end);

////////////////////////////
/////////   Internal methods
///////////////////////////

/* Choose a different side every time a node with 2 children is removed, so
the tree won't become too unbalanced after many removals */
static int choose_smaller(BST *bst)
{
    return bst->node_count % 2;
}

static void a_z_recursive(BST_Node **dest, BST_Node *node, size_t *i)
{
    if(node->smaller != NULL)
        a_z_recursive(dest, node->smaller, i);

    dest[*i] = node;
    *i += 1;

    if(node->greater != NULL)
        a_z_recursive(dest, node->greater, i);
}

/* Return an ordered array containing all nodes starting at a certain point. 
Caller must be sure the tree has at least the root */
static BST_Node **a_z(BST *bst)
{
    //Store pointers to all nodes into an array
    BST_Node **a_z = malloc(bst->node_count * sizeof(BST_Node *));
    if(a_z == NULL)
        return NULL;

    //Keep track of position
    size_t i = 0;
    a_z_recursive(a_z, bst->root, &i);

    return a_z;
}

/* Traverse the tree in order and call some function to work on the node */
static void node_traverse(BST_Node *node, void (*callback)(BST_Node *))

{
    if(node->smaller != NULL)
        node_traverse(node->smaller, callback);

    callback(node);

    if(node->greater != NULL)
        node_traverse(node->greater, callback);
}

//Can't use node_traverse or it will try to read the deallocated node
static void recursive_free(BST_Node *node)
{
    if(node->smaller != NULL)
        recursive_free(node->smaller);

    if(node->greater != NULL)
        recursive_free(node->greater);

    free(node);
}

static void content_traverse(BST_Node *node, void (*callback)(void *))

{
    if(node->smaller != NULL)
        content_traverse(node->smaller, callback);

    callback(node->content);

    if(node->greater != NULL)
        content_traverse(node->greater, callback);
}

static void reverse_content_traverse(BST_Node *node, void (*callback)(void *))

{
    if(node->greater != NULL)
        content_traverse(node->greater, callback);

    callback(node->content);

    if(node->smaller != NULL)
        content_traverse(node->smaller, callback);
}

/* Free all nodes listed in the array */
static void free_all_nodes_from_list(BST_Node **nodes, size_t node_count)
{
    while(node_count-- > 0)
        free(nodes[node_count]);
}

static int insert(BST *bst, BST_Node **where, void *what)
{
    //Allocate new node
    if((*where = malloc(sizeof(BST_Node))) == NULL)
        return BST_NO_MEMORY;

    //Fill
    (*where)->content = what;
    (*where)->smaller = NULL;
    (*where)->greater = NULL;

    ++bst->node_count;

    return BST_SUCCESS;
}

/* Return BST_SUCCESS if found or BST_ERROR if not found */
static BST_Node **find_node_and_its_parent(BST *bst, void *content)
{
    /* Set it up so node will point to the right node, and parent will point to 
    its parent */
    BST_Node *parent;
    BST_Node *node = NULL;
    BST_Node *ite = bst->root;
    int relationship;       

    /* Move until all can be done in a loop, otherwise it will break if the 
    node to be removed is root */   
    relationship = bst->compare(content, ite->content);

    parent = node;
    node = ite;

    if(relationship > 0)
            ite = ite->greater;

    else
    if(relationship < 0)
        ite = ite->smaller;

    //It's root
    else 
        return &bst->root;  


    /* Now it's safe to process */
    while(ite != NULL){

        //Find next path
        relationship = bst->compare(content, ite->content);

        //Update
        parent = node;
        node = ite;

        if(relationship > 0)
            ite = ite->greater;

        else
        if(relationship < 0)
            ite = ite->smaller;

        else {      
            //Found, check which side
            return (parent->greater == node) ? &parent->greater 
                                             : &parent->smaller;
        }
    }

    //Not found
    return NULL;
}

static BST_Node **find_closest_smaller_and_its_parent(BST_Node *node)
{
    /* When it's done ite will be NULL and the other 2 pointers will be set to
    the correct nodes */
    BST_Node *parent = node;
    BST_Node *closest = node->smaller;
    BST_Node *ite = closest->greater;

    while(ite != NULL){
        parent = closest;
        closest = ite;
        ite = ite->greater;     
    }

    return (parent->smaller == closest) ? &parent->smaller : &parent->greater;
}

static BST_Node **find_closest_greater_and_its_parent(BST_Node *node)
{
    BST_Node *parent = node;
    BST_Node *closest = node->greater;
    BST_Node *ite = closest->smaller;

    while(ite != NULL){
        parent = closest;
        closest = ite;
        ite = ite->smaller;     
    }

    return (parent->smaller == closest) ? &parent->smaller : &parent->greater;
}

static void free_and_unlink(BST_Node **node)
{       
    free(*node);
    *node = NULL;
}

static void free_and_link_smaller(BST_Node **node)
{
    void *temp = (*node)->smaller;
    free(*node);
    *node = temp;
}

static void free_and_link_greater(BST_Node **node)
{
    void *temp = (*node)->greater;
    free(*node);
    *node = temp;
}

static void remove_node_with_children(BST *bst, BST_Node *node)
{
    BST_Node **substitute;

    if(choose_smaller(bst))
        substitute = find_closest_smaller_and_its_parent(node);

    else
        substitute = find_closest_greater_and_its_parent(node);

    node->content = (*substitute)->content;
    remove(bst, substitute);
}

static void remove(BST *bst, BST_Node **node)
{
    //Handle removal here if the node is a leaf or has a single child
    if((*node)->smaller)

        //Both left and right, call handler
        if((*node)->greater)
            remove_node_with_children(bst, (*node));

        //Only left
        else 
            free_and_link_smaller(node);

    //Only right child
    else
    if((*node)->greater)
        free_and_link_greater(node);

    //It's a leaf
    else
        free_and_unlink(node);
}

static BST_Node *link_middle(BST_Node **start, BST_Node **end)
{
    if(end - start == 0)
        return NULL;

    //Find middle
    BST_Node **middle = start + (end - start) / 2;

    //Call again for left and right subtrees
    (*middle)->smaller = link_middle(start, middle);
    (*middle)->greater = link_middle(middle + 1, end);

    return *middle;
}

bst.c

#include <stdlib.h>
#include "bst.h"
#include "bst_internal.h"


////////////////////////////
/////////   Public methods
///////////////////////////

//Find where to insert and call insert
int bst_insert(BST *bst, void *content)
{
    //Handle root insertion separately
    if(bst->root == NULL)
        return insert(bst, &bst->root, content);

    //Find the place where the new node is supposed to be.
    BST_Node *ite = bst->root; //Seed
    BST_Node *parent;
    int relationship;

    while(ite != NULL){
        relationship = bst->compare(content, ite->content);
        parent = ite;

        if(relationship > 0)
            ite = ite->greater;

        else
        if(relationship < 0)
            ite = ite->smaller;

        else
            return BST_DUPLICATE;   
    }

    //Found the spot, repeat last decision to see if it went left or right
    return insert(  bst,    (relationship > 0) 
                            ? &parent->greater
                            : &parent->smaller, content ); 
}

//Remove and free node
void bst_remove(BST *bst, void *content)
{
    if(bst->root == NULL)
        return;

    BST_Node **node = find_node_and_its_parent(bst, content);

    if(node == NULL)
        return;

    //Update count
    --bst->node_count;

    remove(bst, node);
}

void *bst_find(BST *bst, void *content)
{
    BST_Node *ite = bst->root;
    int relationship;

    while(ite != NULL){
        relationship = bst->compare(content, ite->content); 

        if(relationship > 0)
            ite = ite->greater;

        else
        if(relationship < 0)
            ite = ite->smaller;

        else
            return ite->content; //Found
    }

    //Didn't find
    return NULL;
}

void bst_free(BST *bst)
{
    if(bst->root == NULL)
        return;

    recursive_free(bst->root);
}

//Callback will be called once per node, from node a to z
void bst_iterate(BST *bst, void (*callback)(void *))
{   
    if(bst->root == NULL)
        return;

    content_traverse(bst->root, callback);
}

//Same as before, but z to a
void bst_iterate_reverse(BST *bst, void (*callback)(void *))
{
    if(bst->root == NULL)
        return;

    reverse_content_traverse(bst->root, callback);
}

//Balance the tree
void bst_balance(BST *bst)
{
    BST_Node **ordered_nodes = a_z(bst); 
    bst->root = link_middle(    ordered_nodes, 
                                ordered_nodes + bst->node_count );
    free(ordered_nodes);
}


////////////////////////////
/////////   Functions to save the tree. They allow to use the tree
//////// to hold data, save it, and load it.
///////////////////////////

// The encoded tree looks like: INDEX + NODE_SIZE + NODE + NODE_SIZE + NODE...

//Use a struct to make it easier to add new fields
typedef struct {
    size_t node_count;
} Index;

/* Calculate total space required for an encoded tree */
static size_t total_space_required_for_encoding
(
    BST_Node **all_nodes,
    size_t node_count,
    ENCODING_SIZE_T (*size_after_encoding)(void *)
)
{
    //Consider the index
    size_t space_required = sizeof(Index);

    //Add NODE_SIZE
    space_required += sizeof(ENCODING_SIZE_T) * node_count;

    //Calculate how much storage each node will need
    ENCODING_SIZE_T node_size;
    while(node_count-- > 0){
        node_size = size_after_encoding(all_nodes[node_count]->content);
        space_required += node_size;

        //Consider padding, so all NODE_SIZEs will be aligned
        space_required += node_size % sizeof(ENCODING_SIZE_T);      
    }

    return space_required;
}

/* Return BST_ERROR if the encoding function signals an error */
static int encode
(
    void *dest, 
    BST_Node **nodes,
    size_t node_count,
    ENCODING_SIZE_T (*size_after_encoding)(void *),
    void *(*encoder)(void *, void *, ENCODING_SIZE_T)
)
{
    //Node size is used when decoding
    ENCODING_SIZE_T *node_size; 

    for(size_t i = 0; i < node_count; ++i){

        //Add the node size
        node_size = dest;
        *node_size = size_after_encoding(nodes[i]->content);

        //Location to copy the content to
        dest = (char *)dest + sizeof(ENCODING_SIZE_T);

        //encoder must fill dest and return one past it
        dest = encoder(dest, nodes[i]->content, *node_size);
        if(dest == NULL)
            return BST_ERROR;

        //Add padding to keep node_size aligned
        dest = (char *)dest + *node_size % sizeof(ENCODING_SIZE_T);
    }

    return BST_SUCCESS;
}

/* Return the total size of encoded tree. container will point to the allocated
memory on success. size_after_encoding must return how much space the element
will require after encoded. encoder will be called like encoder(destination,
element, element_size) and should write the encoded element at destination and
return one past it, if it returns NULL, the encoding will be aborted. */
size_t bst_encode
(
    BST *bst,
    void **container,
    ENCODING_SIZE_T (*size_after_encoding)(void *), //Must not return 0
    void *(*encoder)(void *, void *, ENCODING_SIZE_T) //Must return NULL if it fails
)
{
    //Is there anything?
    if(bst->root == NULL)
        return 0;

    //Get a list of nodes
    BST_Node **nodes = a_z(bst);

    //How much memory will be required?
    size_t space_required = total_space_required_for_encoding
                                (nodes, bst->node_count, size_after_encoding);

    //Try to allocate enough memory
    Index *encoded = malloc(space_required);
    if(encoded == NULL){
        free(nodes);
        return 0;
    }

    //Fill info
    encoded->node_count = bst->node_count;

    if(encode(  encoded + 1, 
                nodes, 
                bst->node_count, 
                size_after_encoding, 
                encoder             )

    == BST_ERROR){
        free(nodes);
        return 0;
    }

    free(nodes);
    *container = encoded;

    return space_required;
}

/* Allocate new nodes, return an array with pointers to all */
static BST_Node **allocate_nodes_and_list(size_t total)
{
    //Allocate list
    BST_Node **nodes = malloc(sizeof(BST_Node *) * total);
    if(nodes == NULL)
        return NULL;

    //Allocate all nodes and add to list
    for(size_t i = 0; i < total; ++i){

        //Create new node, check for errors
        nodes[i] = malloc(sizeof(BST_Node));
        if(nodes[i] == NULL){

            //Abort
            while(i-- > 0)
                free(nodes[i]);

            free(nodes);
            return NULL;        
        }
    }

    return nodes;
}

/* Construct nodes. Destruct all in case of failure */
static int construct_nodes
(
    BST_Node **nodes,
    void *encoded,
    size_t node_count,
    void *(*constructor)(void *, ENCODING_SIZE_T),
    void (destructor)(void *)
)
{
    ENCODING_SIZE_T *node_size;

    for(size_t i = 0; i < node_count; ++i){

        //Get size
        node_size = encoded;

        //Point to content
        encoded = (char *)encoded + sizeof(ENCODING_SIZE_T);

        //Construct
        nodes[i]->content = constructor(encoded, *node_size);

        //Handle error
        if(nodes[i]->content == NULL){

            //Destruct all
            while(i-- > 0)
                destructor(nodes[i]->content);

            return BST_ERROR;                   
        }

        //Update position
        encoded = (char *)encoded + *node_size;

        //Don't forget the padding
        encoded = (char *)encoded + *node_size % sizeof(ENCODING_SIZE_T);       
    }

    return BST_SUCCESS;
}

/* Decode tree. Both constructor and destructor must be set. Constructor is
expected to return a pointer to the unencoded element. Destructor should free
the element created by constructor, it will be called in case of errors. */
int bst_decode
(
    BST *bst,
    void *encoded,
    void *(*constructor)(void *, ENCODING_SIZE_T), //Return NULL if any error
    void (destructor)(void *)
)
{
    Index *index = encoded;

    //Try to allocate all 
    BST_Node **nodes = allocate_nodes_and_list(index->node_count);
    if(nodes == NULL)
        return BST_ERROR;

    if(construct_nodes( nodes, 
                        index + 1, 
                        index->node_count, 
                        constructor, 
                        destructor          )
    == BST_ERROR){
        free_all_nodes_from_list(nodes, index->node_count);
        free(nodes);
        return BST_ERROR;
    }   

    bst->root = link_middle(nodes, nodes + index->node_count);
    bst->node_count = index->node_count;

    free(nodes);

    return BST_SUCCESS;
}

Answer 1

This is simple a first pass with mostly top level thoughts - I dig deeper later.

1. Suggest moving all functions to bst.c including bst_init() size_t bst_encode() bst_decode(). By doing so, ENCODING_SIZE_T MAX_CONTENT_LENGTH BST_Node and the fields of BST may be removed from the application view. The ap does not need to see them.
   [Edit] I now see bst_init() size_t bst_encode() bst_decode() are only prototyped in bst.h - ignore that part of this suggestion. At first blush, the style looked like code, but it is only prototype.

   // BST.h
   typedef struct BST BST;`
   
   // BST.c
   struct BST { 
     BST_Node *root; 
     size_t node_count; 
     int (*compare)(void *, void *);
   };
   

2. I often use the prefix before functions like BST_This() and BST_That(), but I would be consistent with the filename's case. Use BST.c with BST_This() or bst.c with bst_This().

3. Need #include <stdlib.h> for size_t in bst.h. You might have caught that if in bst.c you included the standard headers after "bst.h". Suggest moving those after "bst.h". This prevents requiring users of bst.h to include certain headers files first.

4. Large comment blocks before bst_encode in both bst.h bst.c is a maintenance liability. Suggest only one (in bst.h) a a reference in bst.c to bst.h to show you did not forget.

5. Change compare function type to a const version. qsort()'s type shown for reference.

   int (*compare)(const void *, const void *); 
   // qsort's --> int (*compar)(const void *, const void *);
   

6. Maybe test for compare() against NULL.

7. To expand compare() possibilities, add a compare state variable.

8. Concerning callback, allow that function to return an int and use a non-zero value to short-circuit the progress OR provide a state variable for callback().

More later