encoder and encode/decode test

6 files changed, 641 insertions, 101 deletions

diff --git a/src/Makefile b/src/Makefile
index 94ce7f6..ee7caed 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -1,23 +1,31 @@
 CC = gcc
 CFLAGS = -g -Wall
 
-objects = mytest.o aed.o
+OBJS = aee.o aed.o
 
-mytest: $(objects)
-	$(CC) $(CFLAGS) -o mytest $(objects)
+.PHONY : all clean test
 
-aed.o: aecd.h
+all: libae.a test
 
-.PHONY : clean test vtest
-clean:
-	rm -f mytest ../data/ae_out $(objects)
+test_encode: test_encode.o libae.a
+	$(CC) $(CFLAGS) -o test_encode test_encode.o -L. -lae
+
+test_decode: test_decode.o libae.a
+	$(CC) $(CFLAGS) -o test_decode test_decode.o -L. -lae
+
+libae.a: $(OBJS)
+	$(AR) $(ARFLAGS) $@ $(OBJS)
+	-@ ($(RANLIB) $@ || true) >/dev/null 2>&1
 
-test: mytest
-	./mytest 1 1 < ../data/example_data.szip > ../data/ae_out
-	diff ../data/ae_out ../data/example_data
-	./mytest 9478 16384 < ../data/example_data.szip > ../data/ae_out
-	diff ../data/ae_out ../data/example_data
+aed.o: libae.h
+aee.o: libae.h
+
+clean:
+	rm -f $(OBJS) test_encode.o test_decode.o \
+	test_encode test_decode libae.a \
+	../data/test.ae ../data/test
 
-vtest: mytest
-	valgrind -v ./mytest < ../data/example_data.szip > ../data/ae_out
-	diff ../data/ae_out ../data/example_data
+test: test_encode test_decode
+	./test_encode 1 1 < ../data/example_data > ../data/test.ae
+	./test_decode 1 1 < ../data/test.ae > ../data/test
+	diff ../data/test ../data/example_data
diff --git a/src/aed.c b/src/aed.c
index 02af316..76dd2dc 100644
--- a/src/aed.c
+++ b/src/aed.c
@@ -7,16 +7,16 @@
 #include <inttypes.h>
 #include <string.h>
 
-#include "aecd.h"
+#include "libae.h"
 
-#define REFBLOCK (strm->pp && (strm->total_out / strm->block_size)  \
-                    % strm->segment_size == 0)
 #define SAFE (strm->avail_in >= state->in_blklen        \
               && strm->avail_out >= strm->block_size)
 
 #define ROS 5
 
 typedef struct internal_state {
+    const uint8_t *next_in;
+    uint32_t *next_out;
     int id;            /* option ID */
     uint32_t id_len;   /* bit length of code option identification key */
     int *id_table;     /* table maps IDs to states */
@@ -33,6 +33,7 @@ typedef struct internal_state {
     uint64_t acc;      /* accumulator for currently used bit sequence */
     uint8_t bitp;      /* bit pointer to the next unused bit in accumulator */
     uint32_t fs;       /* last fundamental sequence in accumulator */
+    int ref;           /* 1 if current block has reference sample */
 } decode_state;
 
 /* decoding table for the second-extension option */
@@ -97,7 +98,7 @@ static inline void u_put(ae_streamp strm, uint32_t sample)
         }
         sample = x + D;
     }
-    *strm->next_out++ = state->last_out = sample;
+    *state->next_out++ = state->last_out = sample;
     strm->avail_out--;
     strm->total_out++;
 }
@@ -117,7 +118,7 @@ static inline uint32_t u_get(ae_streamp strm, unsigned int n)
     {
         strm->avail_in--;
         strm->total_in++;
-        state->acc = (state->acc << 8) + *strm->next_in++;
+        state->acc = (state->acc << 8) + *state->next_in++;
         state->bitp += 8;
     }
     state->bitp -= n;
