diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..d948c5e --- /dev/null +++ b/Makefile @@ -0,0 +1,83 @@ +# This Makefile is to build following building blocks +# +AR_NAME := libljmm.a +SO_NAME := libljmm.so +RBTREE_TEST := rbt_test # test program for red-black tree +DEMO_NAME := demo # A demo illustrating how to use the lib being built. + +OPT_FLAGS = -O0 -g -DDEBUG +CFLAGS = -fvisibility=hidden -MMD -Wall $(OPT_FLAGS) +CXXFLAGS = $(CFLAGS) +AR_BUILD_CFLAGS = -DBUILDING_LIB +SO_BUILD_CFLAGS = -DBUILDING_LIB -fPIC + +CC ?= gcc +CXX ?= g++ + +BUILD_AR_DIR = obj/lib +BUILD_SO_DIR = obj/so + +RB_TREE_SRCS = rbtree.c +RB_TEST_SRCS = rb_test.cxx +ALLOC_SRCS = trunk.c page_alloc.c +DEMO_SRCS = demo.c + +C_SRCS = $(RB_TREE_SRCS) $(ALLOC_SRCS) +C_OBJS = ${C_SRCS:%.c=%.o} + +AR_OBJ = $(addprefix obj/lib/, $(C_OBJS)) +SO_OBJ = $(addprefix obj/so/, $(C_OBJS)) + +.PHONY = all clean test + +all: $(AR_NAME) $(SO_NAME) $(RBTREE_TEST) $(DEMO_NAME) + +$(RBTREE_TEST) $(DEMO_NAME) : $(AR_NAME) $(SO_NAME) + +-include ar_dep.txt +-include so_dep.txt + +##################################################################### +# +# Building static lib +# +##################################################################### +# +$(AR_NAME) : $(AR_OBJ) + $(AR) cru $@ $(AR_OBJ) + cat $(BUILD_AR_DIR)/*.d > ar_dep.txt + +$(AR_OBJ) : $(BUILD_AR_DIR)/%.o : %.c + $(CC) -c $(CFLAGS) $(AR_BUILD_CFLAGS) $< -o $@ + +##################################################################### +# +# Building shared lib +# +##################################################################### +$(SO_NAME) : $(SO_OBJ) + $(CC) $(CFLAGS) $(AR_BUILD_CFLAGS) $(SO_OBJ) -shared -o $@ + +$(SO_OBJ) : $(BUILD_SO_DIR)/%.o : %.c + $(CC) -c $(CFLAGS) $(SO_BUILD_CFLAGS) $< -o $@ + +##################################################################### +# +# Building demo program +# +##################################################################### +$(DEMO_NAME) : ${DEMO_SRCS:%.c=%.o} $(AR_NAME) + $(CC) $(filter %.o, $+) -L. -Wl,-static -lljmm -Wl,-Bdynamic -o $@ + +$(RBTREE_TEST) : ${RB_TREE_SRCS:%.c=%.o} ${RB_TEST_SRCS:%.cxx=%.o} + $(CXX) $(filter %.o, $+) -o $@ + +%.o : %.c + $(CC) $(CFLAGS) -c $< + +%.o : %.cxx + $(CXX) $(CXXFLAGS) -c $< + +clean: + rm -f *.o *.d ar_dep.txt so_dep.txt $(BUILD_AR_DIR)/* $(BUILD_SO_DIR)/* + rm -f $(AR_NAME) $(SO_NAME) $(RBTREE_TEST) $(DEMO_NAME) diff --git a/README.md b/README.md index 5289280..aab34ff 100644 --- a/README.md +++ b/README.md @@ -2,3 +2,30 @@ luajit-mm ========= Luajit take full advantage of lower 2G memory on AMD64 platform. + +Rumdimentary implementation. Not yet fully tested. Not yet cleanup the code. +Quite a few optimizatio is not yet implemented. + +Immediate todo +============== + + o.Refine and finish this README. + o.test, add enhancements. + +problem statement: +================== + On Linux/x86-64 platform, Luajit can use no more than 1G memory due to the +combination of bunch of nasty issues. 1G is way too small for server-side application. + + This package is trying to replace mmap/munmap/mremap with hence provide up to +about 2G space. + + +Basic ideas +=========== + o. When a application, which contain luajit, is launched, reserve the the space + from where `sbrk(0)` indidate all the way to 2G. + + o. Perform page allocation on the reserved space. the mmap/munmap/mremap is built + on this page allocation. Currently, we use buddy allocation for page allocation + with some optimizations in an attemp to reduce working set. diff --git a/demo.c b/demo.c new file mode 100644 index 0000000..9fd7b8f --- /dev/null +++ b/demo.c @@ -0,0 +1,53 @@ +#include +#include +#include +#include "lj_mm.h" + +static void* +mmap_wrap(size_t len) { + return lm_mmap(NULL, len, PROT_READ|PROT_WRITE, MAP_32BIT | MAP_PRIVATE, -1, 0); +} + +int +main(int argc, char** argv) { + lm_init(1); + dump_page_alloc(stderr); + + int size1 = 100; + char* p1 = mmap_wrap(size1); + fprintf(stderr, "\nsize=%d, %p\n", size1, p1); + dump_page_alloc(stderr); + + + int size2 = 4097; + char* p2 = mmap_wrap(size2); + fprintf(stderr, "\nsize=%d, %p\n", size2, p2); + dump_page_alloc(stderr); + + int size3 = 4097; + char* p3 = mmap_wrap(size3); + fprintf(stderr, "\nsize=%d %p\n", size3, p3); + dump_page_alloc(stderr); + + int size4 = 4096 * 3; + char* p4 = mmap_wrap(size4); + fprintf(stderr, "\nsize=%d %p\n", size4, p4); + dump_page_alloc(stderr); + + int size5 = 4096 * 2; + char* p5 = mmap_wrap(size5); + fprintf(stderr, "\nsize=%d %p\n", size5, p5); + dump_page_alloc(stderr); + + lm_munmap(p1, size1); + lm_munmap(p2, size2); + lm_munmap(p3, size3); + lm_munmap(p4, size4); + lm_munmap(p5, size5); + + fprintf(stderr, "\n\nAfter delete all allocations\n"); + dump_page_alloc(stderr); + /*lm_fini(); */ + + return 0; +} diff --git a/lj_mm.h b/lj_mm.h new file mode 100644 index 0000000..0aabbca --- /dev/null +++ b/lj_mm.h @@ -0,0 +1,41 @@ +#ifndef _LJ_MM_H_ +#define _LJ_MM_H_ + +#include /* for size_t */ + +#ifdef __cplusplus +extern "C" { +#endif + +#ifdef BUILDING_LIB + #define LJMM_EXPORT __attribute__ ((visibility ("protected"))) +#else + #define LJMM_EXPORT __attribute__ ((visibility ("default"))) +#endif + +/* Inititalize the memory-management system. If auto_fini is set + * (i.e. auto_fini != 0), there is no need to call lm_fini() at exit. + */ +int lm_init(int auto_fini) LJMM_EXPORT; +void lm_fini(void) LJMM_EXPORT; + +/* Same prototype as mmap(2), and munmap(2) */ +void *lm_mmap(void *addr, size_t length, int prot, int flags, + int fd, off_t offset) LJMM_EXPORT; + +int lm_munmap(void *addr, size_t length) LJMM_EXPORT; +void* lm_mremap(void* old_addr, size_t old_size, size_t new_size, int flags) LJMM_EXPORT; + +/* Some "bonus" interface functions */ +void* lm_malloc(size_t sz) LJMM_EXPORT; +int lm_free(void* mem) LJMM_EXPORT; + +#ifdef DEBUG +void dump_page_alloc(FILE*) LJMM_EXPORT; +#endif + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/lm_internal.h b/lm_internal.h new file mode 100644 index 0000000..c8fc1e0 --- /dev/null +++ b/lm_internal.h @@ -0,0 +1,34 @@ +#ifndef _LM_INTERNAL_H_ +#define _LM_INTERNAL_H_ + +#include /* for FILE */ + +#define LJMM_ADDR_UPBOUND ((unsigned int)0x80000000) + +/* "Huge" chunk of memmory. Memmory allocations are to carve blocks + * from the big trunk. + */ +typedef struct { + char* base; /* the starting address of the big trunk */ + char* start; /* the starting address of the usable portion */ + unsigned alloc_size; /* the size of the big trunk */ + unsigned usable_size; /* the size of the usable portion. + * usabe_size = page_num * page_size. + */ + unsigned page_num; /* number of available pages */ + unsigned page_size; /* cache of sysconf(_SC_PAGESIZE); */ +} lm_trunk_t; + +lm_trunk_t* lm_alloc_trunk(void); +void lm_free_trunk(void); +#ifdef DEBUG +void lm_dump_trunk(FILE* f); +#endif + +int lm_init_page_alloc(lm_trunk_t*); +void lm_fini_page_alloc(void); +#ifdef DEBUG +void dump_page_alloc(FILE* f); +#endif + +#endif diff --git a/lm_util.h b/lm_util.h new file mode 100644 index 0000000..2109f77 --- /dev/null +++ b/lm_util.h @@ -0,0 +1,28 @@ +#ifndef _LM_UTIL_H_ +#define _LM_UTIL_H_ + +#ifdef THREAD_SAFE + void entery_mutex(); + void leave_mutex(); + #define ENTER_MUTEX enter_mutex() + #define LEAVE_MUTEX leave_mutex() +#else + #define ENTER_MUTEX + #define LEAVE_MUTEX +#endif + +typedef unsigned int uint; + + +#ifdef DEBUG + // Usage examples: ASSERT(a > b), ASSERT(foo() && "Opps, foo() reutrn 0"); + #define ASSERT(c) if (!(c))\ + { fprintf(stderr, "%s:%d Assert: %s\n", __FILE__, __LINE__, #c); abort(); } +#else + #define ASSERT(c) ((void)0) +#endif + +#define likely(x) __builtin_expect((x),1) +#define unlikely(x) __builtin_expect((x),0) + +#endif diff --git a/obj/lib/empty_file b/obj/lib/empty_file new file mode 100644 index 0000000..e69de29 diff --git a/obj/so/empty_file b/obj/so/empty_file new file mode 100644 index 0000000..e69de29 diff --git a/page_alloc.c b/page_alloc.c new file mode 100644 index 0000000..6e8148c --- /dev/null +++ b/page_alloc.c @@ -0,0 +1,532 @@ +#include /* for intptr_t */ +#include +#include +#include +#include "lm_util.h" +#include "lm_internal.h" +#include "lj_mm.h" +#include "rbtree.h" +#include "lj_mm.h" + +/************************************************************************** + * + * Depicting page + * + ************************************************************************** + */ +typedef struct { + short order; /* the order in the buddy allocation */ + short flags; +} lm_page_t; + +/* A block is a consecutive trunk of memory containing power-of-two pages. A block + * is identified by the id of its first page. + */ +typedef int blk_id_t; + +typedef enum { + PF_LEADER = (1 << 0), /* set if it's the first page of a block */ + PF_ALLOCATED = (1 << 1), /* set if it's "leader" of a allocated block */ + PF_LAST = PF_ALLOCATED, +} page_flag_t; + +static inline int +is_page_leader(lm_page_t * p) { + return p->flags & PF_LEADER; +} + +static inline void +set_page_leader(lm_page_t* p) { + p->flags |= PF_LEADER; +} + +static inline void +reset_page_leader(lm_page_t* p) { + p->flags &= ~PF_LEADER; +} + +static inline int +is_allocated_blk(lm_page_t* p) { + return is_page_leader(p) && (p->flags & PF_ALLOCATED); +} + +static inline void +set_allocated_blk(lm_page_t* p) { + ASSERT(is_page_leader(p)); + p->flags |= PF_ALLOCATED; +} + +static inline void +reset_allocated_blk(lm_page_t* p) { + ASSERT(is_page_leader(p)); + p->flags &= ~PF_ALLOCATED; +} + +static inline int +verify_order(blk_id_t blk_id, int order) { + return (blk_id & ((1<page_num; + int alloc_sz = sizeof(lm_alloc_t) + + sizeof(lm_page_t) * (page_num + 1); + + alloc_info = (lm_alloc_t*) malloc(alloc_sz); + if (!alloc_info) + return 0; + + alloc_info->first_page = trunk->start; + alloc_info->page_num = page_num; + alloc_info->page_size = trunk->page_size; + alloc_info->page_size_log2 = log2_int32(trunk->page_size); + + /* Init the page-info */ + char* p = (char*)(alloc_info + 1); + int align = __alignof__(lm_page_t); + p = (char*)((((intptr_t)p) + align - 1) & ~align); + alloc_info->page_info = (lm_page_t*)p; + + int i; + lm_page_t* pi = alloc_info->page_info; + for (i = 0; i < page_num; i++) { + pi[i].order = INVALID_ORDER; + pi[i].flags = 0; + } + + /* Init the buddy allocator */ + int e; + rb_tree_t* free_blks = &alloc_info->free_blks[0]; + for (i = 0, e = MAX_ORDER; i < e; i++) + rbt_init(&free_blks[i]); + rbt_init(&alloc_info->alloc_blks); + + /* to make subsequent add_block() happy */ + alloc_info->max_order = MAX_ORDER; + + unsigned int bitmask = 0x80000000; + int max_order = 0, order = 31; + int page_id = 0; + + while (bitmask) { + if (page_num & bitmask) { + add_block(page_id, order); + page_id += (1 << order); + + if (max_order == 0) + max_order = order; + } + bitmask = bitmask >> 1; + order --; + } + alloc_info->max_order = max_order; + return 1; +} + +// Return ceil(log2(num)) +static int +ceil_log2_int32 (unsigned num) { + int res = 31 - __builtin_clz(num); + res += (num & (num - 1)) ? 1 : 0; + return res; +} + +static int +log2_int32(unsigned num) { + return 31 - __builtin_clz(num); +} + +/* Add the free block of the given "order" to the buddy system */ +static inline int +add_block(blk_id_t block, int order) { + lm_page_t* page = alloc_info->page_info + block; + + ASSERT(order >= 0 && order <= alloc_info->max_order); + + page->order = order; + set_page_leader(page); + reset_allocated_blk(page); + + return rbt_insert(&alloc_info->free_blks[order], block, 0); +} + +static int +find_block(blk_id_t block, int order, intptr_t* value) { + ASSERT(order >= 0 && order <= alloc_info->max_order && + verify_order(block, order)); + + return rbt_search(&alloc_info->free_blks[order], block, value); +} + +static inline int +remove_block(blk_id_t block, int order) { + lm_page_t* page = alloc_info->page_info + block; + + ASSERT(page->order == order && find_block(block, order, NULL)); + ASSERT(!is_allocated_blk(page) && verify_order(block, order)); + set_allocated_blk(page); + + return rbt_delete(&alloc_info->free_blks[order], block); +} + +void +lm_fini_page_alloc(void) { + if (alloc_info) { + rb_tree_t* free_blks = &alloc_info->free_blks[0]; + int i, e; + for (i = 0, e = MAX_ORDER; i < e; i++) + rbt_fini(free_blks + i); + + rbt_fini(&alloc_info->alloc_blks); + + free(alloc_info); + alloc_info = 0; + } +} + + +static int +free_block(int page_id) { + int del_res = rbt_delete(&alloc_info->alloc_blks, page_id); +#ifdef DEBUG + ASSERT(del_res); +#else + (void)del_res; +#endif + + lm_page_t* pi = alloc_info->page_info; + lm_page_t* page = pi + page_id; + int order = page->order; + ASSERT (find_block(page_id, order, NULL) == 0); + + char* block_addr = alloc_info->first_page + + (page_id << alloc_info->page_size_log2); + size_t block_len = (1<page_size_log2; + madvise(block_addr, block_len, MADV_DONTNEED); + + /* Consolidate adjacent buddies */ + int page_num = alloc_info->page_num; + while (1) { + int buddy_id = page_id ^ (1<= page_num || + pi[buddy_id].order != order || + !is_page_leader(pi + buddy_id) || + is_allocated_blk(pi + buddy_id)) { + break; + } + remove_block(buddy_id, order); + page_id = page_id < buddy_id ? page_id : buddy_id; + order++; + } + + add_block(page_id, order); + return 1; +} + + +/************************************************************************** + * + * Implementation of this package's interface funcs + * + ************************************************************************** + */ +void +lm_fini(void) { + ENTER_MUTEX; + + lm_fini_page_alloc(); + lm_free_trunk(); + + LEAVE_MUTEX; +} + +/* The purpose of this variable is to workaround a link problem: If we were + * directly feeding lm_fini to atexit() in function lm_init(), we would be + * going to see a complaint like this: + * + * "...relocation R_X86_64_PC32 against protected symbol `lm_fini' can not + * be used when making a shared object"... + * + * I think it's perfectly fine using R_X86_64_PC32 as a relocation for + * the protected symbol lm_fini. It seems like it's GNU ld (I'm using 2.24) + * problem. Actually gold linker is able to link successfully. + * + * NOTE: This variable must be visible to other modules, otherwise, with + * higher optimization level, compiler can propagate its initial value (i.e. + * the lm_fini) to where it's referenced. + */ +void (*lm_fini_ptr)() __attribute__((visibility("protected"))) = lm_fini; + +/* Initialize the allocation, return non-zero on success, 0 otherwise. */ +int +lm_init(int auto_fini) { + int res = 1; + + ENTER_MUTEX; + if (auto_fini != 0) { + /* Do not directly feed lm_fini to atexit(), see the comment to + * variable "lm_fini_ptr" for why. + */ + res = atexit(lm_fini_ptr); + + /* Negate the sense of 'success' :-) */ + res = (res == 0) ? 1 : 0; + } + + if (res) + res = lm_init_page_alloc(lm_alloc_trunk()); + + LEAVE_MUTEX; + + return res; +} + +void* +lm_mmap(void *addr, size_t length, int prot, int flags, + int fd, off_t offset) { + void *p = NULL; + + ENTER_MUTEX; + { + if (addr || fd != -1 || !(flags & MAP_32BIT) || !length) { + errno = EINVAL; + } else { + p = lm_malloc(length); + } + } + LEAVE_MUTEX; + + return p ? p : MAP_FAILED; +} + +static int +lm_unmap_helper(void* addr, size_t length) { + long ofst = ((char*)addr) - ((char*)alloc_info->first_page); + if (unlikely (ofst < 0)) + return 0; + + long page_sz = alloc_info->page_size; + if (unlikely ((ofst & (page_sz - 1)))) + return 0; + + long page_id = ofst >> log2_int32(page_sz); + intptr_t map_size; + if (rbt_search(&alloc_info->alloc_blks, page_id, &map_size) == 0) + return 0; + + int page_sz_log2 = alloc_info->page_size_log2; + int map_pages = ((uintptr_t)(page_sz - 1 + map_size)) >> page_sz_log2; + int rel_pages = ((uintptr_t)(page_sz - 1 + length)) >> page_sz_log2; + + if (map_pages != rel_pages) + return 0; + + return free_block(page_id); +} + +int +lm_munmap(void* addr, size_t length) { + int page_sz = alloc_info->page_size; + length = (length + (page_sz - 1)) & ~(page_sz - 1); + int retval = 0; + + ENTER_MUTEX; + { + if (!length || (((uintptr_t)addr) & (page_sz - 1))) { + errno = EINVAL; + retval = -1; + } else { + retval = lm_unmap_helper(addr, length); + /* negate the sense of succ. */ + retval = retval ? 0 : -1; + } + } + LEAVE_MUTEX; + + return retval; +} + +void* +lm_mremap(void* old_addr, size_t old_size, size_t new_size, int flags) { + (void)old_addr; + (void)old_size; + (void)new_size; + (void)flags; + return MAP_FAILED; +} + +int +lm_free(void* mem) { + if (unlikely (!alloc_info)) + return 0; + + long ofst = ((char*)mem) - ((char*)alloc_info->first_page); + if (unlikely (ofst < 0)) + return 0; + + long page_sz = alloc_info->page_size; + if (unlikely ((ofst & (page_sz - 1)))) + return 0; + + long page_id = ofst >> log2_int32(page_sz); + int page_num = alloc_info->page_num; + if (unlikely(page_id >= page_num)) + return 0; + + lm_page_t* pi = alloc_info->page_info; + lm_page_t* page = pi + page_id; + if (unlikely(!is_page_leader(page))) + return 0; + + if (unlikely(!is_allocated_blk(page))) + return 0; + + return free_block(page_id); +} + +void* +lm_malloc(size_t sz) { + if (!alloc_info) { + lm_init(1); + if (!alloc_info) + return NULL; + } + + /* Determine the order of allocation request */ + int req_order = ceil_log2_int32(sz); + req_order -= alloc_info->page_size_log2; + if (req_order < 0) + req_order = 0; + + int max_order = alloc_info->max_order; + if (req_order > max_order) + return 0; + + rb_tree_t* free_blks = alloc_info->free_blks; + + int i, blk_order = -1; + blk_id_t blk_id = -1; + + /* Find the smallest available block big enough to accommodate the + * allocation request. + */ + for (i = req_order; i < max_order; i++) { + rb_tree_t* rbt = free_blks + i; + if (!rbt_is_empty(rbt)) { + blk_id = rbt_get_min(rbt); + blk_order = i; + break; + } + } + + if (blk_id == -1) + return NULL; + + remove_block(blk_id, blk_order); + alloc_info->page_info[blk_id].order = req_order; + + /* The free block may be too big. If this is the case, keep splitting + * the block until it tightly fit the allocation request. + */ + int bo = blk_order; + while (bo > req_order) { + bo --; + int split_block = blk_id + (1 << bo); + add_block(split_block, bo); + } + + int insert_res = rbt_insert(&alloc_info->alloc_blks, blk_id, (intptr_t)sz); + rbt_verify(&alloc_info->alloc_blks); +#ifdef DEBUG + ASSERT(insert_res); +#else + (void)insert_res; +#endif + return alloc_info->first_page + (blk_id << alloc_info->page_size_log2); +} + +/************************************************************************** + * + * Debugging Support & Misc "cold" functions + * + ************************************************************************** + */ +#ifdef DEBUG +void +dump_page_alloc(FILE* f) { + if (!alloc_info) { + fprintf(f, "not initialized yet\n"); + fflush(f); + return; + } + + /* dump the buddy system */ + fprintf (f, "Buddy system: max-order=%d\n", alloc_info->max_order); + + int i, e; + char* page_start_addr = alloc_info->first_page; + int page_sz_log = alloc_info->page_size_log2; + + for (i = 0, e = alloc_info->max_order; i <= e; i++) { + rb_tree_t* free_blks = &alloc_info->free_blks[i]; + + if (rbt_is_empty(free_blks)) + continue; + + fprintf(f, "Order = %3d: ", i); + rb_iter_t iter, iter_e; + for (iter = rbt_iter_begin(free_blks), + iter_e = rbt_iter_end(free_blks); + iter != iter_e; + iter = rbt_iter_inc(free_blks, iter)) { + rb_node_t* node = rbt_iter_deref(iter); + int page_id = node->key; + char* addr = page_start_addr + (page_id << page_sz_log); + fprintf(f, "%d (%p, len=%d), ", page_id, addr, (int)node->value); + } + fputs("\n", f); + } +} +#endif diff --git a/rb_test.cxx b/rb_test.cxx new file mode 100644 index 0000000..c3143d0 --- /dev/null +++ b/rb_test.cxx @@ -0,0 +1,304 @@ +#include +#include +#include "rbtree.h" + +using namespace std; + +#define KEY_VAL_DELTA 123 +#define RN(v) {v + KEY_VAL_DELTA, v, RB_RED} +#define BN(v) {v + KEY_VAL_DELTA, v, RB_BLACK} + +class RB_UNIT_TEST { +public: + // If come across a bug, turn on "dump_tree" to instruct the unit-tester + // to dump the RB-tree before and after the problematic operation. + // + // The "id" is just used to identify the testing case, and hence ease + // trouble-shooting. + RB_UNIT_TEST(int test_id, rb_valcolor_t* nodes, int nd_num, + bool dump_tree = false) : _test_id(test_id), + _dump_tree(dump_tree) { + fprintf(stdout, "Testing unit test %d ...", test_id); + _rbt = rbt_create_manually(nodes, nd_num); + _save_fail_cnt = _fail_cnt; + } + + ~RB_UNIT_TEST() { + if (_rbt) + rbt_destroy(_rbt); + fprintf(stdout, " %s\n", (_save_fail_cnt == _fail_cnt) ? "succ" : "fail"); + } + + static void Reset_Fail_Cnt() { _fail_cnt = 0; } + static int Get_Fail_Cnt() { return _fail_cnt; } + + bool Delete(int val, int expect_ret_val = 1) { + if (DeleteHelper(val, expect_ret_val)) + return true; + _fail_cnt ++; + return false; + } + + bool Insert(int val, int expect_ret_val = 1) { + if (InsertHelper(val, expect_ret_val)) + return true; + _fail_cnt ++; + return false; + } + + bool Search(int key, intptr_t val, bool expect_ret_val = 1) { + if (SearchHelper(key, val, expect_ret_val)) + return true; + _fail_cnt ++; + return false; + } + +private: + bool DeleteHelper(int val, int expect_ret_val) { + if (!_rbt) + return false; + + if (_dump_tree) + Dump_Tree("before_del"); + + if (!rbt_verify(_rbt)) + return false; + + int ret = rbt_delete(_rbt, val); + + if (_dump_tree) + Dump_Tree("after_del"); + + if (ret == expect_ret_val && rbt_verify(_rbt) && VerifyKeyVal()) + return true; + + fprintf(stdout, " fail to delete %d", val); + return false; + } + + bool InsertHelper(int key, int expect_ret_val) { + if (!_rbt) + return false; + + if (_dump_tree) + Dump_Tree("before_insert"); + + if (!rbt_verify(_rbt)) + return false; + + int ret = rbt_insert(_rbt, key, key + KEY_VAL_DELTA); + + if (_dump_tree) + Dump_Tree("after_insert"); + + if (ret == expect_ret_val && rbt_verify(_rbt) && VerifyKeyVal()) + return true; + + fprintf(stdout, " fail to insert %d", key); + return false; + } + + bool SearchHelper(int key, intptr_t val, bool expect_ret_val) { + if (!_rbt) + return false; + + if (_dump_tree) + Dump_Tree("before_search"); + + if (!rbt_verify(_rbt)) + return false; + + intptr_t tmpval; + int ret = rbt_search(_rbt, key, &tmpval); + + if (_dump_tree) + Dump_Tree("after_search"); + + if (ret == expect_ret_val && val == tmpval && rbt_verify(_rbt)) + return true; + + fprintf(stdout, " fail to search %d", key); + + return false; + } + + string GetDumpFileName(const char* op_name) { + char buf[200]; + int len = snprintf(buf, sizeof(buf), "test_%d_%s.dot", _test_id, op_name); + buf[len] = '\0'; + return string(buf); + } + + // In most testing cases, we set "value = key + KEY_VAL_DELTA" (The + // rationale is just to ease testing). This function is to verify the + // key/value relation is broken. + // + bool VerifyKeyVal() { + for (rb_iter_t iter = rbt_iter_begin(_rbt), + iter_e = rbt_iter_begin(_rbt); + iter != iter_e; + iter = rbt_iter_inc(_iter, iter)) { + rb_node_t* nd = rbt_iter_deref(iter); + if (nd->value != nd->key + KEY_VAL_DELTA) + return false; + } + + return true; + } + + void Dump_Tree(const char* op_name) { +#ifdef DEBUG + rbt_dump_dot(_rbt, GetDumpFileName(op_name).c_str()); +#else + (void)op_name; +#endif + } + + int _test_id; + bool _dump_tree; + rb_tree_t* _rbt; + static int _fail_cnt; + int _save_fail_cnt; +}; + +int RB_UNIT_TEST::_fail_cnt = 0; + +bool +unit_test() { + RB_UNIT_TEST::Reset_Fail_Cnt(); + + // test 8. + { + rb_valcolor_t nodes[] = { BN(1), RN(2) }; + RB_UNIT_TEST ut(8, nodes, 2); + ut.Delete(1); + ut.Delete(2); + } + + ///////////////////////////////////////////////////////////////////////// + // + // Insert tests + // + ///////////////////////////////////////////////////////////////////////// + // + + fprintf(stdout, "\n>Testing insert operation...\n"); + // Test 1, Cover the case 1, 2 and 3 of insertion. + // Testing case is from book : + // Thomas H. Cormen et. al, Instruction to Algorithm, 2nd edition, + // page 282. + { + rb_valcolor_t nodes[] = { BN(11), RN(2), BN(14), BN(1), BN(7), + RN(15), RN(5), RN(8) }; + RB_UNIT_TEST ut(1, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Insert(4); + } + + // Test 2. A contrived example for covering case 1', 2' and 3' + // + { + rb_valcolor_t nodes[] = { BN(4), BN(2), RN(10), RN(1), BN(7), + BN(11), RN(6), RN(9) }; + RB_UNIT_TEST ut(2, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Insert(8); + } + + ///////////////////////////////////////////////////////////////////////// + // + // Delete Tests + // + ///////////////////////////////////////////////////////////////////////// + // + fprintf(stdout, "\n>Testing delete operation...\n"); + // test 1 + { + rb_valcolor_t nodes[] = { BN(40), BN(20), BN(60), RN(10), RN(50), RN(70) }; + RB_UNIT_TEST ut(1, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Delete(20); + } + + // test 2 + { + rb_valcolor_t nodes[] = { BN(40), BN(20), BN(60), RN(10), RN(30), + RN(50), RN(70) }; + RB_UNIT_TEST ut(2, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Delete(20); + } + + // test 3 + { + rb_valcolor_t nodes[] = { BN(40), BN(20), BN(60), BN(10), BN(30), + BN(50), BN(70), RN(21) }; + RB_UNIT_TEST ut(3, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Delete(20); + } + + // test 4. Test the case 1 and 2 in rbt_delete_fixup(). + { + rb_valcolor_t nodes[] = { BN(8), RN(2), BN(10), BN(0), BN(4), BN(9), + BN(11), BN(-1), BN(1), BN(3), RN(6), + BN(5), BN(7) }; + RB_UNIT_TEST ut(4, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Delete(2); + } + + // test 5. Test the case 2 and 4 in rbt_delete_fixup(). + { + rb_valcolor_t nodes[] = { BN(80), RN(20), BN(100), BN(00), BN(40), BN(90), + BN(110), BN(-10), BN(10), BN(30), + BN(60), RN(50), RN(70), + BN(49), BN(51), BN(69), BN(71), + BN(-11), BN(-9), BN(9), BN(11), + BN(29), BN(31), + BN(89), BN(91), BN(109), BN(111) }; + RB_UNIT_TEST ut(5, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Delete(20); + } + + // test 6. For case 1' and 2' + { + rb_valcolor_t nodes[] = { BN(4), RN(2), BN(5), BN(1), BN(3) }; + RB_UNIT_TEST ut(6, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Delete(5); + } + + // test 7. For case 2', 3' and 4'. + { + rb_valcolor_t nodes[] = { BN(6), BN(2), BN(8), BN(1), RN(4), BN(7), + BN(9), BN(3), BN(5) }; + RB_UNIT_TEST ut(7, nodes, sizeof(nodes)/sizeof(nodes[0])); + ut.Delete(8); + } + + // test 8. + { + rb_valcolor_t nodes[] = { BN(1), RN(2) }; + RB_UNIT_TEST ut(8, nodes, 2); + ut.Delete(1); + ut.Delete(2); + } + + // test 9. + { + rb_valcolor_t nodes[] = { BN(1) }; + RB_UNIT_TEST ut(8, nodes, 1); + ut.Delete(1); + } + + ///////////////////////////////////////////////////////////////////////// + // + // Search Tests + // + ///////////////////////////////////////////////////////////////////////// + // + + return RB_UNIT_TEST::Get_Fail_Cnt() == 0; +}; + +int +main(int argc, char** argv) { + if (!unit_test()) + return 1; + + return 0; +} diff --git a/rbtree.c b/rbtree.c new file mode 100644 index 0000000..55c78cb --- /dev/null +++ b/rbtree.c @@ -0,0 +1,795 @@ +#ifdef DEBUG +#include +#endif + +#include +#include "rbtree.h" + +#define INVALID_IDX (-1) +#define SENTINEL_IDX 0 +#define likely(x) __builtin_expect((x),1) +#define unlikely(x) __builtin_expect((x),0) + +#define SWAP(a, b) { typeof(a) t = a; a = b; b = t; } + +/**************************************************************************** + * + * Constructors & Destructors + * + **************************************************************************** + */ +int +rbt_init(rb_tree_t* rbt) { + rbt->capacity = 16; + rbt->tree = (rb_node_t*)malloc(rbt->capacity * sizeof(rb_node_t)); + if (rbt->tree == 0) + return 0; + + rbt->root = SENTINEL_IDX; + rbt->node_num = 1; /* sentinel */ + + /* Init the sentinel */ + rb_node_t* s = rbt->tree; + s->left = s->right = s->parent = INVALID_IDX; + s->color = RB_BLACK; + + return 1; +} + +void +rbt_fini(rb_tree_t* rbt) { + if (rbt && rbt->tree) { + free((void*)rbt->tree); + rbt->tree = 0; + rbt->capacity = rbt->node_num = 0; + rbt->root = INVALID_IDX; + } +} + +rb_tree_t* +rbt_create(void) { + rb_tree_t* rbt = (rb_tree_t*)malloc(sizeof(rb_tree_t)); + if (rbt && rbt_init(rbt)) + return rbt; + + return NULL; +} + +void +rbt_destroy(rb_tree_t* rbt) { + if (rbt) { + rbt_fini(rbt); + free((void*)rbt); + } +} + +/**************************************************************************** + * + * Untility functions + * + **************************************************************************** + */ +/* Return 0 if val < node->data, 1 otherwise */ +static inline int +less_than(rb_node_t* node, int key) { + return key < node->key; +} + +static inline int +greater_than(rb_node_t* node, int key) { + return key > node->key; +} + +static void +update_kid(rb_node_t* dad, int kid_was, int kid_is) { + if (dad->left == kid_was) + dad->left = kid_is; + else { + ASSERT(dad->right == kid_was); + dad->right = kid_is; + } +} + +static void +rbt_left_rotate(rb_tree_t* rbt, rb_node_t* node) { + rb_node_t* nd_vect = rbt->tree; + + int node_idx = node - nd_vect; + int kid_idx = node->right; + int par_idx = node->parent; + + rb_node_t* kid = nd_vect + kid_idx; + if (kid->left != INVALID_IDX) { + nd_vect[kid->left].parent = node_idx; + } + + node->right = kid->left; + node->parent = kid_idx; + + kid->left = node_idx; + kid->parent = par_idx; + + if (par_idx != INVALID_IDX) { + rb_node_t* dad = nd_vect + par_idx; + if (dad->left == node_idx) + dad->left = kid_idx; + else { + dad->right = kid_idx; + } + } else { + ASSERT(rbt->root == node_idx); + rbt->root = kid_idx; + } +} + +static void +rbt_right_rotate(rb_tree_t* rbt, rb_node_t* node) { + rb_node_t* nd_vect = rbt->tree; + + int node_idx = node - nd_vect; + int kid_idx = node->left; + int par_idx = node->parent; + + rb_node_t* kid = nd_vect + kid_idx; + if (kid->right != INVALID_IDX) { + nd_vect[kid->right].parent = node_idx; + } + node->left = kid->right; + node->parent = kid_idx; + + kid->right = node_idx; + kid->parent = par_idx; + + if (par_idx != INVALID_IDX) { + rb_node_t* dad = nd_vect + par_idx; + if (dad->left == node_idx) + dad->left = kid_idx; + else { + dad->right = kid_idx; + } + } else { + ASSERT(rbt->root == node_idx); + rbt->root = kid_idx; + } +} + +/* Try to shrink the node vector */ +static int +rbt_try_shrink(rb_tree_t* rbt) { + if (rbt->capacity < 2*rbt->node_num || + rbt->capacity < 32) + return 1; + + int cap = rbt->node_num * 3 / 2; + rbt->tree = (rb_node_t*)realloc(rbt->tree, cap * sizeof(rb_node_t)); + if (rbt->tree == 0) + return 0; + + rbt->capacity = cap; + return 1; +} + +/**************************************************************************** + * + * Binary-search-tree operations + * + **************************************************************************** + */ + +/* Search the binary-search-tree; if found, return the index of node, + * INVALID_IDX otherwise. + */ +inline static int +bst_search(rb_tree_t* rbt, int key) { + rb_node_t* nd_vect = rbt->tree; + rb_node_t* sentinel = rbt->tree; + + rb_node_t* cur = nd_vect + rbt->root; + while (cur != sentinel) { + if (less_than(cur, key)) + cur = nd_vect + cur->left; + else if (greater_than(cur, key)) + cur = nd_vect + cur->right; + else + return cur - nd_vect; + } + + return INVALID_IDX; +} + +/* Insert "val" in the the binary-search-tree. the "bst" stands + * for binary-search-tree. This function has nothing to do with + * rb tree specific features. + * + * Return the index of the node being inserted if it was successful, + * or INVALID_IDX otherwise. + */ +static int +bst_insert(rb_tree_t* t, int key, intptr_t value) { + rb_node_t* nodes = t->tree; + + /* Resize the vector if necessary */ + if (t->capacity <= t->node_num) { + int cap = t->node_num * 3/2; + if (cap <= 16) + cap = 16; + + nodes = t->tree = (rb_node_t*)realloc(nodes, 100* sizeof(rb_node_t)); + t->capacity = cap; + + if (!nodes) + return INVALID_IDX; + } + + /* The tree is empty */ + if (unlikely(t->root == SENTINEL_IDX)) { + int root_id = 1; + t->root = root_id; + t->node_num = 2; + + rb_node_t* root = nodes + root_id; + root->key = key; + root->value = value; + root->left = root->right = SENTINEL_IDX; + root->parent = INVALID_IDX; + root->color = RB_BLACK; + + return root_id; + } + + /* Insert the value in the tree. Care must be taken to avoid inserting a + * value which is already in the tree. + */ + rb_node_t* prev = 0; + rb_node_t* cur = nodes + t->root; + rb_node_t* sentinel = nodes; + + while (cur != sentinel) { + prev = cur; + if (less_than(cur, key)) + cur = nodes + cur->left; + else if (greater_than(cur, key)) + cur = nodes + cur->right; + else { + /* the value is already in the tree */ + return INVALID_IDX; + } + } + + int new_nd_idx = t->node_num++; + rb_node_t* new_nd = nodes + new_nd_idx; + new_nd->key = key; + new_nd->value = value; + new_nd->parent = prev - nodes; + new_nd->left = new_nd->right = SENTINEL_IDX; + new_nd->color = RB_RED; + + if (less_than(prev, key)) + prev->left = new_nd_idx; + else + prev->right = new_nd_idx; + + return new_nd_idx; +} + +/**************************************************************************** + * + * RB-tree operations + * + **************************************************************************** + */ +int +rbt_search(rb_tree_t* rbt, int key, intptr_t* value) { + int nd_idx = bst_search(rbt, key); + if (nd_idx != INVALID_IDX) { + if (value) + *value = rbt->tree[nd_idx].value; + return 1; + } + return 0; +} + +int +rbt_get_min(rb_tree_t* rbt) { + ASSERT(!rbt_is_empty(rbt)); + + rb_node_t* nd_vect = rbt->tree; + rb_node_t* root = nd_vect + rbt->root; + rb_node_t* node = 0; + if (root->left != SENTINEL_IDX) + node = nd_vect + root->left; + else if (root->right != SENTINEL_IDX) + node = nd_vect + root->right; + else + node = root; + + while (node->left != SENTINEL_IDX) { + node = nd_vect + node->left; + } + + return node->key; +} + +int +rbt_insert(rb_tree_t* rbt, int key, intptr_t value) { + /* step 1: insert the key/val pair into the binary-search-tree */ + int nd_idx = bst_insert(rbt, key, value); + if (nd_idx == INVALID_IDX) + return 0; + + /* step 2: Perform color fix up */ + rb_node_t* nd_vect = rbt->tree; + rb_node_t* cur = nd_vect + nd_idx; + + while (cur->parent != INVALID_IDX && nd_vect[cur->parent].color == RB_RED) { + rb_node_t* dad = nd_vect + cur->parent; + rb_node_t* grandpar = nd_vect + dad->parent; + + if (nd_vect + grandpar->left == dad) { + rb_node_t* uncle = nd_vect + grandpar->right; + /* case 1: Both parent and uncle are in red. Just flip the color + * of parent, uncle and grand-parent. + */ + if (uncle->color == RB_RED) { + grandpar->color = RB_RED; + dad->color = uncle->color = RB_BLACK; + cur = grandpar; + continue; + } + + /* case 2: Parent and uncle's color are different (i.e. parent in + * red, uncle in black), and "cur" is parent's *RIGHT* kid. + */ + if (dad->right + nd_vect == cur) { + /* left rotate around parent */ + rbt_left_rotate(rbt, dad); + SWAP(cur, dad); + + /* Fall through to case 3 */ + } + + /* case 3: The condition is the same as case 2, except that 'cur' + * is the *LEFT* kid of the parent. + */ + rbt_right_rotate(rbt, grandpar); + dad->color = RB_BLACK; + grandpar->color = RB_RED; + + break; /* we are done, almost*/ + } else { + rb_node_t* uncle = nd_vect + grandpar->left; + /* case 1': Both parent and uncle are in red. Just flip the color + * of parent, uncle and grand-parent. + */ + if (uncle->color == RB_RED) { + grandpar->color = RB_RED; + dad->color = uncle->color = RB_BLACK; + cur = grandpar; + continue; + } + + + /* case 2': Parent and uncle's color are different (i.e. parent in + * red, uncle in black), and "cur" is parent's *LEFT* kid. + */ + if (dad->left + nd_vect == cur) { + /* left rotate around parent */ + rbt_right_rotate(rbt, dad); + SWAP(cur, dad); + /* Fall through to case 3 */ + } + + /* case 3: The condition is the same as case 2, except that 'cur' + * is the *RIGHT* kid of the parent. + */ + rbt_left_rotate(rbt, grandpar); + dad->color = RB_BLACK; + grandpar->color = RB_RED; + + break; + } + } + + /* make sure the root is in black */ + nd_vect[rbt->root].color = RB_BLACK; + + return 1; +} + +static void +rbt_delete_fixup(rb_tree_t* rbt, int node_idx) { + rb_node_t* nd_vect = rbt->tree; + while (node_idx != rbt->root && nd_vect[node_idx].color == RB_BLACK) { + rb_node_t* node = nd_vect + node_idx; + rb_node_t* dad = nd_vect + node->parent; + + if (dad->left == node_idx) { + int sibling_idx = dad->right; + rb_node_t* sibling = nd_vect + sibling_idx; + + /* case 1: sibling is in red color. Rotate around dad. */ + if (sibling->color == RB_RED) { + sibling->color = RB_BLACK; + dad->color = RB_RED; + rbt_left_rotate(rbt, dad); + + /* Both "current" node and its parent remain unchanged, but + * sibling is changed. + */ + sibling_idx = dad->right; + sibling = nd_vect + sibling_idx; + } + + ASSERT(sibling->color == RB_BLACK); + rb_node_t* slk = nd_vect + sibling->left; + rb_node_t* srk = nd_vect + sibling->right; + + if (slk->color == RB_BLACK && srk->color == RB_BLACK) { + /* case 2: sibling's both kids are in black. Set sibling's + * color to be red. + */ + sibling->color = RB_RED; + node_idx = dad - nd_vect; + } else { + if (srk->color == RB_BLACK) { + /* case 3: sibling's right kid is in black, while the left + * kid in in red. + */ + slk->color = RB_BLACK; + sibling->color = RB_RED; + rbt_right_rotate(rbt, sibling); + + sibling = slk; + sibling_idx = sibling - nd_vect; + } + + /* case 4: sibling's right kid is in red */ + rbt_left_rotate(rbt, dad); + + /* Now dad is still dad, sibling become grand-parent. Propagate + * dad's color to grandpar. + */ + sibling->color = dad->color; + + /* dad and new uncle are in black */ + dad->color = RB_BLACK; + nd_vect[sibling->right].color = RB_BLACK; + + break; + } + continue; + + } else { + int sibling_idx = dad->left; + rb_node_t* sibling = nd_vect + sibling_idx; + + /* case 1': sibling is in red color. Rotate around dad. */ + if (sibling->color == RB_RED) { + sibling->color = RB_BLACK; + dad->color = RB_RED; + rbt_right_rotate(rbt, dad); + + /* Both "current" node and its parent remain unchanged, but + * sibling is changed. + */ + sibling_idx = dad->left; + sibling = nd_vect + sibling_idx; + } + + ASSERT(sibling->color == RB_BLACK); + rb_node_t* slk = nd_vect + sibling->right; + rb_node_t* srk = nd_vect + sibling->left; + + if (slk->color == RB_BLACK && srk->color == RB_BLACK) { + /* case 2': sibling's both kids are in black. Set sibling's + * color to be red. + */ + sibling->color = RB_RED; + node_idx = dad - nd_vect; + } else { + if (srk->color == RB_BLACK) { + /* case 3': sibling's left kid is in black, while the right + * kid in in red. + */ + slk->color = RB_BLACK; + sibling->color = RB_RED; + rbt_left_rotate(rbt, sibling); + + sibling = slk; + sibling_idx = sibling - nd_vect; + } + + /* case 4': sibling's left kid is in red */ + rbt_right_rotate(rbt, dad); + + /* Now dad is still dad, sibling become grand-parent. Propagate + * dad's color to grandpar. + */ + sibling->color = dad->color; + + /* dad and new uncle are in black */ + dad->color = RB_BLACK; + nd_vect[sibling->left].color = RB_BLACK; + + break; + } + continue; + } + } + + nd_vect[node_idx].color = RB_BLACK; +} + +int +rbt_delete(rb_tree_t* rbt, int key) { + /* step 1: find the element to be deleted */ + int nd_idx = bst_search(rbt, key); + if (nd_idx == INVALID_IDX) + return 0; + + /* step 2: delete the element as we normally do with a binary-search tree */ + rb_node_t* nd_vect = rbt->tree; + rb_node_t* node = nd_vect + nd_idx; /* the node being deleted*/ + + int splice_out_idx; + if (node->left == SENTINEL_IDX || node->right == SENTINEL_IDX) { + splice_out_idx = nd_idx; + } else { + /* Get the successor of the node corrponding to nd_idx */ + rb_node_t* succ = nd_vect + node->right; + while (succ->left != SENTINEL_IDX) + succ = nd_vect + succ->left; + + splice_out_idx = succ - nd_vect; + } + + rb_node_t* splice_out = nd_vect + splice_out_idx; + + int so_kid_idx = (splice_out->left != SENTINEL_IDX) ? + splice_out->left : splice_out->right; + + rb_node_t* so_kid = nd_vect + so_kid_idx; + + if (splice_out->parent != INVALID_IDX) { + update_kid(nd_vect + splice_out->parent, splice_out_idx/*was*/, so_kid_idx); + } else { + ASSERT(rbt->root == splice_out_idx); + rbt->root = so_kid_idx; + } + so_kid->parent = splice_out->parent; + + if (splice_out_idx != nd_idx) { + nd_vect[nd_idx].key = splice_out->key; + nd_vect[nd_idx].value = splice_out->value; + } + + /* step 3: color fix up */ + if (splice_out->color == RB_BLACK) { + rbt_delete_fixup(rbt, so_kid_idx); + } + + /* step 4: Misc finialization. + * + * o. Migrate the last element to splice_out to avoid "hole" in + * the node vector. If we perform migration before step 3, care must + * be taken that the last element could be the "so_kid", in that case, + * the "so_kid" need to be re-evaluated. + * o. Misc update. + */ + if (splice_out_idx + 1 != rbt->node_num) { + int last_idx = rbt->node_num - 1; + rb_node_t* last = nd_vect + last_idx; + *splice_out = *last; + if (last->parent != INVALID_IDX) { + update_kid(nd_vect + last->parent, last_idx, splice_out_idx); + } else { + ASSERT(rbt->root == last_idx); + rbt->root = splice_out_idx; + } + nd_vect[last->left].parent = splice_out_idx; + nd_vect[last->right].parent = splice_out_idx; + } + + rbt->node_num--; + nd_vect[rbt->root].color = RB_BLACK; + return rbt_try_shrink(rbt); +} + +/***************************************************************************** + * + * Debugging and Testing Support + * + ***************************************************************************** + */ +#if defined(DEBUG) || defined(ENABLE_TESTING) +rb_tree_t* +rbt_create_manually(rb_valcolor_t* node_info, int len) { + rb_tree_t* rbt = rbt_create(); + if (!rbt) + return NULL; + + int i; + for (i = 0; i < len; i++) { + int nd_idx = bst_insert(rbt, node_info[i].key, node_info[i].value); + if (nd_idx == INVALID_IDX) { + rbt_destroy(rbt); + return NULL; + } + rbt->tree[nd_idx].color = node_info[i].color; + } + + return rbt; +} + +int +rbt_verify(rb_tree_t* rbt) { + /* step 1: Make sure the internal data structure are consistent.*/ + if (rbt->node_num > rbt->capacity) + return 0; + + if (rbt->tree == 0) + return 0; + + /* step 2: Make sure it is a tree */ + int node_num = rbt->node_num; + rb_node_t* nd_vect = rbt->tree; + + int* cnt = (int*)malloc(sizeof(int) * node_num); + int i; + for (i = 0; i < node_num; i++) cnt[i] = 0; + + for (i = SENTINEL_IDX + 1; i < node_num; i++) { + rb_node_t* nd = nd_vect + i; + int kid = nd->left; + if (kid < SENTINEL_IDX || kid >= node_num) + return 0; + cnt[kid]++; + + kid = nd->right; + if (kid < SENTINEL_IDX || kid >= node_num) + return 0; + cnt[kid]++; + + /* Take this opportunity to check if the color is either red or black.*/ + if (nd->color != RB_BLACK && nd->color != RB_RED) + return 0; + + /* make sure the "parent" pointer make sense */ + if (nd->parent != INVALID_IDX) { + rb_node_t* dad = nd_vect + nd->parent; + if (dad->left != i && dad->right != i) + return 0; + } + } + + int root_cnt = 0; + for (i = SENTINEL_IDX + 1; i < node_num; i++) { + int t = cnt[i]; + if (t > 1) { + /* It is DAG, not tree */ + return 0; + } else if (t == 0) { + root_cnt++; + if (i != rbt->root) { + /* we either have multiple roots, or rbt->root is not set + * properly. + */ + free(cnt); + return 0; + } + } + } + free(cnt); + cnt = 0; + if (root_cnt != 1) { + if (root_cnt == 0 && node_num != 1) + return 0; + } + + /* Following is to check if the RB-tree properies are preserved */ + + /* step 3: check if root and leaf are in black color */ + if (nd_vect[rbt->root].color != RB_BLACK || + nd_vect[SENTINEL_IDX].color != RB_BLACK) + return 0; + + /* step 4: Check if there are adjacent red nodes. */ + for (i = SENTINEL_IDX + 1; i < node_num; i++) { + rb_node_t* nd = nd_vect + i; + if (nd->color == RB_RED) { + if (nd_vect[nd->left].color == RB_RED || + nd_vect[nd->right].color == RB_RED) + return 0; + } + } + + /* step 5: check if all paths contain the same number of black nodes.*/ + int len = -1; + rb_node_t* sentinel = rbt->tree; + for (i = SENTINEL_IDX + 1; i < node_num; i++) { + rb_node_t* nd = nd_vect + i; + if (nd->left != SENTINEL_IDX || nd->right != SENTINEL_IDX) { + /* ignore internal node */ + continue; + } + + rb_node_t* cur = nd_vect + rbt->root; + int key = nd->key; + int this_len = 0; + while (cur != sentinel) { + if (cur->color == RB_BLACK) + this_len++; + + if (less_than(cur, key)) { + cur = nd_vect + cur->left; + } else if (greater_than(cur, key)) { + cur = nd_vect + cur->right; + } else { + break; + } + } + + if (len == -1) + len = this_len; + + if (len != this_len) + return 0; + } + + return 1; +} + +int +rbt_dump_dot(rb_tree_t* rbt, const char* name) { + FILE* f = fopen(name, "w"); + if (f == 0) + return 0; + + fprintf(f, "digraph G {\n"); + + int i, e = rbt->node_num; + rb_node_t* nd_vect = rbt->tree; + + for (i = SENTINEL_IDX + 1; i < e; i++) { + rb_node_t* node = nd_vect + i; + fprintf(f, "\t\%d [style=filled, color=%s, fontcolor=white];\n", + node->key, node->color == RB_RED ? "red" : "black"); + } + + for (i = SENTINEL_IDX + 1; i < e; i++) { + rb_node_t* node = nd_vect + i; + if (node->left != SENTINEL_IDX) + fprintf(f, "%d -> %d;\n", node->key, nd_vect[node->left].key); + + if (node->right != SENTINEL_IDX) { + fprintf(f, "%d -> %d [label=r];\n", node->key, + nd_vect[node->right].key); + } + } + + fprintf(f, "}"); + fclose(f); + + return 1; +} + +/* Print the rb tree directly to the console in plain text format*/ +void +rbt_dump_text(rb_tree_t* rbt) { + rb_node_t* nd_vect = rbt->tree; + fprintf(stdout, "RB tree: root id:%d, node_num:%d\n", + rbt->root, rbt->node_num); + + int i, e; + for (i = 0, e = rbt->node_num; i < e; i++) { + rb_node_t* node = nd_vect + i; + fprintf(stdout, + " Node:%d, key:%d, value:%ld, left:%d, right:%d, parent:%d\n", + i, node->key, node->value, node->left, node->right, + node->parent); + } + fprintf(stderr, "\n"); + fflush(stdout); +} + +#endif /*defined(DEBUG) || defined(ENABLE_TESTING)*/ diff --git a/rbtree.h b/rbtree.h new file mode 100644 index 0000000..eb48f2b --- /dev/null +++ b/rbtree.h @@ -0,0 +1,99 @@ +#ifndef _RBTREE_H_ +#define _RBTREE_H_ + +#ifdef __cplusplus +extern "C" { +#endif + +#include /* for intptr_t */ +typedef enum { + RB_BLACK = 0, + RB_RED = 1 +} rb_color_t; + +typedef struct { + intptr_t value; + int key; + int parent; + int left; + int right; + unsigned char color; /* either RB_BLACK or RB_RED */ +} rb_node_t; + +typedef struct { + rb_node_t* tree; + int root; + int node_num; + int capacity; +} rb_tree_t; + +/* Construt/destruct RB tree */ +rb_tree_t* rbt_create(void); +void rbt_destroy(rb_tree_t*); + +/* Alternative way to construct/destrut RB tree -- the instance is already + * allocated (say a global variable), but the content is not yet initialized. + */ +int rbt_init(rb_tree_t*); +void rbt_fini(rb_tree_t*); + +/* RB-tree operations */ +int rbt_insert(rb_tree_t*, int key, intptr_t val); +int rbt_delete(rb_tree_t*, int key); +int rbt_search(rb_tree_t*, int key, intptr_t* val); +int rbt_get_min(rb_tree_t*); +int rbt_get_max(rb_tree_t*); + +#define rbt_is_empty(rbt) ((rbt)->node_num == 1 ? 1 : 0) + +/* RB-tree iterator */ +typedef rb_node_t* rb_iter_t; +#define rbt_iter_begin(rbt) ((rbt)->tree + 1) +#define rbt_iter_end(rbt) ((rbt)->tree + (rbt)->node_num) +#define rbt_iter_inc(rbt, iter) ((iter) + 1) + +#define rbt_iter_deref(iter) ((iter)) + +#if defined(DEBUG) || defined(ENABLE_TESTING) +/* NOTE: If DEBUG is off and ENABLE_TESTING is on, functions enclosed by this + * directive should not be invoked by RB-tree opertaions. The ENABLE_TESTING + * can be used with release-build without degrading the performance. + */ + +/* Create a RB tree manually. The nodes' value, color as well as their + * insertion order are specified by "node_info". The color fix-up is disabled + * each time a node is inserted. It's up to the caller to ensure the final + * resulting tree conform to the RB-tree spec. + * + * The purpose of this function is to constuct a RB tree for testing + * purpose. + */ +typedef struct { + intptr_t value; + int key; + rb_color_t color; +} rb_valcolor_t; +rb_tree_t* rbt_create_manually(rb_valcolor_t* node_info, int len); + +int rbt_verify(rb_tree_t*); +/* dump in dotty format */ +int rbt_dump_dot(rb_tree_t*, const char* file_name); +void rbt_dump_text(rb_tree_t* rbt); + +#endif /*defined(DEBUG) || defined(ENABLE_TESTING)*/ + +#ifdef DEBUG +// Usage examples: ASSERT(a > b), ASSERT(foo() && "Opps, foo() reutrn 0"); +#define ASSERT(c) if (!(c))\ + { fprintf(stderr, "%s:%d Assert: %s\n", __FILE__, __LINE__, #c); abort(); } +#else + +#define ASSERT(c) ((void)0) + +#endif + +#ifdef __cplusplus +} +#endif + +#endif /*_RBTREE_H_*/ diff --git a/trunk.c b/trunk.c new file mode 100644 index 0000000..42c7bc5 --- /dev/null +++ b/trunk.c @@ -0,0 +1,107 @@ +#include +#include +#include +#include /* for bzero() */ +#include "lm_util.h" +#include "lm_internal.h" +#include "lj_mm.h" + +#define LM_ADDR_UPBOUND ((uint)0x80000000) +#define SIZE_1MB ((uint)0x100000) +#define SIZE_1GB ((uint)0x40000000) + +/* If we end up allocating a trunk no larger than this size, we might as well + * pull the plug, and give up for good. + */ +#define MEM_TOO_SMALL (SIZE_1MB * 8) + +static lm_trunk_t big_trunk; + +static int +get_malloc_mmap_max (void) { + /* Per mallopt(3), the default size of max mmap threshold is 65535 */ + size_t mmap_max = 65536; + + const char* envvar = getenv("MALLOC_MMAP_MAX_"); + if (envvar) + mmap_max = atoi(envvar); + + return mmap_max; +} + +lm_trunk_t* +lm_alloc_trunk (void) { + if (big_trunk.base) + return &big_trunk; + + void* cur_brk = sbrk(0); + uint avail = LM_ADDR_UPBOUND - ((intptr_t)cur_brk); + if (avail < MEM_TOO_SMALL) { + /* We can achieve almost nothing with 1MB, might as well bail out. */ + return 0; + } + + /* Disable malloc using mmap */ + int mmap_max = get_malloc_mmap_max (); + if (!mallopt(M_MMAP_MAX, 0)) + return 0; + + intptr_t trunk = (intptr_t)malloc (avail); + + /* Restore the original M_MMAP_MAX*/ + if (!mallopt (M_MMAP_MAX, mmap_max)) { + /* Hey, not all gotos are bad. */ + goto fail_ret; + } + + if (trunk > LM_ADDR_UPBOUND || + (LM_ADDR_UPBOUND - (intptr_t)trunk) <= MEM_TOO_SMALL) { + goto fail_ret; + } + + long page_sz = sysconf(_SC_PAGESIZE); + intptr_t usable_start = (trunk + page_sz - 1) & ~(page_sz - 1); + intptr_t usable_end = (trunk + avail) & ~(page_sz - 1); + if (usable_end > LM_ADDR_UPBOUND) + usable_end = LM_ADDR_UPBOUND; + + if (usable_end - usable_start <= MEM_TOO_SMALL) + goto fail_ret; + + big_trunk.base = (char*)trunk; + big_trunk.start = (char*)usable_start; + big_trunk.alloc_size = avail; + big_trunk.usable_size = usable_end - usable_start; + big_trunk.page_size = page_sz; + big_trunk.page_num = big_trunk.usable_size / page_sz; + ASSERT(big_trunk.page_num * big_trunk.page_size == big_trunk.usable_size); + + return &big_trunk; + +fail_ret: + free ((void*)trunk); + return 0; +} + +void +lm_free_trunk(void) { + ENTER_MUTEX + if (big_trunk.base) { + free(big_trunk.base); + bzero(&big_trunk, sizeof(big_trunk)); + } + + LEAVE_MUTEX +} + +#ifdef DEBUG +void +lm_dump_trunk (FILE* f) { + uint num_page = big_trunk.usable_size / sysconf(_SC_PAGESIZE); + + fprintf(f, "Base:%8p, alloc-size :%fG, usage-start:%8p, usage-size %fG (%u pages)\n", + big_trunk.base, big_trunk.alloc_size / ((float)SIZE_1GB), + big_trunk.start, big_trunk.usable_size / ((float)SIZE_1GB), + num_page); +} +#endif