1// Licensed to the Apache Software Foundation (ASF) under one 2// or more contributor license agreements. See the NOTICE file 3// distributed with this work for additional information 4// regarding copyright ownership. The ASF licenses this file 5// to you under the Apache License, Version 2.0 (the 6// "License"); you may not use this file except in compliance 7// with the License. You may obtain a copy of the License at 8// 9// http://www.apache.org/licenses/LICENSE-2.0 10// 11// Unless required by applicable law or agreed to in writing, software 12// distributed under the License is distributed on an "AS IS" BASIS, 13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14// See the License for the specific language governing permissions and 15// limitations under the License. 16 17package bitutil 18 19import ( 20 "math" 21 "math/bits" 22 "reflect" 23 "unsafe" 24 25 "github.com/apache/arrow/go/v6/arrow/memory" 26) 27 28var ( 29 BitMask = [8]byte{1, 2, 4, 8, 16, 32, 64, 128} 30 FlippedBitMask = [8]byte{254, 253, 251, 247, 239, 223, 191, 127} 31) 32 33// IsMultipleOf8 returns whether v is a multiple of 8. 34func IsMultipleOf8(v int64) bool { return v&7 == 0 } 35 36// IsMultipleOf64 returns whether v is a multiple of 64 37func IsMultipleOf64(v int64) bool { return v&63 == 0 } 38 39func BytesForBits(bits int64) int64 { return (bits + 7) >> 3 } 40 41// NextPowerOf2 rounds x to the next power of two. 42func NextPowerOf2(x int) int { return 1 << uint(bits.Len(uint(x))) } 43 44// CeilByte rounds size to the next multiple of 8. 45func CeilByte(size int) int { return (size + 7) &^ 7 } 46 47// CeilByte64 rounds size to the next multiple of 8. 48func CeilByte64(size int64) int64 { return (size + 7) &^ 7 } 49 50// BitIsSet returns true if the bit at index i in buf is set (1). 51func BitIsSet(buf []byte, i int) bool { return (buf[uint(i)/8] & BitMask[byte(i)%8]) != 0 } 52 53// BitIsNotSet returns true if the bit at index i in buf is not set (0). 54func BitIsNotSet(buf []byte, i int) bool { return (buf[uint(i)/8] & BitMask[byte(i)%8]) == 0 } 55 56// SetBit sets the bit at index i in buf to 1. 57func SetBit(buf []byte, i int) { buf[uint(i)/8] |= BitMask[byte(i)%8] } 58 59// ClearBit sets the bit at index i in buf to 0. 60func ClearBit(buf []byte, i int) { buf[uint(i)/8] &= FlippedBitMask[byte(i)%8] } 61 62// SetBitTo sets the bit at index i in buf to val. 63func SetBitTo(buf []byte, i int, val bool) { 64 if val { 65 SetBit(buf, i) 66 } else { 67 ClearBit(buf, i) 68 } 69} 70 71// CountSetBits counts the number of 1's in buf up to n bits. 72func CountSetBits(buf []byte, offset, n int) int { 73 if offset > 0 { 74 return countSetBitsWithOffset(buf, offset, n) 75 } 76 77 count := 0 78 79 uint64Bytes := n / uint64SizeBits * 8 80 for _, v := range bytesToUint64(buf[:uint64Bytes]) { 81 count += bits.OnesCount64(v) 82 } 83 84 for _, v := range buf[uint64Bytes : n/8] { 85 count += bits.OnesCount8(v) 86 } 87 88 // tail bits 89 for i := n &^ 0x7; i < n; i++ { 90 if BitIsSet(buf, i) { 91 count++ 92 } 93 } 94 95 return count 96} 97 98func countSetBitsWithOffset(buf []byte, offset, n int) int { 99 count := 0 100 101 beg := offset 102 end := offset + n 103 104 begU8 := roundUp(beg, uint64SizeBits) 105 106 init := min(n, begU8-beg) 107 for i := offset; i < beg+init; i++ { 108 if BitIsSet(buf, i) { 109 count++ 110 } 111 } 112 113 nU64 := (n - init) / uint64SizeBits 114 begU64 := begU8 / uint64SizeBits 115 endU64 := begU64 + nU64 116 bufU64 := bytesToUint64(buf) 117 if begU64 < len(bufU64) { 118 for _, v := range bufU64[begU64:endU64] { 119 count += bits.OnesCount64(v) 120 } 121 } 122 123 // FIXME: use a fallback to bits.OnesCount8 124 // before counting the tail bits. 125 126 tail := beg + init + nU64*uint64SizeBits 127 for i := tail; i < end; i++ { 128 if BitIsSet(buf, i) { 129 count++ 130 } 131 } 132 133 return count 134} 135 136func roundUp(v, f int) int { 137 return (v + (f - 1)) / f * f 138} 139 140func min(a, b int) int { 141 if a < b { 142 return a 143 } 144 return b 145} 146 147const ( 148 uint64SizeBytes = int(unsafe.Sizeof(uint64(0))) 149 uint64SizeBits = uint64SizeBytes * 8 150) 151 152func bytesToUint64(b []byte) []uint64 { 153 h := (*reflect.SliceHeader)(unsafe.Pointer(&b)) 154 155 var res []uint64 156 s := (*reflect.SliceHeader)(unsafe.Pointer(&res)) 157 s.Data = h.Data 158 s.Len = h.Len / uint64SizeBytes 159 s.Cap = h.Cap / uint64SizeBytes 160 161 return res 162} 163 164var ( 165 // PrecedingBitmask is a convenience set of values as bitmasks for checking 166 // prefix bits of a byte 167 PrecedingBitmask = [8]byte{0, 1, 3, 7, 15, 31, 63, 127} 168 // TrailingBitmask is the bitwise complement version of kPrecedingBitmask 169 TrailingBitmask = [8]byte{255, 254, 252, 248, 240, 224, 192, 128} 170) 171 172// SetBitsTo is a convenience function to quickly set or unset all the bits 173// in a bitmap starting at startOffset for length bits. 174func SetBitsTo(bits []byte, startOffset, length int64, areSet bool) { 175 if length == 0 { 176 return 177 } 178 179 beg := startOffset 180 end := startOffset + length 181 var fill uint8 = 0 182 if areSet { 183 fill = math.MaxUint8 184 } 185 186 byteBeg := beg / 8 187 byteEnd := end/8 + 1 188 189 // don't modify bits before the startOffset by using this mask 190 firstByteMask := PrecedingBitmask[beg%8] 191 // don't modify bits past the length by using this mask 192 lastByteMask := TrailingBitmask[end%8] 193 194 if byteEnd == byteBeg+1 { 195 // set bits within a single byte 196 onlyByteMask := firstByteMask 197 if end%8 != 0 { 198 onlyByteMask = firstByteMask | lastByteMask 199 } 200 201 bits[byteBeg] &= onlyByteMask 202 bits[byteBeg] |= fill &^ onlyByteMask 203 return 204 } 205 206 // set/clear trailing bits of first byte 207 bits[byteBeg] &= firstByteMask 208 bits[byteBeg] |= fill &^ firstByteMask 209 210 if byteEnd-byteBeg > 2 { 211 memory.Set(bits[byteBeg+1:byteEnd-1], fill) 212 } 213 214 if end%8 == 0 { 215 return 216 } 217 218 bits[byteEnd-1] &= lastByteMask 219 bits[byteEnd-1] |= fill &^ lastByteMask 220} 221