blob: e366851f1c601349209f7a4900d57c2cec027f41 [file] [log] [blame]
// Copyright 2016 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//go:generate go run ./tool/gen.go
package bitfield
import (
"bytes"
"encoding"
"encoding/binary"
"fmt"
"math"
"go.chromium.org/luci/gae/service/datastore"
)
// BitField is a luci/gae-serializable bit field implementation. It should
// be nice and fast for non-AppEngine use as well.
//
// You should construct a new one with bf.Make, rather than by direct
// construction.
type BitField struct {
// data is the actual bits data. rightmost bit of 0th byte is the first bit.
data []byte
size uint32
}
var _ interface {
// Compatible with luci/gae/service/datastore
datastore.PropertyConverter
// Compatible with gob
encoding.BinaryMarshaler
encoding.BinaryUnmarshaler
} = (*BitField)(nil)
// MarshalBinary implements encoding.BinaryMarshaller
func (bf *BitField) MarshalBinary() ([]byte, error) {
ret := make([]byte, binary.MaxVarintLen32+len(bf.data))
n := binary.PutUvarint(ret, uint64(bf.size))
n += copy(ret[n:], bf.data)
return ret[:n], nil
}
// UnmarshalBinary implements encoding.BinaryUnmarshaler
func (bf *BitField) UnmarshalBinary(bs []byte) error {
buf := bytes.NewBuffer(bs)
n, err := binary.ReadUvarint(buf)
if err != nil {
return err
}
if n > math.MaxUint32 {
return fmt.Errorf("encoded number is out of range: %d > %d", n, uint32(math.MaxUint32))
}
size := uint32(n)
if size == 0 {
bf.size = 0
bf.data = nil
return nil
}
data := buf.Bytes()
if uint32(len(data)) != (size+8-1)>>3 {
return fmt.Errorf("mismatched size (%d) v. byte count (%d)", size, len(data))
}
bf.size = size
bf.data = data
return nil
}
// ToProperty implements datastore.PropertyConverter
func (bf *BitField) ToProperty() (datastore.Property, error) {
bs, err := bf.MarshalBinary()
return datastore.MkPropertyNI(bs), err
}
// FromProperty implements datastore.PropertyConverter
func (bf *BitField) FromProperty(p datastore.Property) error {
bin, ok := p.Value().([]byte)
if !ok {
return fmt.Errorf("expected %s, got %s", datastore.PTBytes, p.Type())
}
return bf.UnmarshalBinary(bin)
}
// Reset resets this BitField to the 'empty' (size-0) state.
func (bf *BitField) Reset() {
bf.data = nil
bf.size = 0
}
// Make creates a new BitField.
func Make(size uint32) BitField {
if size == 0 {
return BitField{}
}
return BitField{
data: make([]byte, (size+8-1)>>3),
size: size,
}
}
// Size returns the number of bits which this BitField can hold.
func (bf BitField) Size() uint32 {
return bf.size
}
// Set turns the given bit to true, regardless of its previous value. Will panic
// if idx >= Size().
func (bf BitField) Set(idx uint32) {
if idx >= bf.size {
panic(fmt.Errorf("cannot set bit %d in BitField of size %d", idx, bf.size))
}
bf.data[idx>>3] |= 1 << (idx % 8)
}
// Clear turns the given bit to false, regardless of its previous value. Will
// panic if idx >= Size().
func (bf BitField) Clear(idx uint32) {
if idx >= bf.size {
panic(fmt.Errorf("cannot clear bit %d in BitField of size %d", idx, bf.size))
}
bf.data[idx>>3] &^= 1 << (idx % 8)
}
// All returns true iff all of the bits are equal to `val`.
func (bf BitField) All(val bool) bool {
targ := bf.size
if !val {
targ = 0
}
return bf.CountSet() == targ
}
// CountSet returns the number of true bits.
func (bf BitField) CountSet() (ret uint32) {
for _, b := range bf.data {
if b != 0 {
ret += uint32(bitCount[b])
}
}
return ret
}
// IsSet returns the value of a given bit.
func (bf BitField) IsSet(idx uint32) bool {
return idx < bf.size && ((bf.data[idx>>3] & (1 << (idx % 8))) != 0)
}