btree.cpp

WEB上实现的局域网FTP搜索程序

#include <unistd.h>
#include <sys/mman.h>
#include "conf.h"

// #undef _POSIX_MAPPED_FILES

#ifdef _POSIX_MAPPED_FILES
#define MMAP_VERSION "using a mmap()ed temporary file"
#define INITMMAP() do_mmap_init()
#define ALLOC(x) do_mmap_alloc(x)
#define FREE(x,y) do_mmap_free(x,y)
#define MMAP_PAGESIZE (128*1024)
#else
#define MMAP_VERSION "using a malloc()ed temporary storage"
#define INITMMAP()
#define ALLOC(x) malloc(x)
#define FREE(x,y) free(x)
#endif

TripleList **TriIndex;

char *BTreeMax;
BTreeItem *Bhead=NULL;
TripleList *Thead=NULL,*Tptr=NULL,*Tptr2=NULL;

char genbuf[MAX];
char *ptr;

long triplecount;
long nodecount;
long diskoffset;
long refcount;

#ifdef _POSIX_MAPPED_FILES
off_t  mmap_offset;
int    mmap_fd;
void   *mmap_ptr;
size_t mmap_left;

void do_mmap_init(void)
{
    sprintf(genbuf,"%s", MMAPTMPFILE);
    unlink(genbuf);
    if ((mmap_fd = open(genbuf, O_RDWR|O_CREAT|O_TRUNC, 0644)) == -1)
    {
	perror("Cannot open mmap a file for mmap()ed access");
	exit(55);
    }
    unlink(genbuf);
    mmap_offset = 0;
    mmap_left = 0;
}

char dummy_mmap_block[MMAP_PAGESIZE];

void* do_mmap_alloc(size_t size)
{
    void *ptr;

    if (size > mmap_left)
    {
	if (write(mmap_fd, dummy_mmap_block, MMAP_PAGESIZE) < MMAP_PAGESIZE) 
	{
	    perror("Out of disk space in Btree");
	    exit(112);
	}
	mmap_ptr = mmap(NULL, MMAP_PAGESIZE, PROT_READ|PROT_WRITE, MAP_SHARED, mmap_fd, mmap_offset);
	if ((long)mmap_ptr == -1) 
	{
	    perror("mmap");
	    exit(111);
	}
	mmap_offset += MMAP_PAGESIZE;
	mmap_left = MMAP_PAGESIZE;
    }
    ptr = mmap_ptr;
    mmap_ptr = (void*)((long)mmap_ptr + size);
    mmap_left -= size;
    return ptr;
}

int do_mmap_free(void *ptr, size_t size)
{
    return 0;
}
#endif

void AnnounceMethod()
{
    printf("Btree - creating indexes %s\n", MMAP_VERSION);
    INITMMAP();
}

void addB(TripleList *Tptr, BTreeItem *item, BTreeItem **prev, char **max)
{
    int i;
    if(!item) 
    {
	nodecount++;
	item = (BTreeItem*)ALLOC(sizeof(BTreeItem));
	*prev = item;
	memset(item, 0, sizeof( BTreeItem));
	item->key = Tptr->key;
	item->lastfileno = (BTcnts)(-1);
	*max = item->key;
	return;
    }
    if (!strncmp(Tptr->key, item->key, KEYSIZE)) return;
    for (i = 0; i < B-1; i++)
    {
	if (!item->refs[i] || strncmp(Tptr->key, item->tops[i], KEYSIZE) <= 0)
	{
	    if (strncmp(Tptr->key, *max, KEYSIZE) > 0) 
		*max = Tptr->key;
	    addB(Tptr, item->refs[i], &(item->refs[i]), &(item->tops[i]));
	    break;
	}
    }
    if(i == B-1)
    {
	if (strncmp(Tptr->key, *max, KEYSIZE) > 0) 
	    *max = Tptr->key;
	addB(Tptr, item->refs[B-1], &(item->refs[B-1]), &(item->tops[B-1]));
    }
}

/* hi is not a valid index !! */
void CreateBtree(long lo, long hi)
{
    double dist,off;

    if (lo >= hi) return;
    dist = (double)(hi - lo) / B;
    off = lo + dist;
    while ((long)off < hi)
    {
	addB(TriIndex[(long)off], Bhead, &Bhead, &BTreeMax);
	off += dist;
    }
    CreateBtree(lo, (long)(lo+dist));
    off = lo + dist;
    while (1)
    {
	CreateBtree((long)(off+1), (long)(off+dist));
	off += dist;
	if ((long)(off+dist)>hi) break;
    }
    CreateBtree((long)(off+1), (long)hi);
}

void LoadTriples(void)
{
    long i;

    triplecount=0L;
    while(1)
    {
	fgets(genbuf, MAX, stdin);
	if (feof(stdin)) break;
	genbuf[KEYSIZE] = '\0';
	if ((ptr=strchr(genbuf,CR)) != NULL) *ptr = '\0';
	if ((ptr=strchr(genbuf,LF)) != NULL) *ptr = '\0';
	Tptr2 = new TripleList;
	Tptr2->key = strdup(genbuf);
	Tptr2->next = NULL;
	if (!Tptr) 
	    Thead = Tptr2;
	else 
	    Tptr->next = Tptr2; 
	Tptr = Tptr2;
	triplecount++;
    }
    Tptr = Thead;
    TriIndex = new (TripleList*)[triplecount];
    for (i = 0L; i < triplecount; i++)
    {
	TriIndex[i] = Tptr;
	Tptr = Tptr->next;
    }
}

