From b973939a1967bf4cc9391709eeb930043f4d661b Mon Sep 17 00:00:00 2001
From: Jean-Marc Valin <jeanmarcv@google.com>
Date: Tue, 10 Dec 2024 15:55:51 -0500
Subject: [PATCH] Fix e7e763fa: vary allocation bands

---
 celt/celt_decoder.c |  4 ++--
 celt/celt_encoder.c |  2 +-
 celt/rate.c         | 24 ++++++++++++------------
 celt/rate.h         |  2 +-
 4 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/celt/celt_decoder.c b/celt/celt_decoder.c
index f1d8b50c0..d6eb05b0f 100644
--- a/celt/celt_decoder.c
+++ b/celt/celt_decoder.c
@@ -1040,7 +1040,7 @@ int celt_decode_with_ec_dred(CELTDecoder * OPUS_RESTRICT st, const unsigned char
    opus_int32 qext_bits;
    ec_dec ext_dec;
    int qext_bytes=0;
-   int qext_end;
+   int qext_end=0;
    VARDECL(int, extra_quant);
    VARDECL(int, extra_pulses);
    const CELTMode *qext_mode = NULL;
@@ -1342,7 +1342,7 @@ int celt_decode_with_ec_dred(CELTDecoder * OPUS_RESTRICT st, const unsigned char
    ALLOC(extra_quant, nbEBands+NB_QEXT_BANDS, int);
    ALLOC(extra_pulses, nbEBands+NB_QEXT_BANDS, int);
    qext_bits = ((opus_int32)qext_bytes*8<<BITRES) - ec_tell_frac(dec) - 1;
-   clt_compute_extra_allocation(mode, qext_mode, start, end, NULL, NULL,
+   clt_compute_extra_allocation(mode, qext_mode, start, end, qext_end, NULL, NULL,
          qext_bits, extra_pulses, extra_quant, C, LM, &ext_dec, 0, 0, 0);
    if (qext_bytes > 0) {
       unquant_fine_energy(mode, start, end, oldBandE, fine_quant, extra_quant, &ext_dec, C);
diff --git a/celt/celt_encoder.c b/celt/celt_encoder.c
index 4f5a6d149..37a2dd862 100644
--- a/celt/celt_encoder.c
+++ b/celt/celt_encoder.c
@@ -2525,7 +2525,7 @@ int celt_encode_with_ec(CELTEncoder * OPUS_RESTRICT st, const opus_res * pcm, in
    ALLOC(error_bak, C*nbEBands, celt_glog);
 
    qext_bits = ((opus_int32)qext_bytes*8<<BITRES) - ec_tell_frac(enc) - 1;
-   clt_compute_extra_allocation(mode, qext_mode, start, end, bandLogE, qext_bandLogE,
+   clt_compute_extra_allocation(mode, qext_mode, start, end, qext_end, bandLogE, qext_bandLogE,
          qext_bits, extra_pulses, extra_quant, C, LM, &ext_enc, 1, tone_freq, toneishness);
    OPUS_COPY(error_bak, error, C*nbEBands);
    if (qext_bytes > 0) {
diff --git a/celt/rate.c b/celt/rate.c
index e420b1633..bc1b2ef75 100644
--- a/celt/rate.c
+++ b/celt/rate.c
@@ -687,7 +687,7 @@ static int ec_dec_depth(ec_dec *dec, opus_int32 cap, opus_int32 *last) {
    return depth;
 }
 
-void clt_compute_extra_allocation(const CELTMode *m, const CELTMode *qext_mode, int start, int end, const celt_glog *bandLogE, const celt_glog *qext_bandLogE,
+void clt_compute_extra_allocation(const CELTMode *m, const CELTMode *qext_mode, int start, int end, int qext_end, const celt_glog *bandLogE, const celt_glog *qext_bandLogE,
       opus_int32 total, int *extra_pulses, int *extra_equant, int C, int LM, ec_ctx *ec, int encode, opus_val16 tone_freq, opus_val32 toneishness)
 {
    int i;
@@ -702,8 +702,8 @@ void clt_compute_extra_allocation(const CELTMode *m, const CELTMode *qext_mode,
    if (qext_mode != NULL) {
       celt_assert(start==0);
       celt_assert(end==m->nbEBands);
-      tot_bands = end + NB_QEXT_BANDS;
-      tot_samples = qext_mode->eBands[NB_QEXT_BANDS]*C<<LM;
+      tot_bands = end + qext_end;
+      tot_samples = qext_mode->eBands[qext_end]*C<<LM;
    } else {
       tot_bands = end;
       tot_samples = (m->eBands[end]-m->eBands[start])*C<<LM;
@@ -711,10 +711,10 @@ void clt_compute_extra_allocation(const CELTMode *m, const CELTMode *qext_mode,
    ALLOC(cap, tot_bands, opus_int32);
    for (i=start;i<end;i++) cap[i] = 12;
    if (qext_mode != NULL) {
-      for (i=0;i<NB_QEXT_BANDS;i++) cap[end+i] = 14;
+      for (i=0;i<qext_end;i++) cap[end+i] = 14;
    }
    if (total <= 0) {
-      for (i=start;i<m->nbEBands+NB_QEXT_BANDS;i++) {
+      for (i=start;i<m->nbEBands+qext_end;i++) {
          extra_pulses[i] = extra_equant[i] = 0;
       }
       return;
@@ -745,21 +745,21 @@ void clt_compute_extra_allocation(const CELTMode *m, const CELTMode *qext_mode,
       if (qext_mode != NULL) {
          opus_val16 min_depth = 0;
          /* If we have enough bits, give at least 1 bit of depth to all higher bands. */
-         if (total >= 3*C*(qext_mode->eBands[NB_QEXT_BANDS]-qext_mode->eBands[start])<<LM<<BITRES && (toneishness < QCONST32(.98f, 29) || tone_freq > 1.33f))
+         if (total >= 3*C*(qext_mode->eBands[qext_end]-qext_mode->eBands[start])<<LM<<BITRES && (toneishness < QCONST32(.98f, 29) || tone_freq > 1.33f))
             min_depth = QCONST16(1.f, 10);
-         for (i=0;i<NB_QEXT_BANDS;i++) {
+         for (i=0;i<qext_end;i++) {
             Ncoef[end+i] = (qext_mode->eBands[i+1]-qext_mode->eBands[i])*C<<LM;
             min[end+i] = min_depth;
          }
-         for (i=0;i<NB_QEXT_BANDS;i++) {
+         for (i=0;i<qext_end;i++) {
             flatE[end+i] = PSHR32(qext_bandLogE[i] - GCONST(0.0625f)*qext_mode->logN[i] + SHL32(eMeans[i],DB_SHIFT-4) - GCONST(.0062f)*(end+i+5)*(end+i+5), DB_SHIFT-10);
          }
          if (C==2) {
-            for (i=0;i<NB_QEXT_BANDS;i++) {
-               flatE[end+i] = MAXG(flatE[end+i], PSHR32(qext_bandLogE[NB_QEXT_BANDS+i] - GCONST(0.0625f)*qext_mode->logN[i] + SHL32(eMeans[i],DB_SHIFT-4) - GCONST(.0062f)*(end+i+5)*(end+i+5), DB_SHIFT-10));
+            for (i=0;i<qext_end;i++) {
+               flatE[end+i] = MAXG(flatE[end+i], PSHR32(qext_bandLogE[qext_end+i] - GCONST(0.0625f)*qext_mode->logN[i] + SHL32(eMeans[i],DB_SHIFT-4) - GCONST(.0062f)*(end+i+5)*(end+i+5), DB_SHIFT-10));
             }
          }
-         for (i=0;i<NB_QEXT_BANDS;i++) flatE[end+i] = flatE[end+i] + QCONST16(3.f, 10) + QCONST16(.2f, 10)*i;
+         for (i=0;i<qext_end;i++) flatE[end+i] = flatE[end+i] + QCONST16(3.f, 10) + QCONST16(.2f, 10)*i;
       }
       /* Approximate fill level assuming all bands contribute fully. */
       sum = 0;
@@ -799,7 +799,7 @@ void clt_compute_extra_allocation(const CELTMode *m, const CELTMode *qext_mode,
       extra_pulses[i] = ((((m->eBands[i+1]-m->eBands[i])<<LM)-1)*C * depth[i] * (1<<BITRES) + 2)>>2;
    }
    if (qext_mode) {
-      for (i=0;i<NB_QEXT_BANDS;i++) {
+      for (i=0;i<qext_end;i++) {
          extra_equant[end+i] = (depth[end+i]+3)>>2;
          extra_pulses[end+i] = ((((qext_mode->eBands[i+1]-qext_mode->eBands[i])<<LM)-1)*C * depth[end+i] * (1<<BITRES) + 2)>>2;
       }
diff --git a/celt/rate.h b/celt/rate.h
index 09c8081a2..659c80f3a 100644
--- a/celt/rate.h
+++ b/celt/rate.h
@@ -98,7 +98,7 @@ static OPUS_INLINE int pulses2bits(const CELTMode *m, int band, int LM, int puls
 int clt_compute_allocation(const CELTMode *m, int start, int end, const int *offsets, const int *cap, int alloc_trim, int *intensity, int *dual_stereo,
       opus_int32 total, opus_int32 *balance, int *pulses, int *ebits, int *fine_priority, int C, int LM, ec_ctx *ec, int encode, int prev, int signalBandwidth);
 
-void clt_compute_extra_allocation(const CELTMode *m, const CELTMode *qext_mode, int start, int end, const celt_glog *bandLogE, const celt_glog *qext_bandLogE,
+void clt_compute_extra_allocation(const CELTMode *m, const CELTMode *qext_mode, int start, int end, int qext_end, const celt_glog *bandLogE, const celt_glog *qext_bandLogE,
       opus_int32 total, int *extra_pulses, int *extra_equant, int C, int LM, ec_ctx *ec, int encode, opus_val16 tone_freq, opus_val32 toneishness);
 
 #endif