basis/compression/lzw/lzw.factor

   1 ! Copyright (C) 2009 Doug Coleman.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors alien.accessors byte-arrays combinators
   4 constructors destructors fry io io.binary kernel locals macros
   5 math math.ranges multiline sequences sequences.private ;
   6 IN: bitstreams
   7
   8 QUALIFIED-WITH: bitstreams bs
   9
  10 CONSTANT: clear-code 256
  11 CONSTANT: end-of-information 257
  12
  13 TUPLE: lzw input output end-of-input? table count k omega omega-k #bits
  14 code old-code ;
  15
  16 SYMBOL: table-full
  17
  18 ERROR: index-too-big n ;
  19
  20 : lzw-bit-width ( n -- n' )
  21     {
  22         { [ dup 510 <= ] [ drop 9 ] }
  23         { [ dup 1022 <= ] [ drop 10 ] }
  24         { [ dup 2046 <= ] [ drop 11 ] }
  25         { [ dup 4094 <= ] [ drop 12 ] }
  26         [ drop table-full ]
  27     } cond ;
  28
  29 : lzw-bit-width-compress ( lzw -- n )
  30     count>> lzw-bit-width ;
  31
  32 : lzw-bit-width-uncompress ( lzw -- n )
  33     table>> length lzw-bit-width ;
  34
  35 : initial-compress-table ( -- assoc )
  36     258 iota [ [ 1vector ] keep ] H{ } map>assoc ;
  37
  38 : initial-uncompress-table ( -- seq )
  39     258 iota [ 1vector ] V{ } map-as ;
  40
  41 : reset-lzw ( lzw -- lzw )
  42     257 >>count
  43     V{ } clone >>omega
  44     V{ } clone >>omega-k
  45     9 >>#bits ;
  46
  47 : reset-lzw-compress ( lzw -- lzw )
  48     f >>k
  49     initial-compress-table >>table reset-lzw ;
  50
  51 : reset-lzw-uncompress ( lzw -- lzw )
  52     initial-uncompress-table >>table reset-lzw ;
  53
  54 : <lzw-compress> ( input -- obj )
  55     lzw new
  56         swap >>input
  57         ! binary <byte-writer> <bitstream-writer> >>output
  58         V{ } clone >>output ! TODO
  59         reset-lzw-compress ;
  60
  61 : <lzw-uncompress> ( input -- obj )
  62     lzw new
  63         swap >>input
  64         BV{ } clone >>output
  65         reset-lzw-uncompress ;
  66
  67 : push-k ( lzw -- lzw )
  68     [ ]
  69     [ k>> ]
  70     [ omega>> clone [ push ] keep ] tri >>omega-k ;
  71
  72 : omega-k-in-table? ( lzw -- ? )
  73     [ omega-k>> ] [ table>> ] bi key? ;
  74
  75 ERROR: not-in-table value ;
  76
  77 : write-output ( lzw -- )
  78     [
  79         [ omega>> ] [ table>> ] bi ?at [ not-in-table ] unless
  80     ] [
  81         [ lzw-bit-width-compress ]
  82         [ output>> bs:poke ] bi
  83     ] bi ;
  84
  85 : omega-k>omega ( lzw -- lzw )
  86     dup omega-k>> clone >>omega ;
  87
  88 : k>omega ( lzw -- lzw )
  89     dup k>> 1vector >>omega ;
  90
  91 : add-omega-k ( lzw -- )
  92     [ [ 1+ ] change-count count>> ]
  93     [ omega-k>> clone ]
  94     [ table>> ] tri set-at ;
  95
  96 : lzw-compress-char ( lzw k -- )
  97     >>k push-k dup omega-k-in-table? [
  98         omega-k>omega drop
  99     ] [
 100         [ write-output ]
 101         [ add-omega-k ]
 102         [ k>omega drop ] tri
 103     ] if ;
 104
 105 : (lzw-compress-chars) ( lzw -- )
 106     dup lzw-bit-width-compress table-full = [
 107         drop
 108     ] [
 109         dup input>> stream-read1
 110         [ [ lzw-compress-char ] [ drop (lzw-compress-chars) ] 2bi ]
 111         [ t >>end-of-input? drop ] if*
 112     ] if ;
 113
 114 : lzw-compress-chars ( lzw -- )
 115     {
 116         ! [ [ clear-code lzw-compress-char ] [ drop ] bi ] ! reset-lzw-compress drop ] bi ]
 117         [
 118             [ clear-code ] dip
 119             [ lzw-bit-width-compress ]
 120             [ output>> bs:poke ] bi
 121         ]
 122         [ (lzw-compress-chars) ]
 123         [
 124             [ k>> ]
 125             [ lzw-bit-width-compress ]
 126             [ output>> bs:poke ] tri
 127         ]
 128         [
 129             [ end-of-information ] dip
 130             [ lzw-bit-width-compress ]
 131             [ output>> bs:poke ] bi
 132         ]
 133         [ ]
 134     } cleave dup end-of-input?>> [ drop ] [ lzw-compress-chars ] if ;
 135
 136 : lzw-compress ( byte-array -- seq )
 137     binary <byte-reader> <lzw-compress>
 138     [ lzw-compress-chars ] [ output>> stream>> ] bi ;
 139
 140 : lookup-old-code ( lzw -- vector )
 141     [ old-code>> ] [ table>> ] bi nth ;
 142
 143 : lookup-code ( lzw -- vector )
 144     [ code>> ] [ table>> ] bi nth ;
 145
 146 : code-in-table? ( lzw -- ? )
 147     [ code>> ] [ table>> length ] bi < ;
 148
 149 : code>old-code ( lzw -- lzw )
 150     dup code>> >>old-code ;
 151
 152 : write-code ( lzw -- )
 153     [ lookup-code ] [ output>> ] bi push-all ;
 154
 155 : add-to-table ( seq lzw -- ) table>> push ;
 156
 157 : lzw-read ( lzw -- lzw n )
 158     [ ] [ lzw-bit-width-uncompress ] [ input>> ] tri bs:peek ;
 159
 160 DEFER: lzw-uncompress-char
 161 : handle-clear-code ( lzw -- )
 162     reset-lzw-uncompress
 163     lzw-read dup end-of-information = [
 164         2drop
 165     ] [
 166         >>code
 167         [ write-code ]
 168         [ code>old-code ] bi
 169         lzw-uncompress-char
 170     ] if ;
 171
 172 : handle-uncompress-code ( lzw -- lzw )
 173     dup code-in-table? [
 174         [ write-code ]
 175         [
 176             [
 177                 [ lookup-old-code ]
 178                 [ lookup-code first ] bi suffix
 179             ] [ add-to-table ] bi
 180         ] [ code>old-code ] tri
 181     ] [
 182         [
 183             [ lookup-old-code dup first suffix ] keep
 184             [ output>> push-all ] [ add-to-table ] 2bi
 185         ] [ code>old-code ] bi
 186     ] if ;
 187
 188 : lzw-uncompress-char ( lzw -- )
 189     lzw-read [
 190         >>code
 191         dup code>> end-of-information = [
 192             drop
 193         ] [
 194             dup code>> clear-code = [
 195                 handle-clear-code
 196             ] [
 197                 handle-uncompress-code
 198                 lzw-uncompress-char
 199             ] if
 200         ] if
 201     ] [
 202         drop
 203     ] if* ;
 204
 205 : lzw-uncompress ( seq -- byte-array )
 206     <lsb0-bitstream>
 207     ! binary <byte-reader> ! <bitstream-reader>
 208     <lzw-uncompress> [ lzw-uncompress-char ] [ output>> ] bi ;