Playing with arith

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <limits.h>
#include <assert.h>

struct bitstream
{
   uint8_t *data;
   size_t read_count;
   size_t write_count;
   size_t cap;
};

static void bitstream_set(struct bitstream *stream, size_t index, uint8_t bit)
{
   bit &= 0x01;
   unsigned bitshift = 7 - (index % 8);
   uint8_t bitmask = ~(1 << bitshift);
   stream->data[index / 8] &= bitmask;
   stream->data[index / 8] |= (bit << bitshift);
}

static bool bitstream_equal(const struct bitstream *a, const struct bitstream *b)
{
   if (a->write_count != b->write_count)
      return false;

   if (memcmp(a->data, b->data, a->write_count / 8) != 0)
      return false;

   if (a->write_count % 8)
   {
      uint8_t mask = 0xff << (8 - (a->write_count % 8));
      if ((a->data[a->write_count / 8] ^ b->data[b->write_count / 8]) & mask)
         return false;
   }

   return true;
}

static void bitstream_write(struct bitstream *stream, uint8_t bit)
{
   if (stream->write_count >= CHAR_BIT * stream->cap)
   {
      stream->cap = stream->cap * 2 + 1;
      stream->data = realloc(stream->data, stream->cap);
   }

   bitstream_set(stream, stream->write_count++, bit);
}

static void bitstream_write_byte(struct bitstream *stream, uint8_t byte)
{
   for (int shift = 7; shift >= 0; shift--)
      bitstream_write(stream, byte >> shift);
}

static uint8_t bitstream_get(const struct bitstream *stream, size_t index)
{
   unsigned bitshift = 7 - (index % 8);
   return (stream->data[index / 8] >> bitshift) & 0x01;
}

static bool bitstream_read(struct bitstream *stream, uint8_t *bit)
{
   if (stream->read_count >= stream->write_count)
      return false;

   *bit = bitstream_get(stream, stream->read_count++);
   return true;
}

static size_t bitstream_size(const struct bitstream *stream)
{
   return stream->write_count;
}

static void bitstream_free(struct bitstream *stream)
{
   free(stream->data);
   memset(stream, 0, sizeof(*stream));
}

static void bitstream_print(const struct bitstream *stream)
{
   for (size_t i = 0; i < stream->write_count; i++)
      printf("%u", bitstream_get(stream, i));
   printf("\n");
}

#define FIXED_BITS 30
#define ONE (1 << FIXED_BITS)
#define HALF (ONE >> 1)
#define QUANT_BITS 3
#define QUANT_BITS_ROUND (QUANT_BITS - 1)
#define PROB_QUANT(q) ((((prob_t)(q) + (1 << QUANT_BITS_ROUND)) >> QUANT_BITS) << QUANT_BITS)
#define PROB_MUL(a, b) ((prob_t)(((prob_high_t)(a) * (prob_high_t)(b) + HALF) >> FIXED_BITS))
#define PROB_TEST_FACTOR UINT64_C(15)
#define PROB_QUANT_HARD_BITS 8
#define PROB_QUANT_HARD_BITS_ROUND (PROB_QUANT_HARD_BITS - 1)
#define PROB_QUANT_HARD(q) ((((prob_t)(q) + (1 << PROB_QUANT_HARD_BITS_ROUND)) >> PROB_QUANT_HARD_BITS) << PROB_QUANT_HARD_BITS)
typedef uint64_t prob_high_t;
typedef uint32_t prob_t;

static inline prob_t readjust_prob_mid(prob_t prob_zero, prob_t prob_end, prob_t prob_bit)
{
   assert(prob_end - prob_zero >= (2 << QUANT_BITS));
   prob_t prob_mid = prob_zero + PROB_MUL(prob_end - prob_zero, prob_bit);
   prob_mid = PROB_QUANT(prob_mid);

   if (prob_mid == prob_zero)
      prob_mid += 1 << QUANT_BITS;
   else if (prob_mid == prob_end)
      prob_mid += 1 << QUANT_BITS;

   //printf("Prob adjustment [%u, %u, %u]\n", prob_zero, prob_mid, prob_end);
   return prob_mid;
}

