trie.h

// SPDX-License-Identifier: MIT
// Copyright (c) 2021 Michael Bäck <mhcoma@gmail.com>

#ifndef __CCTL_TRIE_H__
#define __CCTL_TRIE_H__

#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "cctl.h"

// Trie Iterator

/// `trie_iterator(T)` ⇒ `T_trie_iterator`
/// 
/// Generates the iterator type name for `trie(T)`.
#define trie_iterator(T) cctl_join(T, trie_iterator)

/// `trie_iterator_func(FUNC, T)` ⇒ `T_trie_iterator_FUNC`
/// 
/// Generates a mangled function identifier for `trie_iterator(T)`.
#define trie_iterator_func(FUNC, T) cctl_join(trie_iterator(T), FUNC)

/// `trie_iterator_struct(T)` ⇒ `struct T_trie_iterator_struct`
/// 
/// Generates the underlying structure tag for `trie_iterator(T)`.
#define trie_iterator_struct(T) cctl_join(trie_iterator(T), struct)

/// ```c
/// T* trie_iterator_get(T, trie_iterator(T)* p_it)
/// ```
/// Returns a pointer to the value stored in the node currently pointed to by the iterator.
#define trie_iterator_get(T, p_it) ((p_it)->p_node ? &((p_it)->p_node->data) : NULL)

/// ```c
/// bool trie_iterator_is_valid(T, trie_iterator(T)* p_it)
/// ```
/// Checks whether the iterator currently points to a valid trie key node.
#define trie_iterator_is_valid(T, p_it) ((p_it)->p_node != NULL)

/// ```c
/// bool trie_iterator_is_equal(T, trie_iterator(T)* p_it1, trie_iterator(T)* p_it2)
/// ```
/// Checks whether two iterators point to the same trie node.
#define trie_iterator_is_equal(T, p_it1, p_it2) ((p_it1)->p_node == (p_it2)->p_node)

/// ```c
/// void trie_iterator_next(T, trie_iterator(T)* p_it)
/// ```
/// Advances the iterator forward to the next lexicographical key node in the trie.
#define trie_iterator_next(T, p_it) trie_iterator_func(next, T)(p_it)

/// ```c
/// void trie_iterator_prev(T, trie_iterator(T)* p_it)
/// ```
/// Moves the iterator backward to the previous lexicographical key node in the trie.
#define trie_iterator_prev(T, p_it) trie_iterator_func(prev, T)(p_it)

/// ```c
/// char* trie_iterator_get_key(T, trie_iterator(T)* p_it)
/// ```
/// Reconstructs and returns the null-terminated string key for the current iterator node.
/// 
/// The returned string is dynamically allocated; the caller must `free` it.
#define trie_iterator_get_key(T, p_it) trie_func(get_key, T)((p_it)->p_node)

/// Declares trie iterator struct and function prototypes in a header file.
#define trie_iterator_imp_h(T) \
	typedef struct trie_iterator_struct(T) trie_iterator(T); \
	\
	struct trie_iterator_struct(T) { \
		trie(T)* p_node; \
	}; \
	\
	void trie_iterator_func(next, T)(trie_iterator(T)* p_it); \
	void trie_iterator_func(prev, T)(trie_iterator(T)* p_it);

