github.com/cnotch/ipchub@v1.1.0/utils/bits/reader.go (about)

     1  // Copyright (c) 2019,CAOHONGJU All rights reserved.
     2  // Use of this source code is governed by a MIT-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package bits
     6  
     7  const uintBitsCount = int(32 << (^uint(0) >> 63))
     8  
     9  // Reader .
    10  type Reader struct {
    11  	buf    []byte
    12  	offset int // bit base
    13  }
    14  
    15  // NewReader retruns a new Reader.
    16  func NewReader(buf []byte) *Reader {
    17  	return &Reader{
    18  		buf: buf,
    19  	}
    20  }
    21  
    22  // Skip skip n bits.
    23  func (r *Reader) Skip(n int) {
    24  	if n <= 0 {
    25  		return
    26  	}
    27  	_ = r.buf[(r.offset+n-1)>>3] // bounds check hint to compiler; see golang.org/issue/14808
    28  	r.offset += n
    29  }
    30  
    31  // Peek peek the uint32 of n bits.
    32  func (r *Reader) Peek(n int) uint64 {
    33  	clone := *r
    34  	return clone.readUint64(n, 64)
    35  }
    36  
    37  // Read read the uint32 of n bits.
    38  func (r *Reader) Read(n int) uint32 {
    39  	return uint32(r.readUint64(n, 32))
    40  }
    41  
    42  // ReadBit read a bit.
    43  func (r *Reader) ReadBit() uint8 {
    44  	_ = r.buf[(r.offset+1-1)>>3] // bounds check hint to compiler; see golang.org/issue/14808
    45  
    46  	tmp := (r.buf[r.offset>>3] >> (7 - r.offset&0x7)) & 1
    47  	r.offset++
    48  	return tmp
    49  }
    50  
    51  // ReadUe .
    52  func (r *Reader) ReadUe() (res uint32) {
    53  	i := 0
    54  	for {
    55  		if bit := r.ReadBit(); !(bit == 0 && i < 32) {
    56  			break
    57  		}
    58  		i++
    59  	}
    60  
    61  	res = r.Read(i)
    62  	res += (1 << uint(i)) - 1
    63  	return
    64  }
    65  
    66  // ReadSe .
    67  func (r *Reader) ReadSe() (res int32) {
    68  	ui32 := r.ReadUe()
    69  	if ui32&0x01 != 0 {
    70  		res = (int32(res) + 1) / 2
    71  	} else {
    72  		res = -int32(res) / 2
    73  	}
    74  	return
    75  }
    76  
    77  // ==== shortcut methods
    78  
    79  // ReadBool read one bit bool.
    80  func (r *Reader) ReadBool() bool { return bool(r.ReadBit() == 1) }
    81  
    82  // ReadUint read the uint of n bits.
    83  func (r *Reader) ReadUint(n int) uint { return uint(r.readUint64(n, uintBitsCount)) }
    84  
    85  // ReadUint8 read the uint8 of n bits.
    86  func (r *Reader) ReadUint8(n int) uint8 { return uint8(r.readUint64(n, 8)) }
    87  
    88  // ReadUint16 read the uint16 of n bits.
    89  func (r *Reader) ReadUint16(n int) uint16 { return uint16(r.readUint64(n, 16)) }
    90  
    91  // ReadUint32 read the uint32 of n bits.
    92  func (r *Reader) ReadUint32(n int) uint32 { return uint32(r.readUint64(n, 32)) }
    93  
    94  // ReadUint64 read the uint64 of n bits.
    95  func (r *Reader) ReadUint64(n int) uint64 { return r.readUint64(n, 64) }
    96  
    97  // ReadInt read the int of n bits.
    98  func (r *Reader) ReadInt(n int) int { return int(r.readUint64(n, uintBitsCount)) }
    99  
   100  // ReadInt8 read the int8 of n bits.
   101  func (r *Reader) ReadInt8(n int) int8 { return int8(r.readUint64(n, 8)) }
   102  
   103  // ReadInt16 read the int16 of n bits.
   104  func (r *Reader) ReadInt16(n int) int16 { return int16(r.readUint64(n, 16)) }
   105  
   106  // ReadInt32 read the int32 of n bits.
   107  func (r *Reader) ReadInt32(n int) int32 { return int32(r.readUint64(n, 32)) }
   108  
   109  // ReadInt64 read the int64 of n bits.
   110  func (r *Reader) ReadInt64(n int) int64 { return int64(r.readUint64(n, 64)) }
   111  
   112  // ReadUe8 read the UE GolombCode of uint8.
   113  func (r *Reader) ReadUe8() uint8 { return uint8(r.ReadUe()) }
   114  
   115  // ReadUe16 read the UE GolombCode of uint16.
   116  func (r *Reader) ReadUe16() uint16 { return uint16(r.ReadUe()) }
   117  
   118  // ReadSe8 read the SE of int8.
   119  func (r *Reader) ReadSe8() int8 { return int8(r.ReadSe()) }
   120  
   121  // ReadSe16 read the SE of int16.
   122  func (r *Reader) ReadSe16() int16 { return int16(r.ReadSe()) }
   123  
   124  // Offset returns the offset of bits.
   125  func (r *Reader) Offset() int {
   126  	return r.offset
   127  }
   128  
   129  // BitsLeft returns the number of left bits.
   130  func (r *Reader) BitsLeft() int {
   131  	return len(r.buf)<<3 - r.offset
   132  }
   133  
   134  // BytesLeft returns the left byte slice.
   135  func (r *Reader) BytesLeft() []byte {
   136  	return r.buf[r.offset>>3:]
   137  }
   138  
   139  var bitsMask = [9]byte{
   140  	0x00,
   141  	0x01, 0x03, 0x07, 0x0f,
   142  	0x1f, 0x3f, 0x7f, 0xff,
   143  }
   144  
   145  // readUint64 read the uint64 of n bits.
   146  func (r *Reader) readUint64(n, max int) uint64 {
   147  	if n <= 0 || n > max {
   148  		return 0
   149  	}
   150  
   151  	_ = r.buf[(r.offset+n-1)>>3] // bounds check hint to compiler; see golang.org/issue/14808
   152  
   153  	idx := r.offset >> 3
   154  	validBits := 8 - r.offset&0x7
   155  	r.offset += n
   156  
   157  	var tmp uint64
   158  	for n >= validBits {
   159  		n -= validBits
   160  		tmp |= uint64(r.buf[idx]&bitsMask[validBits]) << n
   161  		idx++
   162  		validBits = 8
   163  	}
   164  
   165  	if n > 0 {
   166  		tmp |= uint64((r.buf[idx] >> (validBits - n)) & bitsMask[n])
   167  	}
   168  	return tmp
   169  }