Bitwise Arithmetic Coder

/*
Copyright (C) 2013 by Brendan G Bohannon
Email: [email protected]
Copying: http://pastebin.com/iJxtZHm6
 */

/*

typedef struct BGBBTJ_ArithContext_s BGBBTJ_ArithContext;

struct BGBBTJ_ArithContext_s
{
byte *cs;   //current pos in bitstream (input)
byte *ct;   //current pos in bitstream (output)
byte *cse;  //current end pos in bitstream (input)
byte *cte;  //current end pos in bitstream (output)

u32 rmin;       //window lower range
u32 rmax;       //window upper range
u32 rval;       //window decode value

int wctx;
byte *model;    //pairs of counts
int ctxbits;    //context bits
int ctxmask;    //context mask
};

*/

/*
Bitwise Arithmetic Coder
 */

#include <bgbbtj.h>

#define BGBBTJ_BITARITH_LOWER       0x00000000
#define BGBBTJ_BITARITH_UPPER       0xFFFFFFFF
#define BGBBTJ_BITARITH_NEUTRAL     0x8000
#define BGBBTJ_BITARITH_WBITS       16
#define BGBBTJ_BITARITH_WMASK       ((1<<BGBBTJ_BITARITH_WBITS)-1)

#define BGBBTJ_BITARITH_DATABITS    8   //W: arithmetic coder bits/symbol
//#define BGBBTJ_BITARITH_CTXBITS       16
//#define BGBBTJ_BITARITH_CTXBITS       14
#define BGBBTJ_BITARITH_CTXBITS     13
// #define BGBBTJ_BITARITH_CTXBITS      12
// #define BGBBTJ_BITARITH_CTXBITS      10
// #define BGBBTJ_BITARITH_CTXBITS      8
#define BGBBTJ_BITARITH_CTXMASK     ((1<<BGBBTJ_BITARITH_CTXBITS)-1)

u16 bgbbtj_bitarith_divtab[65536];  //division table
u32 bgbbtj_bitarith_divtab2[65536]; //division table (32 bit)

void BGBBTJ_BitArith_Init()
{
    static int init=0;
    int i, j, k;

    if(init)return;
    init=1;

    for(i=0; i<256; i++)
        for(j=0; j<256; j++)
    {
        k=((i+1)<<16)/(i+j+2);
        bgbbtj_bitarith_divtab[i+(j<<8)]=k;
        bgbbtj_bitarith_divtab2[i+(j<<8)]=k<<16;
    }
}

byte *BGBBTJ_EmitVLI(byte *ct, int val)
{
    if(val<0)
    {
        *ct++=0;
        return(ct);
    }

    if(val<128)
    {
        *ct++=val;
        return(ct);
    }

    if(val<16384)
    {
        *ct++=0x80|(val>>8);
        *ct++=val&0xFF;
        return(ct);
    }

    if(val<2097152)
    {
        *ct++=0xC0|(val>>16);
        *ct++=(val>>8)&0xFF;
        *ct++=val&0xFF;
        return(ct);
    }

    if(val<268435456)
    {
        *ct++=0xE0|(val>>24);
        *ct++=(val>>16)&0xFF;
        *ct++=(val>>8)&0xFF;
        *ct++=val&0xFF;
        return(ct);
    }

    *ct++=0xF0;
    *ct++=(val>>24)&0xFF;
    *ct++=(val>>16)&0xFF;
    *ct++=(val>>8)&0xFF;
    *ct++=val&0xFF;
    return(ct);
}

byte *BGBBTJ_EmitSVLI(byte *ct, int val)
{
    return(BGBBTJ_EmitVLI(ct,
        (val>=0)?(val<<1):(((-val)<<1)-1)));
}

