Initial commit: Lux math library

set/bits.go

@@ -0,0 +1,102 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"encoding/hex"
+	"math/big"
+	"math/bits"
+)
+
+// Bits is a bit-set backed by a big.Int
+// Holds values ranging from [0, INT_MAX] (arch-dependent)
+// Trying to use negative values will result in a panic.
+// This implementation is NOT thread-safe.
+type Bits struct {
+	bits *big.Int
+}
+
+// NewBits returns a new instance of Bits with [bits] set to 1.
+//
+// Invariants:
+// 1. Negative bits will cause a panic.
+// 2. Duplicate bits are allowed but will cause a no-op.
+func NewBits(bits ...int) Bits {
+	b := Bits{new(big.Int)}
+	for _, bit := range bits {
+		b.Add(bit)
+	}
+	return b
+}
+
+// Add sets the [i]'th bit to 1
+func (b Bits) Add(i int) {
+	b.bits.SetBit(b.bits, i, 1)
+}
+
+// Union performs the set union with another set.
+// This adds all elements in [other] to [b]
+func (b Bits) Union(other Bits) {
+	b.bits.Or(b.bits, other.bits)
+}
+
+// Intersection performs the set intersection with another set
+// This sets [b] to include only elements in both [b] and [other]
+func (b Bits) Intersection(other Bits) {
+	b.bits.And(b.bits, other.bits)
+}
+
+// Difference removes all the elements in [other] from this set
+func (b Bits) Difference(other Bits) {
+	b.bits.AndNot(b.bits, other.bits)
+}
+
+// Remove sets the [i]'th bit to 0
+func (b Bits) Remove(i int) {
+	b.bits.SetBit(b.bits, i, 0)
+}
+
+// Clear empties out the bitset
+func (b Bits) Clear() {
+	b.bits.SetUint64(0)
+}
+
+// Contains returns true if the [i]'th bit is 1, and false otherwise
+func (b Bits) Contains(i int) bool {
+	return b.bits.Bit(i) == 1
+}
+
+// BitLen returns the bit length of this bitset
+func (b Bits) BitLen() int {
+	return b.bits.BitLen()
+}
+
+// Len returns the amount of 1's in the bitset
+//
+// This is typically referred to as the "Hamming Weight"
+// of a set of bits.
+func (b Bits) Len() int {
+	result := 0
+	for _, word := range b.bits.Bits() {
+		result += bits.OnesCount(uint(word))
+	}
+	return result
+}
+
+// Returns the byte representation of this bitset
+func (b Bits) Bytes() []byte {
+	return b.bits.Bytes()
+}
+
+// Inverse of Bits.Bytes()
+func BitsFromBytes(bytes []byte) Bits {
+	return Bits{
+		bits: new(big.Int).SetBytes(bytes),
+	}
+}
+
+// String returns the hex representation of this bitset
+func (b Bits) String() string {
+	return hex.EncodeToString(b.bits.Bytes())
+}

set/bits_64.go

@@ -0,0 +1,57 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"fmt"
+	"math/bits"
+)
+
+// Bits64 is a set that can contain uints in the range [0, 64). All functions
+// are O(1). The zero value is the empty set.
+type Bits64 uint64
+
+// Add [i] to the set of ints
+func (b *Bits64) Add(i uint) {
+	*b |= 1 << i
+}
+
+// Union adds all the elements in [s] to this set
+func (b *Bits64) Union(s Bits64) {
+	*b |= s
+}
+
+// Intersection takes the intersection of [s] with this set
+func (b *Bits64) Intersection(s Bits64) {
+	*b &= s
+}
+
+// Difference removes all the elements in [s] from this set
+func (b *Bits64) Difference(s Bits64) {
+	*b &^= s
+}
+
+// Remove [i] from the set of ints
+func (b *Bits64) Remove(i uint) {
+	*b &^= 1 << i
+}
+
+// Clear removes all elements from this set
+func (b *Bits64) Clear() {
+	*b = 0
+}
+
+// Contains returns true if [i] was previously added to this set
+func (b Bits64) Contains(i uint) bool {
+	return b&(1<<i) != 0
+}
+
+// Len returns the number of elements in this set
+func (b Bits64) Len() int {
+	return bits.OnesCount64(uint64(b))
+}
+
+func (b Bits64) String() string {
+	return fmt.Sprintf("%016x", uint64(b))
+}

set/bits_64_test.go

