github.com/weaviate/sroar@v0.0.0-20230210105426-26108af5465d/iterator_test.go (about)

     1  /*
     2   * Copyright 2021 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package sroar
    18  
    19  import (
    20  	"math/rand"
    21  	"sort"
    22  	"testing"
    23  
    24  	"github.com/stretchr/testify/require"
    25  )
    26  
    27  func TestIteratorBasic(t *testing.T) {
    28  	n := uint64(1e5)
    29  	bm := NewBitmap()
    30  	for i := uint64(1); i <= n; i++ {
    31  		bm.Set(uint64(i))
    32  	}
    33  
    34  	it := bm.NewIterator()
    35  	for i := uint64(1); i <= n; i++ {
    36  		v := it.Next()
    37  		require.Equal(t, i, v)
    38  	}
    39  	v := it.Next()
    40  	require.Equal(t, uint64(0), v)
    41  }
    42  
    43  func TestIteratorRanges(t *testing.T) {
    44  	n := uint64(1e5)
    45  	bm := NewBitmap()
    46  	for i := uint64(1); i <= n; i++ {
    47  		bm.Set(uint64(i))
    48  	}
    49  
    50  	iters := bm.NewRangeIterators(8)
    51  	cnt := uint64(1)
    52  	for idx := 0; idx < 8; idx++ {
    53  		it := iters[idx]
    54  		for v := it.Next(); v > 0; v = it.Next() {
    55  			require.Equal(t, cnt, v)
    56  			cnt++
    57  		}
    58  	}
    59  }
    60  
    61  func TestIteratorRandom(t *testing.T) {
    62  	n := uint64(1e6)
    63  	bm := NewBitmap()
    64  	mp := make(map[uint64]struct{})
    65  	var arr []uint64
    66  	for i := uint64(1); i <= n; i++ {
    67  		v := uint64(rand.Intn(int(n) * 5))
    68  		if v == 0 {
    69  			continue
    70  		}
    71  		if _, ok := mp[v]; ok {
    72  			continue
    73  		}
    74  		mp[v] = struct{}{}
    75  		arr = append(arr, v)
    76  		bm.Set(uint64(v))
    77  	}
    78  
    79  	sort.Slice(arr, func(i, j int) bool {
    80  		return arr[i] < arr[j]
    81  	})
    82  
    83  	it := bm.NewIterator()
    84  	v := it.Next()
    85  	for i := uint64(0); i < uint64(len(arr)); i++ {
    86  		require.Equal(t, arr[i], v)
    87  		v = it.Next()
    88  	}
    89  }
    90  
    91  func TestIteratorWithRemoveKeys(t *testing.T) {
    92  	b := NewBitmap()
    93  	N := uint64(1e6)
    94  	for i := uint64(0); i < N; i++ {
    95  		b.Set(i)
    96  	}
    97  
    98  	b.RemoveRange(0, N)
    99  	it := b.NewIterator()
   100  
   101  	cnt := 0
   102  	for it.Next() > 0 {
   103  		cnt++
   104  	}
   105  	require.Equal(t, 0, cnt)
   106  }
   107  
   108  func TestManyIterator(t *testing.T) {
   109  	b := NewBitmap()
   110  	for i := 0; i < int(1e6); i++ {
   111  		b.Set(uint64(i))
   112  	}
   113  
   114  	mi := b.ManyIterator()
   115  	buf := make([]uint64, 1000)
   116  
   117  	i := 0
   118  	for {
   119  		got := mi.NextMany(buf)
   120  		if got == 0 {
   121  			break
   122  		}
   123  		require.Equal(t, 1000, got)
   124  		require.Equal(t, uint64(i*1000), buf[0])
   125  		i++
   126  	}
   127  }
   128  
   129  func BenchmarkIterator(b *testing.B) {
   130  	bm := NewBitmap()
   131  	for i := 0; i < int(1e5); i++ {
   132  		bm.Set(uint64(i))
   133  	}
   134  
   135  	b.ResetTimer()
   136  	for i := 0; i < b.N; i++ {
   137  		it := bm.NewIterator()
   138  		for it.Next() > 0 {
   139  		}
   140  	}
   141  }