github.com/cockroachdb/cockroachdb-parser@v0.23.3-0.20240213214944-911057d40c9a/pkg/util/intsets/bitmap.go (about)

     1  // Copyright 2022 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package intsets
    12  
    13  import "math/bits"
    14  
    15  // smallCutoff is the size of the small bitmap.
    16  // Note: this can be set to a smaller value, e.g. for testing.
    17  const smallCutoff = 128
    18  
    19  // bitmap implements a bitmap of size smallCutoff.
    20  type bitmap struct {
    21  	// We don't use an array because that makes Go always keep the struct on the
    22  	// stack (see https://github.com/golang/go/issues/24416).
    23  	lo, hi uint64
    24  }
    25  
    26  func (v bitmap) IsSet(i int) bool {
    27  	w := v.lo
    28  	if i >= 64 {
    29  		w = v.hi
    30  	}
    31  	return w&(1<<uint64(i&63)) != 0
    32  }
    33  
    34  func (v *bitmap) Set(i int) {
    35  	if i < 64 {
    36  		v.lo |= (1 << uint64(i))
    37  	} else {
    38  		v.hi |= (1 << uint64(i&63))
    39  	}
    40  }
    41  
    42  func (v *bitmap) Unset(i int) {
    43  	if i < 64 {
    44  		v.lo &= ^(1 << uint64(i))
    45  	} else {
    46  		v.hi &= ^(1 << uint64(i&63))
    47  	}
    48  }
    49  
    50  func (v *bitmap) SetRange(from, to int) {
    51  	mask := func(from, to int) uint64 {
    52  		return (1<<(to-from+1) - 1) << from
    53  	}
    54  	switch {
    55  	case to < 64:
    56  		v.lo |= mask(from, to)
    57  	case from >= 64:
    58  		v.hi |= mask(from&63, to&63)
    59  	default:
    60  		v.lo |= mask(from, 63)
    61  		v.hi |= mask(0, to&63)
    62  	}
    63  }
    64  
    65  func (v *bitmap) UnionWith(other bitmap) {
    66  	v.lo |= other.lo
    67  	v.hi |= other.hi
    68  }
    69  
    70  func (v *bitmap) IntersectionWith(other bitmap) {
    71  	v.lo &= other.lo
    72  	v.hi &= other.hi
    73  }
    74  
    75  func (v bitmap) Intersects(other bitmap) bool {
    76  	return ((v.lo & other.lo) | (v.hi & other.hi)) != 0
    77  }
    78  
    79  func (v *bitmap) DifferenceWith(other bitmap) {
    80  	v.lo &^= other.lo
    81  	v.hi &^= other.hi
    82  }
    83  
    84  func (v bitmap) SubsetOf(other bitmap) bool {
    85  	return (v.lo&other.lo == v.lo) && (v.hi&other.hi == v.hi)
    86  }
    87  
    88  func (v bitmap) OnesCount() int {
    89  	return bits.OnesCount64(v.lo) + bits.OnesCount64(v.hi)
    90  }
    91  
    92  func (v bitmap) Next(startVal int) (nextVal int, ok bool) {
    93  	if startVal < 64 {
    94  		if ntz := bits.TrailingZeros64(v.lo >> uint64(startVal)); ntz < 64 {
    95  			// Found next element in the low word.
    96  			return startVal + ntz, true
    97  		}
    98  		startVal = 64
    99  	}
   100  	// Check high word.
   101  	if ntz := bits.TrailingZeros64(v.hi >> uint64(startVal&63)); ntz < 64 {
   102  		return startVal + ntz, true
   103  	}
   104  	return -1, false
   105  }