dummy bitset in golang

bitset.go

package main

import (
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"log"
)

type BitSet struct {
	data []uint64
	size uint
}

func New(s uint) *BitSet {
	return &BitSet{
		data: make([]uint64, s/64+1),
		size: s/64 + 1,
	}
}

func (b *BitSet) Get(i uint) bool {
	v := b.data[i/64]
	return v>>((i%64)-1)&1 == 1
}

func (b *BitSet) Set(i uint) bool {
	if b.Get(i) {
		return false
	} else {
		v := b.data[i/64]
		offset := (i % 64) - 1
		v |= 1 << offset
		b.data[i/64] = v
		return true
	}
}

func (b BitSet) String() string {
	header := fmt.Sprintf("Size [%d] Data", b.size)
	strBuf := make([]string, b.size)
	for i := 0; i < int(b.size); i++ {
		strBuf[i] = fmt.Sprintf("%b", b.data[i])
	}
	return fmt.Sprintf("%s %s", header, strings.Join(strBuf, ","))
}

func (b BitSet) MarshalJSON() ([]byte, error) {
	r := make([]byte, b.size*8)
	for i, v := range b.data {
		binary.BigEndian.PutUint64(r[i*8:], v)
	}
	return json.Marshal(r)
}

func (b *BitSet) UnmarshalJSON(data []byte) error {
	if b == nil {
		return errors.New("json.*BitSet: UnmarshalJSON on nil pointer")
	}
	src := make([]byte, len(data))
	err := json.Unmarshal(data, &src)
	if err != nil {
		return err
	}
	bSize := len(src) / 8
	if bSize > 0 && (len(src)%8) > 0 {
		return errors.New("json.*BitSet: UnmarshalJSON on size illegal")
	}
	b.size = uint(bSize)
	b.data = make([]uint64, bSize)
	for i := 0; i < bSize; i++ {
		b.data[i] = binary.BigEndian.Uint64(src[i*8:])
	}
	return nil
}

func TestBitSet() {
	bs := New(65)
	set := []uint{1, 2, 3, 4, 65}
	for _, v := range set {
		if bs.Set(v) == false {
			log.Fatal("Set Cannot be false %v", v)
		}
		if bs.Set(v) {
			log.Fatal("Set Cannot be true %v", v)
		}
		if !bs.Get(v) {
			log.Fatal("Get Cannot be false %v", v)
		}
	}
}

func TestMarshal(bs *BitSet) ([]byte, error) {
	return json.Marshal(bs)
}

func TestUnmarshal(data []byte) (*BitSet, error) {
	bs := BitSet{}
	err := json.Unmarshal(data, &bs)
	return &bs, err
}

func main() {
	TestBitSet()

	bs := New(65)
	set := []uint{1, 2, 3, 4, 65}
	for _, v := range set {
		bs.Set(v)
	}
	b, err := TestMarshal(bs)
	if err != nil {
		log.Fatal("Marshal BitSet error [%v]", err)
	}
	bs2, err := TestUnmarshal(b)
	if err != nil {
		log.Fatal("Unmarshal BitSet error [%v]", err)
	}
	fmt.Println(bs, bs2)
}