TAD: embedded zero tree coding (basically 1D EZBC)

video_encoder/encoder_tad.c

@@ -854,6 +854,287 @@ void tad32_free_statistics(void) {
     stats_initialized = 0;
 }
 
+//=============================================================================
+// Binary Tree EZBC (1D Variant for TAD)
+//=============================================================================
+
+#include <stdbool.h>
+
+// Bitstream writer for EZBC
+typedef struct {
+    uint8_t *data;
+    size_t capacity;
+    size_t byte_pos;
+    uint8_t bit_pos;  // 0-7, current bit position in current byte
+} tad_bitstream_t;
+
+// Block structure for 1D binary tree
+typedef struct {
+    int start;    // Start index in 1D array
+    int length;   // Block length
+} tad_block_t;
+
+// Queue for block processing
+typedef struct {
+    tad_block_t *blocks;
+    size_t count;
+    size_t capacity;
+} tad_block_queue_t;
+
+// Track coefficient state for refinement
+typedef struct {
+    bool significant;     // Has been marked significant
+    int first_bitplane;   // Bitplane where it became significant
+} tad_coeff_state_t;
+
+// Bitstream operations
+static void tad_bitstream_init(tad_bitstream_t *bs, size_t initial_capacity) {
+    bs->capacity = initial_capacity;
+    bs->data = calloc(1, initial_capacity);
+    bs->byte_pos = 0;
+    bs->bit_pos = 0;
+}
+
+static void tad_bitstream_write_bit(tad_bitstream_t *bs, int bit) {
+    // Grow if needed
+    if (bs->byte_pos >= bs->capacity) {
+        bs->capacity *= 2;
+        bs->data = realloc(bs->data, bs->capacity);
+        // Clear new memory
+        memset(bs->data + bs->byte_pos, 0, bs->capacity - bs->byte_pos);
+    }
+
+    if (bit) {
+        bs->data[bs->byte_pos] |= (1 << bs->bit_pos);
+    }
+
+    bs->bit_pos++;
+    if (bs->bit_pos == 8) {
+        bs->bit_pos = 0;
+        bs->byte_pos++;
+    }
+}
+
+static void tad_bitstream_write_bits(tad_bitstream_t *bs, uint32_t value, int num_bits) {
+    for (int i = 0; i < num_bits; i++) {
+        tad_bitstream_write_bit(bs, (value >> i) & 1);
+    }
+}
+
+static size_t tad_bitstream_size(tad_bitstream_t *bs) {
+    return bs->byte_pos + (bs->bit_pos > 0 ? 1 : 0);
+}
+
+static void tad_bitstream_free(tad_bitstream_t *bs) {
+    free(bs->data);
+}
+
+// Block queue operations
+static void tad_queue_init(tad_block_queue_t *q) {
+    q->capacity = 1024;
+    q->blocks = malloc(q->capacity * sizeof(tad_block_t));
+    q->count = 0;
+}
+
+static void tad_queue_push(tad_block_queue_t *q, tad_block_t block) {
+    if (q->count >= q->capacity) {
+        q->capacity *= 2;
+        q->blocks = realloc(q->blocks, q->capacity * sizeof(tad_block_t));
+    }
+    q->blocks[q->count++] = block;
+}
+
+static void tad_queue_free(tad_block_queue_t *q) {
+    free(q->blocks);
+}
+
+// Check if all coefficients in block have |coeff| < threshold
+static bool tad_is_zero_block(int8_t *coeffs, const tad_block_t *block, int threshold) {
+    for (int i = block->start; i < block->start + block->length; i++) {
+        if (abs(coeffs[i]) >= threshold) {
+            return false;
+        }
+    }
+    return true;
+}
+
+// Find maximum absolute coefficient value
+static int tad_find_max_abs(int8_t *coeffs, size_t count) {
+    int max_abs = 0;
+    for (size_t i = 0; i < count; i++) {
+        int abs_val = abs(coeffs[i]);
+        if (abs_val > max_abs) {
+            max_abs = abs_val;
+        }
+    }
+    return max_abs;
+}
+
+// Get MSB position (bitplane number)
+static int tad_get_msb_bitplane(int value) {
+    if (value == 0) return 0;
+    int bitplane = 0;
+    while (value > 1) {
+        value >>= 1;
+        bitplane++;
+    }
+    return bitplane;
+}
+
+// Context for recursive EZBC processing
+typedef struct {
+    tad_bitstream_t *bs;
+    int8_t *coeffs;
+    tad_coeff_state_t *states;
+    int length;
+    int bitplane;
+    int threshold;
+    tad_block_queue_t *next_insignificant;
+    tad_block_queue_t *next_significant;
+    int *sign_count;
+} tad_ezbc_context_t;
+
+// Recursively process a significant block - subdivide until size 1
+static void tad_process_significant_block_recursive(tad_ezbc_context_t *ctx, tad_block_t block) {
+    // If size 1: emit sign bit and add to significant queue
+    if (block.length == 1) {
+        int idx = block.start;
+        tad_bitstream_write_bit(ctx->bs, ctx->coeffs[idx] < 0 ? 1 : 0);
+        (*ctx->sign_count)++;
+        ctx->states[idx].significant = true;
+        ctx->states[idx].first_bitplane = ctx->bitplane;
+        tad_queue_push(ctx->next_significant, block);
+        return;
+    }
+
+    // Block is > 1: subdivide into left and right halves
+    int mid = block.length / 2;
+    if (mid == 0) mid = 1;
+
+    // Process left child
+    tad_block_t left = {block.start, mid};
+    if (!tad_is_zero_block(ctx->coeffs, &left, ctx->threshold)) {
+        tad_bitstream_write_bit(ctx->bs, 1);  // Significant
+        tad_process_significant_block_recursive(ctx, left);
+    } else {
+        tad_bitstream_write_bit(ctx->bs, 0);  // Insignificant
+        tad_queue_push(ctx->next_insignificant, left);
+    }
+
+    // Process right child (if exists)
+    if (block.length > mid) {
+        tad_block_t right = {block.start + mid, block.length - mid};
+        if (!tad_is_zero_block(ctx->coeffs, &right, ctx->threshold)) {
+            tad_bitstream_write_bit(ctx->bs, 1);
+            tad_process_significant_block_recursive(ctx, right);
+        } else {
+            tad_bitstream_write_bit(ctx->bs, 0);
+            tad_queue_push(ctx->next_insignificant, right);
+        }
+    }
+}
+
+// Binary tree EZBC encoding for a single channel (1D variant)
+static size_t tad_encode_channel_ezbc(int8_t *coeffs, size_t count, uint8_t **output) {
+    tad_bitstream_t bs;
+    tad_bitstream_init(&bs, count / 4);  // Initial guess
+
+    // Track coefficient significance
+    tad_coeff_state_t *states = calloc(count, sizeof(tad_coeff_state_t));
+
+    // Find maximum value to determine MSB bitplane
+    int max_abs = tad_find_max_abs(coeffs, count);
+    int msb_bitplane = tad_get_msb_bitplane(max_abs);
+
+    // Write header: MSB bitplane and length
+    tad_bitstream_write_bits(&bs, msb_bitplane, 8);
+    tad_bitstream_write_bits(&bs, (uint32_t)count, 16);
+
+    // Initialize queues
+    tad_block_queue_t insignificant_queue, next_insignificant;
+    tad_block_queue_t significant_queue, next_significant;
+
+    tad_queue_init(&insignificant_queue);
+    tad_queue_init(&next_insignificant);
+    tad_queue_init(&significant_queue);
+    tad_queue_init(&next_significant);
+
+    // Start with root block as insignificant
+    tad_block_t root = {0, (int)count};
+    tad_queue_push(&insignificant_queue, root);
+
+    // Process bitplanes from MSB to LSB
+    for (int bitplane = msb_bitplane; bitplane >= 0; bitplane--) {
+        int threshold = 1 << bitplane;
+
+        // Process insignificant blocks - check if they become significant
+        for (size_t i = 0; i < insignificant_queue.count; i++) {
+            tad_block_t block = insignificant_queue.blocks[i];
+
+            if (tad_is_zero_block(coeffs, &block, threshold)) {
+                // Still insignificant: emit 0
+                tad_bitstream_write_bit(&bs, 0);
+                // Keep in insignificant queue for next bitplane
+                tad_queue_push(&next_insignificant, block);
+            } else {
+                // Became significant: emit 1
+                tad_bitstream_write_bit(&bs, 1);
+
+                // Use recursive subdivision
+                int sign_count = 0;
+                tad_ezbc_context_t ctx = {
+                    .bs = &bs,
+                    .coeffs = coeffs,
+                    .states = states,
+                    .length = (int)count,
+                    .bitplane = bitplane,
+                    .threshold = threshold,
+                    .next_insignificant = &next_insignificant,
+                    .next_significant = &next_significant,
+                    .sign_count = &sign_count
+                };
+                tad_process_significant_block_recursive(&ctx, block);
+            }
+        }
+
+        // Refinement pass: emit next bit for already-significant coefficients
+        for (size_t i = 0; i < significant_queue.count; i++) {
+            tad_block_t block = significant_queue.blocks[i];
+            int idx = block.start;
+
+            // Emit refinement bit (bit at position 'bitplane')
+            int bit = (abs(coeffs[idx]) >> bitplane) & 1;
+            tad_bitstream_write_bit(&bs, bit);
+        }
+
+        // Swap queues for next bitplane
+        tad_block_queue_t temp_insig = insignificant_queue;
+        insignificant_queue = next_insignificant;
+        next_insignificant = temp_insig;
+        next_insignificant.count = 0;  // Clear for reuse
+
+        tad_block_queue_t temp_sig = significant_queue;
+        significant_queue = next_significant;
+        next_significant = temp_sig;
+        next_significant.count = 0;  // Clear for reuse
+    }
+
+    // Cleanup queues
+    tad_queue_free(&insignificant_queue);
+    tad_queue_free(&next_insignificant);
+    tad_queue_free(&significant_queue);
+    tad_queue_free(&next_significant);
+    free(states);
+
+    // Copy bitstream to output
+    size_t output_size = tad_bitstream_size(&bs);
+    *output = malloc(output_size);
+    memcpy(*output, bs.data, output_size);
+    tad_bitstream_free(&bs);
+
+    return output_size;
+}
+
 //=============================================================================
 // Public API: Chunk Encoding
 //=============================================================================
