basis/compression/huffman/huffman.factor

   1 ! Copyright (C) 2009 Marc Fauconneau.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors arrays assocs fry
   4 hashtables io kernel locals math math.order math.parser
   5 math.ranges multiline sequences bitstreams bit-arrays ;
   6 IN: compression.huffman
   7
   8 QUALIFIED-WITH: bitstreams bs
   9
  10 <PRIVATE
  11
  12 TUPLE: huffman-code
  13     { value fixnum }
  14     { size fixnum }
  15     { code fixnum } ;
  16
  17 : <huffman-code> ( -- huffman-code )
  18     0 0 0 huffman-code boa ; inline
  19
  20 : next-size ( huffman-code -- )
  21     [ 1 + ] change-size
  22     [ 2 * ] change-code drop ; inline
  23
  24 : next-code ( huffman-code -- )
  25     [ 1 + ] change-code drop ; inline
  26
  27 :: all-patterns ( huffman-code n -- seq )
  28     n log2 huffman-code size>> - :> free-bits
  29     free-bits 0 >
  30     [ free-bits 2^ <iota> [ huffman-code code>> free-bits 2^ * + ] map ]
  31     [ huffman-code code>> free-bits neg 2^ /i 1array ] if ;
  32
  33 :: huffman-each ( ... tdesc quot: ( ... huffman-code -- ... ) -- ... )
  34     <huffman-code> :> code
  35     tdesc
  36     [
  37         code next-size
  38         [ code value<< code clone quot call code next-code ] each
  39     ] each ; inline
  40
  41 : update-reverse-table ( huffman-code n table -- )
  42     [ drop all-patterns ]
  43     [ nip '[ _ swap _ set-at ] each ] 3bi ;
  44
  45 :: reverse-table ( tdesc n -- rtable )
  46    n f <array> <enum> :> table
  47    tdesc [ n table update-reverse-table ] huffman-each
  48    table seq>> ;
  49
  50 PRIVATE>
  51
  52 TUPLE: huffman-decoder
  53     { bs bit-reader }
  54     { tdesc array }
  55     { rtable array }
  56     { bits/level fixnum } ;
  57
  58 : <huffman-decoder> ( bs tdesc -- huffman-decoder )
  59     huffman-decoder new
  60         swap >>tdesc
  61         swap >>bs
  62         16 >>bits/level
  63         dup [ tdesc>> ] [ bits/level>> 2^ ] bi reverse-table >>rtable ; inline
  64
  65 : read1-huff ( huffman-decoder -- elt )
  66     16 over [ bs>> bs:peek ] [ rtable>> nth ] bi
  67     [ size>> swap bs>> bs:seek ] [ value>> ] bi ; inline
  68
  69 : reverse-bits ( value bits -- value' )
  70     [ integer>bit-array ] dip
  71     f pad-tail reverse bit-array>integer ; inline
  72
  73 : read1-huff2 ( huffman-decoder -- elt )
  74     16 over [ bs>> bs:peek 16 reverse-bits ] [ rtable>> nth ] bi
  75     [ size>> swap bs>> bs:seek ] [ value>> ] bi ; inline