@@ -0,0 +1,117 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBits64(t *testing.T) {
+	require := require.New(t)
+
+	var bs1 Bits64
+	require.Empty(bs1)
+
+	bs1.Add(5)
+	require.Equal(1, bs1.Len())
+	require.True(bs1.Contains(5))
+
+	bs1.Add(10)
+	require.Equal(2, bs1.Len())
+	require.True(bs1.Contains(5))
+	require.True(bs1.Contains(10))
+
+	bs1.Add(10)
+	require.Equal(2, bs1.Len())
+	require.True(bs1.Contains(5))
+	require.True(bs1.Contains(10))
+
+	var bs2 Bits64
+	require.Empty(bs2)
+
+	bs2.Add(0)
+	require.Equal(1, bs2.Len())
+	require.True(bs2.Contains(0))
+
+	bs2.Union(bs1)
+	require.Equal(2, bs1.Len())
+	require.True(bs1.Contains(5))
+	require.True(bs1.Contains(10))
+	require.Equal(3, bs2.Len())
+	require.True(bs2.Contains(0))
+	require.True(bs2.Contains(5))
+	require.True(bs2.Contains(10))
+
+	bs1.Clear()
+	require.Empty(bs1)
+	require.Equal(3, bs2.Len())
+	require.True(bs2.Contains(0))
+	require.True(bs2.Contains(5))
+	require.True(bs2.Contains(10))
+
+	bs1.Add(63)
+	require.Equal(1, bs1.Len())
+	require.True(bs1.Contains(63))
+
+	bs1.Add(1)
+	require.Equal(2, bs1.Len())
+	require.True(bs1.Contains(1))
+	require.True(bs1.Contains(63))
+
+	bs1.Remove(63)
+	require.Equal(1, bs1.Len())
+	require.True(bs1.Contains(1))
+
+	var bs3 Bits64
+	require.Empty(bs3)
+
+	bs3.Add(0)
+	bs3.Add(2)
+	bs3.Add(5)
+
+	var bs4 Bits64
+	require.Empty(bs4)
+
+	bs4.Add(2)
+	bs4.Add(5)
+
+	bs3.Intersection(bs4)
+
+	require.Equal(2, bs3.Len())
+	require.True(bs3.Contains(2))
+	require.True(bs3.Contains(5))
+	require.Equal(2, bs4.Len())
+	require.True(bs4.Contains(2))
+	require.True(bs4.Contains(5))
+
+	var bs5 Bits64
+	require.Empty(bs5)
+
+	bs5.Add(7)
+	bs5.Add(11)
+	bs5.Add(9)
+
+	var bs6 Bits64
+	require.Empty(bs6)
+
+	bs6.Add(9)
+	bs6.Add(11)
+
+	bs5.Difference(bs6)
+	require.Equal(1, bs5.Len())
+	require.True(bs5.Contains(7))
+	require.Equal(2, bs6.Len())
+	require.True(bs6.Contains(9))
+	require.True(bs6.Contains(11))
+}
+
+func TestBits64String(t *testing.T) {
+	var bs Bits64
+
+	bs.Add(17)
+
+	require.Equal(t, "0000000000020000", bs.String())
+}

set/bits_test.go

