cloudflare
/
luajit-mm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

rudimentary implementation luajit mem-managment for taking full advantage of lower 2G memory

master
Shuxin Yang 2014-08-07 11:05:29 -07:00
parent 31fab6612f
commit 9e6507bd22
13 changed files with 2103 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

83
Makefile Normal file
View File

@ -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)

27
README.md
View File

@ -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.

53
demo.c Normal file
View File

@ -0,0 +1,53 @@
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#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;
}

41
lj_mm.h Normal file
View File

@ -0,0 +1,41 @@
#ifndef _LJ_MM_H_
#define _LJ_MM_H_
#include <stdlib.h> /* 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

34
lm_internal.h Normal file
View File

@ -0,0 +1,34 @@
#ifndef _LM_INTERNAL_H_
#define _LM_INTERNAL_H_
#include <stdio.h> /* 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

28
lm_util.h Normal file
View File

@ -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

0
obj/lib/empty_file Normal file
View File

0
obj/so/empty_file Normal file
View File

532
page_alloc.c Normal file
View File

@ -0,0 +1,532 @@
#include <stdint.h> /* for intptr_t */
#include <unistd.h>
#include <errno.h>
#include <sys/mman.h>
#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<<order) - 1)) == 0;
}
/**************************************************************************
*
* Buddy allocation stuff.
*
**************************************************************************
*/
/* Forward Decl */
static int add_block(blk_id_t block, int order);
/* We could have up to 1M pages (4G/4k). Hence 20 */
#define MAX_ORDER 20
#define INVALID_ORDER (-1)
typedef struct {
char* first_page; /* The starting address of the first page */
lm_page_t* page_info;
/* Free blocks of the same order are ordered by a RB tree. */
rb_tree_t free_blks[MAX_ORDER];
rb_tree_t alloc_blks;
int max_order;
int page_num; /* This many pages in total */
int page_size; /* The size of page in byte, normally 4k*/
int page_size_log2; /* log2(page_size)*/
} lm_alloc_t;
static lm_alloc_t* alloc_info;
static int log2_int32(unsigned num);
/* Initialize the page allocator, return 0 on success, 1 otherwise. */
int
lm_init_page_alloc(lm_trunk_t* trunk) {
if (!trunk) {
/* Trunk is not yet allocated */
return 0;
}
if (alloc_info) {
/* This function was succesfully invoked before */
return 1;
}
int page_num = trunk->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<<order) << alloc_info->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<<order);
if (buddy_id >= 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

304
rb_test.cxx Normal file
View File

@ -0,0 +1,304 @@
#include <stdio.h>
#include <string>
#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;
}

795
rbtree.c Normal file
View File

@ -0,0 +1,795 @@
#ifdef DEBUG
#include <stdio.h>
#endif
#include <stdlib.h>
#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)*/

99
rbtree.h Normal file
View File

@ -0,0 +1,99 @@
#ifndef _RBTREE_H_
#define _RBTREE_H_
#ifdef __cplusplus
extern "C" {
#endif
#include <stdint.h> /* 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_*/

107
trunk.c Normal file
View File

@ -0,0 +1,107 @@
#include <unistd.h>
#include <stdlib.h>
#include <malloc.h>
#include <strings.h> /* 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