@@ -143,25 +144,23 @@ static inline uint32_t u_get_fs(ae_streamp strm)
 
 static inline void fast_split(ae_streamp strm)
 {
-    int i, start, k;
+    int i, k;
     decode_state *state;
 
     state = strm->state;
-    start = 0;
     k = state->id - 1;
 
-    if (REFBLOCK)
+    if (state->ref)
     {
-        start = 1;
         u_put(strm, u_get(strm, strm->bit_per_sample));
     }
 
-    for (i = start; i < strm->block_size; i++)
+    for (i = state->ref; i < strm->block_size; i++)
     {
         state->block[i] = u_get_fs(strm) << k;
     }
 
-    for (i = start; i < strm->block_size; i++)
+    for (i = state->ref; i < strm->block_size; i++)
     {
         state->block[i] += u_get(strm, k);
         u_put(strm, state->block[i]);
@@ -181,7 +180,7 @@ static inline void fast_se(ae_streamp strm)
     int i;
     uint32_t gamma, beta, ms, delta1;
 
-    i = REFBLOCK? 1: 0;
+    i = strm->state->ref;
 
     while (i < strm->bit_per_sample)
     {
@@ -266,17 +265,17 @@ int ae_decode_init(ae_streamp strm)
     return AE_OK;
 }
 
-#define ASK(n)                                          \
-    do {                                                \
-        while (state->bitp < (unsigned)(n))             \
-        {                                               \
-            if (strm->avail_in == 0) goto req_buffer;   \
-            strm->avail_in--;                           \
-            strm->total_in++;                           \
-            state->acc <<= 8;                           \
-            state->acc |= (uint64_t)(*strm->next_in++); \
-            state->bitp += 8;                           \
-        }                                               \
+#define ASK(n)                                           \
+    do {                                                 \
+        while (state->bitp < (unsigned)(n))              \
+        {                                                \
+            if (strm->avail_in == 0) goto req_buffer;    \
+            strm->avail_in--;                            \
+            strm->total_in++;                            \
+            state->acc <<= 8;                            \
+            state->acc |= (uint64_t)(*state->next_in++); \
+            state->bitp += 8;                            \
+        }                                                \
     } while (0)
 
 #define GET(n)                                                  \
@@ -322,8 +321,6 @@ int ae_decode(ae_streamp strm, int flush)
        Can work with one byte input und one sample output buffers. If
        enough buffer space is available, then faster implementations
        of the states are called. Inspired by zlib.
-
-       TODO: Flush modes like in zlib
     */
 
     size_t zero_blocks;
@@ -332,15 +329,23 @@ int ae_decode(ae_streamp strm, int flush)
     decode_state *state;
 
     state = strm->state;
+    state->next_in = strm->next_in;
+    state->next_out = strm->next_out;
 
     for (;;)
     {
         switch(state->mode)
         {
         case M_ID:
-            ASK(3);
-            state->id = GET(3);
-            DROP(3);
+            if (strm->pp 
+                && (strm->total_out / strm->block_size) % strm->segment_size == 0)
+                state->ref = 1;
+            else
+                state->ref = 0;
+
+            ASK(state->id_len);
+            state->id = GET(state->id_len);
+            DROP(state->id_len);
             state->mode = state->id_table[state->id];
             break;
 
@@ -352,7 +357,7 @@ int ae_decode(ae_streamp strm, int flush)
                 break;
             }
 
-            if (REFBLOCK)
+            if (state->ref)
             {
                 COPYSAMPLE();
                 state->n = strm->block_size - 1;
@@ -366,11 +371,10 @@ int ae_decode(ae_streamp strm, int flush)
             state->mode = M_SPLIT_FS;
 
         case M_SPLIT_FS:
-            k = state->id - 1;
             do
             {
                 ASKFS();
-                state->block[state->i] = GETFS() << k;
+                state->block[state->i] = GETFS();
                 DROPFS();
             }
             while(--state->i);
@@ -383,7 +387,7 @@ int ae_decode(ae_streamp strm, int flush)
             do
             {
                 ASK(k);
-                PUT(state->block[state->i] + GET(k));
+                PUT((state->block[state->i] << k) + GET(k));
                 DROP(k);
             }
             while(--state->i);
@@ -398,7 +402,7 @@ int ae_decode(ae_streamp strm, int flush)
             state->mode = M_LOW_ENTROPY_REF;
 
         case M_LOW_ENTROPY_REF:
-            if (REFBLOCK)
+            if (state->ref)
                 COPYSAMPLE();
 
             if(state->id == 1)
@@ -422,7 +426,7 @@ int ae_decode(ae_streamp strm, int flush)
             }
 
 
-            if (REFBLOCK)
+            if (state->ref)
                 state->i = zero_blocks * strm->block_size - 1;
             else
                 state->i = zero_blocks * strm->block_size;
@@ -453,7 +457,7 @@ int ae_decode(ae_streamp strm, int flush)
             }
 
             state->mode = M_SE_DECODE;
-            state->i = REFBLOCK? 1: 0;
+            state->i = state->ref;
 
         case M_SE_DECODE:
             while(state->i < strm->bit_per_sample)
@@ -503,5 +507,7 @@ int ae_decode(ae_streamp strm, int flush)
     }
 
 req_buffer:
+    strm->next_in = state->next_in;
+    strm->next_out = state->next_out;
     return AE_OK;
 }
diff --git a/src/aee.c b/src/aee.c
new file mode 100644
index 0000000..9223640
--- /dev/null
+++ b/src/aee.c
@@ -0,0 +1,431 @@
+/* Adaptive Entropy Encoder            */
+/* CCSDS 121.0-B-1 and CCSDS 120.0-G-2 */
+
+#include <stdio.h>
+#include <unistd.h>
+#include <stdlib.h>
+#include <inttypes.h>
+#include <string.h>
+
+#include "libae.h"
+
+#define ROS 5
+
+#define MIN(a, b) (((a) < (b))? (a): (b))
+
+enum
+{
+    M_NEW_BLOCK,
+    M_GET_BLOCK,
+    M_CHECK_ZERO_BLOCK,
+    M_SELECT_CODE_OPTION,
+    M_ENCODE_SPLIT,
+    M_FLUSH_BLOCK,
+    M_FLUSH_BLOCK_LOOP,
+    M_ENCODE_UNCOMP,
+    M_ENCODE_SE,
+    M_ENCODE_ZERO,
+};
+
+typedef struct internal_state {
+    const uint32_t *next_in;
+    uint8_t *next_out;
+    uint32_t id_len;        /* bit length of code option identification key */
+    uint32_t last_in;       /* previous input for preprocessing */
+    int64_t xmin;           /* minimum integer for preprocessing */
+    int64_t xmax;           /* maximum integer for preprocessing */
+    int mode;               /* current mode of FSM */
+    size_t i;               /* counter for samples */
+    uint32_t *block_in;     /* input block buffer */
+    uint8_t *block_out;     /* output block buffer */
+    uint8_t *bp_out;        /* pointer to current output */
+    uint8_t bitp;           /* bit pointer to the next unused bit in accumulator */
+    uint8_t block_deferred; /* there is a block in the input buffer
+                               but we first have to emit a zero block */
+    uint8_t ref;            /* current buffer has a reference sample */
+    uint8_t zero_ref;       /* current zero block has a reference sample */
+    uint32_t zero_ref_sample; /* reference sample of zero block */
+    size_t zero_blocks;     /* number of contiguous zero blocks */
+} encode_state;
+
+int ae_encode_init(ae_streamp strm)
+{
+    int blklen;
+    encode_state *state;
+
+    /* Some sanity checks */
+    if (strm->bit_per_sample > 32 || strm->bit_per_sample == 0)
+    {
+        return AE_ERRNO;
+    }
+
+    /* Internal state for encoder */
+    state = (encode_state *) malloc(sizeof(encode_state));
+    if (state == NULL)
+    {
+        return AE_MEM_ERROR;
+    }
+    strm->state = state;
+
+    if (16 < strm->bit_per_sample)
+        state->id_len = 5;
+    else if (8 < strm->bit_per_sample)
+        state->id_len = 4;
+    else
+        state->id_len = 3;
+
+    state->block_in = (uint32_t *)malloc(strm->block_size * sizeof(uint32_t));
+    if (state->block_in == NULL)
+    {
+        return AE_MEM_ERROR;
+    }
+
+    blklen = (strm->block_size * strm->bit_per_sample
+              + state->id_len) / 8 + 1;
+
+    state->block_out = (uint8_t *)malloc(blklen);
+    if (state->block_out == NULL)
+    {
+        return AE_MEM_ERROR;
+    }
+    state->bp_out = state->block_out;
+    state->bitp = 8;
+
+    strm->total_in = 0;
+    strm->total_out = 0;
+    state->xmin = 0;
+    state->xmax = (1ULL << strm->bit_per_sample) - 1;
+
+    state->mode = M_NEW_BLOCK;
+
+    state->block_deferred = 0;
+    state->zero_ref = 0;
+    state->ref = 0;
+
+    return AE_OK;
+}
+
+#define EMIT(d, n)                                              \
+    do                                                          \
+    {                                                           \
+        int bits = (n);                                         \
+        uint32_t data = (d);                                    \
+        while(bits)                                             \
+        {                                                       \
+            data &= ((1UL << bits) - 1);                        \
+            if (bits <= state->bitp)                            \
+            {                                                   \
+                data <<= state->bitp - bits;                    \
+                *state->bp_out += data;                         \
+                state->bitp -= bits;                            \
+                bits = 0;                                       \
+            }                                                   \
+            else                                                \
+            {                                                   \
+                *state->bp_out += data >> (bits - state->bitp); \
+                bits -= state->bitp;                            \
+                *++state->bp_out = 0;                           \
+                state->bitp = 8;                                \
+            }                                                   \
+        }                                                       \
+    }                                                           \
+    while (0)
+
+
+#define EMITFS(d)                                               \
+    do                                                          \
+    {                                                           \
+        EMIT(0, d);                                             \
+        EMIT(1, 1);                                             \
+    }                                                           \
+    while (0)
+
+int ae_encode(ae_streamp strm, int flush)
+{
+    /**
+       Finite-state machine implementation of the adaptive entropy
+       encoder.
+    */
+
+    int i, j, zb;
+    int k_len[strm->bit_per_sample - 2];
+    int k, k_min, se_len, blk_head;
+    uint32_t d;
+    int64_t theta, Delta;
+
+    encode_state *state;
+
+    state = strm->state;
+    state->next_in = strm->next_in;
+    state->next_out = strm->next_out;
+
+    for (;;)
+    {
+        switch(state->mode)
+        {
+        case M_NEW_BLOCK:
+            if (state->zero_blocks == 0)
+            {
+                /* copy leftover from last block */
+                *state->block_out = *state->bp_out;
+                state->bp_out = state->block_out;
+            }
+
+            if(state->block_deferred)
+            {
+                state->block_deferred = 0;
+                state->mode = M_SELECT_CODE_OPTION;
+                break;
+            }
+
+            state->i = 0;
+            state->mode = M_GET_BLOCK;
+
+        case M_GET_BLOCK:
+            do
+            {
+                if (strm->avail_in == 0)
+                {
+                    if (flush == AE_FLUSH)
+                    {
+                        if (state->i > 0)
+                        {
+                            /* pad block with zeros if we have
+                               a partial block */
+                            state->block_in[state->i] = 0;
+                        }
+                        else
+                        {
+                            /* Pad last output byte with 1 bits
+                               if user wants to flush, i.e. we got
+                               all input there is.
+                            */
+                            EMIT(0xff, state->bitp);
+                            *state->next_out++ = *state->bp_out;
+                            strm->avail_out--;
+                            strm->total_out++;
+                        }
+                    }
+                    goto req_buffer;
+                }
+                else
+                {
+                    state->block_in[state->i] = *state->next_in++;
+                    strm->avail_in--;
+                    strm->total_in++;
+                }
+            }
+            while (++state->i < strm->block_size);
+
+            /* preprocess block if needed */
+            if (strm->pp)
+            {
+                /* If this is the first block in a segment
+                   then we need to insert a reference sample.
+                */
+                if((strm->total_in / strm->block_size) % strm->segment_size == 1)
+                {
+                    state->ref = 1;
+                    state->last_in = state->block_in[0];
+                }
+                else
+                {
+                    state->ref = 0;
+                }
+
+                for (i = state->ref; i < strm->block_size; i++)
+                {
+                    theta = MIN(state->last_in - state->xmin,
+                                state->xmax - state->last_in);
+                    Delta = (long long)state->block_in[i] - (long long)state->last_in;
+
+                    if (0 <= Delta && Delta <= theta)
+                        d = 2 * Delta;
+                    else if (-theta <= Delta && Delta < 0)
+                        d = 2 * llabs(Delta) - 1;
+                    else
+                        d = theta + llabs(Delta);
+
+                    state->last_in = state->block_in[i];
+                    state->block_in[i] = d;
+                }
+            }
+            state->mode = M_CHECK_ZERO_BLOCK;
+
+        case M_CHECK_ZERO_BLOCK:
+            /* Check zero block */
+            zb = 1;
+            for (i = state->ref; i < strm->block_size && zb; i++)
+                if (state->block_in[i] != 0) zb = 0;
+
+            if (zb)
+            {
+                /* remember ref on first zero block */
+                if (state->zero_blocks == 0)
+                {
+                    state->zero_ref = state->ref;
+                    state->zero_ref_sample = state->block_in[0];
+                }
+
+                state->zero_blocks++;
+
+                if ((strm->total_in / strm->block_size)
+                    % strm->segment_size == 0)
+                {
+                    if (state->zero_blocks > ROS)
+                        state->zero_blocks = ROS;
+                    state->mode = M_ENCODE_ZERO;
+                    break;
+                }
+                state->mode = M_NEW_BLOCK;
+                break;
+            }
+            else if (state->zero_blocks)
+            {
+                state->mode = M_ENCODE_ZERO;
+                /* The current block isn't zero but we have to
+                   emit a previous zero block first. The
+                   current block has to be handled later.
+                */
+                state->block_deferred = 1;
+                break;
+            }
+
+            state->mode = M_SELECT_CODE_OPTION;
+
+        case M_SELECT_CODE_OPTION:
+            /* Encoded block always starts with ID and possibly
+               a reference sample. */
+            blk_head = state->id_len;
+            if (state->ref)
+                blk_head += strm->bit_per_sample;
+
+            for (j = 0; j < strm->bit_per_sample - 2; j++)
+                k_len[j] = blk_head;
+
+            /* Count bits for sample splitting options */
+            for (i = state->ref; i < strm->block_size; i++)
+                for (j = 0; j < strm->bit_per_sample - 2; j++)
+                    k_len[j] += (state->block_in[i] >> j) + 1 + j;
+
+            /* Baseline is the size of an uncompressed block */
+            k_min = state->id_len + strm->block_size * strm->bit_per_sample;
+            k = strm->bit_per_sample;
+
+            /* See if splitting option is better */
+            for (j = 0; j < strm->bit_per_sample - 2; j++)
+            {
+                if (k_len[j] < k_min)
+                {
+                    k = j;
+                    k_min = k_len[j];
+                }
+            }
+
+            /* Count bits for 2nd extension */
+            se_len = blk_head + 1;
+
+            for (i = 0; i < strm->block_size && k_min > se_len; i+= 2)
+            {
+                d = state->block_in[i] + state->block_in[i + 1];
+                se_len += d * (d + 1) / 2 + state->block_in[i + 1];
+            }
+
+            /* Decide which option to use */
+            if (k_min <= se_len)
+            {
+                if (k == strm->bit_per_sample)
+                {
+                    state->mode = M_ENCODE_UNCOMP;
+                    break;
+                }
+                else
+                {
+                    state->mode = M_ENCODE_SPLIT;
+                }
+            }
+            else
+            {
+                state->mode = M_ENCODE_SE;
+                break;
+            }
+
+        case M_ENCODE_SPLIT:
+            EMIT(k + 1, state->id_len);
+            if (state->ref)
+                EMIT(state->block_in[0], strm->bit_per_sample);
+
+            for (i = state->ref; i < strm->block_size; i++)
+                EMITFS(state->block_in[i] >> k);
+
+            for (i = state->ref; i < strm->block_size; i++)
+                EMIT(state->block_in[i], k);
+
+            state->mode = M_FLUSH_BLOCK;
+
+        case M_FLUSH_BLOCK:
+            if (strm->avail_in == 0 && flush == AE_FLUSH)
+            {
+                /* pad last byte with 1 bits */
+                EMIT(0xff, state->bitp);
+            }
+            state->i = 0;
+            state->mode = M_FLUSH_BLOCK_LOOP;
+
+        case M_FLUSH_BLOCK_LOOP:
+            while(state->block_out + state->i < state->bp_out)
+            {
+                if (strm->avail_out == 0)
+                    goto req_buffer;
+
+                *state->next_out++ = state->block_out[state->i];
+                strm->avail_out--;
+                strm->total_out++;
+                state->i++;
+            }
+            state->mode = M_NEW_BLOCK;
+            break;
+
+        case M_ENCODE_UNCOMP:
+            EMIT(0x1f, state->id_len);
+            for (i = 0; i < strm->block_size; i++)
+                EMIT(state->block_in[i], strm->bit_per_sample);
+
+            state->mode = M_FLUSH_BLOCK;
+            break;
+
+        case M_ENCODE_SE:
+            EMIT(1, state->id_len + 1);
+            if (state->ref)
+                EMIT(state->block_in[0], strm->bit_per_sample);
+
+            for (i = 0; i < strm->block_size; i+= 2)
+            {
+                d = state->block_in[i] + state->block_in[i + 1];
+                EMITFS(d * (d + 1) / 2 + state->block_in[i + 1]);
+            }
+
+            state->mode = M_FLUSH_BLOCK;
+            break;
+
+        case M_ENCODE_ZERO:
+            EMIT(0, state->id_len + 1);
+            if (state->zero_ref)
+            {
+                EMIT(state->zero_ref_sample, strm->bit_per_sample);
+            }
+            EMITFS(state->zero_blocks - 1);
+            state->zero_blocks = 0;
+            state->mode = M_FLUSH_BLOCK;
+            break;
+
+        default:
+            return AE_STREAM_ERROR;
+        }
+    }
+
+req_buffer:
+    strm->next_in = state->next_in;
+    strm->next_out = state->next_out;
+    return AE_OK;
+}
diff --git a/src/aecd.h b/src/libae.h
index ecc7163..bbbad81 100644
--- a/src/aecd.h
+++ b/src/libae.h
@@ -1,5 +1,5 @@
-#ifndef AELIB_H
-#define AELIB_H
+#ifndef LIBAE_H
+#define LIBAE_H
 
 #include <inttypes.h>
 
@@ -7,11 +7,11 @@ struct internal_state;
 
 typedef struct _ae_stream
 {
-    const uint8_t *next_in;
+    const void *next_in;
     size_t avail_in;  /* number of bytes available at next_in */
     size_t total_in;  /* total number of input bytes read so far */
 
-    uint32_t *next_out;
+    void *next_out;
     size_t avail_out; /* remaining free space at next_out */
     size_t total_out; /* total number of bytes output so far */
 
@@ -34,15 +34,12 @@ typedef ae_stream *ae_streamp;
 #define AE_MEM_ERROR    (-4)
 
 #define AE_NO_FLUSH      0
-#define AE_PARTIAL_FLUSH 1
-#define AE_SYNC_FLUSH    2
-#define AE_FULL_FLUSH    3
-#define AE_FINISH        4
-#define AE_BLOCK         5
-#define AE_TREES         6
+#define AE_FLUSH         1
 
 int ae_decode_init(ae_streamp strm);
-
 int ae_decode(ae_streamp strm, int flush);
 
-#endif /* AELIB_H */
+int ae_encode_init(ae_streamp strm);
+int ae_encode(ae_streamp strm, int flush);
+
+#endif /* LIBAE_H */
diff --git a/src/mytest.c b/src/test_decode.c
index d029a36..5325727 100644
--- a/src/mytest.c
+++ b/src/test_decode.c
@@ -2,44 +2,35 @@
 #include <unistd.h>
 #include <stdlib.h>
 #include <inttypes.h>
-#include "aecd.h"
+#include "libae.h"
 
 #define CHUNK_OUT 1
 #define CHUNK_IN 1
-#define ALL_IN 9478
 
 int main(int argc, char *argv[])
 {
     ae_stream strm;
-    int c, i, n, status, todo;
     uint8_t *in;
     uint32_t *out;
-    size_t total_out;
-    int chunk_in, chunk_out;
-
-    chunk_in = CHUNK_IN;
-    chunk_out = CHUNK_OUT;
+    int chunk_in, chunk_out, i, c, total_out, status;
+    int input_avail, output_avail;
 
     if (argc == 3)
     {
         chunk_in = atoi(argv[1]);
         chunk_out = atoi(argv[2]);
     }
+    else
+    {
+        chunk_in = CHUNK_IN;
+        chunk_out = CHUNK_OUT;
+    }
 
-    fprintf(stderr, "chunk_in: %i\nchunk_out: %i\n", chunk_in, chunk_out);
-    in = (uint8_t *)malloc(ALL_IN);
+    in = (uint8_t *)malloc(chunk_in);
     out = (uint32_t *)malloc(chunk_out * sizeof(uint32_t));
     if (in == NULL || out == NULL)
         return 1;
 
-    n = 0;
-    while ((c = getc(stdin)) != EOF)
-    {
-        *in++ = c;
-        n++;
-    }
-    in -= n;
-
     strm.bit_per_sample = 8;
     strm.block_size = 8;
     strm.segment_size = 2;
@@ -48,33 +39,32 @@ int main(int argc, char *argv[])
     if (ae_decode_init(&strm) != AE_OK)
         return 1;
 
-    strm.next_in = in;
-    strm.avail_in = chunk_in;
-    strm.next_out = out;
-    strm.avail_out = chunk_out;
-    todo = 1;
     total_out = 0;
+    strm.avail_in = 0;
+    strm.avail_out = chunk_out;
+    strm.next_out = out;
+
+    input_avail = 1;
+    output_avail = 1;
 
-    while(todo)
+    while(input_avail || output_avail)
     {
-        todo = 0;
-        if ((status = ae_decode(&strm, 0)) != AE_OK)
+        if (strm.avail_in == 0)
         {
-            fprintf(stderr, "error is %i\n", status);
-            return 1;
+            i = 0;
+            while(i < chunk_in && (c = getc(stdin)) != EOF)
+                in[i++] = c;
+            strm.avail_in = i;
+
+            strm.next_in = in;
+            if (c == EOF)
+                input_avail = 0;
         }
-//        fprintf(stderr, "avail in %li total in %li avail out %li total out %lx\n", strm.avail_in, strm.total_in, strm.avail_out, strm.total_out);
 
-        if (strm.avail_in == 0 && strm.total_in < ALL_IN)
+        if ((status = ae_decode(&strm, AE_NO_FLUSH)) != AE_OK)
         {
-            in += chunk_in;
-
-            strm.next_in = in;
-            if (ALL_IN - strm.total_in < chunk_in)
-                strm.avail_in = ALL_IN - strm.total_in;
-            else
-                strm.avail_in = chunk_in;
-            todo = 1;
+            fprintf(stderr, "error is %i\n", status);
+            return 1;
         }
 
         if (strm.total_out - total_out > 0)
@@ -84,10 +74,15 @@ int main(int argc, char *argv[])
                 putc(out[i], stdout);
             }
             total_out = strm.total_out;
+            output_avail = 1;
             strm.next_out = out;
             strm.avail_out = chunk_out;
-            todo = 1;
         }
+        else
+        {
+            output_avail = 0;
+        }
+
     }
 
     return 0;
diff --git a/src/test_encode.c b/src/test_encode.c
new file mode 100644
index 0000000..029dca8
--- /dev/null
+++ b/src/test_encode.c
@@ -0,0 +1,103 @@
+#include <stdio.h>
+#include <unistd.h>
+#include <stdlib.h>
+#include <inttypes.h>
+#include "libae.h"
+
+#define CHUNK_OUT 0x4000
+#define CHUNK_IN 1024
+#define ALL_IN 9478
+
+int main(int argc, char *argv[])
+{
+    ae_stream strm;
+    uint32_t *in;
+    uint8_t *out;
+    int chunk_in, chunk_out, i, c, total_out, status;
+    int input_avail, output_avail;
+
+    if (argc == 3)
+    {
+        chunk_in = atoi(argv[1]);
+        chunk_out = atoi(argv[2]);
+    }
+    else
+    {
+        chunk_in = CHUNK_IN;
+        chunk_out = CHUNK_OUT;
+    }
+
+    out = (uint8_t *)malloc(chunk_out);
+    in = (uint32_t *)malloc(chunk_in * sizeof(uint32_t));
+    if (in == NULL || out == NULL)
+        return 1;
+
+    strm.bit_per_sample = 8;
+    strm.block_size = 8;
+    strm.segment_size = 2;
+    strm.pp = 1;
+
+    if (ae_encode_init(&strm) != AE_OK)
+        return 1;
+
+    total_out = 0;
+    strm.avail_in = 0;
+    strm.avail_out = chunk_out;
+    strm.next_out = out;
+
+    input_avail = 1;
+    output_avail = 1;
+
+    while(input_avail || output_avail)
+    {
+        if (strm.avail_in == 0)
+        {
+            i = 0;
+            while(i < chunk_in && (c = getc(stdin)) != EOF)
+                in[i++] = c;
+            strm.avail_in = i;
+
+            strm.next_in = in;
+            if (c == EOF)
+                input_avail = 0;
+        }
+
+        if ((status = ae_encode(&strm, AE_NO_FLUSH)) != AE_OK)
+        {
+            fprintf(stderr, "error is %i\n", status);
+            return 1;
+        }
+
+        if (strm.total_out - total_out > 0)
+        {
+            for (i=0; i < strm.total_out - total_out; i++)
+            {
+                putc(out[i], stdout);
+            }
+            total_out = strm.total_out;
+            output_avail = 1;
+            strm.next_out = out;
+            strm.avail_out = chunk_out;
+        }
+        else
+        {
+            output_avail = 0;
+        }
+
+    }
+
+    if ((status = ae_encode(&strm, AE_FLUSH)) != AE_OK)
+    {
+        fprintf(stderr, "error is %i\n", status);
+        return 1;
+    }
+
+    if (strm.total_out - total_out > 0)
+    {
+        for (i=0; i < strm.total_out - total_out; i++)
+        {
+            putc(out[i], stdout);
+        }
+    }
+    return 0;
+}