1 files changed, 109 insertions, 0 deletions
diff --git a/internal/player/roomset.go b/internal/player/roomset.go
new file mode 100644
index 0000000..3565bb6
--- /dev/null
+++ b/internal/player/roomset.go
@@ -0,0 +1,109 @@
+package player
+
+import (
+	"encoding/base64"
+	"math/bits"
+
+	"gopkg.in/yaml.v3"
+)
+
+type RoomSet struct {
+	bits []byte
+}
+
+func (rs *RoomSet) Add(roomID int) {
+	if roomID < 0 {
+		return
+	}
+	byteIdx := roomID / 8
+	if byteIdx >= len(rs.bits) {
+		nb := make([]byte, byteIdx+1)
+		copy(nb, rs.bits)
+		rs.bits = nb
+	}
+	rs.bits[byteIdx] |= 1 << (roomID % 8)
+}
+
+func (rs *RoomSet) Has(roomID int) bool {
+	if roomID < 0 {
+		return false
+	}
+	byteIdx := roomID / 8
+	return byteIdx < len(rs.bits) && rs.bits[byteIdx]&(1<<(roomID%8)) != 0
+}
+
+func (rs *RoomSet) Len() int {
+	n := 0
+	for _, b := range rs.bits {
+		n += bits.OnesCount8(b)
+	}
+	return n
+}
+
+func (rs RoomSet) IsZero() bool {
+	return len(rs.bits) == 0
+}
+
+func (rs *RoomSet) Keys() []int {
+	var keys []int
+	for byteIdx, b := range rs.bits {
+		if b == 0 {
+			continue
+		}
+		base := byteIdx * 8
+		for bitIdx := 0; bitIdx < 8; bitIdx++ {
+			if b&(1<<bitIdx) != 0 {
+				keys = append(keys, base+bitIdx)
+			}
+		}
+	}
+	return keys
+}
+
+func (rs *RoomSet) Rekey(idMap map[int]int) (RoomSet, bool) {
+	var next RoomSet
+	changed := false
+	for _, k := range rs.Keys() {
+		newK, ok := idMap[k]
+		if !ok {
+			newK = k
+		} else if newK != k {
+			changed = true
+		}
+		next.Add(newK)
+	}
+	return next, changed
+}
+
+func NewRoomSetFrom(ids ...int) RoomSet {
+	var rs RoomSet
+	for _, id := range ids {
+		rs.Add(id)
+	}
+	return rs
+}
+
+func (rs RoomSet) MarshalYAML() (any, error) {
+	return base64.StdEncoding.EncodeToString(rs.bits), nil
+}
+
+func (rs *RoomSet) UnmarshalYAML(value *yaml.Node) error {
+	if value.Kind == yaml.ScalarNode && (value.Tag == "!!null" || value.Value == "null") {
+		rs.bits = nil
+		return nil
+	}
+	var s string
+	if err := value.Decode(&s); err != nil {
+		return err
+	}
+	if s == "" {
+		rs.bits = nil
+		return nil
+	}
+	b, err := base64.StdEncoding.DecodeString(s)
+	if err != nil {
+		return err
+	}
+	rs.bits = b
+	return nil
+}