Use raw bits for cubic codebook

xiph · Jan 30, 2025 · a9d0953 · a9d0953
1 parent 01e2eb5
commit a9d0953
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 17 deletions.
diff --git a/celt/bands.c b/celt/bands.c
@@ -1165,8 +1165,8 @@ static unsigned quant_partition(struct band_ctx *ctx, celt_norm *X,
             extra_bits = IMAX(extra_bits-1, 0);
          }
          extra_bits = IMIN(14, extra_bits);
-         if (encode) cm = cubic_quant(X, N, 1<<extra_bits, B, ctx->ext_ec, gain, ctx->resynth);
-         else cm = cubic_unquant(X, N, 1<<extra_bits, B, ctx->ext_ec, gain);
+         if (encode) cm = cubic_quant(X, N, extra_bits, B, ctx->ext_ec, gain, ctx->resynth);
+         else cm = cubic_unquant(X, N, extra_bits, B, ctx->ext_ec, gain);
 #endif
       } else {
          /* If there's no pulse, fill the band anyway */
@@ -1216,18 +1216,20 @@ static unsigned quant_partition(struct band_ctx *ctx, celt_norm *X,
 #ifdef ENABLE_QEXT
 static unsigned cubic_quant_partition(struct band_ctx *ctx, celt_norm *X, int N, int b, int B, ec_ctx *ec, int LM, opus_val32 gain, int resynth, int encode)
 {
-   int res;
    celt_assert(LM>=0);
    ctx->remaining_bits = ctx->ec->storage*8*8 - ec_tell_frac(ctx->ec);
-   res = IMIN((b+N/2)/8/(N-1),  (ctx->remaining_bits-(ctx->m->logN[ctx->i]+8+8*LM))/8/(N-1));
-   res = IMIN(14, IMAX(0, res));
-   if (LM==0 || res<=2) {
-      int K;
-      K=1<<res;
-      if (B!=1) K=IMAX(1, K-1);
-      if (encode) return cubic_quant(X, N, K, B, ec, gain, resynth);
-      else return cubic_unquant(X, N, K, B, ec, gain);
-      /* Using odd K on transients to avoid adding pre-echo. */
+   b = IMIN(b, ctx->remaining_bits);
+   if (LM==0 || b<=2*N<<BITRES) {
+      int K, res, ret;
+      /* Resolution left after taking into account coding the cube face. */
+      res = (b-(1<<BITRES)-ctx->m->logN[ctx->i]-(LM<<BITRES)-1)/(N-1);
+      res = IMIN(14<<BITRES, IMAX(0, res));
+      int res2 = IMAX(0, (ctx->remaining_bits-(1<<BITRES)-ctx->m->logN[ctx->i]-(LM<<BITRES)-1)/(N-1));
+      res=IMIN(((res)>>BITRES), (IMIN(14*8, res2)>>BITRES));
+      if (encode) ret = cubic_quant(X, N, res, B, ec, gain, resynth);
+      else ret = cubic_unquant(X, N, res, B, ec, gain);
+      ctx->remaining_bits = ctx->ec->storage*8*8 - ec_tell_frac(ctx->ec);
+      return ret;
    } else {
       celt_norm *Y;
       opus_int32 itheta_q30;
@@ -1243,7 +1245,7 @@ static unsigned cubic_quant_partition(struct band_ctx *ctx, celt_norm *X, int N,
       Y = X+N;
       LM -= 1;
       B = (B+1)>>1;
-      theta_res = IMIN(16, res+1);
+      theta_res = IMIN(16, (b>>BITRES)/(N0-1) + 1);
       if (encode) {
          itheta_q30 = stereo_itheta(X, Y, 0, N, ctx->arch);
          qtheta = (itheta_q30+(1<<(29-theta_res)))>>(30-theta_res);

diff --git a/celt/vq.c b/celt/vq.c
@@ -768,15 +768,19 @@ static void cubic_synthesis(celt_norm *X, int *iy, int N, int K, int face, int s
 #endif
 }
 
-unsigned cubic_quant(celt_norm *X, int N, int K, int B, ec_enc *enc, opus_val32 gain, int resynth) {
+unsigned cubic_quant(celt_norm *X, int N, int res, int B, ec_enc *enc, opus_val32 gain, int resynth) {
    int i;
    int face=0;
+   int K;
    VARDECL(int, iy);
    celt_norm faceval=-1;
    float norm;
    int sign;
    SAVE_STACK;
    ALLOC(iy, N, int);
+   K = 1<<res;
+   /* Using odd K on transients to avoid adding pre-echo. */
+   if (B!=1) K=IMAX(1, K-1);
    if (K==1) {
       if (resynth) OPUS_CLEAR(X, N);
       return 0;
@@ -793,7 +797,7 @@ unsigned cubic_quant(celt_norm *X, int N, int K, int B, ec_enc *enc, opus_val32
    norm = .5f*K/(faceval+EPSILON);
    for (i=0;i<N;i++) {
       iy[i] = IMIN(K-1, (int)floor((X[i]+faceval)*norm));
-      if (i != face) ec_enc_uint(enc, iy[i], K);
+      if (i != face) ec_enc_bits(enc, iy[i], res);
    }
    if (resynth) {
       cubic_synthesis(X, iy, N, K, face, sign, gain);
@@ -802,21 +806,25 @@ unsigned cubic_quant(celt_norm *X, int N, int K, int B, ec_enc *enc, opus_val32
    return (1<<B)-1;
 }
 
-unsigned cubic_unquant(celt_norm *X, int N, int K, int B, ec_dec *dec, opus_val32 gain) {
+unsigned cubic_unquant(celt_norm *X, int N, int res, int B, ec_dec *dec, opus_val32 gain) {
    int i;
    int face;
    int sign;
+   int K;
    VARDECL(int, iy);
    SAVE_STACK;
    ALLOC(iy, N, int);
+   K = 1<<res;
+   /* Using odd K on transients to avoid adding pre-echo. */
+   if (B!=1) K=IMAX(1, K-1);
    if (K==1) {
       OPUS_CLEAR(X, N);
       return 0;
    }
    face = ec_dec_uint(dec, N);
    sign = ec_dec_bits(dec, 1);
    for (i=0;i<N;i++) {
-      if (i != face) iy[i] = ec_dec_uint(dec, K);
+      if (i != face) iy[i] = ec_dec_bits(dec, res);
    }
    iy[face]=0;
    cubic_synthesis(X, iy, N, K, face, sign, gain);