Playing with binary search trees.

tree.c

#include <stdio.h>
#include <stdlib.h>
#include <readline/readline.h>
#include <readline/history.h>

#define USE_VT100 1

#if USE_VT100
#define BOX_DRAWINGS_LIGHT_VERTICAL "\x1B(0\x78\x1B(B"
#define BOX_DRAWINGS_LIGHT_DOWN_AND_HORIZONTAL "\x1B(0\x77\x1B(B"
#define BOX_DRAWINGS_LIGHT_UP_AND_RIGHT "\x1B(0\x6D\x1B(B"
#else
#define BOX_DRAWINGS_LIGHT_VERTICAL "\xE2\x94\x82"
#define BOX_DRAWINGS_LIGHT_DOWN_AND_HORIZONTAL "\xE2\x94\xAC"
#define BOX_DRAWINGS_LIGHT_UP_AND_RIGHT "\xE2\x94\x94"
#endif

typedef struct Node {
	struct Node *parent;
	struct Node *left;
	struct Node *right;
	int value;
} Node;

typedef struct BacktrackNode {
	struct BacktrackNode *prev;
	int length;
	int display;
} BacktrackNode;

void
tree_print_fill(BacktrackNode *backtrack) {
	if (backtrack) {
		tree_print_fill(backtrack->prev);
		printf("%*c %s ",
		       backtrack->length, ' ',
		       backtrack->display ? BOX_DRAWINGS_LIGHT_VERTICAL : " ");
	}
}

void
tree_print_node(Node *node, BacktrackNode *prev) {
	BacktrackNode backtrack;
	char buf[12];

	backtrack.length = snprintf(buf, sizeof buf, "%d", node->value);
	backtrack.prev = prev;
	backtrack.display = 1;

	printf("%s "BOX_DRAWINGS_LIGHT_DOWN_AND_HORIZONTAL" ", buf);

	if (node->right) {
		tree_print_node(node->right, &backtrack);
	} else {
		printf("NULL\n");
	}

	tree_print_fill(prev);
	printf("%*c "BOX_DRAWINGS_LIGHT_VERTICAL"\n", backtrack.length, ' ');
	tree_print_fill(prev);
	printf("%*c "BOX_DRAWINGS_LIGHT_UP_AND_RIGHT" ", backtrack.length, ' ');

	backtrack.display = 0;

	if (node->left) {
		tree_print_node(node->left, &backtrack);
	} else {
		printf("NULL\n");
	}
}

void
tree_print(Node *root) {
	if (root != NULL) {
		tree_print_node(root, NULL);
	}
}

void
tree_rotate_left(Node **root) {
	Node *x = *root;
	Node *y = x->right;

	*root = y;
	y->parent = x->parent;

	x->right = y->left;
	if (x->right != NULL) {
		x->right->parent = x;
	}

	y->left = x;
	x->parent = y;
}

void
tree_rotate_right(Node **root) {
	Node *x = *root;
	Node *y = x->left;

	*root = y;
	y->parent = x->parent;

	x->left = y->right;
	if (x->left != NULL) {
		x->left->parent = x;
	}

	y->right = x;
	x->parent = y;
}

Node **
tree_find(Node **root, int value, Node **parent) {
	Node *node;

	if (*root == NULL) {
		return root;
	}
	node = *root;

	if (value == node->value) {
		return root;
	}

	if (parent != NULL) {
		*parent = node;
	}
	if (value < node->value) {
		return tree_find(&node->left, value, parent);
	} else {
		return tree_find(&node->right, value, parent);
	}
}

Node *
tree_minimum(Node *node) {
	while (node->left != NULL) {
		node = node->left;
	}
	return node;
}

Node *
tree_successor(Node *node) {
	Node *y;

	if (node->right != NULL) {
		return tree_minimum(node->right);
	}

	y = node->parent;
	while (y != NULL && node == y->right) {
		node = y;
		y = y->parent;
	}

	return y;
}

int
tree_insert(Node **root, int value) {
	Node *parent = NULL;

	root = tree_find(root, value, &parent);

	if (*root == NULL) {
		Node *node = *root = calloc(1, sizeof *node);
		node->value = value;
		node->parent = parent;
		return 1;
	}

	/* No duplicate nodes. */
	return 0;
}

int
tree_remove(Node **root, int value) {
	Node *node;

	root = tree_find(root, value, NULL);
	if (*root == NULL) {
		return 0;
	}

	node = *root;
	if (node->left == NULL) {
		if (node->right == NULL) {
			*root = NULL;
			free(node);
		} else {
			*root = node->right;
			(*root)->parent = node->parent;
			free(node);
		}
	} else {
		if (node->right == NULL) {
			*root = node->left;
			(*root)->parent = node->parent;
			free(node);
		} else {
			Node *successor = tree_successor(node);

			/* successor->left must be NULL */
			if (successor == successor->parent->left) {
				successor->parent->left = successor->right;
			} else {
				successor->parent->right = successor->right;
			}
			if (successor->right != NULL) {
				successor->right->parent = successor->parent;
			}

			successor->right = node->right;
			if (successor->right != NULL) {
				successor->right->parent = successor;
			}

			successor->left = node->left;
			if (successor->left != NULL) {
				successor->left->parent = successor;
			}

			*root = successor;
			free(node);
		}
	}

	return 1;
}

