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