@@ -0,0 +1,508 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_Bits_New(t *testing.T) {
+	tests := []struct {
+		name   string
+		bits   []int
+		length int
+	}{
+		{
+			name:   "empty",
+			bits:   []int{},
+			length: 0,
+		},
+		{
+			name:   "populated",
+			bits:   []int{0, 9, 99, 999, 9999},
+			length: 10_000,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+			b := NewBits(test.bits...)
+
+			for _, bit := range test.bits {
+				require.True(b.Contains(bit))
+			}
+
+			require.Equal(test.length, b.BitLen())
+		})
+	}
+}
+
+func Test_Bits_AddRemove(t *testing.T) {
+	tests := []struct {
+		name             string
+		toAdd            []int
+		toRemove         []int
+		expectedElements []int
+		expectedLen      int
+	}{
+		{
+			name:             "empty sets",
+			toAdd:            []int{},
+			toRemove:         []int{},
+			expectedElements: []int{}, // []
+			expectedLen:      0,
+		},
+		{
+			name:             "add only",
+			toAdd:            []int{0, 1, 2},
+			toRemove:         []int{},
+			expectedElements: []int{0, 1, 2}, // [1, 1, 1]
+			expectedLen:      3,
+		},
+		{
+			name:             "remove left-most",
+			toAdd:            []int{0, 1, 2},
+			toRemove:         []int{0},
+			expectedElements: []int{1, 2}, // [1, 1, 0]
+			expectedLen:      3,
+		},
+		{
+			name:             "remove middle",
+			toAdd:            []int{0, 1, 2},
+			toRemove:         []int{1},
+			expectedElements: []int{2, 0}, // [1, 0, 1]
+			expectedLen:      3,
+		},
+		{
+			name:             "remove right-most",
+			toAdd:            []int{0, 1, 2},
+			toRemove:         []int{2},
+			expectedElements: []int{0, 1}, // [1, 1]
+			expectedLen:      2,
+		},
+		{
+			name:             "remove all",
+			toAdd:            []int{0, 1, 2},
+			toRemove:         []int{0, 1, 2},
+			expectedElements: []int{}, // [1, 1, 1]
+			expectedLen:      0,
+		},
+		{
+			name:             "remove reverse-order",
+			toAdd:            []int{0, 1, 2},
+			toRemove:         []int{2, 1, 0},
+			expectedElements: []int{}, // []
+			expectedLen:      0,
+		},
+		{
+			name:             "remove non-existent elements",
+			toAdd:            []int{0, 1, 2},
+			toRemove:         []int{3, 4, 5},
+			expectedElements: []int{0, 1, 2}, // [1, 1, 1]
+			expectedLen:      3,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+			b := NewBits()
+
+			for _, add := range test.toAdd {
+				b.Add(add)
+			}
+
+			for _, remove := range test.toRemove {
+				b.Remove(remove)
+			}
+
+			for _, element := range test.expectedElements {
+				require.True(b.Contains(element))
+			}
+
+			require.Equal(test.expectedLen, b.BitLen())
+		})
+	}
+}
+
+func Test_Bits_Union(t *testing.T) {
+	tests := []struct {
+		name        string
+		left        []int
+		right       []int
+		expected    []int
+		expectedLen int
+	}{
+		{
+			name:        "empty sets",
+			left:        []int{},
+			right:       []int{},
+			expected:    []int{}, // []
+			expectedLen: 0,
+		},
+		{
+			name:        "left and right are same",
+			left:        []int{2, 1, 0},
+			right:       []int{2, 1, 0},
+			expected:    []int{2, 1, 0}, // [1, 1, 1]
+			expectedLen: 3,
+		},
+		{
+			name:        "left and no right",
+			left:        []int{2, 1, 0},
+			right:       []int{},
+			expected:    []int{2, 1, 0}, // [1, 1, 1]
+			expectedLen: 3,
+		},
+		{
+			name:        "right and no left",
+			left:        []int{},
+			right:       []int{2, 1, 0},
+			expected:    []int{2, 1, 0}, // [1, 1, 1]
+			expectedLen: 3,
+		},
+		{
+			name:        "left and right overlap",
+			left:        []int{2, 1},
+			right:       []int{1, 0},
+			expected:    []int{2, 1, 0}, // [1, 1, 1]
+			expectedLen: 3,
+		},
+		{
+			name:        "left and right overlap different sizes",
+			left:        []int{5, 3, 1},
+			right:       []int{8, 6, 4, 2, 0},
+			expected:    []int{8, 6, 5, 4, 3, 2, 1, 0}, // [1, 0, 1, 1, 1, 1, 1, 1, 1]
+			expectedLen: 9,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+			b := NewBits()
+
+			for _, add := range test.left {
+				b.Add(add)
+			}
+			for _, add := range test.right {
+				b.Add(add)
+			}
+
+			for _, element := range test.expected {
+				require.True(b.Contains(element))
+			}
+
+			require.Equal(test.expectedLen, b.BitLen())
+		})
+	}
+}
+
+func Test_Bits_Intersection(t *testing.T) {
+	tests := []struct {
+		name        string
+		left        []int
+		right       []int
+		expected    []int
+		expectedLen int
+	}{
+		{
+			name:        "empty sets",
+			left:        []int{},
+			right:       []int{},
+			expected:    []int{}, // []
+			expectedLen: 0,
+		},
+		{
+			name:        "left and right are same",
+			left:        []int{2, 1, 0},
+			right:       []int{2, 1, 0},
+			expected:    []int{2, 1, 0}, // [1, 1, 1]
+			expectedLen: 3,
+		},
+		{
+			name:        "left and no right",
+			left:        []int{2, 1, 0},
+			right:       []int{},
+			expected:    []int{}, // []
+			expectedLen: 0,
+		},
+		{
+			name:        "right and no left",
+			left:        []int{},
+			right:       []int{2, 1, 0},
+			expected:    []int{}, // []
+			expectedLen: 0,
+		},
+		{
+			name:        "left and right overlap",
+			left:        []int{2, 1},
+			right:       []int{1, 0},
+			expected:    []int{1}, // [1, 0]
+			expectedLen: 2,
+		},
+		{
+			name:        "left and right overlap different sizes",
+			left:        []int{5, 3, 1},
+			right:       []int{8, 6, 4, 2, 0},
+			expected:    []int{}, // []
+			expectedLen: 0,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+			left := NewBits()
+			right := NewBits()
+			for _, add := range test.left {
+				left.Add(add)
+			}
+			for _, add := range test.right {
+				right.Add(add)
+			}
+
+			left.Intersection(right)
+
+			expected := NewBits()
+			for _, element := range test.expected {
+				expected.Add(element)
+			}
+
+			require.ElementsMatch(left.bits.Bits(), expected.bits.Bits())
+		})
+	}
+}
+
+func Test_Bits_Difference(t *testing.T) {
+	tests := []struct {
+		name        string
+		left        []int
+		right       []int
+		expected    []int
+		expectedLen int
+	}{
+		{
+			name:        "empty sets",
+			left:        []int{},
+			right:       []int{},
+			expected:    []int{}, // []
+			expectedLen: 0,
+		},
+		{
+			name:        "left and right are same",
+			left:        []int{2, 1, 0},
+			right:       []int{2, 1, 0},
+			expected:    []int{}, // []
+			expectedLen: 0,
+		},
+		{
+			name:        "left and no right",
+			left:        []int{2, 1, 0},
+			right:       []int{},
+			expected:    []int{2, 1, 0}, // [1, 1, 1]
+			expectedLen: 3,
+		},
+		{
+			name:        "right and no left",
+			left:        []int{},
+			right:       []int{2, 1, 0},
+			expected:    []int{}, // []
+			expectedLen: 3,
+		},
+		{
+			name:        "left and right overlap",
+			left:        []int{2, 1},
+			right:       []int{1, 0},
+			expected:    []int{2}, // [1, 0, 0]
+			expectedLen: 3,
+		},
+		{
+			name:        "left and right overlap different sizes",
+			left:        []int{5, 3, 1},
+			right:       []int{8, 6, 4, 2, 0},
+			expected:    []int{5, 3, 1}, // [1, 0, 1, 0, 1, 0]
+			expectedLen: 6,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+			left := NewBits()
+			right := NewBits()
+			for _, add := range test.left {
+				left.Add(add)
+			}
+			for _, add := range test.right {
+				right.Add(add)
+			}
+
+			left.Difference(right)
+
+			expected := NewBits()
+			for _, element := range test.expected {
+				expected.Add(element)
+			}
+
+			require.ElementsMatch(left.bits.Bits(), expected.bits.Bits())
+		})
+	}
+}
+
+func Test_Bits_Clear(t *testing.T) {
+	tests := []struct {
+		name   string
+		bitset []int
+	}{
+		{
+			name:   "empty",
+			bitset: []int{}, // []
+		},
+		{
+			name:   "populated",
+			bitset: []int{5, 4, 3, 2, 1}, // [1, 1, 1, 1, 1]
+		},
+		{
+			name:   "populated - big",
+			bitset: []int{255}, // [1, 0...]
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+			b := NewBits()
+
+			for bit := range test.bitset {
+				b.Add(bit)
+			}
+
+			b.Clear()
+
+			require.Zero(b.BitLen())
+		})
+	}
+}
+
+func Test_Bits_String(t *testing.T) {
+	tests := []struct {
+		name     string
+		bitset   []int
+		expected string
+	}{
+		{
+			name:     "empty",
+			bitset:   []int{},
+			expected: "", // []
+		},
+		{
+			name:     "populated",
+			bitset:   []int{7, 6, 5, 4, 3, 2, 1, 0}, // [1, 1, 1, 1, 1, 1, 1, 1]
+			expected: "ff",
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+			b := NewBits()
+
+			for _, bit := range test.bitset {
+				b.Add(bit)
+			}
+
+			require.Equal(test.expected, b.String())
+		})
+	}
+}
+
+func Test_Bits_Len(t *testing.T) {
+	tests := []struct {
+		name     string
+		bitset   []int
+		expected int
+	}{
+		{
+			name:     "empty",
+			bitset:   []int{}, // []
+			expected: 0,
+		},
+		{
+			name:     "populated - more than one word",
+			bitset:   []int{255}, // [1, 0...]
+			expected: 1,
+		},
+		{
+			name:     "populated - all ones",
+			bitset:   []int{5, 4, 3, 2, 1, 0}, // [1, 1, 1, 1, 1, 1]
+			expected: 6,
+		},
+		{
+			name:     "populated - trailing zeroes",
+			bitset:   []int{5, 4, 3}, // [1, 1, 1, 0, 0, 0]
+			expected: 3,
+		},
+		{
+			name:     "populated - interwoven 1",
+			bitset:   []int{4, 2, 0}, // [1, 0, 1, 0, 1]
+			expected: 3,
+		},
+		{
+			name:     "populated - interwoven 2",
+			bitset:   []int{3, 1}, // [1, 0, 1, 0]
+			expected: 2,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require := require.New(t)
+			b := NewBits()
+
+			for _, bit := range test.bitset {
+				b.Add(bit)
+			}
+
+			require.Equal(test.expected, b.Len())
+		})
+	}
+}
+
+func Test_Bits_Bytes(t *testing.T) {
+	type test struct {
+		name string
+		elts []int
+	}
+
+	tests := []test{
+		{
+			name: "empty",
+			elts: []int{},
+		},
+		{
+			name: "single; element > 63",
+			elts: []int{1337},
+		},
+		{
+			name: "multiple",
+			elts: []int{1, 2, 3},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+
+			b := NewBits(tt.elts...)
+			bytes := b.Bytes()
+			fromBytes := BitsFromBytes(bytes)
+
+			require.Equal(len(tt.elts), fromBytes.Len())
+			for _, elt := range tt.elts {
+				require.True(fromBytes.Contains(elt))
+			}
+		})
+	}
+}

