github.com/unicornultrafoundation/go-u2u@v1.0.0-rc1.0.20240205080301-e74a83d3fadc/utils/bits/bits.go (about)

     1  package bits
     2  
     3  type (
     4  	// Array is a bitmap array
     5  	Array struct {
     6  		Bytes []byte
     7  	}
     8  
     9  	// Writer of numbers to Array.
    10  	Writer struct {
    11  		*Array
    12  		bitOffset int
    13  	}
    14  
    15  	// Reader of numbers from Array.
    16  	Reader struct {
    17  		*Array
    18  		byteOffset int
    19  		bitOffset  int
    20  	}
    21  )
    22  
    23  // NewWriter is a bitmap writer
    24  func NewWriter(arr *Array) *Writer {
    25  	return &Writer{
    26  		Array: arr,
    27  	}
    28  }
    29  
    30  // NewReader is a bitmap reader
    31  func NewReader(arr *Array) *Reader {
    32  	return &Reader{
    33  		Array: arr,
    34  	}
    35  }
    36  
    37  func (a *Writer) byteBitsFree() int {
    38  	return 8 - a.bitOffset
    39  }
    40  
    41  func (a *Writer) writeIntoLastByte(v uint) {
    42  	a.Bytes[len(a.Bytes)-1] |= byte(v << a.bitOffset)
    43  }
    44  
    45  func zeroTopByteBits(v uint, bits int) uint {
    46  	mask := uint(0xff) >> bits
    47  	return v & mask
    48  }
    49  
    50  // Write bits of a number into array.
    51  func (a *Writer) Write(bits int, v uint) {
    52  	if a.bitOffset == 0 {
    53  		a.Bytes = append(a.Bytes, byte(0))
    54  	}
    55  	free := a.byteBitsFree()
    56  	if bits <= free {
    57  		toWrite := bits
    58  		// appending v to the bit array
    59  		a.writeIntoLastByte(v)
    60  		// increment offsets
    61  		if toWrite == free {
    62  			a.bitOffset = 0
    63  		} else {
    64  			a.bitOffset += toWrite
    65  		}
    66  	} else {
    67  		toWrite := free
    68  		clear := a.bitOffset // 8 - free
    69  		// zeroing top `clear` bits and appending result to the bit array
    70  		a.writeIntoLastByte(zeroTopByteBits(v, clear))
    71  		// increment offsets
    72  		a.bitOffset = 0
    73  		a.Write(bits-toWrite, v>>toWrite)
    74  	}
    75  }
    76  
    77  func (a *Reader) byteBitsFree() int {
    78  	return 8 - a.bitOffset
    79  }
    80  
    81  func (a *Reader) Read(bits int) (v uint) {
    82  	// perform all the checks in the same function to make CPU branch predictor work better
    83  	if bits == 0 {
    84  		return 0
    85  	}
    86  	/*if bits > a.NonReadBits() {
    87  		panic(io.ErrUnexpectedEOF)
    88  	}*/
    89  
    90  	free := a.byteBitsFree()
    91  	if bits <= free {
    92  		toRead := bits
    93  		clear := 8 - (a.bitOffset + toRead)
    94  		v = zeroTopByteBits(uint(a.Bytes[a.byteOffset]), clear) >> a.bitOffset
    95  		// increment offsets
    96  		if toRead == free {
    97  			a.bitOffset = 0
    98  			a.byteOffset++
    99  		} else {
   100  			a.bitOffset += toRead
   101  		}
   102  	} else {
   103  		toRead := free
   104  		v = uint(a.Bytes[a.byteOffset]) >> a.bitOffset
   105  		// increment offsets
   106  		a.bitOffset = 0
   107  		a.byteOffset++
   108  		// read rest
   109  		rest := a.Read(bits - toRead)
   110  		v |= rest << toRead
   111  	}
   112  	return
   113  }
   114  
   115  func (a *Reader) View(bits int) (v uint) {
   116  	cp := *a
   117  	cpp := &cp
   118  	return cpp.Read(bits)
   119  }
   120  
   121  // NonReadBytes returns a number of non-consumed bytes
   122  func (a *Reader) NonReadBytes() int {
   123  	return len(a.Bytes) - a.byteOffset
   124  }
   125  
   126  // NonReadBytes returns a number of non-consumed bits
   127  func (a *Reader) NonReadBits() int {
   128  	//return a.nonReadBits
   129  	return a.NonReadBytes()*8 - a.bitOffset
   130  }