byte *BGBBTJ_DecodeVLI(byte *cs, int *rval)
{
    int i;

    if(!cs)return(NULL);

    i=*cs++;
    if(!(i&0x80))
    {
        *rval=i; return(cs);
    }else if((i&0xC0)==0x80)
    {
        i=((i&0x3F)<<8)|(*cs++);
        *rval=i; return(cs);
    }else if((i&0xE0)==0xC0)
    {
        i=((i&0x1F)<<8)|(*cs++);
        i=(i<<8)|(*cs++);
        *rval=i; return(cs);
    }else if((i&0xF0)==0xE0)
    {
        i=((i&0x0F)<<8)|(*cs++);
        i=(i<<8)|(*cs++);
        i=(i<<8)|(*cs++);
        *rval=i; return(cs);
    }else if((i&0xF8)==0xF0)
    {
        i=((i&0x0F)<<8)|(*cs++);
        i=(i<<8)|(*cs++);
        i=(i<<8)|(*cs++);
        i=(i<<8)|(*cs++);
        *rval=i; return(cs);
    }

//  *(int *)-1=-1;
    *rval=0;
//  return(cs);
    return(NULL);
}

byte *BGBBTJ_DecodeSVLI(byte *cs, int *rval)
{
    int i;

    cs=BGBBTJ_DecodeVLI(cs, &i);

    i=(i&1)?(-((i+1)>>1)):(i>>1);

    *rval=i;
    return(cs);
}

int BGBBTJ_BitArith_InputByte(BGBBTJ_ArithContext *ctx)
{
    if(ctx->cs>ctx->cse)
        return(0);

    return(*ctx->cs++);
}

void BGBBTJ_BitArith_OutputByte(BGBBTJ_ArithContext *ctx, int i)
{
    if(ctx->ct>ctx->cte)
        return;

    *ctx->ct++=i;
}

void BGBBTJ_BitArith_OutputBit(BGBBTJ_ArithContext *ctx,
    int i, u32 w)
{
    u32 r, r2, v;
    int j;

    r=ctx->rmax-ctx->rmin;
    v=ctx->rmin+((((u64)r)*((u64)(w<<(32-BGBBTJ_BITARITH_WBITS))))>>32);

    if(i)ctx->rmin=v+1;
    else ctx->rmax=v;

    //while bits
//  while((ctx->rmin&0xFF000000)==(ctx->rmax&0xFF000000))
    while(!((ctx->rmin^ctx->rmax)>>24))
    {
        BGBBTJ_BitArith_OutputByte(ctx, ctx->rmin>>24);
        ctx->rmin<<=8;
        ctx->rmax<<=8;
        ctx->rmax|=255;
    }
}

int BGBBTJ_BitArith_InputBit(BGBBTJ_ArithContext *ctx, u32 w)
{
    u32 r, r2, v, i;

    r=ctx->rmax-ctx->rmin;
    v=ctx->rmin+((((u64)r)*((u64)(w<<(32-BGBBTJ_BITARITH_WBITS))))>>32);

    i=(ctx->rval>v);
    if(i)ctx->rmin=v+1;
    else ctx->rmax=v;

    //while bits
//  while((ctx->rmin&0xFF000000)==(ctx->rmax&0xFF000000))
    while(!((ctx->rmin^ctx->rmax)>>24))
    {
        ctx->rmin<<=8;
        ctx->rmax<<=8;
        ctx->rmax|=255;

        ctx->rval<<=8;
        ctx->rval|=BGBBTJ_BitArith_InputByte(ctx);
    }

    return(i);
}

int BGBBTJ_BitArith_InputModelBit(BGBBTJ_ArithContext *ctx,
    byte *model, int wctx)
{
    u32 r, v, i, w;
    int wc;

    r=ctx->rmax-ctx->rmin;
    w=bgbbtj_bitarith_divtab2[((u16 *)model)[wctx]];
    v=ctx->rmin+(((u64)r*(u64)w)>>32);
    wc=wctx<<1;

    i=(ctx->rval>v);
    if(i)
    {
        ctx->rmin=v+1;

        if((++model[wc|1])&0x80)
        {
            model[wc|0]>>=1;
            model[wc|1]>>=1;
        }
    }else
    {
        ctx->rmax=v;

        if((++model[wc])&0x80)
        {
            model[wc|0]>>=1;
            model[wc|1]>>=1;
        }
    }

    //while bits
//  while((ctx->rmin&0xFF000000)==(ctx->rmax&0xFF000000))
//  while(!((ctx->rmin^ctx->rmax)&0xFF000000))
    while(!((ctx->rmin^ctx->rmax)>>24))
    {
        ctx->rmin<<=8;
        ctx->rmax<<=8;
        ctx->rmax|=255;

        ctx->rval<<=8;
        ctx->rval|=BGBBTJ_BitArith_InputByte(ctx);
    }

    return(i);
}

