Propagate balance from compute_allocation() to quant_all_bands().

Instead of just dumping excess bits into the first band after
 allocation, use them to initialize the rebalancing loop in
 quant_all_bands().
This allows these bits to be redistributed over several bands, like
 normal.

5 files changed, 18 insertions, 19 deletions

diff --git a/libcelt/bands.c b/libcelt/bands.c
index 929eab9..3888978 100644
--- a/libcelt/bands.c
+++ b/libcelt/bands.c
@@ -1139,10 +1139,9 @@ static unsigned quant_band(int encode, const CELTMode *m, int i, celt_norm *X, c
 void quant_all_bands(int encode, const CELTMode *m, int start, int end,
       celt_norm *_X, celt_norm *_Y, unsigned char *collapse_masks, const celt_ener *bandE, int *pulses,
       int shortBlocks, int spread, int dual_stereo, int intensity, int *tf_res, int resynth,
-      celt_int32 total_bits, void *ec, int LM, int codedBands, ec_uint32 *seed)
+      celt_int32 total_bits, celt_int32 balance, void *ec, int LM, int codedBands, ec_uint32 *seed)
 {
    int i;
-   celt_int32 balance;
    celt_int32 remaining_bits;
    const celt_int16 * restrict eBands = m->eBands;
    celt_norm * restrict norm, * restrict norm2;
@@ -1162,7 +1161,6 @@ void quant_all_bands(int encode, const CELTMode *m, int start, int end,
    norm = _norm;
    norm2 = norm + M*eBands[m->nbEBands];
 
-   balance = 0;
    lowband_offset = 0;
    for (i=start;i<end;i++)
    {
diff --git a/libcelt/bands.h b/libcelt/bands.h
index d1cfe6e..3e6bf74 100644
--- a/libcelt/bands.h
+++ b/libcelt/bands.h
@@ -88,7 +88,7 @@ void haar1(celt_norm *X, int N0, int stride);
 void quant_all_bands(int encode, const CELTMode *m, int start, int end,
       celt_norm * X, celt_norm * Y, unsigned char *collapse_masks, const celt_ener *bandE, int *pulses,
       int time_domain, int fold, int dual_stereo, int intensity, int *tf_res, int resynth,
-      celt_int32 total_bits, void *enc, int M, int codedBands, ec_uint32 *seed);
+      celt_int32 total_bits, celt_int32 balance, void *enc, int M, int codedBands, ec_uint32 *seed);
 
 
 void stereo_decision(const CELTMode *m, celt_norm * restrict X, int *stereo_mode, int len, int M);
diff --git a/libcelt/celt.c b/libcelt/celt.c
index 053b180..c6c8124 100644
--- a/libcelt/celt.c
+++ b/libcelt/celt.c
@@ -890,6 +890,7 @@ int celt_encode_with_ec_float(CELTEncoder * restrict st, const celt_sig * pcm, i
    celt_int32 vbr_rate;
    celt_int32 total_bits;
    celt_int32 total_boost;
+   celt_int32 balance;
    celt_int32 tell;
    int prefilter_tapset=0;
    int pf_on;
@@ -1398,8 +1399,8 @@ int celt_encode_with_ec_float(CELTEncoder * restrict st, const celt_sig * pcm, i
    anti_collapse_rsv = isTransient&&LM>=2&&bits>=(LM+2<<BITRES) ? (1<<BITRES) : 0;
    bits -= anti_collapse_rsv;
    codedBands = compute_allocation(st->mode, st->start, st->end, offsets, cap,
-         alloc_trim, &intensity, &dual_stereo, bits, pulses, fine_quant,
-         fine_priority, C, LM, enc, 1, st->lastCodedBands);
+         alloc_trim, &intensity, &dual_stereo, bits, &balance, pulses,
+         fine_quant, fine_priority, C, LM, enc, 1, st->lastCodedBands);
    st->lastCodedBands = codedBands;
 
    quant_fine_energy(st->mode, st->start, st->end, oldBandE, error, fine_quant, enc, C);
@@ -1419,7 +1420,7 @@ int celt_encode_with_ec_float(CELTEncoder * restrict st, const celt_sig * pcm, i
    ALLOC(collapse_masks, st->mode->nbEBands, unsigned char);
    quant_all_bands(1, st->mode, st->start, st->end, X, C==2 ? X+N : NULL, collapse_masks,
          bandE, pulses, shortBlocks, st->spread_decision, dual_stereo, intensity, tf_res, resynth,
-         nbCompressedBytes*(8<<BITRES)-anti_collapse_rsv, enc, LM, codedBands, &st->rng);
+         nbCompressedBytes*(8<<BITRES)-anti_collapse_rsv, balance, enc, LM, codedBands, &st->rng);
 
    if (anti_collapse_rsv > 0)
    {
@@ -2125,6 +2126,7 @@ int celt_decode_with_ec_float(CELTDecoder * restrict st, const unsigned char *da
    int intensity=0;
    int dual_stereo=0;
    celt_int32 total_bits;
+   celt_int32 balance;
    celt_int32 tell;
    int dynalloc_logp;
    int postfilter_tapset;
@@ -2319,8 +2321,8 @@ int celt_decode_with_ec_float(CELTDecoder * restrict st, const unsigned char *da
    anti_collapse_rsv = isTransient&&LM>=2&&bits>=(LM+2<<BITRES) ? (1<<BITRES) : 0;
    bits -= anti_collapse_rsv;
    codedBands = compute_allocation(st->mode, st->start, st->end, offsets, cap,
-         alloc_trim, &intensity, &dual_stereo, bits, pulses, fine_quant,
-         fine_priority, C, LM, dec, 0, 0);
+         alloc_trim, &intensity, &dual_stereo, bits, &balance, pulses,
+         fine_quant, fine_priority, C, LM, dec, 0, 0);
    
    unquant_fine_energy(st->mode, st->start, st->end, oldBandE, fine_quant, dec, C);
 
@@ -2328,7 +2330,7 @@ int celt_decode_with_ec_float(CELTDecoder * restrict st, const unsigned char *da
    ALLOC(collapse_masks, st->mode->nbEBands, unsigned char);
    quant_all_bands(0, st->mode, st->start, st->end, X, C==2 ? X+N : NULL, collapse_masks,
          NULL, pulses, shortBlocks, spread_decision, dual_stereo, intensity, tf_res, 1,
-         len*(8<<BITRES)-anti_collapse_rsv, dec, LM, codedBands, &st->rng);
+         len*(8<<BITRES)-anti_collapse_rsv, balance, dec, LM, codedBands, &st->rng);
 
    if (anti_collapse_rsv > 0)
    {
diff --git a/libcelt/rate.c b/libcelt/rate.c
index 6023459..68845d2 100644
--- a/libcelt/rate.c
+++ b/libcelt/rate.c
@@ -254,8 +254,8 @@ void compute_pulse_cache(CELTMode *m, int LM)
 #define ALLOC_STEPS 6
 
 static inline int interp_bits2pulses(const CELTMode *m, int start, int end, int skip_start,
-      const int *bits1, const int *bits2, const int *thresh, const int *cap, int total, int skip_rsv,
-      int *intensity, int intensity_rsv, int *dual_stereo, int dual_stereo_rsv, int *bits,
+      const int *bits1, const int *bits2, const int *thresh, const int *cap, int total, celt_int32 *_balance,
+      int skip_rsv, int *intensity, int intensity_rsv, int *dual_stereo, int dual_stereo_rsv, int *bits,
       int *ebits, int *fine_priority, int _C, int LM, void *ec, int encode, int prev)
 {
    int psum;
@@ -512,10 +512,9 @@ static inline int interp_bits2pulses(const CELTMode *m, int start, int end, int
       celt_assert(bits[j] >= 0);
       celt_assert(ebits[j] >= 0);
    }
-   /* Sweep any bits over the caps into the first band.
-      They'll be reallocated by the normal rebalancing code, which gives
-       them the best chance to be used _somewhere_. */
-   bits[start]+=balance;
+   /* Save any remaining bits over the cap for the rebalancing in
+       quant_all_bands(). */
+   *_balance = balance;
 
    /* The skipped bands use all their bits for fine energy. */
    for (;j<end;j++)
@@ -530,7 +529,7 @@ static inline int interp_bits2pulses(const CELTMode *m, int start, int end, int
 }
 
 int compute_allocation(const CELTMode *m, int start, int end, const int *offsets, const int *cap, int alloc_trim, int *intensity, int *dual_stereo,
-      int total, int *pulses, int *ebits, int *fine_priority, int _C, int LM, void *ec, int encode, int prev)
+      int total, celt_int32 *balance, int *pulses, int *ebits, int *fine_priority, int _C, int LM, void *ec, int encode, int prev)
 {
    int lo, hi, len, j;
    const int C = CHANNELS(_C);
@@ -632,7 +631,7 @@ int compute_allocation(const CELTMode *m, int start, int end, const int *offsets
       bits2[j] -= bits1[j];
    }
    codedBands = interp_bits2pulses(m, start, end, skip_start, bits1, bits2, thresh, cap,
-         total, skip_rsv, intensity, intensity_rsv, dual_stereo, dual_stereo_rsv,
+         total, balance, skip_rsv, intensity, intensity_rsv, dual_stereo, dual_stereo_rsv,
          pulses, ebits, fine_priority, C, LM, ec, encode, prev);
    RESTORE_STACK;
    return codedBands;
diff --git a/libcelt/rate.h b/libcelt/rate.h
index 617e852..eabba93 100644
--- a/libcelt/rate.h
+++ b/libcelt/rate.h
@@ -106,7 +106,7 @@ celt_int16 **compute_alloc_cache(CELTMode *m, int M);
  @return Total number of bits allocated
 */
 int compute_allocation(const CELTMode *m, int start, int end, const int *offsets, const int *cap, int alloc_trim, int *intensity, int *dual_stero,
-      int total, int *pulses, int *ebits, int *fine_priority, int _C, int LM, void *ec, int encode, int prev);
+      int total, celt_int32 *balance, int *pulses, int *ebits, int *fine_priority, int _C, int LM, void *ec, int encode, int prev);
 
 
 #endif