basis/compression/inflate/inflate.factor

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