/// Implements trie iterator functions in a source file.
#define trie_iterator_imp_c(T) \
	void trie_iterator_func(next, T)(trie_iterator(T)* p_it) { \
		if (!p_it || !(p_it->p_node)) return; \
		trie(T)* p_curr = p_it->p_node; \
		size_t next_child_index = 0; \
		while (true) { \
			bool found_child = false; \
			for (size_t i = next_child_index; i < 256; i++) { \
				if (p_curr->ap_children[i]) { \
					p_curr = p_curr->ap_children[i]; \
					if (p_curr->existence) { \
						p_it->p_node = p_curr; \
						return; \
					} \
					next_child_index = 0; \
					found_child = true; \
					break; \
				} \
			} \
			if (found_child) continue; \
			if (p_curr->p_parent == NULL) { \
				p_it->p_node = NULL; \
				return; \
			} \
			next_child_index = p_curr->curr_char + 1; \
			p_curr = p_curr->p_parent; \
		} \
	} \
	\
	void trie_iterator_func(prev, T)(trie_iterator(T)* p_it) { \
		if (!p_it || !(p_it->p_node)) return; \
		trie(T)* p_curr = p_it->p_node; \
		while (p_curr->p_parent != NULL) { \
			uint8_t curr_char = p_curr->curr_char; \
			trie(T)* p_parent = p_curr->p_parent; \
			bool found_sibling = false; \
			for (size_t i = curr_char; i > 0; i--) { \
				if (p_parent->ap_children[i - 1]) { \
					p_curr = p_parent->ap_children[i - 1]; \
					found_sibling = true; \
					break; \
				} \
			} \
			if (found_sibling) { \
				while (true) { \
					bool found_child = false; \
					for (size_t i = 256; i > 0; i--) { \
						if (p_curr->ap_children[i - 1]) { \
							p_curr = p_curr->ap_children[i - 1]; \
							found_child = true; \
							break; \
						} \
					} \
					if (!found_child) { \
						p_it->p_node = p_curr; \
						return; \
					} \
				} \
			} \
			p_curr = p_parent; \
			if (p_curr->existence) { \
				p_it->p_node = p_curr; \
				return; \
			} \
		} \
		p_it->p_node = NULL; \
	}


// Trie

/// `trie(T)` ⇒ `T_trie`
/// 
/// Generates the Radix-256 Trie container type name for `T`.
#define trie(T) cctl_join(T, trie)

/// `trie_pair(T)` ⇒ `T_trie_pair`
/// 
/// Generates the key-value pair type name for `trie(T)`.
#define trie_pair(T) cctl_join(T, trie_pair)

/// `trie_pair_struct(T)` ⇒ `struct T_trie_pair_struct`
/// 
/// Generates the underlying structure tag for `trie_pair(T)`.
#define trie_pair_struct(T) cctl_join(trie_pair(T), struct)

/// `trie_func(FUNC, T)` ⇒ `T_trie_FUNC`
/// 
/// Generates a mangled function identifier for `trie(T)`.
#define trie_func(FUNC, T) cctl_join(trie(T), FUNC)

/// `trie_struct(T)` ⇒ `struct T_trie_struct`
/// 
/// Generates the underlying structure tag for `trie(T)`.
#define trie_struct(T) cctl_join(trie(T), struct)

/// ```c
/// void trie_init(T, trie(T)* p_t)
/// ```
/// Initializes the trie root node.
/// 
/// `p_t` is a pointer to the root trie instance.
#define trie_init(T, p_t) trie_func(init, T)(p_t)

/// ```c
/// trie(T) trie_new(T)
/// ```
/// Constructs and returns an initialized empty trie.
#define trie_new(T) trie_func(new, T)()

/// ```c
/// trie(T) trie_from_items(T, ...)
/// ```
/// Constructs and returns an initialized trie containing all `{key, value}` pairs provided in `...`.
#define trie_from_items(T, ...) \
	trie_func(from_items, T)((const trie_pair(T)[]) { __VA_ARGS__ }, sizeof((const trie_pair(T)[]) { __VA_ARGS__ }) / sizeof(trie_pair(T)))

/// ```c
/// bool trie_extend_items(T, trie(T)* p_t, ...)
/// ```
/// Inserts multiple `{key, value}` pairs provided in `...` into the trie.
#define trie_extend_items(T, p_t, ...) \
	trie_func(extend_items, T)(p_t, (const trie_pair(T)[]) { __VA_ARGS__ }, sizeof((const trie_pair(T)[]) { __VA_ARGS__ }) / sizeof(trie_pair(T)))


/// ```c
/// void trie_free(T, trie(T)* p_t)
/// ```
/// Recursively frees all descendant nodes in the trie.
#define trie_free(T, p_t) trie_func(free, T)(p_t)