int
tree_check(Node *root) {
	int found_error = 0;

	if (root->left) {
		if (root->left->value > root->value) {
			printf("ERROR: Left child %d is greater than parent %d.\n",
			       root->left->value, root->value);
			found_error = 1;
		}
		if (root->left->parent != root) {
			printf("ERROR: Parent pointer of %d is not set to %d.\n",
			       root->left->value, root->value);
			found_error = 1;
		}
		if (tree_check(root->left)) {
			found_error = 1;
		}
	}

	if (root->right) {
		if (root->right->value < root->value) {
			printf("ERROR: Right child %d is less than parent %d.\n",
			       root->right->value, root->value);
			found_error = 1;
		}
		if (root->right->parent != root) {
			printf("ERROR: Parent pointer of %d is not set to %d.\n",
			       root->right->value, root->value);
			found_error = 1;
		}
		if (tree_check(root->right)) {
			found_error = 1;
		}
	}

	return found_error;
}

void
tree_free(Node **root) {
	if (*root == NULL) {
		return;
	}

	tree_free(&(*root)->left);
	tree_free(&(*root)->right);

	free(*root);
	*root = NULL;
}

int
main(int argc, char **argv) {
	Node *root = NULL;
	char *line;

	while ((line = readline("> ")) != NULL) {
		char *cmd = strtok(line, " ");

		if (cmd == NULL) {
			goto next;
		} else if (strcmp(cmd, "insert") == 0) {
			char *str_value;

			while ((str_value = strtok(NULL, " ")) != NULL) {
				char *endptr;
				int value;

				value = strtol(str_value, &endptr, 10);
				if (*endptr != '\0') {
					printf("Invalid integer '%s'.\n", str_value);
					goto next;
				}

				if (!tree_insert(&root, value)) {
					printf("Cowardly refusing to insert %d twice.\n", value);
				}
			}
		} else if (strcmp(cmd, "remove") == 0) {
			char *str_value;

			while ((str_value = strtok(NULL, " ")) != NULL) {
				char *endptr;
				int value;

				value = strtol(str_value, &endptr, 10);
				if (*endptr != '\0') {
					printf("Invalid integer '%s'.\n", str_value);
					goto next;
				}

				if (!tree_remove(&root, value)) {
					printf("No such element: %d\n", value);
				}
			}
		} else if (strcmp(cmd, "rotate") == 0) {
			char *str_value = strtok(NULL, " ");
			char *direction = strtok(NULL, " ");
			char *endptr;
			int value;
			Node **node;

			if (str_value == NULL || direction == NULL) {
				printf("Usage: rotate <value> left|right\n");
				goto next;
			}

			value = strtol(str_value, &endptr, 10);
			if (*endptr != '\0') {
				printf("Invalid integer '%s'.\n", str_value);
				goto next;
			}

			node = tree_find(&root, value, NULL);
			if (*node == NULL) {
				printf("No such node.\n");
				goto next;
			}

			if (strcmp(direction, "left") == 0) {
				if ((*node)->right == NULL) {
					printf("Node does not have a right child to pivot.\n");
				} else {
					tree_rotate_left(node);
				}
			} else if (strcmp(direction, "right") == 0) {
				if ((*node)->left == NULL) {
					printf("Node does not have a left child to pivot.\n");
				} else {
					tree_rotate_right(node);
				}
			} else {
				printf("Invalid direction.\n");
			}
		} else if (strcmp(cmd, "successor") == 0) {
			char *str_value = strtok(NULL, " ");
			char *endptr;
			int value;
			Node **node;
			Node *successor;

			if (str_value == NULL) {
				printf("Usage: successor <value>\n");
				goto next;
			}

			value = strtol(str_value, &endptr, 10);
			if (*endptr != '\0') {
				printf("Invalid integer '%s'.\n", str_value);
				goto next;
			}

			node = tree_find(&root, value, NULL);
			if (*node == NULL) {
				printf("No such node.\n");
				goto next;
			}

			successor = tree_successor(*node);
			if (successor == NULL) {
				printf("%d has no successor.\n", value);
			} else {
				printf("%d\n", successor->value);
			}
		} else if (strcmp(cmd, "print") == 0) {
			tree_print(root);
		} else if (strcmp(cmd, "quit") == 0) {
			free(line);
			break;
		} else {
			printf("Unknown command.\n");
		}

	next:
		free(line);

		if (root != NULL && tree_check(root)) {
			tree_print(root);
		}
	}

	tree_free(&root);
	return 0;
}