factor/basis/bitstreams/bitstreams.factor

97 lines
2.5 KiB
Factor

! Copyright (C) 2009 Doug Coleman.
! See http://factorcode.org/license.txt for BSD license.
USING: accessors byte-arrays destructors fry io kernel locals
math sequences ;
IN: bitstreams
TUPLE: bitstream stream end-of-stream? current-bits #bits disposed ;
TUPLE: bitstream-reader < bitstream ;
: reset-bitstream ( stream -- stream )
0 >>#bits 0 >>current-bits ; inline
: new-bitstream ( stream class -- bitstream )
new
swap >>stream
reset-bitstream ; inline
M: bitstream-reader dispose ( stream -- )
stream>> dispose ;
: <bitstream-reader> ( stream -- bitstream )
bitstream-reader new-bitstream ; inline
: read-next-byte ( bitstream -- bitstream )
dup stream>> stream-read1 [
>>current-bits 8 >>#bits
] [
0 >>#bits
t >>end-of-stream?
] if* ;
: maybe-read-next-byte ( bitstream -- bitstream )
dup #bits>> 0 = [ read-next-byte ] when ; inline
: shift-one-bit ( bitstream -- n )
[ current-bits>> ] [ #bits>> ] bi 1- neg shift 1 bitand ; inline
: next-bit ( bitstream -- n/f ? )
maybe-read-next-byte
dup end-of-stream?>> [
drop f
] [
[ shift-one-bit ]
[ [ 1- ] change-#bits maybe-read-next-byte drop ] bi
] if dup >boolean ;
: read-bit ( bitstream -- n ? )
dup #bits>> 1 = [
[ current-bits>> 1 bitand ]
[ read-next-byte drop ] bi t
] [
next-bit
] if ; inline
: bits>integer ( seq -- n )
0 [ [ 1 shift ] dip bitor ] reduce ; inline
: read-bits ( width bitstream -- n width ? )
[
'[ _ read-bit drop ] replicate
[ f = ] trim-tail
[ bits>integer ] [ length ] bi
] 2keep drop over = ;
TUPLE: bitstream-writer < bitstream ;
: <bitstream-writer> ( stream -- bitstream )
bitstream-writer new-bitstream ; inline
: write-bit ( n bitstream -- )
[ 1 shift bitor ] change-current-bits
[ 1+ ] change-#bits
dup #bits>> 8 = [
[ [ current-bits>> ] [ stream>> stream-write1 ] bi ]
[ reset-bitstream drop ] bi
] [
drop
] if ; inline
ERROR: invalid-bit-width n ;
:: write-bits ( n width bitstream -- )
n 0 < [ n invalid-bit-width ] when
n 0 = [
width [ 0 bitstream write-bit ] times
] [
width n log2 1+ dup :> n-length - [ 0 bitstream write-bit ] times
n-length [
n-length swap - 1- neg n swap shift 1 bitand
bitstream write-bit
] each
] if ;
: flush-bits ( bitstream -- ) stream>> stream-flush ;
: bitstream-output ( bitstream -- bytes ) stream>> >byte-array ;