/// ```c
/// T* trie_insert(T, trie(T)* p_t, const char* p_key, T item)
/// ```
/// Inserts a key string `p_key` associated with `item` into the trie.
/// 
/// Returns a pointer to the stored item, or `NULL` on allocation failure.
#define trie_insert(T, p_t, p_key, item) trie_func(insert, T)(p_t, p_key, item)

/// ```c
/// void trie_remove(T, trie(T)* p_t, const char* p_key)
/// ```
/// Removes the key string `p_key` from the trie and cleans up unused branch nodes.
#define trie_remove(T, p_t, p_key) trie_func(remove, T)(p_t, p_key)

/// ```c
/// T* trie_find(T, trie(T)* p_t, const char* p_key)
/// ```
/// Searches for `p_key` in the trie and returns a pointer to its mapped value, or `NULL`.
#define trie_find(T, p_t, p_key) trie_func(find, T)(p_t, p_key)

/// ```c
/// char* trie_get_key(T, trie(T)* p_node)
/// ```
/// Reconstructs the string key leading from root to `p_node`.
/// 
/// The returned string is dynamically allocated; the caller must `free` it.
#define trie_get_key(T, p_node) trie_func(get_key, T)(p_node)

/// ```c
/// void trie_clear(T, trie(T)* p_t)
/// ```
/// Clears all entries from the trie and resets it to empty.
#define trie_clear(T, p_t) trie_func(clear, T)(p_t)

/// ```c
/// bool trie_contains(T, trie(T)* p_t, const char* p_key)
/// ```
/// Checks whether the trie contains an entry with `p_key`.
#define trie_contains(T, p_t, p_key) (trie_find(T, p_t, p_key) != NULL)

/// ```c
/// T trie_get_or_default(T, trie(T)* p_t, const char* p_key, T default_value)
/// ```
/// Returns the mapped value for `p_key`, or `default_value` if not found.
#define trie_get_or_default(T, p_t, p_key, default_value) trie_func(get_or_default, T)(p_t, p_key, default_value)

/// ```c
/// char* trie_min_key(T, trie(T)* p_t)
/// ```
/// Returns the smallest lexicographical key in the trie. The returned string must be freed by the caller, or `NULL` if empty.
#define trie_min_key(T, p_t) trie_func(min_key, T)(p_t)

/// ```c
/// char* trie_max_key(T, trie(T)* p_t)
/// ```
/// Returns the largest lexicographical key in the trie. The returned string must be freed by the caller, or `NULL` if empty.
#define trie_max_key(T, p_t) trie_func(max_key, T)(p_t)

/// ```c
/// bool trie_is_empty(T, trie(T)* p_t)
/// ```
/// Checks whether the trie contains no keys.
#define trie_is_empty(T, p_t) ((p_t)->count == 0)

/// ```c
/// size_t trie_size(T, trie(T)* p_t)
/// ```
/// Returns the number of keys stored in the trie as a read-only rvalue.
#define trie_size(T, p_t) ((size_t)((p_t)->count))

/// ```c
/// trie_iterator(T) trie_lower_bound(T, trie(T)* p_t, const char* p_key)
/// ```
/// Returns an iterator pointing to the first key node with key >= `p_key`.
#define trie_lower_bound(T, p_t, p_key) trie_func(lower_bound, T)(p_t, p_key)

/// ```c
/// trie_iterator(T) trie_upper_bound(T, trie(T)* p_t, const char* p_key)
/// ```
/// Returns an iterator pointing to the first key node with key > `p_key`.
#define trie_upper_bound(T, p_t, p_key) trie_func(upper_bound, T)(p_t, p_key)

/// ```c
/// trie_iterator(T) trie_begin(T, trie(T)* p_t)
/// ```
/// Returns an iterator pointing to the first lexicographical key node in the trie.
#define trie_begin(T, p_t) trie_func(begin, T)(p_t)

/// ```c
/// trie_iterator(T) trie_end(T, trie(T)* p_t)
/// ```
/// Returns an end iterator pointing to `NULL`.
#define trie_end(T, p_t) trie_func(end, T)(p_t)

