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