basis/compression/inflate/inflate.factor

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