/// ```c
/// trie_iterator(T) trie_rbegin(T, trie(T)* p_t)
/// ```
/// Returns a reverse iterator pointing to the last lexicographical key node in the trie.
#define trie_rbegin(T, p_t) trie_func(rbegin, T)(p_t)

/// ```c
/// trie_iterator(T) trie_rend(T, trie(T)* p_t)
/// ```
/// Returns a reverse end iterator pointing to `NULL`.
#define trie_rend(T, p_t) trie_func(rend, T)(p_t)

/// ```c
/// trie_foreach(T, trie(T)* p_t, IT_NAME)
/// ```
/// Iterates over all key nodes in the trie in ascending lexicographical order.
#define trie_foreach(T, p_t, it) for (trie_iterator(T) it = trie_begin(T, p_t); trie_iterator_is_valid(T, &it); trie_iterator_next(T, &it))

/// ```c
/// trie_rforeach(T, trie(T)* p_t, IT_NAME)
/// ```
/// Iterates over all key nodes in the trie in descending lexicographical order.
#define trie_rforeach(T, p_t, it) for (trie_iterator(T) it = trie_rbegin(T, p_t); trie_iterator_is_valid(T, &it); trie_iterator_prev(T, &it))

/// Forward declares the `trie(T)` container type for prior references.
#define trie_fd(T) \
	typedef struct trie_pair_struct(T) trie_pair(T); \
	typedef struct trie_struct(T) trie(T);

/// Declares trie struct and function prototypes in a header file.
#define trie_imp_h(T) \
	struct trie_pair_struct(T) { \
		const char* p_key; \
		T value; \
	}; \
	\
	trie_iterator_imp_h(T); \
	\
	struct trie_struct(T) { \
		T data; \
		bool existence; \
		uint8_t curr_char; \
		size_t count; \
		trie(T)* p_parent; \
		trie(T)* ap_children[256]; \
	}; \
	\
	void trie_func(init, T)(trie(T)* p_t); \
	trie(T) trie_func(new, T)(void); \
	trie(T) trie_func(from_items, T)(const trie_pair(T)* p_items, size_t count); \
	bool trie_func(extend_items, T)(trie(T)* p_t, const trie_pair(T)* p_items, size_t count); \
	void trie_func(free, T)(trie(T)* p_t); \
	void trie_func(clear, T)(trie(T)* p_t); \
	T* trie_func(insert, T)(trie(T)* p_t, const char* p_key, T item); \
	void trie_func(remove, T)(trie(T)* p_t, const char* p_key); \
	T* trie_func(find, T)(trie(T)* p_t, const char* p_key); \
	T trie_func(get_or_default, T)(trie(T)* p_t, const char* p_key, T default_value); \
	char* trie_func(min_key, T)(trie(T)* p_t); \
	char* trie_func(max_key, T)(trie(T)* p_t); \
	char* trie_func(get_key, T)(trie(T)* p_node); \
	trie_iterator(T) trie_func(lower_bound, T)(trie(T)* p_t, const char* p_key); \
	trie_iterator(T) trie_func(upper_bound, T)(trie(T)* p_t, const char* p_key); \
	trie_iterator(T) trie_func(begin, T)(trie(T)* p_t); \
	trie_iterator(T) trie_func(end, T)(trie(T)* p_t); \
	trie_iterator(T) trie_func(rbegin, T)(trie(T)* p_t); \
	trie_iterator(T) trie_func(rend, T)(trie(T)* p_t);