set/sampleable_set.go

@@ -0,0 +1,235 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"bytes"
+	"encoding/json"
+	"slices"
+
+	"github.com/luxfi/node/utils"
+	"github.com/luxfi/node/utils/sampler"
+	"github.com/luxfi/node/utils/wrappers"
+
+	avajson "github.com/luxfi/node/utils/json"
+)
+
+var _ json.Marshaler = (*Set[int])(nil)
+
+// SampleableSet is a set of elements that supports sampling.
+type SampleableSet[T comparable] struct {
+	// indices maps the element in the set to the index that it appears in
+	// elements.
+	indices  map[T]int
+	elements []T
+}
+
+// OfSampleable returns a Set initialized with [elts]
+func OfSampleable[T comparable](elts ...T) SampleableSet[T] {
+	s := NewSampleableSet[T](len(elts))
+	s.Add(elts...)
+	return s
+}
+
+// Return a new sampleable set with initial capacity [size].
+// More or less than [size] elements can be added to this set.
+// Using NewSampleableSet() rather than SampleableSet[T]{} is just an
+// optimization that can be used if you know how many elements will be put in
+// this set.
+func NewSampleableSet[T comparable](size int) SampleableSet[T] {
+	if size < 0 {
+		return SampleableSet[T]{}
+	}
+	return SampleableSet[T]{
+		indices:  make(map[T]int, size),
+		elements: make([]T, 0, size),
+	}
+}
+
+// Add all the elements to this set.
+// If the element is already in the set, nothing happens.
+func (s *SampleableSet[T]) Add(elements ...T) {
+	s.resize(2 * len(elements))
+	for _, e := range elements {
+		s.add(e)
+	}
+}
+
+// Union adds all the elements from the provided set to this set.
+func (s *SampleableSet[T]) Union(set SampleableSet[T]) {
+	s.resize(2 * set.Len())
+	for _, e := range set.elements {
+		s.add(e)
+	}
+}
+
+// Difference removes all the elements in [set] from [s].
+func (s *SampleableSet[T]) Difference(set SampleableSet[T]) {
+	for _, e := range set.elements {
+		s.remove(e)
+	}
+}
+
+// Contains returns true iff the set contains this element.
+func (s SampleableSet[T]) Contains(e T) bool {
+	_, contains := s.indices[e]
+	return contains
+}
+
+// Overlaps returns true if the intersection of the set is non-empty
+func (s SampleableSet[T]) Overlaps(big SampleableSet[T]) bool {
+	small := s
+	if small.Len() > big.Len() {
+		small, big = big, small
+	}
+
+	for _, e := range small.elements {
+		if _, ok := big.indices[e]; ok {
+			return true
+		}
+	}
+	return false
+}
+
+// Len returns the number of elements in this set.
+func (s SampleableSet[_]) Len() int {
+	return len(s.elements)
+}
+
+// Remove all the given elements from this set.
+// If an element isn't in the set, it's ignored.
+func (s *SampleableSet[T]) Remove(elements ...T) {
+	for _, e := range elements {
+		s.remove(e)
+	}
+}
+
+// Clear empties this set
+func (s *SampleableSet[T]) Clear() {
+	clear(s.indices)
+	for i := range s.elements {
+		s.elements[i] = utils.Zero[T]()
+	}
+	s.elements = s.elements[:0]
+}
+
+// List converts this set into a list
+func (s SampleableSet[T]) List() []T {
+	return slices.Clone(s.elements)
+}
+
+// Equals returns true if the sets contain the same elements
+func (s SampleableSet[T]) Equals(other SampleableSet[T]) bool {
+	if len(s.indices) != len(other.indices) {
+		return false
+	}
+	for k := range s.indices {
+		if _, ok := other.indices[k]; !ok {
+			return false
+		}
+	}
+	return true
+}
+
+func (s SampleableSet[T]) Sample(numToSample int) []T {
+	if numToSample <= 0 {
+		return nil
+	}
+
+	uniform := sampler.NewUniform()
+	uniform.Initialize(uint64(len(s.elements)))
+	indices, _ := uniform.Sample(min(len(s.elements), numToSample))
+	elements := make([]T, len(indices))
+	for i, index := range indices {
+		elements[i] = s.elements[index]
+	}
+	return elements
+}
+
+func (s *SampleableSet[T]) UnmarshalJSON(b []byte) error {
+	str := string(b)
+	if str == avajson.Null {
+		return nil
+	}
+	var elements []T
+	if err := json.Unmarshal(b, &elements); err != nil {
+		return err
+	}
+	s.Clear()
+	s.Add(elements...)
+	return nil
+}
+
+func (s *SampleableSet[_]) MarshalJSON() ([]byte, error) {
+	var (
+		elementBytes = make([][]byte, len(s.elements))
+		err          error
+	)
+	for i, e := range s.elements {
+		elementBytes[i], err = json.Marshal(e)
+		if err != nil {
+			return nil, err
+		}
+	}
+	// Sort for determinism
+	slices.SortFunc(elementBytes, bytes.Compare)
+
+	// Build the JSON
+	var (
+		jsonBuf = bytes.Buffer{}
+		errs    = wrappers.Errs{}
+	)
+	_, err = jsonBuf.WriteString("[")
+	errs.Add(err)
+	for i, elt := range elementBytes {
+		_, err := jsonBuf.Write(elt)
+		errs.Add(err)
+		if i != len(elementBytes)-1 {
+			_, err := jsonBuf.WriteString(",")
+			errs.Add(err)
+		}
+	}
+	_, err = jsonBuf.WriteString("]")
+	errs.Add(err)
+
+	return jsonBuf.Bytes(), errs.Err
+}
+
+func (s *SampleableSet[T]) resize(size int) {
+	if s.elements == nil {
+		if minSetSize > size {
+			size = minSetSize
+		}
+		s.indices = make(map[T]int, size)
+	}
+}
+
+func (s *SampleableSet[T]) add(e T) {
+	_, ok := s.indices[e]
+	if ok {
+		return
+	}
+
+	s.indices[e] = len(s.elements)
+	s.elements = append(s.elements, e)
+}
+
+func (s *SampleableSet[T]) remove(e T) {
+	indexToRemove, ok := s.indices[e]
+	if !ok {
+		return
+	}
+
+	lastIndex := len(s.elements) - 1
+	if indexToRemove != lastIndex {
+		lastElement := s.elements[lastIndex]
+
+		s.indices[lastElement] = indexToRemove
+		s.elements[indexToRemove] = lastElement
+	}
+
+	delete(s.indices, e)
+	s.elements[lastIndex] = utils.Zero[T]()
+	s.elements = s.elements[:lastIndex]
+}