void BGBBTJ_BitArith_FlushWBits(BGBBTJ_ArithContext *ctx)
{
    while((ctx->rmin!=BGBBTJ_BITARITH_LOWER) ||
        (ctx->rmax!=BGBBTJ_BITARITH_UPPER))
    {
        BGBBTJ_BitArith_OutputByte(ctx, ctx->rmin>>24);

        ctx->rmin<<=8;
        ctx->rmax<<=8;
        ctx->rmax|=255;
    }
}


int BGBBTJ_BitArith_SetupEncode(BGBBTJ_ArithContext *ctx, byte *out, int sz)
{
    ctx->ct=out;
    ctx->cte=out+sz;
    ctx->rmin=BGBBTJ_BITARITH_LOWER;
    ctx->rmax=BGBBTJ_BITARITH_UPPER;
}

int BGBBTJ_BitArith_SetupDecode(BGBBTJ_ArithContext *ctx, byte *in, int sz)
{
    int i;

    ctx->cs=in;
    ctx->cse=in+sz;
    ctx->rmin=BGBBTJ_BITARITH_LOWER;
    ctx->rmax=BGBBTJ_BITARITH_UPPER;

    ctx->rval=BGBBTJ_BITARITH_LOWER;
    for(i=0; i<4; i++)
        ctx->rval=(ctx->rval<<8)|BGBBTJ_BitArith_InputByte(ctx);
}

//arithmetic frontend/modeler
#if 1
int BGBBTJ_BitArith_Predict(
    BGBBTJ_ArithContext *ctx, byte *model, int wctx)
{
    return(bgbbtj_bitarith_divtab[((u16 *)ctx->model)[wctx]]);

//  return(bgbbtj_bitarith_divtab[
//      model[(ctx<<1)|0]|
//      (model[(ctx<<1)|1]<<8)]);
}
#endif

#if 0
int BGBBTJ_BitArith_Predict(BGBBTJ_ArithContext *ctx, byte *model, int wctx)
{
    u32 i, j, k;

    wctx<<=1;
    i=model[wctx|0];
    j=model[wctx|1];

    k=((i+1)<<BGBBTJ_BITARITH_WBITS)/(i+j+2);

//  k=(k<WARITH2_CLAMP)?WARITH2_CLAMP:
//      ((k>(4096-WARITH2_CLAMP))?(4096-WARITH2_CLAMP):k);
//  k=BGBBTJ_BITARITH_NEUTRAL;

    return(k);
}
#endif

int BGBBTJ_BitArith_Update(
    BGBBTJ_ArithContext *ctx, byte *model, int wctx, int v)
{
    int i;

    i=(wctx<<1);
    ctx->model[i|v]++;
    if(ctx->model[i|v]&0x80)
    {
        ctx->model[i|0]>>=1;
        ctx->model[i|1]>>=1;
    }
}

void BGBBTJ_BitArith_EncodeSymbol(BGBBTJ_ArithContext *ctx, int v)
{
    int i, j, k;

    i=BGBBTJ_BITARITH_DATABITS;
    while(i--)
    {
        j=(v>>i)&1;

        k=BGBBTJ_BitArith_Predict(ctx, ctx->model, ctx->wctx);
        BGBBTJ_BitArith_Update(ctx, ctx->model, ctx->wctx, j);
//      k=BGBBTJ_BITARITH_NEUTRAL;

        BGBBTJ_BitArith_OutputBit(ctx, j, k);
//      ctx->wctx=((ctx->wctx<<1)|j)&BGBBTJ_BITARITH_CTXMASK;
        ctx->wctx=((ctx->wctx<<1)|j)&ctx->ctxmask;
    }
}

