! Copyright (C) 2009 Marc Fauconneau.
! See http://factorcode.org/license.txt for BSD license.
USING: accessors alien.accessors byte-arrays combinators
constructors fry io io.binary kernel locals macros math
math.parser math.ranges multiline sequences ;
IN: alt.bitstreams

TUPLE: bitstream
    { bytes byte-array }
    { byte-pos fixnum initial: 0 }
    { bit-pos fixnum initial: 0 } ;

TUPLE: msb0-bitstream < bitstream ;
TUPLE: lsb0-bitstream < bitstream ;

CONSTRUCTOR: msb0-bitstream ( bytes -- bs ) ;
CONSTRUCTOR: lsb0-bitstream ( bytes -- bs ) ;


! interface

GENERIC: peek ( n bitstream -- value )
GENERIC: poke ( value n bitstream -- )

: seek ( n bitstream -- )
    {
        [ byte-pos>> 8 * ] 
        [ bit-pos>> + + 8 /mod ] 
        [ (>>bit-pos) ] 
        [ (>>byte-pos) ]
    } cleave ; inline

: read ( n bitstream -- value )
    [ peek ] [ seek ] 2bi ; inline

: write ( value n bitstream -- )
    [ poke ] [ seek ] 2bi ; inline


! reading

<PRIVATE

MACRO: multi-alien-unsigned-1 ( seq -- quot ) 
    [ '[ _ + alien-unsigned-1 ] ] map 2cleave>quot ;

GENERIC: fetch3-le-unsafe ( n byte-array -- value )
GENERIC: fetch3-be-unsafe ( n byte-array -- value )

: fetch3-unsafe ( byte-array n offsets -- value ) 
    multi-alien-unsigned-1 8 2^ * + 8 2^ * + ; inline

M: byte-array fetch3-le-unsafe ( n byte-array -- value ) 
    swap { 0 1 2 } fetch3-unsafe ; inline
M: byte-array fetch3-be-unsafe ( n byte-array -- value ) 
    swap { 2 1 0 } fetch3-unsafe ; inline

: fetch3 ( n byte-array -- value ) 
    [ 3 [0,b) [ + ] with map ] dip [ nth ] curry map ;
    
: fetch3-le ( n byte-array -- value ) fetch3 le> ;
: fetch3-be ( n byte-array -- value ) fetch3 be> ;
    
GENERIC: peek16 ( n bitstream -- value )

M:: lsb0-bitstream peek16 ( n bs -- v )
    bs byte-pos>> bs bytes>> fetch3-le
    bs bit-pos>> 2^ /i
    n 2^ mod ;

M:: msb0-bitstream peek16 ( n bs -- v )
    bs byte-pos>> bs bytes>> fetch3-be
    24 n bs bit-pos>> + - 2^ /i
    n 2^ mod ;

PRIVATE>

M: lsb0-bitstream peek ( n bs -- v ) peek16 ;
M: msb0-bitstream peek ( n bs -- v ) peek16 ;


! writing

<PRIVATE

: or-bytes ( bytes bitstream -- )
    dup '[ _ byte-pos>> + _ bytes>> [ bitor ] change-nth ] each-index ;
    
PRIVATE>
    
! %unimplemented...


! utilities

: bits>string ( value n -- str ) 
    swap [ >bin ] [ CHAR: 0 pad-head ] bi* ;

: bits. ( value n -- ) bits>string print ;