aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--src/huffman/huffmanencoder.nim10
-rw-r--r--tests/thuffman.nim12
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
19import huffmantree 19import huffmantree
20 20
21type HuffmanEncoder*[T, U: SomeUnsignedInt] = object 21type HuffmanEncoder*[T, U: SomeUnsignedInt] = object
22 codebook: TableRef[T, U] 22 codebook: TableRef[T, tuple[bitLength: int, value: U]]
23 23
24proc buildCodebook*[T, U](tree: HuffmanTreeNode[T], codeType: typedesc[U]): TableRef[T, U] = 24proc buildCodebook*[T, U](tree: HuffmanTreeNode[T], codeType: typedesc[U]): TableRef[T, tuple[bitLength: int, value: U]] =
25 var codebook = newTable[T, U]() 25 var codebook = newTable[T, tuple[bitLength: int, value: U]]()
26 proc addCode(node: HuffmanTreeNode[T], path: U, depth: int) = 26 proc addCode(node: HuffmanTreeNode[T], path: U, depth: int) =
27 case node.kind: 27 case node.kind:
28 of branch: 28 of branch:
29 addCode(node.left, path, depth + 1) 29 addCode(node.left, path, depth + 1)
30 addCode(node.right, path or (1.U shl depth), depth + 1) 30 addCode(node.right, path or (1.U shl depth), depth + 1)
31 of leaf: 31 of leaf:
32 codebook[node.value] = path 32 codebook[node.value] = (depth, path)
33 addCode(tree, 0.U, 0) 33 addCode(tree, 0.U, 0)
34 codebook 34 codebook
35 35
36proc encoder*[T, U](tree: HuffmanTreeNode[T], codeType: typedesc[U]): HuffmanEncoder[T, U] = 36proc encoder*[T, U](tree: HuffmanTreeNode[T], codeType: typedesc[U]): HuffmanEncoder[T, U] =
37 HuffmanEncoder[T, U](codebook: buildCodebook(tree, codeType)) 37 HuffmanEncoder[T, U](codebook: buildCodebook(tree, codeType))
38 38
39proc encode*[T, U](decoder: HuffmanEncoder[T, U], value: T): U = 39proc encode*[T, U](decoder: HuffmanEncoder[T, U], value: T): tuple[bitLength: int, value: U] =
40 decoder.codebook[value] 40 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":
90 test "buildCodebook": 90 test "buildCodebook":
91 let codebook = buildCodebook(tree, uint) 91 let codebook = buildCodebook(tree, uint)
92 check codebook.len == 3 92 check codebook.len == 3
93 check codebook[1'u] == 0b0 93 check codebook[1'u] == (1, 0b0'u)
94 check codebook[2'u] == 0b01 94 check codebook[2'u] == (2, 0b01'u)
95 check codebook[3'u] == 0b11 95 check codebook[3'u] == (2, 0b11'u)
96 96
97 test "encode": 97 test "encode":
98 let encoder = tree.encoder(uint) 98 let encoder = tree.encoder(uint)
99 check encoder.encode(1'u) == 0b0 99 check encoder.encode(1'u) == (1, 0b0'u)
100 check encoder.encode(2'u) == 0b01 100 check encoder.encode(2'u) == (2, 0b01'u)
101 check encoder.encode(3'u) == 0b11 101 check encoder.encode(3'u) == (2, 0b11'u)
102 102
103suite "huffdecoder": 103suite "huffdecoder":
104 test "decode": 104 test "decode":