readjpg.c

这是一个同样来自贝尔实验室的和UNIX有着渊源的操作系统, 其简洁的设计和实现易于我们学习和理解

#include <u.h>
#include <libc.h>
#include <bio.h>
#include <draw.h>
#include "imagefile.h"

enum {
	/* Constants, all preceded by byte 0xFF */
	SOF	=0xC0,	/* Start of Frame */
	SOF2=0xC2,	/* Start of Frame; progressive Huffman */
	JPG	=0xC8,	/* Reserved for JPEG extensions */
	DHT	=0xC4,	/* Define Huffman Tables */
	DAC	=0xCC,	/* Arithmetic coding conditioning */
	RST	=0xD0,	/* Restart interval termination */
	RST7	=0xD7,	/* Restart interval termination (highest value) */
	SOI	=0xD8,	/* Start of Image */
	EOI	=0xD9,	/* End of Image */
	SOS	=0xDA,	/* Start of Scan */
	DQT	=0xDB,	/* Define quantization tables */
	DNL	=0xDC,	/* Define number of lines */
	DRI	=0xDD,	/* Define restart interval */
	DHP	=0xDE,	/* Define hierarchical progression */
	EXP	=0xDF,	/* Expand reference components */
	APPn	=0xE0,	/* Reserved for application segments */
	JPGn	=0xF0,	/* Reserved for JPEG extensions */
	COM	=0xFE,	/* Comment */

	CLAMPOFF	= 300,
	NCLAMP		= CLAMPOFF+700
};

typedef struct Framecomp Framecomp;
typedef struct Header Header;
typedef struct Huffman Huffman;

struct Framecomp	/* Frame component specifier from SOF marker */
{
	int	C;
	int	H;
	int	V;
	int	Tq;
};

struct Huffman
{
	int	*size;	/* malloc'ed */
	int	*code;	/* malloc'ed */
	int	*val;		/* malloc'ed */
	int	mincode[17];
	int	maxcode[17];
	int	valptr[17];
	/* fast lookup */
	int	value[256];
	int	shift[256];
};


struct Header
{
	Biobuf	*fd;
	char		err[256];
	jmp_buf	errlab;
	/* variables in i/o routines */
	int		sr;	/* shift register, right aligned */
	int		cnt;	/* # bits in right part of sr */
	uchar	*buf;
	int		nbuf;
	int		peek;

	int		Nf;

	Framecomp	comp[3];
	uchar	mode;
	int		X;
	int		Y;
	int		qt[4][64];		/* quantization tables */
	Huffman	dcht[4];
	Huffman	acht[4];
	int		**data[3];
	int		ndata[3];
	
	uchar	*sf;	/* start of frame; do better later */
	uchar	*ss;	/* start of scan; do better later */
	int		ri;	/* restart interval */

	/* progressive scan */
	Rawimage *image;
	Rawimage **array;
	int		*dccoeff[3];
	int		**accoeff[3];	/* only need 8 bits plus quantization */
	int		naccoeff[3];
	int		nblock[3];
	int		nacross;
	int		ndown;
	int		Hmax;
	int		Vmax;
};

static	uchar	clamp[NCLAMP];

static	Rawimage	*readslave(Header*, int);
static	int			readsegment(Header*, int*);
static	void			quanttables(Header*, uchar*, int);
static	void			huffmantables(Header*, uchar*, int);
static	void			soiheader(Header*);
static	int			nextbyte(Header*, int);
static	int			int2(uchar*, int);
static	void			nibbles(int, int*, int*);
static	int			receive(Header*, int);
static	int			receiveEOB(Header*, int);
static	int			receivebit(Header*);
static	void			restart(Header*, int);
static	int			decode(Header*, Huffman*);
static	Rawimage*	baselinescan(Header*, int);
static	void			progressivescan(Header*, int);
static	Rawimage*	progressiveIDCT(Header*, int);
static	void			idct(int*);
static	void			colormap1(Header*, int, Rawimage*, int*, int, int);
static	void			colormapall1(Header*, int, Rawimage*, int*, int*, int*, int, int);
static	void			colormap(Header*, int, Rawimage*, int**, int**, int**, int, int, int, int, int*, int*);
static	void			jpgerror(Header*, char*, ...);