struct BTreeItem *LocateBItem(char *key, BTreeItem *item)
{
    int i;

    if (!strncasecmp(key, item->key, KEYSIZE)) return item;

    for (i = 0; i < B-1; i++) 
    {
	if(!item->refs[i]) 
	    return NULL;
	else 
	    if(strncasecmp(key,item->tops[i],KEYSIZE)<=0) return LocateBItem(key, item->refs[i]);
    }
    if (!item->refs[B-1]) return NULL;
    else return LocateBItem(key, item->refs[B-1]);
}

void AddRef(BTreeItem *MyTriple, long ref)
{
    BTreeInfo *lastinfo = MyTriple->lastinfo;
    BTreeHost *lasthost = MyTriple->lasthost;

    MyTriple->infocnt++;
    if (!lastinfo || lastinfo->cnt == BTREENUMINFO)
    {
	lastinfo = (BTreeInfo*)ALLOC(sizeof(BTreeInfo));
	lastinfo->cnt = 0;
	lastinfo->next = NULL;
	if (MyTriple->lastinfo == NULL)
	    MyTriple->info = lastinfo;
	else
	    MyTriple->lastinfo->next = lastinfo;
	MyTriple->lastinfo = lastinfo;
    }
    lastinfo->refs[lastinfo->cnt] = ref;
    lastinfo->cnt++;

    lasthost->refs[lasthost->cnt-1].len++;
}

void AddHostRef(BTreeItem *MyTriple, int fileno)
{
    BTreeHost *lasthost = MyTriple->lasthost;

    if (lasthost) lasthost->cnt++;
    if (!lasthost || lasthost->cnt > BTREENUMHOST)
    {
	if (lasthost) lasthost->cnt--;
	lasthost = (BTreeHost*)ALLOC(sizeof(BTreeHost));
	lasthost->cnt = 1;
	lasthost->next = NULL;
	if (MyTriple->host == NULL)
	    MyTriple->host = lasthost;
	else
	    MyTriple->lasthost->next = lasthost;
	MyTriple->lasthost = lasthost;
    }
    lasthost->refs[lasthost->cnt-1].host = fileno;
    lasthost->refs[lasthost->cnt-1].offset = MyTriple->infocnt;
    lasthost->refs[lasthost->cnt-1].len = 0;
}

short NumHostRefs(BTreeHost *ptr)
{
    if (!ptr) return 0;
    return ptr->cnt + NumHostRefs(ptr->next);
}

void IndexHostOffset(BTreeItem *Bptr, long *offset)
{
    int i;

    Bptr->diskoffset = *offset;
    *offset += sizeof(DiskBTreeHead) +
	NumHostRefs(Bptr->host) * sizeof(HostRef);
    for (i=0; i<B; i++) 
    {
	if (Bptr->refs[i]) IndexHostOffset(Bptr->refs[i], offset);
    }
}

void WriteRefs(BTreeInfo *ptr, FILE *ref)
{
    BTreeInfo *next;
    while (ptr)
    {
	next = ptr->next;
	fwrite(ptr->refs, sizeof(long), ptr->cnt, ref);
	refcount += ptr->cnt;
	FREE(ptr, sizeof(BTreeInfo));
	ptr = next;
    }
}

void WriteHostRefs(BTreeHost *ptr, FILE *idx)
{
    BTreeHost *next;
    int i;
    while (ptr)
    {
	next = ptr->next;
	for (i = 0; i < ptr->cnt; i++) ptr->refs[i].offset += refcount;
	fwrite(ptr->refs, sizeof(HostRef), ptr->cnt, idx);
	FREE(ptr, sizeof(BTreeHost));
	ptr = next;
    }
}

void SaveBHostItem(BTreeItem *Bptr, FILE *idx, FILE *ref, DiskBTreeHead *head)
{
    int i;

    if(ftell(idx) != Bptr->diskoffset)
    {
	printf("ERROR !!! Offsets does not match for `%s' (%ld != %ld)\n",
	       Bptr->key, ftell(idx), Bptr->diskoffset);
	exit(110);
    }
    memcpy(head->key, Bptr->key, KEYSIZE);
    for(i=0; i<B; i++)
    {
	if(Bptr->tops[i]) 
	{
	    memcpy(head->treerefs[i].top, Bptr->tops[i], KEYSIZE);
	    head->treerefs[i].offset = Bptr->refs[i]->diskoffset;
	}
	else memcpy(head->treerefs[i].top, "\0\0\0", KEYSIZE);
    }
    head->numhosts = NumHostRefs(Bptr->host);
    head->numrefs = Bptr->infocnt;

    fwrite(head, sizeof(DiskBTreeHead), 1, idx);
    WriteHostRefs(Bptr->host, idx);
    WriteRefs(Bptr->info, ref);

    for(i=0; i<B; i++) 
    {
	if(Bptr->refs[i]) 
	    SaveBHostItem(Bptr->refs[i], idx, ref, head);
    }
    FREE(Bptr, sizeof(BTreeItem));
}

void SaveHostIndex(char *IndexName, char* RefName)
{
    FILE *idx, *ref;
    DiskBTreeHead head;

    refcount = 0;
    if (!(idx = fopen(IndexName,"w"))) 
    {
	perror("Cannot open Index");
	exit(109);
    }
    if (!(ref = fopen(RefName,"w"))) 
    {
	perror("Cannot open Ref");
	exit(109);
    }
    SaveBHostItem(Bhead, idx, ref, &head);
    fclose(idx);
    fclose(ref);
}