int BGBBTJ_BitArith_DecodeSymbol(BGBBTJ_ArithContext *ctx)
{
    int i, j, k, v;

//  v=0;
    i=BGBBTJ_BITARITH_DATABITS;
    while(i--)
    {
//      k=BGBBTJ_BitArith_Predict(ctx->model, ctx->wctx);
//      k=BGBBTJ_BITARITH_NEUTRAL;
//      j=BGBBTJ_BitArith_InputBit(k);
//      BGBBTJ_BitArith_Update(ctx->model, ctx->wctx, j);

//      ctx->wctx=v&((1<<WARITH2_CTXBITS)-1);
        j=BGBBTJ_BitArith_InputModelBit(ctx, ctx->model, ctx->wctx);

//      ctx->wctx=((ctx->wctx<<1)|j)&((1<<WARITH2_CTXBITS)-1);
//      ctx->wctx=((ctx->wctx<<1)|j)&BGBBTJ_BITARITH_CTXMASK;
        ctx->wctx=((ctx->wctx<<1)|j)&ctx->ctxmask;
    }

//  return(ctx->wctx&0xFF);
    return(ctx->wctx&((1<<BGBBTJ_BITARITH_DATABITS)-1));
//  return(v);
}

void BGBBTJ_BitArith_SetupContextBits(
    BGBBTJ_ArithContext *ctx, int bits)
{
    int i, j;

    if(ctx->model && (ctx->ctxbits==bits))
    {
        j=1<<bits;
        for(i=0; i<j*2; i++)ctx->model[i]=0;
        ctx->wctx=0;
        return;
    }

    if(ctx->model)
        { gcfree(ctx->model); }

    j=1<<bits;
    ctx->model=gcalloc(j*2*sizeof(byte));
    ctx->ctxbits=bits;
    ctx->ctxmask=(1<<bits)-1;
    ctx->wctx=0;
}

BGBBTJ_API int BGBBTJ_BitArith_EncodeDataCtx(
    BGBBTJ_ArithContext *ctx, byte *ibuf, int isz, byte *obuf, int osz)
{
    int i, j, k, l, ll;
    byte *cs, *ct, *cse, *cte, *s2;

    if(isz<=0)return(0);

    BGBBTJ_BitArith_Init();

#if 0
//  j=1<<WARITH2_CTXBITS;
    j=1<<BGBBTJ_BITARITH_CTXBITS;
//  ctx->model=malloc(j*2*sizeof(byte));
    for(i=0; i<j*2; i++)ctx->model[i]=0;
    ctx->wctx=0;
#endif

    BGBBTJ_BitArith_SetupContextBits(ctx,
        BGBBTJ_BITARITH_CTXBITS);

    ct=obuf; cte=obuf+osz;

    //Emit Header
    *ct++=0xBA;     //magic byte
    *ct++=0x00|(BGBBTJ_BITARITH_CTXBITS-12)&7;

    ct=BGBBTJ_EmitVLI(ct, isz);

    BGBBTJ_BitArith_SetupEncode(ctx, ct, cte-ct);
//  BGBBTJ_BitArith_SetupEncode(ctx, obuf, osz);

    cs=ibuf;
    cse=ibuf+isz;
    i=0;

    while(cs<cse)
    {
        j=*cs++;

//      if(i<16)printf("%02X ", j);

        BGBBTJ_BitArith_EncodeSymbol(ctx, j);
        i++;
    }
//  printf("\n");
//  BGBBTJ_BitArith_EncodeSymbol(ctx, 258);
    BGBBTJ_BitArith_FlushWBits(ctx);

#if 0
    for(i=0; i<16; i++)
        printf("%02X ", out[i]);
    printf("\n");
#endif

#if 0
//test decode
//  for(i=0; i<1024; i++)ctx->model[i]=0;
//  j=1<<WARITH2_CTXBITS;
    j=1<<BGBBTJ_BITARITH_CTXBITS;
    for(i=0; i<j*2; i++)ctx->model[i]=0;
    ctx->wctx=0;

    BGBBTJ_BitArith_SetupDecode(ctx, obuf, osz);

    cs=ibuf;
    ce=ibuf+isz;
    i=0;

    while(cs<ce)
    {
        j=*cs++;

        k=BGBBTJ_BitArith_DecodeSymbol(ctx);

        if(i<16)printf("%02X ", k);
        if(j!=k)
        {
            printf("break at %d\n", i);
            break;
        }
        i++;
    }

    printf("\n");
#endif

    i=ctx->ct-obuf;
    return(i);
}