static	char		readerr[] = "ReadJPG: read error: %r";
static	char		memerr[] = "ReadJPG: malloc failed: %r";

static	int zig[64] = {
	0, 1, 8, 16, 9, 2, 3, 10, 17, /* 0-7 */
	24, 32, 25, 18, 11, 4, 5, /* 8-15 */
	12, 19, 26, 33, 40, 48, 41, 34, /* 16-23 */
	27, 20, 13, 6, 7, 14, 21, 28, /* 24-31 */
	35, 42, 49, 56, 57, 50, 43, 36, /* 32-39 */
	29, 22, 15, 23, 30, 37, 44, 51, /* 40-47 */
	58, 59, 52, 45, 38, 31, 39, 46, /* 48-55 */
	53, 60, 61, 54, 47, 55, 62, 63 /* 56-63 */
};

static
void
jpginit(void)
{
	int k;
	static int inited;

	if(inited)
		return;
	inited = 1;
	for(k=0; k<CLAMPOFF; k++)
		clamp[k] = 0;
	for(; k<CLAMPOFF+256; k++)
		clamp[k] = k-CLAMPOFF;
	for(; k<NCLAMP; k++)
		clamp[k] = 255;
}

static
void*
jpgmalloc(Header *h, int n, int clear)
{
	void *p;

	p = malloc(n);
	if(p == nil)
		jpgerror(h, memerr);
	if(clear)
		memset(p, 0, n);
	return p;
}

static
void
clear(void **p)
{
	if(*p){
		free(*p);
		*p = nil;
	}
}

static
void
jpgfreeall(Header *h, int freeimage)
{
	int i, j;

	clear(&h->buf);
	if(h->dccoeff[0])
		for(i=0; i<3; i++)
			clear(&h->dccoeff[i]);
	if(h->accoeff[0])
		for(i=0; i<3; i++){
			if(h->accoeff[i])
				for(j=0; j<h->naccoeff[i]; j++)
					clear(&h->accoeff[i][j]);
			clear(&h->accoeff[i]);
		}
	for(i=0; i<4; i++){
		clear(&h->dcht[i].size);
		clear(&h->acht[i].size);
		clear(&h->dcht[i].code);
		clear(&h->acht[i].code);
		clear(&h->dcht[i].val);
		clear(&h->acht[i].val);
	}
	if(h->data[0])
		for(i=0; i<3; i++){
			if(h->data[i])
				for(j=0; j<h->ndata[i]; j++)
					clear(&h->data[i][j]);
			clear(&h->data[i]);
		}
	if(freeimage && h->image!=nil){
		clear(&h->array);
		clear(&h->image->cmap);
		for(i=0; i<3; i++)
			clear(&h->image->chans[i]);
		clear(&h->image);
	}
}

static
void
jpgerror(Header *h, char *fmt, ...)
{
	va_list arg;

	va_start(arg, fmt);
	vseprint(h->err, h->err+sizeof h->err, fmt, arg);
	va_end(arg);

	werrstr(h->err);
	jpgfreeall(h, 1);
	longjmp(h->errlab, 1);
}

