basis/compression/lzw/lzw.factor

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