From 58e4346410587d19e41c11c88585807eeb693510 Mon Sep 17 00:00:00 2001
From: Koji Arai <jca02266@gmail.com>
Date: Wed, 27 Aug 2008 21:54:01 +0900
Subject: [PATCH] should check the Kraft's inequality for tree

---
 src/pm2.c     |  2 +-
 src/pm2tree.c | 51 ++++++++++++++++++++++++++++++++++++++++-----------
 src/pm2tree.h |  2 +-
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/pm2.c b/src/pm2.c
index be74f5a..070ce63 100644
--- a/src/pm2.c
+++ b/src/pm2.c
@@ -41,7 +41,7 @@ decode_c_pm2(void)
     }
 
     while (decode_count >= nextcount) {
-        /* Actually it will never loop, because count doesn't grow that fast.
+        /* Actually it will never loop, because decode_count doesn't grow that fast.
            However, this is the way LHA does it.
            Probably other encoding methods can have repeats larger than 256 bytes.
            Note: LHA puts this code in decode_p...
diff --git a/src/pm2tree.c b/src/pm2tree.c
index cae183a..6d66c37 100644
--- a/src/pm2tree.c
+++ b/src/pm2tree.c
@@ -35,7 +35,7 @@ maketree1()
             x = getbits(nbits);
             table1[i] = (x == 0 ? 0 : x - 1 + mindepth);
         }
-        tree_rebuild(&tree1, tree1bound, mindepth, table1);
+        tree_rebuild(&tree1, tree1bound, mindepth, 31, table1);
     }
 }
 
@@ -69,7 +69,7 @@ maketree2(int par_b) /* in use: 5 <= par_b <= 8 */
     }
     else if (count > 1) {
         mindepth = 1;
-        tree_rebuild(&tree2, 8, mindepth, table2);
+        tree_rebuild(&tree2, 8, mindepth, 7, table2);
     }
     // Note: count == 0 is possible!
     //       Excluding that possibility was a bug in version 1.
@@ -97,12 +97,42 @@ void
 tree_rebuild(struct tree *t,
              unsigned char bound,
              unsigned char mindepth,
+             unsigned char maxdepth,
              unsigned char *table)
 {
-    unsigned char *parentarr, d;
+    unsigned char parentarr[32], d;
     int i, curr, empty, n;
 
-    parentarr = (unsigned char *)xmalloc(bound);
+    /* validate table */
+    {
+        unsigned int count[32];
+        double total;
+
+        memset(count, 0, sizeof(count));
+        for (i = 0; i < bound; i++) {
+            if (table[i] > maxdepth) {
+                error("Bad table");
+                exit(1);
+            }
+            count[table[i]]++;
+        }
+        total = 0.0;
+        for (i = 1; i <= maxdepth; i++) {
+            int max_leaves = (1<<i);
+            if (count[i] > max_leaves) {
+                error("Bad table");
+                exit(1);
+            }
+            total += 1.0/max_leaves * count[i];
+        }
+        if (total != 1.0) {
+            /* check the Kraft's inequality */
+            error("Bad table");
+            exit(1);
+        }
+    }
+
+    /* initialize tree */
     t->root = 0;
     for (i = 0; i < bound; i++) {
         t->leftarr[i] = 0;
@@ -110,6 +140,7 @@ tree_rebuild(struct tree *t,
         parentarr[i] = 0;
     }
 
+    /* build tree */
     for (i = 0; i < mindepth - 1; i++) {
         t->leftarr[i] = i + 1;
         parentarr[i + 1] = i;
@@ -117,7 +148,7 @@ tree_rebuild(struct tree *t,
 
     curr = mindepth - 1;
     empty = mindepth;
-    for (d = mindepth; TRUE; d++) {
+    for (d = mindepth; d <= maxdepth; d++) {
         for (i = 0; i < bound; i++) {
             if (table[i] != d)
                 continue;
@@ -131,7 +162,6 @@ tree_rebuild(struct tree *t,
             n = 0;
             while (t->rightarr[curr] != 0) {
                 if (curr == 0) {        /* root? -> done */
-                    free(parentarr);
                     return;
                 }
                 curr = parentarr[curr];
@@ -156,13 +186,12 @@ tree_rebuild(struct tree *t,
         else
             t->rightarr[curr] = empty;
 
-        if (empty >= bound) {
-            error("bad archive");
-            exit(1);
-        }
-
         parentarr[empty] = curr;
         curr = empty;
         empty++;
     }
+
+    /* unreachable */
+    error("bad table");
+    exit(1);
 }
diff --git a/src/pm2tree.h b/src/pm2tree.h
index afd1b1d..3f68e3f 100644
--- a/src/pm2tree.h
+++ b/src/pm2tree.h
@@ -13,4 +13,4 @@ void maketree2(int par_b);
 int tree_get(struct tree *t);
 void tree_setsingle(struct tree *t, unsigned char value);
 void tree_rebuild(struct tree *t, unsigned char bound,
-                  unsigned char mindepth, unsigned char *table);
+                  unsigned char mindepth, unsigned char maxdepth, unsigned char *table);
-- 
2.11.0