#include "wyrm-huff.h"
#include <limits.h>
#include <assert.h>

<Data structures>
<Splay the tree for each transmitted character>
<Create and initialise the context data structure>
<Compression step for the next byte>
<Expansion step for the next byte>
<Compress an object representation>
<Expand an object representation>
<Scramble the compression tree>
<Tcl interface to the compressor>
			
		

#ifndef WYRM_HUFFMAN_H
#define WYRM_HUFFMAN_H

	#include "wyrmwif.h"
	
	<Opaque data structures>

	<Declare wyrm_huffCreate>
	<Declare wyrm_huffCompress>
	<Declare wyrm_huffCompress>
	<Declare wyrm_huffExpand>
	<Declare wyrm_huffCompressObj>
	<Declare wyrm_huffExpandObj>
	<Declare wyrm_huffPrime>
	<Declare wyrm_huffDelete>
	<Declare wyrm_huffCommandInit>

#endif

		

enum WyrmHuffConstants {
	wyrm_huffEOS = 256,
	wyrm_huffLongsize = 32,
	wyrm_huffMaxStates = 64
};

typedef struct WyrmHuffTree WyrmHuffTree,*pWyrmHuffTree;
typedef long WyrmHuffBuffer;

typedef WyrmHuffBuffer (*rWyrmHuffReceive)(ptr,bool *ofl);
typedef void (*rWyrmHuffTransmit)(ptr,WyrmHuffBuffer);


		

enum {
	maxchar = wyrm_huffEOS,
	succmax = maxchar+1,
	twicemax = 2*maxchar+1,
	root = 1
};

typedef short codetype; /* 0..maxchar */
typedef short upindex; /* 1..maxchar */
typedef short downindex; /* 1..twicemax */
typedef short stateindex; /* 0..wyrm_huffMaxStates-1 */

struct WyrmHuffTree {
	stateindex state;
	stateindex nstates;
	ptr context;
	rWyrmHuffReceive receive;
	rWyrmHuffTransmit transmit;
	WyrmHuffBuffer buffer;
	short bp;
	bool ofl;
	downindex left[wyrm_huffMaxStates][maxchar+1];
	downindex right[wyrm_huffMaxStates][maxchar+1];
	upindex up[wyrm_huffMaxStates][twicemax+1];
};

		

assert(sizeof(short)>sizeof(byte));
assert(SHRT_MAX>=twicemax);
assert(SHRT_MAX>=wyrm_huffMaxStates);
assert(sizeof(WyrmHuffBuffer)*CHAR_BIT==wyrm_huffLongsize);

		

static void splay(WyrmHuffTree *C,codetype plain) {
	downindex a,b;
	upindex c,d;
	upindex *up = C->up[C->state];
	downindex *left = C->left[C->state];
	downindex *right = C->right[C->state];

	a = plain+succmax;
	do {
		c = up[a];
		if (c!=root) {
			d = up[c]; b = left[d];
			if (c==b) {
				b = right[d]; right[d] = a;
			}else {
				left[d] = a;
			}
			if (a==left[c]) {
				left[c] = b;
			}else {
				right[c] = b;
			}
			up[a] = d; up[b] = c;
			a = d;
		}else {
			a = c;
		}
	} while (a!=root);
	C->state = (plain) & C->nstates;
}

		

pWyrmHuffTree wyrm_huffCreate(int nstates,ptr context,rWyrmHuffReceive receive,rWyrmHuffTransmit transmit);
		

pWyrmHuffTree wyrm_huffCreate(int nstates,ptr context,rWyrmHuffReceive receive,rWyrmHuffTransmit transmit) {
	stateindex s;
	downindex i;
	upindex j;
	pWyrmHuffTree C = heap(WyrmHuffTree);
	C->state = 0; C->context = context; C->receive = receive; C->transmit = transmit;
	C->buffer = C->bp = 0; C->ofl = 0;
	for (s=0; s<wyrm_huffMaxStates; s++) {
		for (i=2; i<=twicemax; i++) {C->up[s][i] = i/2;}
		for (j=1; j<=maxchar; j++) {C->left[s][j] = 2*j; C->right[s][j] = 2*j+1;}
	}
	for (s=1; s<wyrm_huffMaxStates && s<nstates; s+=s) ;
	C->nstates = s-1;
	return C;
}

		

#define wyrm_huffDelete(C) (Tcl_Free((chars)(C)))
		

void wyrm_huffPrime(pWyrmHuffTree C,chars);
		

void wyrm_huffPrime(pWyrmHuffTree C,chars prime) {
	while (*prime) splay(C,*prime++);
}

		

void wyrm_huffCompress(pWyrmHuffTree C,short plain);
		

void wyrm_huffCompress(WyrmHuffTree *C,short plain) {
	short sp;
	bool stack[maxchar+1];
	downindex a;
	upindex *up = C->up[C->state];
	downindex *right = C->right[C->state];

	if (plain<0) plain &= ((1<<CHAR_BIT)-1);
	a = plain+succmax;
	sp = 1;
	do {
		stack[sp] = right[up[a]]==a; sp++;
		a = up[a];
	} while (a!=root);
	do {
		sp--;
		if (!C->transmit) abort();
		C->buffer = (C->buffer<<1) | (stack[sp]&1);
		C->bp += 1;
		if (C->bp==wyrm_huffLongsize) {
			(C->transmit)(C->context,htonl(C->buffer));
			C->buffer = C->bp = 0;
		}
	} while (sp>1);
	splay(C,plain);
	if (plain==wyrm_huffEOS) {
		if (!C->transmit) abort();
		if (C->bp>0) {
			C->buffer <<= wyrm_huffLongsize-C->bp;
			(C->transmit)(C->context,htonl(C->buffer));
			C->buffer = C->bp = 0;
		}
	}
}

		