Rawimage**
Breadjpg(Biobuf *b, int colorspace)
{
	Rawimage *r, **array;
	Header *h;
	char buf[ERRMAX];

	buf[0] = '\0';
	if(colorspace!=CYCbCr && colorspace!=CRGB){
		errstr(buf, sizeof buf);	/* throw it away */
		werrstr("ReadJPG: unknown color space");
		return nil;
	}
	jpginit();
	h = malloc(sizeof(Header));
	array = malloc(sizeof(Header));
	if(h==nil || array==nil){
		free(h);
		free(array);
		return nil;
	}
	h->array = array;
	memset(h, 0, sizeof(Header));
	h->fd = b;
	errstr(buf, sizeof buf);	/* throw it away */
	if(setjmp(h->errlab))
		r = nil;
	else
		r = readslave(h, colorspace);
	jpgfreeall(h, 0);
	free(h);
	array[0] = r;
	array[1] = nil;
	return array;
}

Rawimage**
readjpg(int fd, int colorspace)
{
	Rawimage** a;
	Biobuf b;

	if(Binit(&b, fd, OREAD) < 0)
		return nil;
	a = Breadjpg(&b, colorspace);
	Bterm(&b);
	return a;
}

static
Rawimage*
readslave(Header *header, int colorspace)
{
	Rawimage *image;
	int nseg, i, H, V, m, n;
	uchar *b;

	soiheader(header);
	nseg = 0;
	image = nil;

	header->buf = jpgmalloc(header, 4096, 0);
	header->nbuf = 4096;
	while(header->err[0] == '\0'){
		nseg++;
		n = readsegment(header, &m);
		b = header->buf;
		switch(m){
		case -1:
			return image;

		case APPn+0:
			if(nseg==1 && strncmp((char*)b, "JFIF", 4)==0)  /* JFIF header; check version */
				if(b[5]>1 || b[6]>2)
					sprint(header->err, "ReadJPG: can't handle JFIF version %d.%2d", b[5], b[6]);
			break;

		case APPn+1: case APPn+2: case APPn+3: case APPn+4: case APPn+5:
		case APPn+6: case APPn+7: case APPn+8: case APPn+9: case APPn+10:
		case APPn+11: case APPn+12: case APPn+13: case APPn+14: case APPn+15:
			break;

		case DQT:
			quanttables(header, b, n);
			break;

		case SOF:
		case SOF2:
			header->Y = int2(b, 1);
			header->X = int2(b, 3);
			header->Nf =b[5];
			for(i=0; i<header->Nf; i++){
				header->comp[i].C = b[6+3*i+0];
				nibbles(b[6+3*i+1], &H, &V);
				if(H<=0 || V<=0)
					jpgerror(header, "non-positive sampling factor (Hsamp or Vsamp)");
				header->comp[i].H = H;
				header->comp[i].V = V;
				header->comp[i].Tq = b[6+3*i+2];
			}
			header->mode = m;
			header->sf = b;
			break;

		case  SOS:
			header->ss = b;
			switch(header->mode){
			case SOF:
				image = baselinescan(header, colorspace);
				break;
			case SOF2:
				progressivescan(header, colorspace);
				break;
			default:
				sprint(header->err, "unrecognized or unspecified encoding %d", header->mode);
				break;
			}
			break;

		case  DHT:
			huffmantables(header, b, n);
			break;

		case  DRI:
			header->ri = int2(b, 0);
			break;

		case  COM:
			break;

		case EOI:
			if(header->mode == SOF2)
				image = progressiveIDCT(header, colorspace);
			return image;

		default:
			sprint(header->err, "ReadJPG: unknown marker %.2x", m);
			break;
		}
	}
	return image;
}

/* readsegment is called after reading scan, which can have */
/* read ahead a byte.  so we must check peek here */
static
int
readbyte(Header *h)
{
	uchar x;

	if(h->peek >= 0){
		x = h->peek;
		h->peek = -1;
	}else if(Bread(h->fd, &x, 1) != 1)
		jpgerror(h, readerr);
	return x;
}

static
int
marker(Header *h)
{
	int c;

	while((c=readbyte(h)) == 0)
		fprint(2, "ReadJPG: skipping zero byte at offset %lld\n", Boffset(h->fd));
	if(c != 0xFF)
		jpgerror(h, "ReadJPG: expecting marker; found 0x%x at offset %lld\n", c, Boffset(h->fd));
	while(c == 0xFF)
		c = readbyte(h);
	return c;
}

