1 files changed, 25 insertions, 8 deletions

diff --git a/libavcodec/aacenc_tns.c b/libavcodec/aacenc_tns.c
index da353d65c5..8e0a00fc78 100644
--- a/libavcodec/aacenc_tns.c
+++ b/libavcodec/aacenc_tns.c
@@ -31,17 +31,34 @@
 #include "aacenc_utils.h"
 #include "aacenc_quantization.h"
 
+/*
+ * Shifts the values as well if compression is possible.
+ */
+static inline int compress_coeffs(int *coef, int order, int c_bits)
+{
+    int i, res = 0;
+    const int low_idx   = c_bits ?  4 : 2;
+    const int shift_val = c_bits ?  8 : 4;
+    const int high_idx  = c_bits ? 11 : 5;
+    for (i = 0; i < order; i++)
+        if (coef[i] < low_idx && coef[i] > high_idx)
+            res++;
+    if (res == order)
+        for (i = 0; i < order; i++)
+            coef[i] -= (coef[i] > high_idx) ? shift_val : 0;
+    return res == order;
+}
+
 /**
  * Encode TNS data.
  * Coefficient compression saves a single bit per coefficient.
  */
 void ff_aac_encode_tns_info(AACEncContext *s, SingleChannelElement *sce)
 {
-    uint8_t u_coef;
-    const uint8_t coef_res = TNS_Q_BITS == 4;
     int i, w, filt, coef_len, coef_compress = 0;
     const int is8 = sce->ics.window_sequence[0] == EIGHT_SHORT_SEQUENCE;
     TemporalNoiseShaping *tns = &sce->tns;
+    const int c_bits = is8 ? TNS_Q_BITS_SHORT == 4 : TNS_Q_BITS == 4;
 
     if (!sce->tns.present)
         return;
@@ -49,18 +66,18 @@ void ff_aac_encode_tns_info(AACEncContext *s, SingleChannelElement *sce)
     for (i = 0; i < sce->ics.num_windows; i++) {
         put_bits(&s->pb, 2 - is8, sce->tns.n_filt[i]);
         if (tns->n_filt[i]) {
-            put_bits(&s->pb, 1, coef_res);
+            put_bits(&s->pb, 1, c_bits);
             for (filt = 0; filt < tns->n_filt[i]; filt++) {
                 put_bits(&s->pb, 6 - 2 * is8, tns->length[i][filt]);
                 put_bits(&s->pb, 5 - 2 * is8, tns->order[i][filt]);
                 if (tns->order[i][filt]) {
+                    coef_compress = compress_coeffs(tns->coef_idx[i][filt],
+                                                    tns->order[i][filt], c_bits);
                     put_bits(&s->pb, 1, !!tns->direction[i][filt]);
                     put_bits(&s->pb, 1, !!coef_compress);
-                    coef_len = coef_res + 3 - coef_compress;
-                    for (w = 0; w < tns->order[i][filt]; w++) {
-                        u_coef = (tns->coef_idx[i][filt][w])&(~(~0<<coef_len));
-                        put_bits(&s->pb, coef_len, u_coef);
-                    }
+                    coef_len = c_bits + 3 - coef_compress;
+                    for (w = 0; w < tns->order[i][filt]; w++)
+                        put_bits(&s->pb, coef_len, tns->coef_idx[i][filt][w]);
                 }
             }
         }