basis/bitstreams/bitstreams.factor

   1 ! Copyright (C) 2009 Doug Coleman.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors alien.accessors assocs byte-arrays combinators
   4 destructors fry io io.binary io.encodings.binary io.streams.byte-array
   5 kernel locals macros math math.ranges multiline sequences
   6 sequences.private vectors byte-vectors combinators.short-circuit
   7 math.bitwise ;
   8 IN: bitstreams
   9
  10 TUPLE: widthed { bits integer read-only } { #bits integer read-only } ;
  11
  12 ERROR: invalid-widthed bits #bits ;
  13
  14 : check-widthed ( bits #bits -- bits #bits )
  15     dup 0 < [ invalid-widthed ] when
  16     2dup { [ nip 0 = ] [ drop 0 = not ] } 2&& [ invalid-widthed ] when
  17     over 0 = [
  18         2dup [ dup 0 < [ neg ] when log2 1 + ] dip > [ invalid-widthed ] when
  19     ] unless ;
  20
  21 : <widthed> ( bits #bits -- widthed )
  22     check-widthed
  23     widthed boa ;
  24
  25 : zero-widthed ( -- widthed ) 0 0 <widthed> ;
  26 : zero-widthed? ( widthed -- ? ) zero-widthed = ;
  27
  28 TUPLE: bit-reader
  29     { bytes byte-array }
  30     { byte-pos array-capacity initial: 0 }
  31     { bit-pos array-capacity initial: 0 } ;
  32
  33 TUPLE: bit-writer
  34     { bytes byte-vector }
  35     { widthed widthed } ;
  36
  37 TUPLE: msb0-bit-reader < bit-reader ;
  38 TUPLE: lsb0-bit-reader < bit-reader ;
  39
  40 : <msb0-bit-reader> ( bytes -- bs )
  41     msb0-bit-reader new swap >>bytes ; inline
  42
  43 : <lsb0-bit-reader> ( bytes -- bs )
  44     lsb0-bit-reader new swap >>bytes ; inline
  45
  46 TUPLE: msb0-bit-writer < bit-writer ;
  47 TUPLE: lsb0-bit-writer < bit-writer ;
  48
  49 : new-bit-writer ( class -- bs )
  50     new
  51         BV{ } clone >>bytes
  52         0 0 <widthed> >>widthed ; inline
  53
  54 : <msb0-bit-writer> ( -- bs )
  55     msb0-bit-writer new-bit-writer ;
  56
  57 : <lsb0-bit-writer> ( -- bs )
  58     lsb0-bit-writer new-bit-writer ;
  59
  60 GENERIC: peek ( n bitstream -- value )
  61 GENERIC: poke ( value n bitstream -- )
  62
  63 : get-abp ( bitstream -- abp )
  64     [ byte-pos>> 8 * ] [ bit-pos>> + ] bi ; inline
  65
  66 : set-abp ( abp bitstream -- )
  67     [ 8 /mod ] dip [ (>>bit-pos) ] [ (>>byte-pos) ] bi ; inline
  68
  69 : seek ( n bitstream -- )
  70     [ get-abp + ] [ set-abp ] bi ; inline
  71
  72 : (align) ( n m -- n' )
  73     [ /mod 0 > [ 1 + ] when ] [ * ] bi ; inline
  74
  75 : align ( n bitstream -- )
  76     [ get-abp swap (align) ] [ set-abp ] bi ; inline
  77
  78 : read ( n bitstream -- value )
  79     [ peek ] [ seek ] 2bi ; inline
  80
  81 <PRIVATE
  82
  83 ERROR: not-enough-bits widthed n ;
  84
  85 : widthed-bits ( widthed n -- bits )
  86     dup 0 < [ not-enough-bits ] when
  87     2dup [ #bits>> ] dip < [ not-enough-bits ] when
  88     [ [ bits>> ] [ #bits>> ] bi ] dip
  89     [ - neg shift ] keep <widthed> ;
  90
  91 : split-widthed ( widthed n -- widthed1 widthed2 )
  92     2dup [ #bits>> ] dip < [
  93         drop zero-widthed
  94     ] [
  95         [ widthed-bits ]
  96         [ [ [ bits>> ] [ #bits>> ] bi ] dip - [ bits ] keep <widthed> ] 2bi
  97     ] if ;
  98
  99 : widthed>bytes ( widthed -- bytes widthed )
 100     [ 8 split-widthed dup zero-widthed? not ]
 101     [ swap bits>> ] B{ } produce-as nip swap ;
 102
 103 :: |widthed ( widthed1 widthed2 -- widthed3 )
 104     widthed1 bits>> :> bits1
 105     widthed1 #bits>> :> #bits1
 106     widthed2 bits>> :> bits2
 107     widthed2 #bits>> :> #bits2
 108     bits1 #bits2 shift bits2 bitor
 109     #bits1 #bits2 + <widthed> ;
 110
 111 PRIVATE>
 112
 113 M:: lsb0-bit-writer poke ( value n bs -- )
 114     value n <widthed> :> widthed
 115     widthed
 116     bs widthed>> #bits>> 8 swap - split-widthed :> remainder :> byte
 117     byte bs widthed>> |widthed :> new-byte
 118     new-byte #bits>> 8 = [
 119         new-byte bits>> bs bytes>> push
 120         zero-widthed bs (>>widthed)
 121         remainder widthed>bytes
 122         [ bs bytes>> push-all ] [ bs (>>widthed) ] bi*
 123     ] [
 124         byte bs (>>widthed)
 125     ] if ;
 126
 127 : enough-bits? ( n bs -- ? )
 128     [ bytes>> length ]
 129     [ byte-pos>> - 8 * ]
 130     [ bit-pos>> - ] tri <= ;
 131
 132 ERROR: not-enough-bits n bit-reader ;
 133
 134 : #bits>#bytes ( #bits -- #bytes )
 135     8 /mod 0 = [ 1 + ] unless ; inline
 136
 137 :: subseq>bits-le ( bignum n bs -- bits )
 138     bignum bs bit-pos>> neg shift n bits ;
 139
 140 :: subseq>bits-be ( bignum n bs -- bits )
 141     bignum
 142     8 bs bit-pos>> - n - 8 mod dup 0 < [ 8 + ] when
 143     neg shift n bits ;
 144
 145 :: adjust-bits ( n bs -- )
 146     n 8 /mod :> #bits :> #bytes
 147     bs [ #bytes + ] change-byte-pos
 148     bit-pos>> #bits + dup 8 >= [
 149         8 - bs (>>bit-pos)
 150         bs [ 1 + ] change-byte-pos drop
 151     ] [
 152         bs (>>bit-pos)
 153     ] if ;
 154
 155 :: (peek) ( n bs endian> subseq-endian -- bits )
 156     n bs enough-bits? [ n bs not-enough-bits ] unless
 157     bs [ byte-pos>> ] [ bit-pos>> n + ] bi #bits>#bytes dupd +
 158     bs bytes>> subseq endian> execute( seq -- x ) :> bignum
 159     bignum n bs subseq-endian execute( bignum n bs -- bits ) ;
 160
 161 M: lsb0-bit-reader peek ( n bs -- bits ) \ le> \ subseq>bits-le (peek) ;
 162
 163 M: msb0-bit-reader peek ( n bs -- bits ) \ be> \ subseq>bits-be (peek) ;
 164
 165 :: bit-writer-bytes ( writer -- bytes )
 166     writer widthed>> #bits>> :> n
 167     n 0 = [
 168         writer widthed>> bits>> 8 n - shift
 169         writer bytes>> swap push
 170     ] unless
 171     writer bytes>> ;
 172
 173 :: byte-array-n>seq ( byte-array n -- seq )
 174     byte-array length 8 * n / iota
 175     byte-array <msb0-bit-reader> '[
 176         drop n _ read
 177     ] { } map-as ;