homebox/backend/internal/repo/id_set.go

package repo

import "github.com/google/uuid"

// HasID is an interface to entities that have an ID uuid.UUID field and a GetID() method.
// This interface is fulfilled by all entities generated by entgo.io/ent via a custom template
type HasID interface {
	GetID() uuid.UUID
}

// IDSet is a utility set-like type for working with sets of uuid.UUIDs within a repository
// instance. Most useful for comparing lists of UUIDs for processing relationship
// IDs and remove/adding relationships as required.
//
// # See how ItemRepo uses it to manage the Labels-To-Items relationship
//
// NOTE: may be worth moving this to a more generic package/set implementation
// or use a 3rd party set library, but this is good enough for now
type IDSet struct {
	mp map[uuid.UUID]struct{}
}

func NewIDSet(l int) *IDSet {
	return &IDSet{
		mp: make(map[uuid.UUID]struct{}, l),
	}
}

func EntitiesToIDSet[T HasID](entities []T) *IDSet {
	s := NewIDSet(len(entities))
	for _, e := range entities {
		s.Add(e.GetID())
	}
	return s
}

func (t *IDSet) Slice() []uuid.UUID {
	s := make([]uuid.UUID, 0, len(t.mp))
	for k := range t.mp {
		s = append(s, k)
	}
	return s
}

func (t *IDSet) Add(ids ...uuid.UUID) {
	for _, id := range ids {
		t.mp[id] = struct{}{}
	}
}

func (t *IDSet) Has(id uuid.UUID) bool {
	_, ok := t.mp[id]
	return ok
}

func (t *IDSet) Len() int {
	return len(t.mp)
}

func (t *IDSet) Remove(id uuid.UUID) {
	delete(t.mp, id)
}