Obj wyrm_huffCompressObj(Obj x,int nstates,chars prime);
		

typedef struct {
	int m,n; bytes p;
} CompressContext;

static void transmitToObj(ptr o,WyrmHuffBuffer b) {
	CompressContext *context = o; WyrmHuffBuffer B = b;
	if (context->n+sizeof B>=context->m) {
		context->m = 2*(context->n+sizeof B); context->p = reheap(context->m,byte,context->p);
	}
	memcpy(context->p+context->n,&B,sizeof B); context->n += sizeof B;
}

Obj wyrm_huffCompressObj(Obj x,int nstates,chars prime) {
	Obj outp; CompressContext context;
	pWyrmHuffTree C = wyrm_huffCreate(nstates,&context,0,transmitToObj);
	int n; bytes s = Tcl_GetByteArrayFromObj(x,&n);
	context.m = context.n = 0; context.p = 0;
	if (prime) wyrm_huffPrime(C,prime);
	while (n-->0) wyrm_huffCompress(C,*s++);
	wyrm_huffCompress(C,wyrm_huffEOS);
	wyrm_huffDelete(C);
	outp = incr(Tcl_NewByteArrayObj(context.p,context.n));
	dispose(context.p);
	return outp;
}

		

short wyrm_huffExpand(pWyrmHuffTree C);
		

short wyrm_huffExpand(WyrmHuffTree *C) {
	downindex *left = C->left[C->state];
	downindex *right = C->right[C->state];
	downindex a;

	a = root;
	do {
		int receive;
		if (!C->receive) abort();
		if (C->bp==0) {
			C->buffer = ntohl((C->receive)(C->context,&C->ofl));
			C->bp = wyrm_huffLongsize;
			if (C->ofl) return wyrm_huffEOS;
		}
		C->bp -= 1;
		receive = (C->buffer>>C->bp) & 1;
		if (receive==0) {
			a = left[a];
		}else {
			a = right[a];
		}
	} while (a<=maxchar);
	a -= succmax;
	splay(C,a);
	return a;
}

		

Obj wyrm_huffExpandObj(Obj x,int nstates,chars prime);
		

typedef struct {Obj inp; int off;} ExpandContext;

static WyrmHuffBuffer receiveFromObj(ptr e,bool *ofl) {
	ExpandContext *E = e; int n; bytes s = Tcl_GetByteArrayFromObj(E->inp,&n);
	WyrmHuffBuffer B;
	n -= E->off; s += E->off; if (n>sizeof B) n = sizeof B; E->off += n;
	*ofl = n<=0;
	memset(&B,0,sizeof(B)); memcpy(&B,s,n); return B;
}

Obj wyrm_huffExpandObj(Obj x,int nstates,chars prime) {
	ExpandContext E; int m=0,n=0; bytes p=0; Obj obj; short c;
	pWyrmHuffTree C = wyrm_huffCreate(nstates,&E,receiveFromObj,0);
	if (prime) wyrm_huffPrime(C,prime);
	E.inp = x; E.off = 0;
	while ((c=wyrm_huffExpand(C))!=wyrm_huffEOS) {
		if (n+1>=m) {m = 2*(n+1); p = reheap(m,byte,p);}
		p[n++] = c;
	}
	wyrm_huffDelete(C);
	obj = incr(Tcl_NewByteArrayObj(p,n));
	dispose(p);
	return obj;
}

		

int wyrm_huffCommandInit(Intr intr);
		

static int wyrm_huffCommand(ClientData clientData,Intr intr,int N,Tcl_Obj *const P[]) {
	chars p,prime; Obj R; int nstates = 16;
	Tcl_ResetResult(intr);
	#ifdef TESTING
		if (streq(Tcl_GetStringFromObj(P[1],0),"splash")) {
			Tcl_SetResult(intr,splashgif,TCL_STATIC);
			return TCL_OK;
		}
	#endif
	if (N>=3 && streq(Tcl_GetStringFromObj(P[1],0),"-n")) {
		if (Tcl_GetIntFromObj(intr,P[2],&nstates)!=TCL_OK)
			return TCL_ERROR;
		P+=2; N-=2;
	}
	if (N!=3 && N!=4) {
		usage:
			Tcl_AppendResult(intr,"usage: huffman compress|expand <string> [<prime>]",0);
			return TCL_ERROR;
	}
	prime = N==4 ? Tcl_GetStringFromObj(P[3],0) : 0;
	p = Tcl_GetStringFromObj(P[1],0);
	if (stribegins("c",p))
		R = wyrm_huffCompressObj(P[2],nstates,prime);
	else if (stribegins("e",p))
		R = wyrm_huffExpandObj(P[2],nstates,prime);
	else {
		Tcl_AppendResult(intr,"expected 'compress' or 'expand': ",p,"\n",0); goto usage;
	}
	Tcl_SetObjResult(intr,R); decr(R); return TCL_OK;
}


int wyrm_huffCommandInit(Intr intr) {
	char package[] = "namespace eval ::wyrm {namespace export huffman}\n";
	<Assert types are large enough>
	Tcl_CreateObjCommand(intr,"::wyrm::huffman",wyrm_huffCommand,0,0);
	return Tcl_Eval(intr,package);
}