srs/trunk/3rdparty/srs-bench/vendor/github.com/yapingcat/gomedia/codec/bitstream.go

package codec

import (
    "encoding/binary"
)

var BitMask [8]byte = [8]byte{0x01, 0x03, 0x07, 0x0F, 0x1F, 0x3F, 0x7F, 0xFF}

type BitStream struct {
    bits        []byte
    bytesOffset int
    bitsOffset  int
    bitsmark    int
    bytemark    int
}

func NewBitStream(buf []byte) *BitStream {
    return &BitStream{
        bits:        buf,
        bytesOffset: 0,
        bitsOffset:  0,
        bitsmark:    0,
        bytemark:    0,
    }
}

func (bs *BitStream) Uint8(n int) uint8 {
    return uint8(bs.GetBits(n))
}

func (bs *BitStream) Uint16(n int) uint16 {
    return uint16(bs.GetBits(n))
}

func (bs *BitStream) Uint32(n int) uint32 {
    return uint32(bs.GetBits(n))
}

func (bs *BitStream) GetBytes(n int) []byte {
    if bs.bytesOffset+n > len(bs.bits) {
        panic("OUT OF RANGE")
    }
    if bs.bitsOffset != 0 {
        panic("invaild operation")
    }
    data := make([]byte, n)
    copy(data, bs.bits[bs.bytesOffset:bs.bytesOffset+n])
    bs.bytesOffset += n
    return data
}

//n <= 64
func (bs *BitStream) GetBits(n int) uint64 {
    if bs.bytesOffset >= len(bs.bits) {
        panic("OUT OF RANGE")
    }
    var ret uint64 = 0
    if 8-bs.bitsOffset >= n {
        ret = uint64((bs.bits[bs.bytesOffset] >> (8 - bs.bitsOffset - n)) & BitMask[n-1])
        bs.bitsOffset += n
        if bs.bitsOffset == 8 {
            bs.bytesOffset++
            bs.bitsOffset = 0
        }
    } else {
        ret = uint64(bs.bits[bs.bytesOffset] & BitMask[8-bs.bitsOffset-1])
        bs.bytesOffset++
        n -= 8 - bs.bitsOffset
        bs.bitsOffset = 0
        for n > 0 {
            if bs.bytesOffset >= len(bs.bits) {
                panic("OUT OF RANGE")
            }
            if n >= 8 {
                ret = ret<<8 | uint64(bs.bits[bs.bytesOffset])
                bs.bytesOffset++
                n -= 8
            } else {
                ret = (ret << n) | uint64((bs.bits[bs.bytesOffset]>>(8-n))&BitMask[n-1])
                bs.bitsOffset = n
                break
            }
        }
    }
    return ret
}

func (bs *BitStream) GetBit() uint8 {
    if bs.bytesOffset >= len(bs.bits) {
        panic("OUT OF RANGE")
    }
    ret := bs.bits[bs.bytesOffset] >> (7 - bs.bitsOffset) & 0x01
    bs.bitsOffset++
    if bs.bitsOffset >= 8 {
        bs.bytesOffset++
        bs.bitsOffset = 0
    }
    return ret
}

func (bs *BitStream) SkipBits(n int) {
    bytecount := n / 8
    bitscount := n % 8
    bs.bytesOffset += bytecount
    if bs.bitsOffset+bitscount < 8 {
        bs.bitsOffset += bitscount
    } else {
        bs.bytesOffset += 1
        bs.bitsOffset += bitscount - 8
    }
}

func (bs *BitStream) Markdot() {
    bs.bitsmark = bs.bitsOffset
    bs.bytemark = bs.bytesOffset
}

func (bs *BitStream) DistanceFromMarkDot() int {
    bytecount := bs.bytesOffset - bs.bytemark - 1
    bitscount := bs.bitsOffset + (8 - bs.bitsmark)
    return bytecount*8 + bitscount
}

func (bs *BitStream) RemainBytes() int {
    if bs.bitsOffset > 0 {
        return len(bs.bits) - bs.bytesOffset - 1
    } else {
        return len(bs.bits) - bs.bytesOffset
    }
}

func (bs *BitStream) RemainBits() int {
    if bs.bitsOffset > 0 {
        return bs.RemainBytes()*8 + 8 - bs.bitsOffset
    } else {
        return bs.RemainBytes() * 8
    }

}

func (bs *BitStream) Bits() []byte {
    return bs.bits
}

func (bs *BitStream) RemainData() []byte {
    return bs.bits[bs.bytesOffset:]
}

//无符号哥伦布熵编码
func (bs *BitStream) ReadUE() uint64 {
    leadingZeroBits := 0
    for bs.GetBit() == 0 {
        leadingZeroBits++
    }
    if leadingZeroBits == 0 {
        return 0
    }
    info := bs.GetBits(leadingZeroBits)
    return uint64(1)<<leadingZeroBits - 1 + info
}

//有符号哥伦布熵编码
func (bs *BitStream) ReadSE() int64 {
    v := bs.ReadUE()
    if v%2 == 0 {
        return -1 * int64(v/2)
    } else {
        return int64(v+1) / 2
    }
}

func (bs *BitStream) ByteOffset() int {
    return bs.bytesOffset
}

func (bs *BitStream) UnRead(n int) {
    if n-bs.bitsOffset <= 0 {
        bs.bitsOffset -= n
    } else {
        least := n - bs.bitsOffset
        for least >= 8 {
            bs.bytesOffset--
            least -= 8
        }
        if least > 0 {
            bs.bytesOffset--
            bs.bitsOffset = 8 - least
        }
    }
}