set/sampleable_set_test.go

@@ -0,0 +1,174 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"encoding/json"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSampleableSet(t *testing.T) {
+	require := require.New(t)
+	id1 := 1
+
+	s := SampleableSet[int]{}
+
+	s.Add(id1)
+	require.True(s.Contains(id1))
+
+	s.Remove(id1)
+	require.False(s.Contains(id1))
+
+	s.Add(id1)
+	require.True(s.Contains(id1))
+	require.Len(s.List(), 1)
+	require.Equal(id1, s.List()[0])
+
+	s.Clear()
+	require.False(s.Contains(id1))
+
+	s.Add(id1)
+
+	s2 := SampleableSet[int]{}
+
+	require.False(s.Overlaps(s2))
+
+	s2.Union(s)
+	require.True(s2.Contains(id1))
+	require.True(s.Overlaps(s2))
+
+	s2.Difference(s)
+	require.False(s2.Contains(id1))
+	require.False(s.Overlaps(s2))
+}
+
+func TestSampleableSetClear(t *testing.T) {
+	require := require.New(t)
+
+	set := SampleableSet[int]{}
+	for i := 0; i < 25; i++ {
+		set.Add(i)
+	}
+	set.Clear()
+	require.Zero(set.Len())
+	set.Add(1337)
+	require.Equal(1, set.Len())
+}
+
+func TestSampleableSetMarshalJSON(t *testing.T) {
+	require := require.New(t)
+	set := SampleableSet[int]{}
+	{
+		asJSON, err := set.MarshalJSON()
+		require.NoError(err)
+		require.Equal("[]", string(asJSON))
+	}
+	id1, id2 := 1, 2
+	id1JSON, err := json.Marshal(id1)
+	require.NoError(err)
+	id2JSON, err := json.Marshal(id2)
+	require.NoError(err)
+	set.Add(id1)
+	{
+		asJSON, err := set.MarshalJSON()
+		require.NoError(err)
+		require.Equal(fmt.Sprintf("[%s]", string(id1JSON)), string(asJSON))
+	}
+	set.Add(id2)
+	{
+		asJSON, err := set.MarshalJSON()
+		require.NoError(err)
+		require.Equal(fmt.Sprintf("[%s,%s]", string(id1JSON), string(id2JSON)), string(asJSON))
+	}
+}
+
+func TestSampleableSetUnmarshalJSON(t *testing.T) {
+	require := require.New(t)
+	set := SampleableSet[int]{}
+	{
+		require.NoError(set.UnmarshalJSON([]byte("[]")))
+		require.Zero(set.Len())
+	}
+	id1, id2 := 1, 2
+	id1JSON, err := json.Marshal(id1)
+	require.NoError(err)
+	id2JSON, err := json.Marshal(id2)
+	require.NoError(err)
+	{
+		require.NoError(set.UnmarshalJSON([]byte(fmt.Sprintf("[%s]", string(id1JSON)))))
+		require.Equal(1, set.Len())
+		require.True(set.Contains(id1))
+	}
+	{
+		require.NoError(set.UnmarshalJSON([]byte(fmt.Sprintf("[%s,%s]", string(id1JSON), string(id2JSON)))))
+		require.Equal(2, set.Len())
+		require.True(set.Contains(id1))
+		require.True(set.Contains(id2))
+	}
+	{
+		require.NoError(set.UnmarshalJSON([]byte(fmt.Sprintf("[%d,%d,%d]", 3, 4, 5))))
+		require.Equal(3, set.Len())
+		require.True(set.Contains(3))
+		require.True(set.Contains(4))
+		require.True(set.Contains(5))
+	}
+	{
+		require.NoError(set.UnmarshalJSON([]byte(fmt.Sprintf("[%d,%d,%d, %d]", 3, 4, 5, 3))))
+		require.Equal(3, set.Len())
+		require.True(set.Contains(3))
+		require.True(set.Contains(4))
+		require.True(set.Contains(5))
+	}
+	{
+		set1 := SampleableSet[int]{}
+		set2 := SampleableSet[int]{}
+		require.NoError(set1.UnmarshalJSON([]byte(fmt.Sprintf("[%s,%s]", string(id1JSON), string(id2JSON)))))
+		require.NoError(set2.UnmarshalJSON([]byte(fmt.Sprintf("[%s,%s]", string(id2JSON), string(id1JSON)))))
+		require.True(set1.Equals(set2))
+	}
+}
+
+func TestOfSampleable(t *testing.T) {
+	tests := []struct {
+		name     string
+		elements []int
+		expected []int
+	}{
+		{
+			name:     "nil",
+			elements: nil,
+			expected: []int{},
+		},
+		{
+			name:     "empty",
+			elements: []int{},
+			expected: []int{},
+		},
+		{
+			name:     "unique elements",
+			elements: []int{1, 2, 3},
+			expected: []int{1, 2, 3},
+		},
+		{
+			name:     "duplicate elements",
+			elements: []int{1, 2, 3, 1, 2, 3},
+			expected: []int{1, 2, 3},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+
+			s := OfSampleable(tt.elements...)
+
+			require.Equal(len(tt.expected), s.Len())
+			for _, expected := range tt.expected {
+				require.True(s.Contains(expected))
+			}
+		})
+	}
+}

