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