# Paste: huffman

Author: Ed factor Sat, 14 Feb 2009 01:31:34
Plain Text |
```USING: accessors arrays assocs bit-arrays
combinators.short-circuit lists fry kernel make math math.order
sequences sorting summary ;
IN: huffman-coding

TUPLE: huffman-code tree-root dictionary ;

<PRIVATE

TUPLE: huffman-leaf symbol probability ;
TUPLE: huffman-node probability left right ;

C: <huffman-leaf> huffman-leaf
C: <huffman-node> huffman-node

: combine-nodes ( a b -- c )
[ [ probability>> ] bi@ + ] 2keep
<huffman-node> ;

: (huffman-sorted-insert) ( list node -- )
2dup [ cadr probability>> ] [ probability>> ] bi* >
[ over cdr cons >>cdr drop ]
[ [ cdr ] dip (huffman-sorted-insert) ] if ;

: huffman-sorted-insert ( list node -- list )
2dup {
[ drop nil? ]
[ [ car probability>> ] [ probability>> ] bi* >= ]
} 2||
[ swap cons ] [ dupd (huffman-sorted-insert) ] if ; inline

: build-huffman-tree ( list -- tree )
dup cdr nil? not [
[ cdr cdr ] [ 2car ] bi
combine-nodes huffman-sorted-insert
build-huffman-tree
] [ car ] if ;

GENERIC# (build-dictionary) 1 ( internal-tree accum -- )

M: huffman-leaf (build-dictionary)
[ symbol>> ] dip >bit-array 2array , ;

M: huffman-node (build-dictionary)
[ [ left>> ] [ right>> ] bi ] dip
dup clone
[ dup t swap push ] [ dup f swap push ] bi*
swapd [ (build-dictionary) ] 2bi@ ;

: build-dictionary ( internal-tree -- dictionary )
[ V{ } (build-dictionary) ] { } make H{ } assoc-like ;

PRIVATE>

: <huffman-code> ( assoc -- huffman-code )
[ <huffman-leaf> ] { } assoc>map
[ [ probability>> ] bi@ <=> ] sort seq>list
build-huffman-tree dup build-dictionary
huffman-code boa ;

<PRIVATE

ERROR: illegal-huffman-symbol symbol ;

M: illegal-huffman-symbol summary
drop "You tried to encode a symbol which is not in the huffman-code" ;

PRIVATE>

: huffman-encode ( huffman-code seq -- code )
swap dictionary>> '[
dup _ at* [ nip ]
[ drop illegal-huffman-symbol ] if
] { } map-as concat ;

<PRIVATE

TUPLE: huffman-decoder root current ;

: <huffman-decoder> ( huffman-code -- huffman-decoder )
tree-root>> dup huffman-decoder boa ;

GENERIC: update-decoder ( decoder tree -- decoder )

M: huffman-node update-decoder
>>current ;

M: huffman-leaf update-decoder
symbol>> , dup root>> >>current ;

: decode-step ( huffman-decoder bit -- huffman-decoder )
[ dup current>> ] dip
[ left>> ] [ right>> ] if
update-decoder ;

PRIVATE>

: huffman-decode-as ( huffman-code code exemplar -- seq )
[
[ <huffman-decoder> ] dip
[ decode-step ]
[ each ]
] dip make nip ;

: huffman-decode ( huffman-code code -- array )
{ } huffman-decode-as ; inline

<PRIVATE

: get-statistics ( seq -- hash )
H{ } clone [ '[ _ inc-at ] each ] keep ;

PRIVATE>

: huffman-count-and-encode ( seq -- huffman-code code )
[ get-statistics <huffman-code> dup ] keep
huffman-encode ;```