set: Switch to a common set implementation

chat/command.go

@@ -8,7 +8,7 @@ import (
 	"strings"
 
 	"github.com/shazow/ssh-chat/chat/message"
-	"github.com/shazow/ssh-chat/common"
+	"github.com/shazow/ssh-chat/set"
 )
 
 // The error returned when an invalid command is issued.
@@ -250,8 +250,9 @@ func InitCommands(c *Commands) {
 			id := strings.TrimSpace(strings.TrimLeft(msg.Body(), "/ignore"))
 			if id == "" {
 				var names []string
-				msg.From().Ignored.Each(func(i common.Identified) {
+				msg.From().Ignored.Each(func(_ string, item set.Item) error {
-					names = append(names, i.Id())
+					names = append(names, item.Key())
+					return nil
 				})
 
 				var systemMsg string

chat/message/user.go

@@ -9,7 +9,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/shazow/ssh-chat/common"
+	"github.com/shazow/ssh-chat/set"
 )
 
 const messageBuffer = 5
@@ -25,7 +25,7 @@ type User struct {
 	joined   time.Time
 	msg      chan Message
 	done     chan struct{}
-	Ignored  *common.IdSet
+	Ignored  *set.Set
 
 	screen    io.WriteCloser
 	closeOnce sync.Once
@@ -42,7 +42,7 @@ func NewUser(identity Identifier) *User {
 		joined:     time.Now(),
 		msg:        make(chan Message, messageBuffer),
 		done:       make(chan struct{}),
-		Ignored:    common.NewIdSet(),
+		Ignored:    set.New(),
 	}
 	u.SetColorIdx(rand.Int())
 
@@ -189,20 +189,20 @@ func (u *User) Send(m Message) error {
 	return nil
 }
 
-func (u *User) Ignore(identified common.Identified) error {
+func (u *User) Ignore(other Identifier) error {
-	if identified == nil {
+	if other == nil {
 		return errors.New("user is nil.")
 	}
 
-	if identified.Id() == u.Id() {
+	if other.Id() == u.Id() {
 		return errors.New("cannot ignore self.")
 	}
 
-	if u.Ignored.In(identified) {
+	if u.Ignored.In(other.Id()) {
 		return errors.New("user already ignored.")
 	}
 
-	u.Ignored.Add(identified)
+	u.Ignored.Add(set.Itemize(other.Id(), other))
 	return nil
 }
 
@@ -211,12 +211,7 @@ func (u *User) Unignore(id string) error {
 		return errors.New("user is nil.")
 	}
 
-	identified, err := u.Ignored.Get(id)
+	return u.Ignored.Remove(id)
-	if err != nil {
-		return err
-	}
-
-	return u.Ignored.Remove(identified)
 }
 
 // Container for per-user configurations.

chat/room.go

@@ -7,7 +7,7 @@ import (
 	"sync"
 
 	"github.com/shazow/ssh-chat/chat/message"
-	"github.com/shazow/ssh-chat/common"
+	"github.com/shazow/ssh-chat/set"
 )
 
 const historyLen = 20
@@ -35,8 +35,8 @@ type Room struct {
 	closed    bool
 	closeOnce sync.Once
 
-	Members *common.IdSet
+	Members *set.Set
-	Ops     *common.IdSet
+	Ops     *set.Set
 }
 
 // NewRoom creates a new room.
@@ -48,8 +48,8 @@ func NewRoom() *Room {
 		history:   message.NewHistory(historyLen),
 		commands:  *defaultCommands,
 
-		Members: common.NewIdSet(),
+		Members: set.New(),
-		Ops:     common.NewIdSet(),
+		Ops:     set.New(),
 	}
 }
 
@@ -62,8 +62,9 @@ func (r *Room) SetCommands(commands Commands) {
 func (r *Room) Close() {
 	r.closeOnce.Do(func() {
 		r.closed = true
-		r.Members.Each(func(m common.Identified) {
+		r.Members.Each(func(_ string, item set.Item) error {
-			m.(*Member).Close()
+			item.Value().(*Member).Close()
+			return nil
 		})
 		r.Members.Clear()
 		close(r.broadcast)
@@ -96,10 +97,10 @@ func (r *Room) HandleMsg(m message.Message) {
 		}
 
 		r.history.Add(m)
-		r.Members.Each(func(u common.Identified) {
+		r.Members.Each(func(_ string, item set.Item) (err error) {
-			user := u.(*Member).User
+			user := item.Value().(*Member).User
 
-			if fromMsg != nil && user.Ignored.In(fromMsg.From()) {
+			if fromMsg != nil && user.Ignored.In(fromMsg.From().Id()) {
 				// Skip because ignored
 				return
 			}
@@ -115,6 +116,7 @@ func (r *Room) HandleMsg(m message.Message) {
 				}
 			}
 			user.Send(m)
+			return
 		})
 	}
 }
@@ -145,40 +147,40 @@ func (r *Room) Join(u *message.User) (*Member, error) {
 	if u.Id() == "" {
 		return nil, ErrInvalidName
 	}
-	member := Member{u}
+	member := &Member{u}
-	err := r.Members.Add(&member)
+	err := r.Members.Add(set.Itemize(u.Id(), member))
 	if err != nil {
 		return nil, err
 	}
 	r.History(u)
 	s := fmt.Sprintf("%s joined. (Connected: %d)", u.Name(), r.Members.Len())
 	r.Send(message.NewAnnounceMsg(s))
-	return &member, nil
+	return member, nil
 }
 
 // Leave the room as a user, will announce. Mostly used during setup.
 func (r *Room) Leave(u message.Identifier) error {
-	err := r.Members.Remove(u)
+	err := r.Members.Remove(u.Id())
 	if err != nil {
 		return err
 	}
-	r.Ops.Remove(u)
+	r.Ops.Remove(u.Id())
 	s := fmt.Sprintf("%s left.", u.Name())
 	r.Send(message.NewAnnounceMsg(s))
 	return nil
 }
 
 // Rename member with a new identity. This will not call rename on the member.
-func (r *Room) Rename(oldId string, identity message.Identifier) error {
+func (r *Room) Rename(oldId string, u message.Identifier) error {
-	if identity.Id() == "" {
+	if u.Id() == "" {
 		return ErrInvalidName
 	}
-	err := r.Members.Replace(oldId, identity)
+	err := r.Members.Replace(oldId, set.Itemize(u.Id(), u))
 	if err != nil {
 		return err
 	}
 
-	s := fmt.Sprintf("%s is now known as %s.", oldId, identity.Id())
+	s := fmt.Sprintf("%s is now known as %s.", oldId, u.Id())
 	r.Send(message.NewAnnounceMsg(s))
 	return nil
 }
@@ -202,12 +204,12 @@ func (r *Room) MemberById(id string) (*Member, bool) {
 	if err != nil {
 		return nil, false
 	}
-	return m.(*Member), true
+	return m.Value().(*Member), true
 }
 
 // IsOp returns whether a user is an operator in this room.
 func (r *Room) IsOp(u *message.User) bool {
-	return r.Ops.In(u)
+	return r.Ops.In(u.Id())
 }
 
 // Topic of the room.
@@ -223,10 +225,10 @@ func (r *Room) SetTopic(s string) {
 // NamesPrefix lists all members' names with a given prefix, used to query
 // for autocompletion purposes.
 func (r *Room) NamesPrefix(prefix string) []string {
-	members := r.Members.ListPrefix(prefix)
+	items := r.Members.ListPrefix(prefix)
-	names := make([]string, len(members))
+	names := make([]string, len(items))
-	for i, u := range members {
+	for i, item := range items {
-		names[i] = u.(*Member).User.Name()
+		names[i] = item.Value().(*Member).User.Name()
 	}
 	return names
 }

chat/set_test.go

@@ -4,35 +4,35 @@ import (
 	"testing"
 
 	"github.com/shazow/ssh-chat/chat/message"
-	"github.com/shazow/ssh-chat/common"
+	"github.com/shazow/ssh-chat/set"
 )
 
 func TestSet(t *testing.T) {
 	var err error
-	s := common.NewIdSet()
+	s := set.New()
 	u := message.NewUser(message.SimpleId("foo"))
 
-	if s.In(u) {
+	if s.In(u.Id()) {
 		t.Errorf("Set should be empty.")
 	}
 
-	err = s.Add(u)
+	err = s.Add(set.Itemize(u.Id(), u))
 	if err != nil {
 		t.Error(err)
 	}
 
-	if !s.In(u) {
+	if !s.In(u.Id()) {
 		t.Errorf("Set should contain user.")
 	}
 
 	u2 := message.NewUser(message.SimpleId("bar"))
-	err = s.Add(u2)
+	err = s.Add(set.Itemize(u2.Id(), u2))
 	if err != nil {
 		t.Error(err)
 	}
 
-	err = s.Add(u2)
+	err = s.Add(set.Itemize(u2.Id(), u2))
-	if err != common.ErrIdTaken {
+	if err != set.ErrCollision {
 		t.Error(err)
 	}
 

common/set.go

@@ -1,158 +0,0 @@
-package common
-
-import (
-	"errors"
-	"strings"
-	"sync"
-)
-
-// The error returned when an added id already exists in the set.
-var ErrIdTaken = errors.New("id already taken")
-
-// The error returned when a requested item does not exist in the set.
-var ErrIdentifiedMissing = errors.New("item does not exist")
-
-// Interface for an item storeable in the set
-type Identified interface {
-	Id() string
-}
-
-// Set with string lookup.
-// TODO: Add trie for efficient prefix lookup?
-type IdSet struct {
-	sync.RWMutex
-	lookup map[string]Identified
-}
-
-// newIdSet creates a new set.
-func NewIdSet() *IdSet {
-	return &IdSet{
-		lookup: map[string]Identified{},
-	}
-}
-
-// Clear removes all items and returns the number removed.
-func (s *IdSet) Clear() int {
-	s.Lock()
-	n := len(s.lookup)
-	s.lookup = map[string]Identified{}
-	s.Unlock()
-	return n
-}
-
-// Len returns the size of the set right now.
-func (s *IdSet) Len() int {
-	s.RLock()
-	defer s.RUnlock()
-	return len(s.lookup)
-}
-
-// In checks if an item exists in this set.
-func (s *IdSet) In(item Identified) bool {
-	id := s.normalize(item.Id())
-	s.RLock()
-	_, ok := s.lookup[id]
-	s.RUnlock()
-	return ok
-}
-
-// Get returns an item with the given Id.
-func (s *IdSet) Get(id string) (Identified, error) {
-	s.RLock()
-	item, ok := s.lookup[s.normalize(id)]
-	s.RUnlock()
-
-	if !ok {
-		return nil, ErrIdentifiedMissing
-	}
-
-	return item, nil
-}
-
-// Add item to this set if it does not exist already.
-func (s *IdSet) Add(item Identified) error {
-	s.Lock()
-	defer s.Unlock()
-
-	id := s.normalize(item.Id())
-	_, found := s.lookup[id]
-	if found {
-		return ErrIdTaken
-	}
-
-	s.lookup[id] = item
-	return nil
-}
-
-// Remove item from this set.
-func (s *IdSet) Remove(item Identified) error {
-	s.Lock()
-	defer s.Unlock()
-	id := s.normalize(item.Id())
-	_, found := s.lookup[id]
-	if !found {
-		return ErrIdentifiedMissing
-	}
-	delete(s.lookup, id)
-	return nil
-}
-
-// Replace item from old id with new Identified.
-// Used for moving the same Identified to a new Id, such as a rename.
-func (s *IdSet) Replace(oldId string, item Identified) error {
-	id := s.normalize(item.Id())
-	oldId = s.normalize(oldId)
-
-	s.Lock()
-	defer s.Unlock()
-
-	// Check if it already exists
-	_, found := s.lookup[id]
-	if found {
-		return ErrIdTaken
-	}
-
-	// Remove oldId
-	_, found = s.lookup[oldId]
-	if !found {
-		return ErrIdentifiedMissing
-	}
-	delete(s.lookup, oldId)
-
-	// Add new Identified
-	s.lookup[id] = item
-
-	return nil
-}
-
-// Each loops over every item while holding a read lock and applies fn to each
-// element.
-func (s *IdSet) Each(fn func(item Identified)) {
-	s.RLock()
-	for _, item := range s.lookup {
-		fn(item)
-	}
-	s.RUnlock()
-}
-
-// ListPrefix returns a list of items with a prefix, case insensitive.
-func (s *IdSet) ListPrefix(prefix string) []Identified {
-	r := []Identified{}
-	prefix = strings.ToLower(prefix)
-
-	s.RLock()
-	defer s.RUnlock()
-
-	for id, item := range s.lookup {
-		if !strings.HasPrefix(string(id), prefix) {
-			continue
-		}
-		r = append(r, item)
-	}
-
-	return r
-}
-
-func (s *IdSet) normalize(id string) string {
-	return strings.ToLower(id)
-}

host.go

@@ -12,6 +12,7 @@ import (
 	"github.com/shazow/rateio"
 	"github.com/shazow/ssh-chat/chat"
 	"github.com/shazow/ssh-chat/chat/message"
+	"github.com/shazow/ssh-chat/set"
 	"github.com/shazow/ssh-chat/sshd"
 )
 
@@ -126,7 +127,7 @@ func (h *Host) Connect(term *sshd.Terminal) {
 
 	// Should the user be op'd on join?
 	if h.isOp(term.Conn) {
-		h.Room.Ops.Add(member)
+		h.Room.Ops.Add(set.Itemize(member.Id(), member))
 	}
 	ratelimit := rateio.NewSimpleLimiter(3, time.Second*3)
 
@@ -493,7 +494,8 @@ func (h *Host) InitCommands(c *chat.Commands) {
 			if !ok {
 				return errors.New("user not found")
 			}
-			room.Ops.Add(member)
+			room.Ops.Add(set.Itemize(member.Id(), member))
+
 			id := member.Identifier.(*Identity)
 			h.auth.Op(id.PublicKey(), until)
 

host_test.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 
 	"github.com/shazow/ssh-chat/chat/message"
+	"github.com/shazow/ssh-chat/set"
 	"github.com/shazow/ssh-chat/sshd"
 	"golang.org/x/crypto/ssh"
 )
@@ -193,7 +194,7 @@ func TestHostKick(t *testing.T) {
 			if member == nil {
 				return errors.New("failed to load MemberById")
 			}
-			host.Room.Ops.Add(member)
+			host.Room.Ops.Add(set.Itemize(member.Id(), member))
 
 			// Block until second client is here
 			connected <- struct{}{}

set/item.go

@@ -8,6 +8,23 @@ type Item interface {
 	Value() interface{}
 }
 
+type item struct {
+	key   string
+	value interface{}
+}
+
+func (item *item) Key() string {
+	return item.key
+}
+
+func (item *item) Value() interface{} {
+	return item.value
+}
+
+func Itemize(key string, value interface{}) Item {
+	return &item{key, value}
+}
+
 type StringItem string
 
 func (item StringItem) Key() string {