static void arith_decompress(struct bitstream *output, struct bitstream *input, size_t size)
{
   prob_t prob_bit = PROB_QUANT_HARD(ONE * PROB_TEST_FACTOR / 256);

   prob_t prob_zero = 0;
   prob_t prob_mid = prob_bit;
   prob_t prob_end = ONE;

   uint8_t bit = 0;
   prob_t bit_start = 0;
   prob_t bit_end = ONE;
   while (bitstream_read(input, &bit))
   {
      if (bit)
         bit_start += (bit_end - bit_start) >> 1;
      else
         bit_end -= (bit_end - bit_start) >> 1;
      //printf("bit range: [%u, %u], prob range: [%u, %u, %u]\n", bit_start, bit_end, prob_zero, prob_mid, prob_end);

      while ((bit_start >= prob_zero && bit_end <= prob_mid) ||
            (bit_start >= prob_mid && bit_end <= prob_end))
      {
         if (bitstream_size(output) >= size)
            return;

         if (bit_start >= prob_mid)
         {
            bitstream_write(output, 1);
            prob_zero = PROB_QUANT(prob_mid);
         }
         else
         {
            bitstream_write(output, 0);
            prob_end = PROB_QUANT(prob_mid);
         }

         prob_mid = readjust_prob_mid(prob_zero, prob_end, prob_bit);

         while (prob_zero >= HALF || prob_end <= HALF)
         {
            if (prob_zero >= HALF)
            {
               prob_zero = PROB_QUANT((prob_zero << 1) - ONE);
               prob_end = PROB_QUANT((prob_end << 1) - ONE);
               prob_mid = readjust_prob_mid(prob_zero, prob_end, prob_bit);
               bit_start = (bit_start << 1) - ONE;
               bit_end = (bit_end << 1) - ONE;
            }
            else
            {
               prob_zero = PROB_QUANT(prob_zero << 1);
               prob_end  = PROB_QUANT(prob_end << 1);
               prob_mid = readjust_prob_mid(prob_zero, prob_end, prob_bit);
               bit_start = bit_start << 1;
               bit_end = bit_end << 1;
            }
         }
      }
   }
}

static void arith_compress(struct bitstream *output, struct bitstream *input)
{
   prob_t prob_bit = PROB_QUANT_HARD(ONE * PROB_TEST_FACTOR / 256);

   // Current probability range.
   prob_t prob_zero = 0;
   prob_t prob_mid  = prob_bit;
   prob_t prob_end  = ONE;

   uint8_t bit = 0;
   while (bitstream_read(input, &bit))
   {
      if (bit)
         prob_zero = PROB_QUANT(prob_mid);
      else
         prob_end = PROB_QUANT(prob_mid);
      prob_mid = readjust_prob_mid(prob_zero, prob_end, prob_bit);

      prob_zero = PROB_QUANT(prob_zero);
      prob_end = PROB_QUANT(prob_end);

      // Check if we can encode bits. We can encode a new bit if the probability range falls perfectly inside a half of the bit range.
      // If so, normalization is performed as well. Prob range is shifted so that bit_range is [0, 1] again.
      while (prob_zero >= HALF || prob_end <= HALF)
      {
         if (prob_zero >= HALF)
         {
            bitstream_write(output, 1);
            //printf("Encoded 1.\n");
            prob_zero = PROB_QUANT((prob_zero << 1) - ONE);
            prob_end  = PROB_QUANT((prob_end << 1) - ONE);
         }
         else
         {
            bitstream_write(output, 0);
            //printf("Encoded 0.\n");
            prob_zero = PROB_QUANT(prob_zero << 1);
            prob_end  = PROB_QUANT(prob_end << 1);
         }

         prob_mid = readjust_prob_mid(prob_zero, prob_end, prob_bit);
      }
   }

   // Right now, the encoded range is outside probability range.
   // Make sure that encoded range is inside probability range so we can terminate the encoding.
   prob_t bit_start = 0;
   prob_t bit_end = ONE;
   while (bit_start < prob_zero || bit_end > prob_end)
   {
      //printf("bit: (%u, %u), prob: (%u, %u).\n", bit_start, bit_end, prob_zero, prob_end);
      prob_t center = (prob_zero + prob_end) >> 1;
      int dist_start  = center - bit_start;
      int dist_end    = bit_end - center;
      if (dist_end < 0)
      {
         bit_start += (bit_end - bit_start) >> 1;
         //printf("Write 1 #1.\n");
         bitstream_write(output, 1);
      }
      else if (dist_start < 0)
      {
         bit_end -= (bit_end - bit_start) >> 1;
         //printf("Write 0 #1.\n");
         bitstream_write(output, 0);
      }
      else if (dist_start <= dist_end)
      {
         bit_end -= (bit_end - bit_start) >> 1;
         //printf("Write 0 #2.\n");
         bitstream_write(output, 0);
      }
      else
      {
         bit_start += (bit_end - bit_start) >> 1;
         //printf("Write 1 #2.\n");
         bitstream_write(output, 1);
      }
   }
}

int main(void)
{
   struct bitstream stream = {0};
   struct bitstream out_stream = {0};
   struct bitstream out_decode = {0};
   for (unsigned i = 0; i < 10000000; i++)
      bitstream_write_byte(&stream, 0x10 + i);

   //printf("Input:\n");
   //bitstream_print(&stream);

   arith_compress(&out_stream, &stream);
   //printf("Encoded:\n");
   //bitstream_print(&out_stream);

   arith_decompress(&out_decode, &out_stream, bitstream_size(&stream));
   //printf("Decoded:\n");
   //bitstream_print(&out_decode);

   if (bitstream_equal(&out_decode, &stream))
      printf("Correctly decoded!\n");
   else
      printf("Incorrectly decoded!\n");

   bitstream_free(&stream);
   bitstream_free(&out_stream);
   bitstream_free(&out_decode);
}