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\r
  79     }\r
  80 \r
  81 CONSTANT: dist-table\r
  82     { 1 2 3 4 \r
  83       5 7 9 13 \r
  84       17 25 33 49\r
  85       65 97 129 193\r
  86       257 385 513 769\r
  87       1025 1537 2049 3073\r
  88       4097 6145 8193 12289\r
  89       16385 24577 }\r
  90 \r
  91 : nth* ( n seq -- elt )\r
  92     [ length 1- swap - ] [ nth ] bi ;\r
  93 \r
  94 :: inflate-lz77 ( seq -- bytes )\r
  95     1000 <byte-vector> :> bytes\r
  96     seq\r
  97     [\r
  98         dup array?\r
  99         [ first2 '[ _ 1- bytes nth* bytes push ] times ]\r
 100         [ bytes push ] if\r
 101     ] each \r
 102     bytes ;\r
 103 \r
 104 :: inflate-dynamic ( bitstream -- bytes )\r
 105     bitstream decode-huffman-tables\r
 106     bitstream '[ _ swap <huffman-decoder> ] map :> tables\r
 107     [\r
 108         tables first read1-huff2\r
 109         dup 256 >\r
 110         [\r
 111             dup 285 = \r
 112             [ ]\r
 113             [ \r
 114                 dup 264 > \r
 115                 [ \r
 116                     dup 261 - 4 /i dup 5 > \r
 117                     [ bad-zlib-data ] when \r
 118                     bitstream bs:read 2array \r
 119                 ]\r
 120                 when \r
 121             ] if\r
 122             ! 5 bitstream read-bits ! distance\r
 123             tables second read1-huff2\r
 124             dup 3 > \r
 125             [ \r
 126                 dup 2 - 2 /i dup 13 >\r
 127                 [ bad-zlib-data ] when\r
 128                 bitstream bs:read 2array\r
 129             ] \r
 130             when\r
 131             2array\r
 132         ]\r
 133         when\r
 134         dup 256 = not\r
 135     ]\r
 136     [ ] produce nip\r
 137     [\r
 138         dup array? [\r
 139             first2\r
 140             [  \r
 141                 dup array? [ first2 ] [ 0 ] if\r
 142                 [ 257 - length-table nth ] [ + ] bi*\r
 143             ] \r
 144             [\r
 145                 dup array? [ first2 ] [ 0 ] if\r
 146                 [ dist-table nth ] [ + ] bi*\r
 147             ] bi*\r
 148             2array\r
 149         ] when\r
 150     ] map ;\r
 151     \r
 152 : inflate-raw ( bitstream -- bytes ) zlib-unimplemented ;\r
 153 : inflate-static ( bitstream -- bytes ) zlib-unimplemented ;\r
 154 \r
 155 :: inflate-loop ( bitstream -- bytes )\r
 156     [ 1 bitstream bs:read 0 = ]\r
 157     [\r
 158         bitstream\r
 159         2 bitstream bs:read ! B\r
 160         { \r
 161             { 0 [ inflate-raw ] }\r
 162             { 1 [ inflate-static ] }\r
 163             { 2 [ inflate-dynamic ] }\r
 164             { 3 [ bad-zlib-data f ] }\r
 165         }\r
 166         case\r
 167     ]\r
 168     [ produce ] keep call suffix concat ;\r
 169     \r
 170   !  [ produce ] keep dip swap suffix\r
 171 \r
 172 :: paeth ( a b c -- p ) \r
 173     a b + c - { a b c } [ [ - abs ] keep 2array ] with map \r
 174     sort-keys first second ;\r
 175     \r
 176 :: png-unfilter-line ( prev curr filter -- curr' )\r
 177     prev :> c\r
 178     prev 3 tail-slice :> b\r
 179     curr :> a\r
 180     curr 3 tail-slice :> x\r
 181     x length [0,b)\r
 182     filter\r
 183     {\r
 184         { 0 [ drop ] }\r
 185         { 1 [ [| n | n x nth n a nth + 256 wrap n x set-nth ] each ] }\r
 186         { 2 [ [| n | n x nth n b nth + 256 wrap n x set-nth ] each ] }\r
 187         { 3 [ [| n | n x nth n a nth n b nth + 2/ + 256 wrap n x set-nth ] each ] }\r
 188         { 4 [ [| n | n x nth n a nth n b nth n c nth paeth + 256 wrap n x set-nth ] each ] }\r
 189         \r
 190     } case \r
 191     curr 3 tail ;\r
 192 \r
 193 PRIVATE>\r
 194 \r
 195 ! for debug -- shows residual values\r
 196 : reverse-png-filter' ( lines -- filtered )\r
 197     [ first ] [ 1 tail ] [ map ] bi-curry@ bi nip\r
 198     concat [ 128 + 256 wrap ] map ;\r
 199     \r
 200 : reverse-png-filter ( lines -- filtered )\r
 201     dup first [ 0 ] replicate prefix\r
 202     [ { 0 0 } prepend  ] map\r
 203     2 clump [\r
 204         first2 dup [ third ] [ 0 2 rot set-nth ] bi png-unfilter-line\r
 205     ] map concat ;\r
 206 \r
 207 : zlib-inflate ( bytes -- bytes )\r
 208     bs:<lsb0-bit-reader>\r
 209     [ check-zlib-header ]\r
 210     [ inflate-loop ] bi\r
 211     inflate-lz77 ;\r