set/set.go

@@ -0,0 +1,194 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"bytes"
+	"encoding/json"
+	"slices"
+
+	"golang.org/x/exp/maps"
+
+	"github.com/luxfi/node/utils"
+	"github.com/luxfi/node/utils/wrappers"
+
+	avajson "github.com/luxfi/node/utils/json"
+)
+
+// The minimum capacity of a set
+const minSetSize = 16
+
+var _ json.Marshaler = (*Set[int])(nil)
+
+// Set is a set of elements.
+type Set[T comparable] map[T]struct{}
+
+// Of returns a Set initialized with [elts]
+func Of[T comparable](elts ...T) Set[T] {
+	s := NewSet[T](len(elts))
+	s.Add(elts...)
+	return s
+}
+
+// Return a new set with initial capacity [size].
+// More or less than [size] elements can be added to this set.
+// Using NewSet() rather than Set[T]{} is just an optimization that can
+// be used if you know how many elements will be put in this set.
+func NewSet[T comparable](size int) Set[T] {
+	if size < 0 {
+		return Set[T]{}
+	}
+	return make(map[T]struct{}, size)
+}
+
+func (s *Set[T]) resize(size int) {
+	if *s == nil {
+		if minSetSize > size {
+			size = minSetSize
+		}
+		*s = make(map[T]struct{}, size)
+	}
+}
+
+// Add all the elements to this set.
+// If the element is already in the set, nothing happens.
+func (s *Set[T]) Add(elts ...T) {
+	s.resize(2 * len(elts))
+	for _, elt := range elts {
+		(*s)[elt] = struct{}{}
+	}
+}
+
+// Union adds all the elements from the provided set to this set.
+func (s *Set[T]) Union(set Set[T]) {
+	s.resize(2 * set.Len())
+	for elt := range set {
+		(*s)[elt] = struct{}{}
+	}
+}
+
+// Difference removes all the elements in [set] from [s].
+func (s *Set[T]) Difference(set Set[T]) {
+	for elt := range set {
+		delete(*s, elt)
+	}
+}
+
+// Contains returns true iff the set contains this element.
+func (s *Set[T]) Contains(elt T) bool {
+	_, contains := (*s)[elt]
+	return contains
+}
+
+// Overlaps returns true if the intersection of the set is non-empty
+func (s *Set[T]) Overlaps(big Set[T]) bool {
+	small := *s
+	if small.Len() > big.Len() {
+		small, big = big, small
+	}
+
+	for elt := range small {
+		if _, ok := big[elt]; ok {
+			return true
+		}
+	}
+	return false
+}
+
+// Len returns the number of elements in this set.
+func (s Set[_]) Len() int {
+	return len(s)
+}
+
+// Remove all the given elements from this set.
+// If an element isn't in the set, it's ignored.
+func (s *Set[T]) Remove(elts ...T) {
+	for _, elt := range elts {
+		delete(*s, elt)
+	}
+}
+
+// Clear empties this set
+func (s *Set[_]) Clear() {
+	clear(*s)
+}
+
+// List converts this set into a list
+func (s Set[T]) List() []T {
+	return maps.Keys(s)
+}
+
+// Equals returns true if the sets contain the same elements
+func (s Set[T]) Equals(other Set[T]) bool {
+	return maps.Equal(s, other)
+}
+
+// Removes and returns an element.
+// If the set is empty, does nothing and returns false.
+func (s *Set[T]) Pop() (T, bool) {
+	for elt := range *s {
+		delete(*s, elt)
+		return elt, true
+	}
+	return utils.Zero[T](), false
+}
+
+func (s *Set[T]) UnmarshalJSON(b []byte) error {
+	str := string(b)
+	if str == avajson.Null {
+		return nil
+	}
+	var elts []T
+	if err := json.Unmarshal(b, &elts); err != nil {
+		return err
+	}
+	s.Clear()
+	s.Add(elts...)
+	return nil
+}
+
+func (s Set[_]) MarshalJSON() ([]byte, error) {
+	var (
+		eltBytes = make([][]byte, len(s))
+		i        int
+		err      error
+	)
+	for elt := range s {
+		eltBytes[i], err = json.Marshal(elt)
+		if err != nil {
+			return nil, err
+		}
+		i++
+	}
+	// Sort for determinism
+	slices.SortFunc(eltBytes, bytes.Compare)
+
+	// Build the JSON
+	var (
+		jsonBuf = bytes.Buffer{}
+		errs    = wrappers.Errs{}
+	)
+	_, err = jsonBuf.WriteString("[")
+	errs.Add(err)
+	for i, elt := range eltBytes {
+		_, err := jsonBuf.Write(elt)
+		errs.Add(err)
+		if i != len(eltBytes)-1 {
+			_, err := jsonBuf.WriteString(",")
+			errs.Add(err)
+		}
+	}
+	_, err = jsonBuf.WriteString("]")
+	errs.Add(err)
+
+	return jsonBuf.Bytes(), errs.Err
+}
+
+// Returns a random element. If the set is empty, returns false
+func (s *Set[T]) Peek() (T, bool) {
+	for elt := range *s {
+		return elt, true
+	}
+	return utils.Zero[T](), false
+}