/// Implements trie functions in a source file.
#define trie_imp_c(T) \
	trie_iterator_imp_c(T); \
	\
	static void cctl_join(trie_free_subtree, T)(trie(T)* p_node) { \
		if (!p_node) return; \
		for (size_t i = 0; i < 256; i++) { \
			if (p_node->ap_children[i]) { \
				cctl_join(trie_free_subtree, T)(p_node->ap_children[i]); \
				free(p_node->ap_children[i]); \
				p_node->ap_children[i] = NULL; \
			} \
		} \
	} \
	\
	void trie_func(init, T)(trie(T)* p_t) { \
		memset(p_t, 0, sizeof(trie(T))); \
	} \
	\
	trie(T) trie_func(new, T)(void) { \
		trie(T) t; \
		trie_func(init, T)(&t); \
		return t; \
	} \
	\
	trie(T) trie_func(from_items, T)(const trie_pair(T)* p_items, size_t count) { \
		trie(T) t; \
		trie_func(init, T)(&t); \
		if (p_items && count > 0) { \
			if (!trie_func(extend_items, T)(&t, p_items, count)) { \
				trie_func(free, T)(&t); \
			} \
		} \
		return t; \
	} \
	\
	bool trie_func(extend_items, T)(trie(T)* p_t, const trie_pair(T)* p_items, size_t count) { \
		if (!p_t || !p_items || count == 0) return true; \
		for (size_t i = 0; i < count; i++) { \
			if (!trie_func(insert, T)(p_t, p_items[i].p_key, p_items[i].value)) return false; \
		} \
		return true; \
	} \
	\
	void trie_func(clear, T)(trie(T)* p_t) { \
		if (!p_t) return; \
		cctl_join(trie_free_subtree, T)(p_t); \
		p_t->existence = false; \
		p_t->count = 0; \
	} \
	\
	void trie_func(free, T)(trie(T)* p_t) { \
		if (!p_t) return; \
		trie_func(clear, T)(p_t); \
	} \
	\
	T* trie_func(insert, T)(trie(T)* p_t, const char* p_key, T item) { \
		if (!p_t || !p_key) return NULL; \
		trie(T)* p_curr = p_t; \
		while (*p_key) { \
			uint8_t index = (uint8_t)*p_key; \
			if (!p_curr->ap_children[index]) { \
				trie(T)* p_new_node = (trie(T)*) malloc(sizeof(trie(T))); \
				if (!p_new_node) return NULL; \
				trie_func(init, T)(p_new_node); \
				p_new_node->curr_char = index; \
				p_new_node->p_parent = p_curr; \
				p_curr->ap_children[index] = p_new_node; \
			} \
			p_curr = p_curr->ap_children[index]; \
			p_key++; \
		} \
		p_curr->data = item; \
		if (!p_curr->existence) { \
			p_curr->existence = true; \
			p_t->count++; \
		} \
		return &p_curr->data; \
	} \
	\
	void trie_func(remove, T)(trie(T)* p_t, const char* p_key) { \
		if (!p_t || !p_key) return; \
		trie(T)* p_curr = p_t; \
		while (*p_key) { \
			uint8_t index = (uint8_t)*p_key; \
			if (!p_curr->ap_children[index]) return; \
			p_curr = p_curr->ap_children[index]; \
			p_key++; \
		} \
		if (!p_curr->existence) return; \
		p_curr->existence = false; \
		p_t->count--; \
		while (p_curr->p_parent != NULL && !p_curr->existence) { \
			bool has_children = false; \
			for (size_t i = 0; i < 256; i++) { \
				if (p_curr->ap_children[i]) { \
					has_children = true; \
					break; \
				} \
			} \
			if (has_children) break; \
			trie(T)* p_parent = p_curr->p_parent; \
			p_parent->ap_children[p_curr->curr_char] = NULL; \
			free(p_curr); \
			p_curr = p_parent; \
		} \
	} \
	\
	T* trie_func(find, T)(trie(T)* p_t, const char* p_key) { \
		if (!p_t || !p_key) return NULL; \
		trie(T)* p_curr = p_t; \
		while (*p_key) { \
			uint8_t index = (uint8_t)*p_key; \
			if (!p_curr->ap_children[index]) return NULL; \
			p_curr = p_curr->ap_children[index]; \
			p_key++; \
		} \
		return p_curr->existence ? &p_curr->data : NULL; \
	} \
	\
	T trie_func(get_or_default, T)(trie(T)* p_t, const char* p_key, T default_value) { \
		T* p_value = trie_func(find, T)(p_t, p_key); \
		return p_value ? *p_value : default_value; \
	} \
	\
	char* trie_func(min_key, T)(trie(T)* p_t) { \
		if (!p_t || p_t->count == 0) return NULL; \
		trie_iterator(T) it = trie_func(begin, T)(p_t); \
		if (!trie_iterator_is_valid(T, &it)) return NULL; \
		return trie_iterator_get_key(T, &it); \
	} \
	\
	char* trie_func(max_key, T)(trie(T)* p_t) { \
		if (!p_t || p_t->count == 0) return NULL; \
		trie_iterator(T) it = trie_func(rbegin, T)(p_t); \
		if (!trie_iterator_is_valid(T, &it)) return NULL; \
		return trie_iterator_get_key(T, &it); \
	} \
	\
	char* trie_func(get_key, T)(trie(T)* p_node) { \
		if (!p_node) return NULL; \
		size_t depth = 0; \
		trie(T)* p_temp = p_node; \
		while (p_temp->p_parent != NULL) { \
			depth++; \
			p_temp = p_temp->p_parent; \
		} \
		char* p_key = (char*) malloc(depth + 1); \
		if (!p_key) return NULL; \
		p_key[depth] = '\0'; \
		p_temp = p_node; \
		for (size_t i = depth; i > 0; i--) { \
			p_key[i - 1] = (char) p_temp->curr_char; \
			p_temp = p_temp->p_parent; \
		} \
		return p_key; \
	} \
	\
	trie_iterator(T) trie_func(lower_bound, T)(trie(T)* p_t, const char* p_key) { \
		trie_iterator(T) it = trie_func(begin, T)(p_t); \
		while (trie_iterator_is_valid(T, &it)) { \
			char* p_curr_key = trie_iterator_get_key(T, &it); \
			if (p_curr_key) { \
				int compare_result = strcmp(p_curr_key, p_key); \
				free(p_curr_key); \
				if (compare_result >= 0) return it; \
			} \
			trie_iterator_next(T, &it); \
		} \
		it.p_node = NULL; \
		return it; \
	} \
	\
	trie_iterator(T) trie_func(upper_bound, T)(trie(T)* p_t, const char* p_key) { \
		trie_iterator(T) it = trie_func(begin, T)(p_t); \
		while (trie_iterator_is_valid(T, &it)) { \
			char* p_curr_key = trie_iterator_get_key(T, &it); \
			if (p_curr_key) { \
				int compare_result = strcmp(p_curr_key, p_key); \
				free(p_curr_key); \
				if (compare_result > 0) return it; \
			} \
			trie_iterator_next(T, &it); \
		} \
		it.p_node = NULL; \
		return it; \
	} \
	\
	trie_iterator(T) trie_func(begin, T)(trie(T)* p_t) { \
		trie_iterator(T) it = { p_t }; \
		if (!p_t->existence) trie_iterator_func(next, T)(&it); \
		return it; \
	} \
	\
	trie_iterator(T) trie_func(end, T)(trie(T)* p_t) { \
		(void) p_t; \
		trie_iterator(T) it = { NULL }; \
		return it; \
	} \
	\
	trie_iterator(T) trie_func(rbegin, T)(trie(T)* p_t) { \
		trie_iterator(T) it = { p_t }; \
		trie(T)* p_curr = p_t; \
		while (true) { \
			bool found_child = false; \
			for (size_t i = 256; i > 0; i--) { \
				if (p_curr->ap_children[i - 1]) { \
					p_curr = p_curr->ap_children[i - 1]; \
					found_child = true; \
					break; \
				} \
			} \
			if (!found_child) break; \
		} \
		it.p_node = p_curr; \
		if (!p_curr->existence) trie_iterator_func(prev, T)(&it); \
		return it; \
	} \
	\
	trie_iterator(T) trie_func(rend, T)(trie(T)* p_t) { \
		(void) p_t; \
		trie_iterator(T) it = { NULL }; \
		return it; \
	}

#endif