static
int
int2(uchar *buf, int n)
{
	return (buf[n]<<8) + buf[n+1];
}

static
void
nibbles(int b, int *p0, int *p1)
{
	*p0 = (b>>4) & 0xF;
	*p1 = b & 0xF;
}

static
void
soiheader(Header *h)
{
	h->peek = -1;
	if(marker(h) != SOI)
		jpgerror(h, "ReadJPG: unrecognized marker in header");
	h->err[0] = '\0';
	h->mode = 0;
	h->ri = 0;
}

static
int
readsegment(Header *h, int *markerp)
{
	int m, n;
	uchar tmp[2];

	m = marker(h);
	switch(m){
	case EOI:
		*markerp = m;
		return 0;
	case 0:
		jpgerror(h, "ReadJPG: expecting marker; saw %.2x at offset %lld", m, Boffset(h->fd));
	}
	if(Bread(h->fd, tmp, 2) != 2)
    Readerr:
		jpgerror(h, readerr);
	n = int2(tmp, 0);
	if(n < 2)
		goto Readerr;
	n -= 2;
	if(n > h->nbuf){
		free(h->buf);
		h->buf = jpgmalloc(h, n+1, 0); /* +1 for sentinel */
		h->nbuf = n;
	}
	if(Bread(h->fd, h->buf, n) != n)
		goto Readerr;
	*markerp = m;
	return n;
}

static
int
huffmantable(Header *h, uchar *b)
{
	Huffman *t;
	int Tc, th, n, nsize, i, j, k, v, cnt, code, si, sr, m;
	int *maxcode;

	nibbles(b[0], &Tc, &th);
	if(Tc > 1)
		jpgerror(h, "ReadJPG: unknown Huffman table class %d", Tc);
	if(th>3 || (h->mode==SOF && th>1))
		jpgerror(h, "ReadJPG: unknown Huffman table index %d", th);
	if(Tc == 0)
		t = &h->dcht[th];
	else
		t = &h->acht[th];

	/* flow chart C-2 */
	nsize = 0;
	for(i=0; i<16; i++)
		nsize += b[1+i];
	t->size = jpgmalloc(h, (nsize+1)*sizeof(int), 1);
	k = 0;
	for(i=1; i<=16; i++){
		n = b[i];
		for(j=0; j<n; j++)
			t->size[k++] = i;
	}
	t->size[k] = 0;

	/* initialize HUFFVAL */
	t->val = jpgmalloc(h, nsize*sizeof(int), 1);
	for(i=0; i<nsize; i++)
		t->val[i] = b[17+i];

	/* flow chart C-3 */
	t->code = jpgmalloc(h, (nsize+1)*sizeof(int), 1);
	k = 0;
	code = 0;
	si = t->size[0];
	for(;;){
		do
			t->code[k++] = code++;
		while(t->size[k] == si);
		if(t->size[k] == 0)
			break;
		do{
			code <<= 1;
			si++;
		}while(t->size[k] != si);
	}

	/* flow chart F-25 */
	i = 0;
	j = 0;
	for(;;){
		for(;;){
			i++;
			if(i > 16)
				goto outF25;
			if(b[i] != 0)
				break;
			t->maxcode[i] = -1;
		}
		t->valptr[i] = j;
		t->mincode[i] = t->code[j];
		j += b[i]-1;
		t->maxcode[i] = t->code[j];
		j++;
	}
outF25:

	/* create byte-indexed fast path tables */
	maxcode = t->maxcode;
	/* stupid startup algorithm: just run machine for each byte value */
	for(v=0; v<256; ){
		cnt = 7;
		m = 1<<7;
		code = 0;
		sr = v;
		i = 1;
		for(;;i++){
			if(sr & m)
				code |= 1;
			if(code <= maxcode[i])