1 ! Copyright (C) 2009, 2020 Marc Fauconneau, Abtin Molavi, and Jacob Fischer.
2 ! See http://factorcode.org/license.txt for BSD license.
3 USING: accessors arrays assocs bit-arrays bitstreams combinators
4 fry hashtables heaps io kernel locals math math.bits math.order
5 math.parser math.ranges multiline namespaces sequences sorting
7 QUALIFIED-WITH: bitstreams bs
8 IN: compression.huffman
20 : <huffman-code> ( -- huffman-code )
21 0 0 0 huffman-code boa ; inline
23 : next-size ( huffman-code -- )
25 [ 2 * ] change-code drop ; inline
27 : next-code ( huffman-code -- )
28 [ 1 + ] change-code drop ; inline
30 :: all-patterns ( huffman-code n -- seq )
31 n log2 huffman-code size>> - :> free-bits
33 [ free-bits 2^ <iota> [ huffman-code code>> free-bits 2^ * + ] map ]
34 [ huffman-code code>> free-bits neg 2^ /i 1array ] if ;
36 :: huffman-each ( ... tdesc quot: ( ... huffman-code -- ... ) -- ... )
37 <huffman-code> :> code
41 [ code value<< code clone quot call code next-code ] each
44 : update-reverse-table ( huffman-code n table -- )
46 [ nip '[ _ swap _ set-at ] each ] 3bi ;
48 :: reverse-table ( tdesc n -- rtable )
49 n f <array> <enumerated> :> table
50 tdesc [ n table update-reverse-table ] huffman-each
54 { code maybe{ fixnum } }
55 { left maybe{ huffman-tree } }
56 { right maybe{ huffman-tree } } ;
58 : <huffman-tree> ( code left right -- huffman-tree )
61 : <huffman-internal> ( left right -- huffman-tree )
62 huffman-tree new swap >>left swap >>right ;
64 : leaf? ( huff-tree -- ? )
65 [ left>> not ] [ right>> not ] bi and ;
67 : gen-leaves ( lit-seq -- leaves )
68 [ huffman-tree new swap >>code ] map ;
70 : build-leaf-table ( leaves -- )
71 dup empty? [ drop ] [ dup first leaf-table get inc-at rest build-leaf-table ] if ;
73 : insert-leaves ( -- ) leaf-table get unzip swap zip node-heap get heap-push-all ;
76 node-heap get heap-pop node-heap get heap-pop swap [ + ] dip pick <huffman-internal> swap node-heap get heap-push drop ;
78 : build-tree ( lit-seq -- heap )
79 gen-leaves build-leaf-table insert-leaves [ node-heap get heap-size 1 > ] [ combine-two ] while node-heap get ;
81 ! Walks down a huffman tree and outputs a dictionary of codes
82 : (generate-codes) ( huff-tree -- code-dict )
84 { [ dup leaf? ] [ code>> ?{ } swap H{ } clone ?set-at ] }
85 { [ dup left>> not ] [ right>> (generate-codes) [ ?{ t } prepend ] assoc-map ] }
86 { [ dup right>> not ] [ left>> (generate-codes) [ ?{ f } prepend ] assoc-map ] }
88 [ left>> (generate-codes) [ ?{ f } prepend ] assoc-map ]
89 [ right>> (generate-codes) [ ?{ t } prepend ] assoc-map ] bi assoc-union!
93 : generate-codes ( lit-seq -- code-dict )
96 [ H{ } clone leaf-table set
97 <min-heap> node-heap set
98 build-tree heap-pop swap (generate-codes) nip ]
102 ! Ordering of codes that is useful for generating canonical codes.
103 ! Sort by length, then lexicographically.
104 :: <==> ( b1 b2 -- <=> )
106 { [ b1 second length b2 second length < ] [ +lt+ ] }
107 { [ b2 second length b1 second length < ] [ +gt+ ] }
108 { [ b1 first b2 first < ] [ +lt+ ] }
109 { [ b2 first b1 first < ] [ +gt+ ] }
113 : sort-values! ( obj -- sortedseq )
114 >alist [ <==> ] sort ;
116 : get-next-code ( code current -- next )
117 [ reverse bit-array>integer 1 + ] [ length ] bi <bits> >bit-array reverse dup length pick length swap - [ f ] replicate append nip ;
119 ! Does most of the work of converting a collection of codes to canonical ones.
120 : (canonize-codes) ( current codes -- codes )
121 dup empty? [ 2drop V{ } clone ] [ dup first pick get-next-code dup pick 1 tail (canonize-codes) ?push 2nip ] if ;
123 ! Basically a wrapper for the above recursive helper
124 : canonize-codes ( codes -- codes )
125 [ V{ } clone ] [ dup first length <bit-array> dup pick 1 tail (canonize-codes) ?push nip reverse ] if-empty ;
127 :: length-limit-codes ( max-len old-codes -- new-codes )
128 old-codes [ length ] assoc-map [ dup length max-len < [ drop max-len ] when ] assoc-map ;
132 TUPLE: huffman-decoder
136 { bits/level fixnum } ;
138 : <huffman-decoder> ( bs tdesc -- huffman-decoder )
143 dup [ tdesc>> ] [ bits/level>> 2^ ] bi reverse-table >>rtable ; inline
145 : read1-huff ( huffman-decoder -- elt )
146 16 over [ bs>> bs:peek ] [ rtable>> nth ] bi
147 [ size>> swap bs>> bs:seek ] [ value>> ] bi ; inline
149 : reverse-bits ( value bits -- value' )
150 [ integer>bit-array ] dip
151 f pad-tail reverse bit-array>integer ; inline
153 : read1-huff2 ( huffman-decoder -- elt )
154 16 over [ bs>> bs:peek 16 reverse-bits ] [ rtable>> nth ] bi
155 [ size>> swap bs>> bs:seek ] [ value>> ] bi ; inline
157 ! Outputs a dictionary of canonical codes
158 : generate-canonical-codes ( lit-seq -- code-dict )
159 generate-codes sort-values! unzip canonize-codes zip ;