basis/compression/inflate/inflate.factor

   1 ! Copyright (C) 2009 Marc Fauconneau.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors arrays assocs byte-vectors combinators
   4 compression.huffman fry hashtables io.binary kernel locals math
   5 math.bitwise math.order math.ranges sequences sorting ;
   6 QUALIFIED-WITH: bitstreams bs
   7 IN: compression.inflate
   8
   9 QUALIFIED-WITH: bitstreams bs
  10
  11 <PRIVATE
  12
  13 : enum>seq ( assoc -- seq )
  14     dup keys [ ] [ max ] map-reduce 1 + f <array>
  15     [ '[ swap _ set-nth ] assoc-each ] keep ;
  16
  17 ERROR: zlib-unimplemented ;
  18 ERROR: bad-zlib-data ;
  19 ERROR: bad-zlib-header ;
  20
  21 :: check-zlib-header ( data -- )
  22     16 data bs:peek 2 >le be> 31 mod    ! checksum
  23     0 assert=
  24     4 data bs:read 8 assert=            ! compression method: deflate
  25     4 data bs:read                      ! log2(max length)-8, 32K max
  26     7 <= [ bad-zlib-header ] unless
  27     5 data bs:seek                      ! drop check bits
  28     1 data bs:read 0 assert=            ! dictionnary - not allowed in png
  29     2 data bs:seek                      ! compression level; ignore
  30     ;
  31
  32 :: default-table ( -- table )
  33     0 <hashtable> :> table
  34     0 143 [a,b] 280 287 [a,b] append 8 table set-at
  35     144 255 [a,b] >array 9 table set-at
  36     256 279 [a,b] >array 7 table set-at
  37     table enum>seq 1 tail ;
  38
  39 CONSTANT: clen-shuffle { 16 17 18 0 8 7 9 6 10 5 11 4 12 3 13 2 14 1 15 }
  40
  41 : get-table ( values size -- table )
  42     16 f <array> clone <enum>
  43     [ '[ _ push-at ] 2each ] keep seq>> 1 tail [ natural-sort ] map ;
  44
  45 :: decode-huffman-tables ( bitstream -- tables )
  46     5 bitstream bs:read 257 +
  47     5 bitstream bs:read 1 +
  48     4 bitstream bs:read 4 +
  49     clen-shuffle swap head
  50     dup [ drop 3 bitstream bs:read ] map
  51     get-table
  52     bitstream swap <huffman-decoder>
  53     [ 2dup + ] dip swap :> k!
  54     '[
  55         _ read1-huff2
  56         {
  57             { [ dup 16 = ] [ 2 bitstream bs:read 3 + 2array ] }
  58             { [ dup 17 = ] [ 3 bitstream bs:read 3 + 2array ] }
  59             { [ dup 18 = ] [ 7 bitstream bs:read 11 + 2array ] }
  60             [ ]
  61         } cond
  62         dup array? [ dup second ] [ 1 ] if
  63         k swap - dup k! 0 >
  64     ]
  65     [ ] produce swap suffix
  66     { } [ dup array? [ dup first 16 = ] [ f ] if [ [ unclip-last ] [ second 1 + swap <repetition> append ] bi* ] [ suffix ] if ] reduce
  67     [ dup array? [ second 0 <repetition> ] [ 1array ] if ] map concat
  68     nip swap cut 2array [ [ length>> [0,b) ] [ ] bi get-table ] map ;
  69
  70 : static-huffman-tables ( -- tables )
  71       0 143 [a,b] [ 8 ] replicate
  72     144 255 [a,b] [ 9 ] replicate append
  73     256 279 [a,b] [ 7 ] replicate append
  74     280 287 [a,b] [ 8 ] replicate append
  75     0 31 [a,b] [ 5 ] replicate
  76     2array [ [ length>> [0,b) ] [ ] bi get-table ] map ;
  77
  78 CONSTANT: length-table
  79     {
  80         3 4 5 6 7 8 9 10
  81         11 13 15 17
  82         19 23 27 31
  83         35 43 51 59
  84         67 83 99 115
  85         131 163 195 227 258
  86     }
  87
  88 CONSTANT: dist-table
  89     {
  90         1 2 3 4
  91         5 7 9 13
  92         17 25 33 49
  93         65 97 129 193
  94         257 385 513 769
  95         1025 1537 2049 3073
  96         4097 6145 8193 12289
  97         16385 24577
  98     }
  99
 100 : nth* ( n seq -- elt )
 101     [ length 1 - swap - ] [ nth ] bi ;
 102
 103 :: inflate-lz77 ( seq -- bytes )
 104     1000 <byte-vector> :> bytes
 105     seq
 106     [
 107         dup array?
 108         [ first2 '[ _ 1 - bytes nth* bytes push ] times ]
 109         [ bytes push ] if
 110     ] each
 111     bytes ;
 112
 113 :: inflate-huffman ( bitstream tables -- bytes )
 114     tables bitstream '[ _ swap <huffman-decoder> ] map :> tables
 115     [
 116         tables first read1-huff2
 117         dup 256 >
 118         [
 119             dup 285 =
 120             [ ]
 121             [
 122                 dup 264 >
 123                 [
 124                     dup 261 - 4 /i dup 5 >
 125                     [ bad-zlib-data ] when
 126                     bitstream bs:read 2array
 127                 ]
 128                 when
 129             ] if
 130             ! 5 bitstream read-bits ! distance
 131             tables second read1-huff2
 132             dup 3 >
 133             [
 134                 dup 2 - 2 /i dup 13 >
 135                 [ bad-zlib-data ] when
 136                 bitstream bs:read 2array
 137             ]
 138             when
 139             2array
 140         ]
 141         when
 142         dup 256 = not
 143     ]
 144     [ ] produce nip
 145     [
 146         dup array? [
 147             first2
 148             [
 149                 dup array? [ first2 ] [ 0 ] if
 150                 [ 257 - length-table nth ] [ + ] bi*
 151             ]
 152             [
 153                 dup array? [ first2 ] [ 0 ] if
 154                 [ dist-table nth ] [ + ] bi*
 155             ] bi*
 156             2array
 157         ] when
 158     ] map ;
 159
 160 :: inflate-raw ( bitstream -- bytes )
 161     8 bitstream bs:align
 162     16 bitstream bs:read :> len
 163     16 bitstream bs:read :> nlen
 164     len nlen + 16 >signed -1 assert= ! len + ~len = -1
 165     bitstream byte-pos>>
 166     bitstream byte-pos>> len +
 167     bitstream bytes>> <slice>
 168     len 8 * bitstream bs:seek ;
 169
 170 : inflate-dynamic ( bitstream -- bytes )
 171     dup decode-huffman-tables inflate-huffman ;
 172
 173 : inflate-static ( bitstream -- bytes )
 174     static-huffman-tables inflate-huffman ;
 175
 176 :: inflate-loop ( bitstream -- bytes )
 177     [ 1 bitstream bs:read 0 = ]
 178     [
 179         bitstream
 180         2 bitstream bs:read
 181         {
 182             { 0 [ inflate-raw ] }
 183             { 1 [ inflate-static ] }
 184             { 2 [ inflate-dynamic ] }
 185             { 3 [ bad-zlib-data f ] }
 186         }
 187         case
 188     ]
 189     [ produce ] keep call suffix concat ;
 190
 191 PRIVATE>
 192
 193 : zlib-inflate ( bytes -- bytes )
 194     bs:<lsb0-bit-reader>
 195     [ check-zlib-header ] [ inflate-loop ] bi
 196     inflate-lz77 ;