set/set_benchmark_test.go

@@ -0,0 +1,39 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"strconv"
+	"testing"
+)
+
+func BenchmarkSetList(b *testing.B) {
+	sizes := []int{5, 25, 100, 100_000} // Test with various sizes
+	for size := range sizes {
+		b.Run(strconv.Itoa(size), func(b *testing.B) {
+			set := Set[int]{}
+			for i := 0; i < size; i++ {
+				set.Add(i)
+			}
+			b.ResetTimer()
+			for n := 0; n < b.N; n++ {
+				set.List()
+			}
+		})
+	}
+}
+
+func BenchmarkSetClear(b *testing.B) {
+	for _, numElts := range []int{10, 25, 50, 100, 250, 500, 1000} {
+		b.Run(strconv.Itoa(numElts), func(b *testing.B) {
+			set := NewSet[int](numElts)
+			for n := 0; n < b.N; n++ {
+				for i := 0; i < numElts; i++ {
+					set.Add(i)
+				}
+				set.Clear()
+			}
+		})
+	}
+}

set/set_test.go

@@ -0,0 +1,228 @@
+// Copyright (C) 2019-2024, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package set
+
+import (
+	"encoding/json"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSet(t *testing.T) {
+	require := require.New(t)
+	id1 := 1
+
+	s := Set[int]{id1: struct{}{}}
+
+	s.Add(id1)
+	require.True(s.Contains(id1))
+
+	s.Remove(id1)
+	require.False(s.Contains(id1))
+
+	s.Add(id1)
+	require.True(s.Contains(id1))
+	require.Len(s.List(), 1)
+	require.Equal(id1, s.List()[0])
+
+	s.Clear()
+	require.False(s.Contains(id1))
+
+	s.Add(id1)
+
+	s2 := Set[int]{}
+
+	require.False(s.Overlaps(s2))
+
+	s2.Union(s)
+	require.True(s2.Contains(id1))
+	require.True(s.Overlaps(s2))
+
+	s2.Difference(s)
+	require.False(s2.Contains(id1))
+	require.False(s.Overlaps(s2))
+}
+
+func TestOf(t *testing.T) {
+	tests := []struct {
+		name     string
+		elements []int
+		expected []int
+	}{
+		{
+			name:     "nil",
+			elements: nil,
+			expected: []int{},
+		},
+		{
+			name:     "empty",
+			elements: []int{},
+			expected: []int{},
+		},
+		{
+			name:     "unique elements",
+			elements: []int{1, 2, 3},
+			expected: []int{1, 2, 3},
+		},
+		{
+			name:     "duplicate elements",
+			elements: []int{1, 2, 3, 1, 2, 3},
+			expected: []int{1, 2, 3},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+
+			s := Of(tt.elements...)
+
+			require.Len(s, len(tt.expected))
+			for _, expected := range tt.expected {
+				require.True(s.Contains(expected))
+			}
+		})
+	}
+}
+
+func TestSetClear(t *testing.T) {
+	require := require.New(t)
+
+	set := Set[int]{}
+	for i := 0; i < 25; i++ {
+		set.Add(i)
+	}
+	set.Clear()
+	require.Empty(set)
+	set.Add(1337)
+	require.Len(set, 1)
+}
+
+func TestSetPop(t *testing.T) {
+	require := require.New(t)
+
+	var s Set[int]
+	_, ok := s.Pop()
+	require.False(ok)
+
+	s = make(Set[int])
+	_, ok = s.Pop()
+	require.False(ok)
+
+	id1, id2 := 0, 1
+	s.Add(id1, id2)
+
+	got, ok := s.Pop()
+	require.True(ok)
+	require.True(got == id1 || got == id2)
+	require.Equal(1, s.Len())
+
+	got, ok = s.Pop()
+	require.True(ok)
+	require.True(got == id1 || got == id2)
+	require.Zero(s.Len())
+
+	_, ok = s.Pop()
+	require.False(ok)
+}
+
+func TestSetMarshalJSON(t *testing.T) {
+	require := require.New(t)
+	set := Set[int]{}
+	{
+		asJSON, err := set.MarshalJSON()
+		require.NoError(err)
+		require.Equal("[]", string(asJSON))
+	}
+	id1, id2 := 1, 2
+	id1JSON, err := json.Marshal(id1)
+	require.NoError(err)
+	id2JSON, err := json.Marshal(id2)
+	require.NoError(err)
+	set.Add(id1)
+	{
+		asJSON, err := set.MarshalJSON()
+		require.NoError(err)
+		require.Equal(fmt.Sprintf("[%s]", string(id1JSON)), string(asJSON))
+	}
+	set.Add(id2)
+	{
+		asJSON, err := set.MarshalJSON()
+		require.NoError(err)
+		require.Equal(fmt.Sprintf("[%s,%s]", string(id1JSON), string(id2JSON)), string(asJSON))
+	}
+}
+
+func TestSetUnmarshalJSON(t *testing.T) {
+	require := require.New(t)
+	set := Set[int]{}
+	{
+		require.NoError(set.UnmarshalJSON([]byte("[]")))
+		require.Empty(set)
+	}
+	id1, id2 := 1, 2
+	id1JSON, err := json.Marshal(id1)
+	require.NoError(err)
+	id2JSON, err := json.Marshal(id2)
+	require.NoError(err)
+	{
+		require.NoError(set.UnmarshalJSON([]byte(fmt.Sprintf("[%s]", string(id1JSON)))))
+		require.Len(set, 1)
+		require.Contains(set, id1)
+	}
+	{
+		require.NoError(set.UnmarshalJSON([]byte(fmt.Sprintf("[%s,%s]", string(id1JSON), string(id2JSON)))))
+		require.Len(set, 2)
+		require.Contains(set, id1)
+		require.Contains(set, id2)
+	}
+	{
+		require.NoError(set.UnmarshalJSON([]byte(fmt.Sprintf("[%d,%d,%d]", 3, 4, 5))))
+		require.Len(set, 3)
+		require.Contains(set, 3)
+		require.Contains(set, 4)
+		require.Contains(set, 5)
+	}
+	{
+		require.NoError(set.UnmarshalJSON([]byte(fmt.Sprintf("[%d,%d,%d, %d]", 3, 4, 5, 3))))
+		require.Len(set, 3)
+		require.Contains(set, 3)
+		require.Contains(set, 4)
+		require.Contains(set, 5)
+	}
+	{
+		set1 := Set[int]{}
+		set2 := Set[int]{}
+		require.NoError(set1.UnmarshalJSON([]byte(fmt.Sprintf("[%s,%s]", string(id1JSON), string(id2JSON)))))
+		require.NoError(set2.UnmarshalJSON([]byte(fmt.Sprintf("[%s,%s]", string(id2JSON), string(id1JSON)))))
+		require.Equal(set1, set2)
+	}
+}
+
+func TestSetReflectJSONMarshal(t *testing.T) {
+	require := require.New(t)
+	set := Set[int]{}
+	{
+		asJSON, err := json.Marshal(set)
+		require.NoError(err)
+		require.Equal("[]", string(asJSON))
+	}
+	id1JSON, err := json.Marshal(1)
+	require.NoError(err)
+	id2JSON, err := json.Marshal(2)
+	require.NoError(err)
+	set.Add(1)
+	{
+		asJSON, err := json.Marshal(set)
+		require.NoError(err)
+		require.Equal(fmt.Sprintf("[%s]", string(id1JSON)), string(asJSON))
+	}
+	set.Add(2)
+	{
+		asJSON, err := json.Marshal(set)
+		require.NoError(err)
+		require.Equal(fmt.Sprintf("[%s,%s]", string(id1JSON), string(id2JSON)), string(asJSON))
+	}
+}