basis/bitstreams/bitstreams.factor

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