basis/compression/huffman/huffman.factor

   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 ranges multiline namespaces sequences sorting
   6 vectors ;
   7 QUALIFIED-WITH: bitstreams bs
   8 IN: compression.huffman
   9
  10 <PRIVATE
  11
  12 SYMBOL: leaf-table
  13 SYMBOL: node-heap
  14
  15 TUPLE: huffman-code
  16     { value fixnum }
  17     { size fixnum }
  18     { code fixnum } ;
  19
  20 : <huffman-code> ( -- huffman-code )
  21     0 0 0 huffman-code boa ; inline
  22
  23 : next-size ( huffman-code -- )
  24     [ 1 + ] change-size
  25     [ 2 * ] change-code drop ; inline
  26
  27 : next-code ( huffman-code -- )
  28     [ 1 + ] change-code drop ; inline
  29
  30 :: all-patterns ( huffman-code n -- seq )
  31     n log2 huffman-code size>> - :> free-bits
  32     free-bits 0 >
  33     [ free-bits 2^ <iota> [ huffman-code code>> free-bits 2^ * + ] map ]
  34     [ huffman-code code>> free-bits neg 2^ /i 1array ] if ;
  35
  36 :: huffman-each ( ... tdesc quot: ( ... huffman-code -- ... ) -- ... )
  37     <huffman-code> :> code
  38     tdesc
  39     [
  40         code next-size
  41         [ code value<< code clone quot call code next-code ] each
  42     ] each ; inline
  43
  44 : update-reverse-table ( huffman-code n table -- )
  45     [ drop all-patterns ]
  46     [ nip '[ _ swap _ set-at ] each ] 3bi ;
  47
  48 :: reverse-table ( tdesc n -- rtable )
  49    n f <array> <enumerated> :> table
  50    tdesc [ n table update-reverse-table ] huffman-each
  51    table seq>> ;
  52
  53 TUPLE: huffman-tree
  54     { code maybe{ fixnum } }
  55     { left maybe{ huffman-tree } }
  56     { right maybe{ huffman-tree } } ;
  57
  58 : <huffman-tree> ( code left right -- huffman-tree )
  59     huffman-tree boa ;
  60
  61 : <huffman-internal> ( left right -- huffman-tree )
  62     huffman-tree new swap >>left swap >>right ;
  63
  64 : leaf? ( huff-tree -- ? )
  65     [ left>> not ] [ right>> not ] bi and ;
  66
  67 : gen-leaves ( lit-seq -- leaves )
  68     [ huffman-tree new swap >>code ] map ;
  69
  70 : build-leaf-table ( leaves --  )
  71     dup empty? [ drop ] [ dup first leaf-table get inc-at rest build-leaf-table ] if ;
  72
  73 : insert-leaves ( -- ) leaf-table get unzip swap zip node-heap get heap-push-all  ;
  74
  75 : combine-two ( -- )
  76     node-heap get heap-pop node-heap get heap-pop swap [ + ] dip pick <huffman-internal> swap node-heap get heap-push drop ;
  77
  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 ;
  80
  81 ! Walks down a huffman tree and outputs a dictionary of codes
  82 : (generate-codes) ( huff-tree -- code-dict )
  83     {
  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 ] }
  87           [
  88             [ left>> (generate-codes) [ ?{ f } prepend ] assoc-map ]
  89             [ right>> (generate-codes) [ ?{ t } prepend ] assoc-map ] bi assoc-union!
  90           ]
  91      } cond ;
  92
  93 : generate-codes ( lit-seq -- code-dict )
  94     [
  95        [ H{ } clone ]
  96        [ H{ } clone leaf-table set
  97         <min-heap> node-heap set
  98         build-tree heap-pop swap (generate-codes) nip ]
  99         if-empty
 100     ] with-scope ;
 101
 102 ! Ordering of codes that is useful for generating canonical codes.
 103 ! Sort by length, then lexicographically.
 104 :: <==> ( b1 b2  -- <=> )
 105     {
 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+ ] }
 110       [ +eq+ ]
 111     } cond ;
 112
 113 : sort-values! ( obj -- sortedseq )
 114     >alist [ <==> ] sort ;
 115
 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 ;
 118
 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 ;
 122
 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 ;
 126
 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 ;
 129
 130 PRIVATE>
 131
 132 TUPLE: huffman-decoder
 133     { bs bit-reader }
 134     { tdesc array }
 135     { rtable array }
 136     { bits/level fixnum } ;
 137
 138 : <huffman-decoder> ( bs tdesc -- huffman-decoder )
 139     huffman-decoder new
 140         swap >>tdesc
 141         swap >>bs
 142         16 >>bits/level
 143         dup [ tdesc>> ] [ bits/level>> 2^ ] bi reverse-table >>rtable ; inline
 144
 145 : read1-huff ( huffman-decoder -- elt )
 146     16 over [ bs>> bs:peek ] [ rtable>> nth ] bi
 147     [ size>> swap bs>> bs:seek ] [ value>> ] bi ; inline
 148
 149 : reverse-bits ( value bits -- value' )
 150     [ integer>bit-array ] dip
 151     f pad-tail reverse bit-array>integer ; inline
 152
 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
 156
 157 ! Outputs a dictionary of canonical codes
 158 : generate-canonical-codes ( lit-seq -- code-dict )
 159     generate-codes sort-values! unzip canonize-codes zip ;