Bag data structure implemented in Go

bag.go

package bag

import (
	"bytes"
	"errors"
	"fmt"
)

type bag struct {
	m    map[interface{}]int
	size int
}

func New(items ...interface{}) *bag {
	b := new(bag)

	b.init()

	for _, item := range items {
		b.Add(item)
	}

	return b
}

func (b *bag) String() string {
	var s bytes.Buffer

	s.WriteString("{ ")

	b.ForEach(func(item interface{}) {
		fmt.Fprintf(&s, "%v ", item)
	})

	s.WriteString("}")

	return s.String()
}

func (b *bag) Add(item interface{}) {
	if b.Has(item) {
		b.m[item] += 1
	} else {
		b.m[item] = 1
	}

	b.size++
}

func (b *bag) Merge(other *bag) {
	for k, v := range other.m {
		if b.Has(k) {
			b.m[k] += v
		} else {
			b.m[k] = v
		}

		b.size += v
	}
}

func (b *bag) Remove(item interface{}) error {
	if b.Has(item) {
		if b.m[item] > 1 {
			b.m[item]--
			b.size--
		} else {
			b.RemoveAll(item)
		}

		return nil
	}

	return errors.New("can not remove item not present in the bag")
}

func (b *bag) RemoveAll(item interface{}) error {
	if b.Has(item) {
		b.size -= b.Count(item)
		delete(b.m, item)
		return nil
	}

	return errors.New("can not remove item not present in the bag")
}

func (b *bag) init() {
	b.m = make(map[interface{}]int)
	b.size = 0
}

func (b *bag) Clear() {
	b.init()
}

func (b *bag) IsEmpty() bool {
	return b.size == 0
}

func (b *bag) Count(item interface{}) int {
	if b.Has(item) {
		return b.m[item]
	}

	return 0
}

func (b *bag) Size() int {
	return b.size
}

func (b *bag) ForEach(f func(interface{})) {
	for k, count := range b.m {
		for i := 0; i < count; i++ {
			f(k)
		}
	}
}

func (b *bag) Has(item interface{}) bool {
	_, present := b.m[item]
	return present
}

bag_test.go

package bag

import "testing"

func TestNew(t *testing.T) {
	b := New(1, 1, 1, 1, 2, 2, 2, 3, 3, 9)

	if c := b.Count(1); c != 4 {
		t.Errorf("%v != %v", c, 4)
	}

	if c := b.Count(2); c != 3 {
		t.Errorf("%v != %v", c, 3)
	}

	if c := b.Count(3); c != 2 {
		t.Errorf("%v != %v", c, 2)
	}

	if c := b.Count(9); c != 1 {
		t.Errorf("%v != %v", c, 1)
	}

	if c := b.Count(100500); c != 0 {
		t.Errorf("%v != %v", c, 0)
	}
}

func TestAdd(t *testing.T) {
	b := New()

	b.Add(1)
	b.Add(1)
	b.Add(2)

	if c := b.Count(1); c != 2 {
		t.Errorf("%v != %v", c, 2)
	}

	if c := b.Count(2); c != 1 {
		t.Errorf("%v != %v", c, 1)
	}

	if c := b.Count(3); c != 0 {
		t.Errorf("%v != %v", c, 0)
	}
}

func TestIsEmpty(t *testing.T) {
	b := New()

	if !b.IsEmpty() {
		t.Error("newly created bag without arguments should be empty")
	}

	b.Add(42)

	if b.IsEmpty() {
		t.Error("bag with contents should not be empty")
	}
}

func TestClear(t *testing.T) {
	b := New(1, 2, 3)
	b.Clear()

	if !b.IsEmpty() {
		t.Error("freshly cleared bag should be empty")
	}
}

func TestRemove(t *testing.T) {
	b := New(1, 2, 2, 3, 4, 3)

	err := b.Remove(2)

	if err != nil {
		t.Error(err)
	}

	if c := b.Count(2); c != 1 {
		t.Errorf("%d != %d", c, 1)
	}
}

func TestRemoveAll(t *testing.T) {
	b := New(1, 2, 2, 2, 3)

	err := b.RemoveAll(2)

	if err != nil {
		t.Error(err)
	}

	if c := b.Count(2); c != 0 {
		t.Errorf("%d != %d", c, 0)
	}

}

func TestSize(t *testing.T) {
	b := New()

	if s := b.Size(); s != 0 {
		t.Error("size of empty bag should be zero")
	}

	b.Add(true)

	if s := b.Size(); s != 1 {
		t.Errorf("%v != %v", s, 1)
	}

	b.Add(true)

	if s := b.Size(); s != 2 {
		t.Errorf("%v != %v", s, 2)
	}

	b.Clear()

	if s := b.Size(); s != 0 {
		t.Error("size of empty bag should be zero")
	}
}

func TestForEach(t *testing.T) {
	b := New()

	for i := 0; i <= 9; i++ {
		b.Add(i)
	}

	iteratedOver := make([]bool, 10)

	for i := range iteratedOver {
		iteratedOver[i] = false
	}

	b.ForEach(func(item interface{}) {
		iteratedOver[item.(int)] = true
	})

	for _, iterated := range iteratedOver {
		if !iterated {
			t.Error("should have iterated over the whole thing")
		}
	}
}

func TestForEachWithDuplicates(t *testing.T) {
	b := New(1, 1, 1, 1, 1, 1)

	onesSeen := 0

	b.ForEach(func(item interface{}) {
		if item.(int) == 1 {
			onesSeen++
		}
	})

	if onesSeen != 6 {
		t.Errorf("%d != %d", onesSeen, 6)
	}
}

func TestHas(t *testing.T) {
	b := New()

	if b.Has(0) {
		t.Error("empty list does not have any elements")
	}

	b.Add(0)

	if !b.Has(0) {
		t.Error("should have item after it was added")
	}

	b.Remove(0)

	if b.Has(0) {
		t.Error("should not have item after it was removed")
	}
}

func TestMerge(t *testing.T) {
	b1 := New(1, 2)
	b2 := New(2, 3)

	b1.Merge(b2)

	if b1.Count(1) != 1 || b1.Count(2) != 2 || b1.Count(3) != 1 {
		t.Error("did not add correctly")
	}

	if s := b1.Size(); s != 4 {
		t.Errorf("%d != %d", s, 4)
	}
}

func TestString(t *testing.T) {
	b := New(1, 1, 1)

	actual := b.String()
	expected := "{ 1 1 1 }"

	if actual != expected {
		t.Errorf("%s != %s", actual, expected)
	}
}