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

     1  package bits
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestBitArrayEmpty(t *testing.T) {
    12  	testBitArray(t, []testWord{}, "empty")
    13  }
    14  
    15  func TestBitArrayB0(t *testing.T) {
    16  	testBitArray(t, []testWord{
    17  		{1, 0b0},
    18  	}, "b0")
    19  }
    20  
    21  func TestBitArrayB1(t *testing.T) {
    22  	testBitArray(t, []testWord{
    23  		{1, 0b1},
    24  	}, "b1")
    25  }
    26  
    27  func TestBitArrayB010101010(t *testing.T) {
    28  	testBitArray(t, []testWord{
    29  		{9, 0b010101010},
    30  	}, "b010101010")
    31  }
    32  
    33  func TestBitArrayV01010101010101010(t *testing.T) {
    34  	testBitArray(t, []testWord{
    35  		{17, 0b01010101010101010},
    36  	}, "b01010101010101010")
    37  }
    38  
    39  func TestBitArrayRand1(t *testing.T) {
    40  	r := rand.New(rand.NewSource(0))
    41  	for i := 0; i < 50; i++ {
    42  		testBitArray(t, genTestWords(r, 24, 1), fmt.Sprintf("1 bit, case#%d", i))
    43  	}
    44  }
    45  
    46  func TestBitArrayRand8(t *testing.T) {
    47  	r := rand.New(rand.NewSource(0))
    48  	for i := 0; i < 50; i++ {
    49  		testBitArray(t, genTestWords(r, 100, 8), fmt.Sprintf("8 bits, case#%d", i))
    50  	}
    51  }
    52  
    53  func TestBitArrayRand17(t *testing.T) {
    54  	r := rand.New(rand.NewSource(0))
    55  	for i := 0; i < 50; i++ {
    56  		testBitArray(t, genTestWords(r, 50, 17), fmt.Sprintf("17 bits, case#%d", i))
    57  	}
    58  }
    59  
    60  func genTestWords(r *rand.Rand, maxCount int, maxBits int) []testWord {
    61  	count := r.Intn(maxCount)
    62  	words := make([]testWord, count)
    63  	for i := range words {
    64  		if maxBits == 1 {
    65  			words[i].bits = 1
    66  		} else {
    67  			words[i].bits = 1 + r.Intn(maxBits-1)
    68  		}
    69  		words[i].v = uint(r.Intn(1 << words[i].bits))
    70  	}
    71  	return words
    72  }
    73  
    74  func bytesToFit(bits int) int {
    75  	if bits%8 == 0 {
    76  		return bits / 8
    77  	}
    78  	return bits/8 + 1
    79  }
    80  
    81  type testWord struct {
    82  	bits int
    83  	v    uint
    84  }
    85  
    86  func testBitArray(t *testing.T, words []testWord, name string) {
    87  	arr := Array{make([]byte, 0, 100)}
    88  	writer := NewWriter(&arr)
    89  	reader := NewReader(&arr)
    90  
    91  	totalBitsWritten := 0
    92  	for _, w := range words {
    93  		writer.Write(w.bits, w.v)
    94  		totalBitsWritten += w.bits
    95  	}
    96  	assert.EqualValuesf(t, bytesToFit(totalBitsWritten), len(arr.Bytes), name)
    97  
    98  	totalBitsRead := 0
    99  	for _, w := range words {
   100  		assert.EqualValuesf(t, bytesToFit(totalBitsWritten)*8-totalBitsRead, reader.NonReadBits(), name)
   101  		assert.EqualValuesf(t, bytesToFit(reader.NonReadBits()), reader.NonReadBytes(), name)
   102  
   103  		v := reader.Read(w.bits)
   104  		assert.EqualValuesf(t, w.v, v, name)
   105  		totalBitsRead += w.bits
   106  
   107  		assert.EqualValuesf(t, bytesToFit(totalBitsWritten)*8-totalBitsRead, reader.NonReadBits(), name)
   108  		assert.EqualValuesf(t, bytesToFit(reader.NonReadBits()), reader.NonReadBytes(), name)
   109  	}
   110  
   111  	// read the tail
   112  	assert.Panicsf(t, func() {
   113  		reader.Read(reader.NonReadBits() + 1)
   114  	}, name)
   115  	zero := reader.Read(reader.NonReadBits())
   116  	assert.EqualValuesf(t, uint(0), zero, name)
   117  	assert.EqualValuesf(t, int(0), reader.NonReadBits(), name)
   118  	assert.EqualValuesf(t, int(0), reader.NonReadBytes(), name)
   119  }
   120  
   121  func BenchmarkArray_write(b *testing.B) {
   122  	for bits := 1; bits <= 9; bits++ {
   123  		b.Run(fmt.Sprintf("%d bits", bits), func(b *testing.B) {
   124  			b.ResetTimer()
   125  
   126  			arr := Array{make([]byte, 0, bytesToFit(bits*b.N))}
   127  			writer := NewWriter(&arr)
   128  			for i := 0; i < b.N; i++ {
   129  				writer.Write(bits, 0xff)
   130  			}
   131  		})
   132  	}
   133  }
   134  
   135  func BenchmarkArray_read(b *testing.B) {
   136  	for bits := 1; bits <= 9; bits++ {
   137  		b.Run(fmt.Sprintf("%d bits", bits), func(b *testing.B) {
   138  			b.ResetTimer()
   139  
   140  			arr := Array{make([]byte, bytesToFit(bits*b.N))}
   141  			writer := NewReader(&arr)
   142  			for i := 0; i < b.N; i++ {
   143  				_ = writer.Read(bits)
   144  			}
   145  		})
   146  	}
   147  }