github.com/outcaste-io/sroar@v0.0.0-20221229172112-1fb64f14314c/iterator.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/bits"
    21  )
    22  
    23  type Iterator struct {
    24  	bm *Bitmap
    25  
    26  	keys   []uint64
    27  	keyIdx int
    28  
    29  	contIdx int
    30  
    31  	bitmapIdx int
    32  	bitset    uint16
    33  }
    34  
    35  func (bm *Bitmap) NewRangeIterators(numRanges int) []*Iterator {
    36  	keyn := bm.keys.numKeys()
    37  	iters := make([]*Iterator, numRanges)
    38  	width := keyn / numRanges
    39  	rem := keyn % numRanges
    40  	cnt := 0
    41  
    42  	// This loop distributes the key equally to the ranges. For example: If numRanges = 3
    43  	// and keyn = 8 then it will be distributes as [3, 3, 2]
    44  	for i := 0; i < numRanges; i++ {
    45  		iters[i] = bm.NewIterator()
    46  		n := width
    47  		if i < rem {
    48  			n = width + 1
    49  		}
    50  		iters[i].keys = iters[i].keys[cnt : cnt+2*n]
    51  		cnt = cnt + 2*n
    52  	}
    53  	return iters
    54  }
    55  
    56  func (bm *Bitmap) NewIterator() *Iterator {
    57  	if bm == nil {
    58  		return nil
    59  	}
    60  	return &Iterator{
    61  		bm:        bm,
    62  		keys:      bm.keys[indexNodeStart : indexNodeStart+bm.keys.numKeys()*2],
    63  		keyIdx:    0,
    64  		contIdx:   -1,
    65  		bitmapIdx: -1,
    66  	}
    67  }
    68  
    69  func (it *Iterator) Next() uint64 {
    70  	if it == nil || len(it.keys) == 0 {
    71  		return 0
    72  	}
    73  
    74  	key := it.keys[it.keyIdx]
    75  	off := it.keys[it.keyIdx+1]
    76  	cont := it.bm.getContainer(off)
    77  	card := getCardinality(cont)
    78  
    79  	// Loop until we find a container on which next operation is possible. When such a container
    80  	// is found, reset the variables responsible for container iteration.
    81  	for card == 0 || it.contIdx+1 >= card {
    82  		if it.keyIdx+2 >= len(it.keys) {
    83  			return 0
    84  		}
    85  		// jump by 2 because key is followed by a value
    86  		it.keyIdx += 2
    87  		it.contIdx = -1
    88  		it.bitmapIdx = -1
    89  		it.bitset = 0
    90  		key = it.keys[it.keyIdx]
    91  		off = it.keys[it.keyIdx+1]
    92  		cont = it.bm.getContainer(off)
    93  		card = getCardinality(cont)
    94  	}
    95  
    96  	//  The above loop assures that we can do next in this container.
    97  	it.contIdx++
    98  	switch cont[indexType] {
    99  	case typeArray:
   100  		return key | uint64(cont[int(startIdx)+it.contIdx])
   101  	case typeBitmap:
   102  		// A bitmap container is an array of uint16s.
   103  		// If the container is bitmap, go to the index which has a non-zero value.
   104  		for it.bitset == 0 && it.bitmapIdx+1 < len(cont[startIdx:]) {
   105  			it.bitmapIdx++
   106  			it.bitset = cont[int(startIdx)+it.bitmapIdx]
   107  		}
   108  		assert(it.bitset > 0)
   109  
   110  		// msbIdx is the index of most-significant bit. In this iteration we choose this set bit
   111  		// and make it zero.
   112  		msbIdx := uint16(bits.LeadingZeros16(it.bitset))
   113  		msb := 1 << (16 - msbIdx - 1)
   114  		it.bitset ^= uint16(msb)
   115  		return key | uint64(it.bitmapIdx*16+int(msbIdx))
   116  	}
   117  	return 0
   118  }
   119  
   120  type ManyItr struct {
   121  	index int
   122  	arr   []uint64
   123  }
   124  
   125  // TODO: See if this is needed, we should remove this
   126  func (r *Bitmap) ManyIterator() *ManyItr {
   127  	return &ManyItr{
   128  		arr: r.ToArray(),
   129  	}
   130  
   131  }
   132  
   133  func (itr *ManyItr) NextMany(buf []uint64) int {
   134  	count := 0
   135  	for i := 0; i < len(buf); i++ {
   136  		if itr.index == len(itr.arr) {
   137  			break
   138  		}
   139  		buf[i] = itr.arr[itr.index]
   140  		itr.index++
   141  		count++
   142  	}
   143  	return count
   144  }