BGBBTJ_API int BGBBTJ_BitArith_DecodeDataCtx(
    BGBBTJ_ArithContext *ctx,
    byte *ibuf, int isz, byte *obuf, int osz)
{
    int i, j, k, l, ll, osz1;
    byte *cs, *ct, *cse, *cte, *s2;

    if(isz<=0)return(0);

    BGBBTJ_BitArith_Init();

//  j=1<<BGBBTJ_BITARITH_CTXBITS;
//  for(i=0; i<j*2; i++)ctx->model[i]=0;
//  ctx->wctx=0;

    cs=ibuf; cse=ibuf+isz;
//  if((cs[0]!=0xBA)||(cs[1]!=0x00))
    if((cs[0]!=0xBA)||(cs[1]&0xF8))
        return(-1);

    BGBBTJ_BitArith_SetupContextBits(ctx, 12+(cs[1]&7));

    cs+=2;

    cs=BGBBTJ_DecodeVLI(cs, &osz1);
    if(osz1>osz)return(-1);

    BGBBTJ_BitArith_SetupDecode(ctx, cs, cse-cs);
//  BGBBTJ_BitArith_SetupDecode(ctx, ibuf, isz);

    ct=obuf;
//  cte=obuf+osz;
    cte=obuf+osz1;
    i=0;

    while(ct<cte)
    {
        j=BGBBTJ_BitArith_DecodeSymbol(ctx);
        *ct++=j;
    }

//  i=ctx->cs-ibuf;
    i=ct-obuf;
    return(i);
}

BGBBTJ_API int BGBBTJ_BitArith_EncodeData(
    byte *ibuf, int isz, byte *obuf, int osz)
{
    BGBBTJ_ArithContext *ctx;
    int i, j;

    ctx=gcalloc(sizeof(BGBBTJ_ArithContext));
    ctx->cs=NULL; ctx->ct=NULL;
//  j=1<<BGBBTJ_BITARITH_CTXBITS;
//  ctx->model=gcalloc(j*2*sizeof(byte));
    i=BGBBTJ_BitArith_EncodeDataCtx(ctx, ibuf, isz, obuf, osz);

    if(ctx->model)gcfree(ctx->model);
    gcfree(ctx);

    return(i);
}

BGBBTJ_API int BGBBTJ_BitArith_DecodeData(
    byte *ibuf, int isz, byte *obuf, int osz)
{
    BGBBTJ_ArithContext *ctx;
    int i, j;

    ctx=gcalloc(sizeof(BGBBTJ_ArithContext));
    ctx->cs=NULL; ctx->ct=NULL;
//  j=1<<BGBBTJ_BITARITH_CTXBITS;
//  ctx->model=gcalloc(j*2*sizeof(byte));
    i=BGBBTJ_BitArith_DecodeDataCtx(ctx, ibuf, isz, obuf, osz);

    if(ctx->model)gcfree(ctx->model);
    gcfree(ctx);

    return(i);
}

BGBBTJ_API int BGBBTJ_BitArith_EncodeTestData(
    byte *ibuf, int isz, byte *obuf, int osz)
{
    BGBBTJ_ArithContext *ctx;
    byte *tbuf;
    int i, j, k;

    ctx=gcalloc(sizeof(BGBBTJ_ArithContext));
    ctx->cs=NULL; ctx->ct=NULL;
//  j=1<<BGBBTJ_BITARITH_CTXBITS;
//  ctx->model=gcalloc(j*2*sizeof(byte));
    i=BGBBTJ_BitArith_EncodeDataCtx(ctx, ibuf, isz, obuf, osz);

    tbuf=malloc(2*isz);
    j=BGBBTJ_BitArith_DecodeDataCtx(ctx, obuf, i, tbuf, 2*isz);

    if(j!=isz)
    {
        printf("BGBBTJ_BitArith_EncodeTestData: Size Check %d->%d\n",
            isz, j);
    }

    for(k=0; k<j; k++)
    {
        if(ibuf[k]!=tbuf[k])
        {
            printf("BGBBTJ_BitArith_EncodeTestData: "
                "Decode Fail At %d\n", k);
            break;
        }
    }

    if(k>=j)
    {
        printf("BGBBTJ_BitArith_EncodeTestData: OK\n");
    }

    free(tbuf);

    if(ctx->model)gcfree(ctx->model);
    gcfree(ctx);

    return(i);
}