set.go

// Copyright 2014 CoreOS, Inc.
//
// 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.

package pkg

import (
	"reflect"
	"sort"
	"sync"
)

type Set interface {
	Add(string)
	Remove(string)
	Contains(string) bool
	Equals(Set) bool
	Length() int
	Values() []string
	Copy() Set
	Sub(Set) Set
}

func NewUnsafeSet(values ...string) *unsafeSet {
	set := &unsafeSet{make(map[string]struct{})}
	for _, v := range values {
		set.Add(v)
	}
	return set
}

func NewThreadsafeSet(values ...string) *tsafeSet {
	us := NewUnsafeSet(values...)
	return &tsafeSet{us, sync.RWMutex{}}
}

type unsafeSet struct {
	d map[string]struct{}
}

// Add adds a new value to the set (no-op if the value is already present)
func (us *unsafeSet) Add(value string) {
	us.d[value] = struct{}{}
}

// Remove removes the given value from the set
func (us *unsafeSet) Remove(value string) {
	delete(us.d, value)
}

// Contains returns whether the set contains the given value
func (us *unsafeSet) Contains(value string) (exists bool) {
	_, exists = us.d[value]
	return
}

// Equals returns whether the contents of two sets are identical
func (us *unsafeSet) Equals(other Set) bool {
	v1 := sort.StringSlice(us.Values())
	v2 := sort.StringSlice(other.Values())
	v1.Sort()
	v2.Sort()
	return reflect.DeepEqual(v1, v2)
}

// Length returns the number of elements in the set
func (us *unsafeSet) Length() int {
	return len(us.d)
}

// Values returns the values of the Set in an unspecified order.
func (us *unsafeSet) Values() (values []string) {
	values = make([]string, 0)
	for val, _ := range us.d {
		values = append(values, val)
	}
	return
}

// Copy creates a new Set containing the values of the first
func (us *unsafeSet) Copy() Set {
	cp := NewUnsafeSet()
	for val, _ := range us.d {
		cp.Add(val)
	}

	return cp
}

// Sub removes all elements in other from the set
func (us *unsafeSet) Sub(other Set) Set {
	oValues := other.Values()
	result := us.Copy().(*unsafeSet)

	for _, val := range oValues {
		if _, ok := result.d[val]; !ok {
			continue
		}
		delete(result.d, val)
	}

	return result
}

type tsafeSet struct {
	us *unsafeSet
	m  sync.RWMutex
}

func (ts *tsafeSet) Add(value string) {
	ts.m.Lock()
	defer ts.m.Unlock()
	ts.us.Add(value)
}

func (ts *tsafeSet) Remove(value string) {
	ts.m.Lock()
	defer ts.m.Unlock()
	ts.us.Remove(value)
}

func (ts *tsafeSet) Contains(value string) (exists bool) {
	ts.m.RLock()
	defer ts.m.RUnlock()
	return ts.us.Contains(value)
}

func (ts *tsafeSet) Equals(other Set) bool {
	ts.m.RLock()
	defer ts.m.RUnlock()
	return ts.us.Equals(other)
}

func (ts *tsafeSet) Length() int {
	ts.m.RLock()
	defer ts.m.RUnlock()
	return ts.us.Length()
}

func (ts *tsafeSet) Values() (values []string) {
	ts.m.RLock()
	defer ts.m.RUnlock()
	return ts.us.Values()
}

func (ts *tsafeSet) Copy() Set {
	ts.m.RLock()
	defer ts.m.RUnlock()
	usResult := ts.us.Copy().(*unsafeSet)
	return &tsafeSet{usResult, sync.RWMutex{}}
}

func (ts *tsafeSet) Sub(other Set) Set {
	ts.m.RLock()
	defer ts.m.RUnlock()
	usResult := ts.us.Sub(other).(*unsafeSet)
	return &tsafeSet{usResult, sync.RWMutex{}}
}