func (bs *BitStream) NextBits(n int) uint64 {
    r := bs.GetBits(n)
    bs.UnRead(n)
    return r
}

func (bs *BitStream) EOS() bool {
    return bs.bytesOffset == len(bs.bits) && bs.bitsOffset == 0
}

type BitStreamWriter struct {
    bits       []byte
    byteoffset int
    bitsoffset int
    bitsmark   int
    bytemark   int
}

func NewBitStreamWriter(n int) *BitStreamWriter {
    return &BitStreamWriter{
        bits:       make([]byte, n),
        byteoffset: 0,
        bitsoffset: 0,
        bitsmark:   0,
        bytemark:   0,
    }
}

func (bsw *BitStreamWriter) expandSpace(n int) {
    if (len(bsw.bits)-bsw.byteoffset-1)*8+8-bsw.bitsoffset < n {
        newlen := 0
        if len(bsw.bits)*8 < n {
            newlen = len(bsw.bits) + n/8 + 1
        } else {
            newlen = len(bsw.bits) * 2
        }
        tmp := make([]byte, newlen)
        copy(tmp, bsw.bits)
        bsw.bits = tmp
    }
}

func (bsw *BitStreamWriter) ByteOffset() int {
    return bsw.byteoffset
}

func (bsw *BitStreamWriter) BitOffset() int {
    return bsw.bitsoffset
}

func (bsw *BitStreamWriter) Markdot() {
    bsw.bitsmark = bsw.bitsoffset
    bsw.bytemark = bsw.byteoffset
}

func (bsw *BitStreamWriter) DistanceFromMarkDot() int {
    bytecount := bsw.byteoffset - bsw.bytemark - 1
    bitscount := bsw.bitsoffset + (8 - bsw.bitsmark)
    return bytecount*8 + bitscount
}

func (bsw *BitStreamWriter) PutByte(v byte) {
    bsw.expandSpace(8)
    if bsw.bitsoffset == 0 {
        bsw.bits[bsw.byteoffset] = v
        bsw.byteoffset++
    } else {
        bsw.bits[bsw.byteoffset] |= v >> byte(bsw.bitsoffset)
        bsw.byteoffset++
        bsw.bits[bsw.byteoffset] = v & BitMask[bsw.bitsoffset-1]
    }
}

func (bsw *BitStreamWriter) PutBytes(v []byte) {
    if bsw.bitsoffset != 0 {
        panic("bsw.bitsoffset > 0")
    }
    bsw.expandSpace(8 * len(v))
    copy(bsw.bits[bsw.byteoffset:], v)
    bsw.byteoffset += len(v)
}

func (bsw *BitStreamWriter) PutRepetValue(v byte, n int) {
    if bsw.bitsoffset != 0 {
        panic("bsw.bitsoffset > 0")
    }
    bsw.expandSpace(8 * n)
    for i := 0; i < n; i++ {
        bsw.bits[bsw.byteoffset] = v
        bsw.byteoffset++
    }
}

func (bsw *BitStreamWriter) PutUint8(v uint8, n int) {
    bsw.PutUint64(uint64(v), n)
}

func (bsw *BitStreamWriter) PutUint16(v uint16, n int) {
    bsw.PutUint64(uint64(v), n)
}

func (bsw *BitStreamWriter) PutUint32(v uint32, n int) {
    bsw.PutUint64(uint64(v), n)
}

func (bsw *BitStreamWriter) PutUint64(v uint64, n int) {
    bsw.expandSpace(n)
    if 8-bsw.bitsoffset >= n {
        bsw.bits[bsw.byteoffset] |= uint8(v) & BitMask[n-1] << (8 - bsw.bitsoffset - n)
        bsw.bitsoffset += n
        if bsw.bitsoffset == 8 {
            bsw.bitsoffset = 0
            bsw.byteoffset++
        }
    } else {
        bsw.bits[bsw.byteoffset] |= uint8(v>>(n-int(8-bsw.bitsoffset))) & BitMask[8-bsw.bitsoffset-1]
        bsw.byteoffset++
        n -= 8 - bsw.bitsoffset
        for n-8 >= 0 {
            bsw.bits[bsw.byteoffset] = uint8(v>>(n-8)) & 0xFF
            bsw.byteoffset++
            n -= 8
        }
        bsw.bitsoffset = n
        if n > 0 {
            bsw.bits[bsw.byteoffset] |= (uint8(v) & BitMask[n-1]) << (8 - n)
        }
    }
}

func (bsw *BitStreamWriter) SetByte(v byte, where int) {
    bsw.bits[where] = v
}

func (bsw *BitStreamWriter) SetUint16(v uint16, where int) {
    binary.BigEndian.PutUint16(bsw.bits[where:where+2], v)
}

func (bsw *BitStreamWriter) Bits() []byte {
    if bsw.byteoffset == len(bsw.bits) {
        return bsw.bits
    }
    if bsw.bitsoffset > 0 {
        return bsw.bits[0 : bsw.byteoffset+1]
    } else {
        return bsw.bits[0:bsw.byteoffset]
    }
}

//用v 填充剩余字节
func (bsw *BitStreamWriter) FillRemainData(v byte) {
    for i := bsw.byteoffset; i < len(bsw.bits); i++ {
        bsw.bits[i] = v
    }
    bsw.byteoffset = len(bsw.bits)
    bsw.bitsoffset = 0
}

func (bsw *BitStreamWriter) Reset() {
    for i := 0; i < len(bsw.bits); i++ {
        bsw.bits[i] = 0
    }
    bsw.bitsmark = 0
    bsw.bytemark = 0
    bsw.bitsoffset = 0
    bsw.byteoffset = 0
}