2 files changed, 11 insertions, 11 deletions

diff --git a/src/huffman/huffmanencoder.nim b/src/huffman/huffmanencoder.nim
index 5ae1a68..09c6701 100644
--- a/src/huffman/huffmanencoder.nim
+++ b/src/huffman/huffmanencoder.nim
@@ -19,22 +19,22 @@ import ../bitio/integers, ../bitio/bitwriter
 import huffmantree
 
 type HuffmanEncoder*[T, U: SomeUnsignedInt] = object
-  codebook: TableRef[T, U]
+  codebook: TableRef[T, tuple[bitLength: int, value: U]]
 
-proc buildCodebook*[T, U](tree: HuffmanTreeNode[T], codeType: typedesc[U]): TableRef[T, U] =
-  var codebook = newTable[T, U]()
+proc buildCodebook*[T, U](tree: HuffmanTreeNode[T], codeType: typedesc[U]): TableRef[T, tuple[bitLength: int, value: U]] =
+  var codebook = newTable[T, tuple[bitLength: int, value: U]]()
   proc addCode(node: HuffmanTreeNode[T], path: U, depth: int) =
     case node.kind:
       of branch:
         addCode(node.left, path, depth + 1)
         addCode(node.right, path or (1.U shl depth), depth + 1)
       of leaf:
-        codebook[node.value] = path
+        codebook[node.value] = (depth, path)
   addCode(tree, 0.U, 0)
   codebook
 
 proc encoder*[T, U](tree: HuffmanTreeNode[T], codeType: typedesc[U]): HuffmanEncoder[T, U] =
   HuffmanEncoder[T, U](codebook: buildCodebook(tree, codeType))
 
-proc encode*[T, U](decoder: HuffmanEncoder[T, U], value: T): U =
+proc encode*[T, U](decoder: HuffmanEncoder[T, U], value: T): tuple[bitLength: int, value: U] =
   decoder.codebook[value]
diff --git a/tests/thuffman.nim b/tests/thuffman.nim
index 0294694..db209f4 100644
--- a/tests/thuffman.nim
+++ b/tests/thuffman.nim
@@ -90,15 +90,15 @@ suite "huffencoder":
   test "buildCodebook":
     let codebook = buildCodebook(tree, uint)
     check codebook.len == 3
-    check codebook[1'u] == 0b0
-    check codebook[2'u] == 0b01
-    check codebook[3'u] == 0b11
+    check codebook[1'u] == (1, 0b0'u)
+    check codebook[2'u] == (2, 0b01'u)
+    check codebook[3'u] == (2, 0b11'u)
 
   test "encode":
     let encoder = tree.encoder(uint)
-    check encoder.encode(1'u) == 0b0
-    check encoder.encode(2'u) == 0b01
-    check encoder.encode(3'u) == 0b11
+    check encoder.encode(1'u) == (1, 0b0'u)
+    check encoder.encode(2'u) == (2, 0b01'u)
+    check encoder.encode(3'u) == (2, 0b11'u)
 
 suite "huffdecoder":
   test "decode":