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