@@ -931,17 +1212,21 @@ size_t tad32_encode_chunk(const float *pcm32_stereo, size_t num_samples,
         accumulate_quantized(quant_side, dwt_levels, num_samples, side_quant_accumulators);
     }
 
-    // Step 5: Encode with twobit-map significance map or raw int8_t storage
-    uint8_t *temp_buffer = malloc(num_samples * 4);  // Generous buffer
-    size_t mid_size, side_size;
+    // Step 5: Encode with binary tree EZBC (1D variant)
+    uint8_t *mid_ezbc = NULL;
+    uint8_t *side_ezbc = NULL;
 
-    // Raw int8_t storage
-    memcpy(temp_buffer, quant_mid, num_samples);
-    mid_size = num_samples;
-    memcpy(temp_buffer + mid_size, quant_side, num_samples);
-    side_size = num_samples;
+    size_t mid_size = tad_encode_channel_ezbc(quant_mid, num_samples, &mid_ezbc);
+    size_t side_size = tad_encode_channel_ezbc(quant_side, num_samples, &side_ezbc);
 
+    // Concatenate EZBC outputs
     size_t uncompressed_size = mid_size + side_size;
+    uint8_t *temp_buffer = malloc(uncompressed_size);
+    memcpy(temp_buffer, mid_ezbc, mid_size);
+    memcpy(temp_buffer + mid_size, side_ezbc, side_size);
+
+    free(mid_ezbc);
+    free(side_ezbc);
 
     // Step 6: Optional Zstd compression
     uint8_t *write_ptr = output;