fix corner cases of symbol table construction logic

diff --git a/src/tANS.c b/src/tANS.c
index c7db81c63545c5fd08b06d251ddc5d1cb28ab736..00a0999de137d6258fc917df854a41caf00b1acd 100644 (file)
--- a/src/tANS.c
+++ b/src/tANS.c
@@ -103,12 +103,14 @@ static int tANS_init_symbol_tbls(struct tANS_symbol_tbl symbol_tbls[static 3], c
                     tbls[j].freq[i] = 1;
             for (i = 0, sum = 0, zeroes = 0; i < N_SYMBOLS; ++i)
                 sum += tbls[j].freq[i];
-            for (i = 0; i < N_SYMBOLS && sum < 1 << LOG2_TBLSZ; ++i)
-                if (!(j == 2 && i == 0))
-                    ++tbls[j].freq[i], ++sum;
-            for (i = 0; i < N_SYMBOLS && sum > 1 << LOG2_TBLSZ; ++i)
-                if (tbls[j].freq[i] > 1)
-                    --tbls[j].freq[i], --sum;
+            while (sum < 1 << LOG2_TBLSZ)
+                for (i = 0; i < N_SYMBOLS && sum < 1 << LOG2_TBLSZ; ++i)
+                    if (!(j == 2 && i == 0))
+                        ++tbls[j].freq[i], ++sum;
+            while (sum > 1 << LOG2_TBLSZ)
+                for (i = 0; i < N_SYMBOLS && sum > 1 << LOG2_TBLSZ; ++i)
+                    if (tbls[j].freq[i] > 1)
+                        --tbls[j].freq[i], --sum;
         }
     }