github.com/mitranim/gg@v0.1.17/coll.go

package gg

import "encoding/json"

/*
Short for "valid primary key". Returns the primary key generated by the given
input, asserts that the key is non-zero, and returns the resulting key.
Used internally by `Coll`.
*/
func ValidPk[Key comparable, Val Pker[Key]](val Val) Key {
	key := val.Pk()
	if IsZero(key) {
		panic(Errf(`unexpected empty key %v in %v`, Type[Key](), Type[Val]()))
	}
	return key
}

/*
Syntactic shortcut for making a `Coll` of the given arguments. Reuses the given
slice as-is with no reallocation.
*/
func CollOf[Key comparable, Val Pker[Key]](src ...Val) Coll[Key, Val] {
	var tar Coll[Key, Val]
	tar.Reset(src...)
	return tar
}

/*
Syntactic shortcut for making a `Coll` from any number of source slices. When
called with exactly one argument, this reuses the given slice as-is with no
reallocation.
*/
func CollFrom[Key comparable, Val Pker[Key], Slice ~[]Val](src ...Slice) Coll[Key, Val] {
	var tar Coll[Key, Val]

	switch len(src) {
	case 1:
		tar.Reset(src[0]...)
	default:
		for _, src := range src {
			tar.Add(src...)
		}
	}

	return tar
}

/*
Short for "collection". Represents an ordered map where keys are automatically
derived from values. Keys must be non-zero. Similarly to a map, this ensures
value uniqueness by primary key, and allows efficient access by key. Unlike a
map, values in this type are ordered and can be iterated cheaply, because they
are stored in a publicly-accessible slice. However, as a tradeoff, this type
does not support deletion.
*/
type Coll[
	Key comparable,
	Val Pker[Key],
] struct {
	Slice []Val `role:"ref"`
	Index map[Key]int
}

// Same as `len(self.Slice)`.
func (self Coll[_, _]) Len() int { return len(self.Slice) }

// True if `.Len` <= 0. Inverse of `.IsNotEmpty`.
func (self Coll[_, _]) IsEmpty() bool { return self.Len() <= 0 }

// True if `.Len` > 0. Inverse of `.IsEmpty`.
func (self Coll[_, _]) IsNotEmpty() bool { return self.Len() > 0 }

// True if the index has the given key.
func (self Coll[Key, _]) Has(key Key) bool {
	return MapHas(self.Index, key)
}

// Returns the value indexed on the given key, or the zero value of that type.
func (self Coll[Key, Val]) Get(key Key) Val {
	return PtrGet(self.Ptr(key))
}

/*
Short for "get required". Returns the value indexed on the given key. Panics if
the value is missing.
*/
func (self Coll[Key, Val]) GetReq(key Key) Val {
	ptr := self.Ptr(key)
	if ptr != nil {
		return *ptr
	}
	panic(errCollMissing[Val](key))
}

/*
Returns the value indexed on the given key and a boolean indicating if the value
was actually present.
*/
func (self Coll[Key, Val]) Got(key Key) (Val, bool) {
	// Note: we must check `ok` because if the entry is missing, `ind` is `0`,
	// which is invalid.
	ind, ok := self.Index[key]
	if ok {
		return Got(self.Slice, ind)
	}
	return Zero[Val](), false
}

/*
Short for "pointer". Returns a pointer to the value indexed on the given key, or
nil if the value is missing.
*/
func (self Coll[Key, Val]) Ptr(key Key) *Val {
	// Note: we must check `ok` because if the entry is missing, `ind` is `0`,
	// which is invalid.
	ind, ok := self.Index[key]
	if ok {
		return GetPtr(self.Slice, ind)
	}
	return nil
}

/*
Short for "pointer required". Returns a non-nil pointer to the value indexed
on the given key, or panics if the value is missing.
*/
func (self Coll[Key, Val]) PtrReq(key Key) *Val {
	ptr := self.Ptr(key)
	if ptr != nil {
		return ptr
	}
	panic(errCollMissing[Val](key))
}

/*
Idempotently adds each given value to both the inner slice and the inner index.
Every value whose key already exists in the index is replaced at the existing
position in the slice.
*/
func (self *Coll[Key, Val]) Add(src ...Val) *Coll[Key, Val] {
	index := MapInit(&self.Index)

	for _, val := range src {
		key := ValidPk[Key](val)
		ind, ok := index[key]
		if ok {
			self.Slice[ind] = val
			continue
		}
		index[key] = AppendIndex(&self.Slice, val)
	}

	return self
}

/*
Same as `Coll.Add`, but panics if any inputs are redundant, as in, their primary
keys are already present in the index.
*/
func (self *Coll[Key, Val]) AddUniq(src ...Val) *Coll[Key, Val] {
	index := MapInit(&self.Index)

	for _, val := range src {
		key := ValidPk[Key](val)
		if MapHas(index, key) {
			panic(Errf(
				`unexpected redundant %v with key %v`,
				Type[Val](), key,
			))
		}
		index[key] = AppendIndex(&self.Slice, val)
	}

	return self
}

/*
Replaces `.Slice` with the given slice and rebuilds `.Index`. Uses the slice
as-is with no reallocation. Callers must be careful to avoid modifying the
source data, which may invalidate the collection's index.
*/
func (self *Coll[Key, Val]) Reset(src ...Val) *Coll[Key, Val] {
	self.Slice = src
	self.Reindex()
	return self
}

// Nullifies both the slice and the index. Does not preserve their capacity.
func (self *Coll[Key, Val]) Clear() *Coll[Key, Val] {
	if self != nil {
		self.Slice = nil
		self.Index = nil
	}
	return self
}

/*
Rebuilds the inner index from the inner slice, without checking the validity of
the existing index. Can be useful for external code that directly modifies the
inner `.Slice`, for example by sorting it. This is NOT used when adding items
via `.Add`, which modifies the index incrementally rather than all-at-once.
*/
func (self *Coll[Key, Val]) Reindex() *Coll[Key, Val] {
	slice := self.Slice
	if len(slice) <= 0 {
		self.Index = nil
		return self
	}

	index := make(map[Key]int, len(slice))
	for ind, val := range slice {
		index[ValidPk[Key](val)] = ind
	}
	self.Index = index

	return self
}

/*
Swaps two elements both in `.Slice` and in `.Index`. Useful for sorting.
`.Index` may be nil, in which case it's unaffected. Slice indices must be
either equal or valid.
*/
func (self Coll[Key, _]) Swap(ind0, ind1 int) {
	if ind0 == ind1 {
		return
	}

	slice := self.Slice
	val0, val1 := slice[ind0], slice[ind1]
	slice[ind0], slice[ind1] = val1, val0

	index := self.Index
	if index != nil {
		index[ValidPk[Key](val0)], index[ValidPk[Key](val1)] = ind1, ind0
	}
}

// Implement `json.Marshaler`. Encodes the inner slice, ignoring the index.
func (self Coll[_, _]) MarshalJSON() ([]byte, error) {
	return json.Marshal(self.Slice)
}

/*
Implement `json.Unmarshaler`. Decodes the input into the inner slice and
rebuilds the index.
*/
func (self *Coll[_, _]) UnmarshalJSON(src []byte) error {
	err := json.Unmarshal(src, &self.Slice)
	self.Reindex()
	return err
}

// Free cast to equivalent `LazyColl`.
func (self *Coll[Key, Val]) LazyColl() *LazyColl[Key, Val] {
	return (*LazyColl[Key, Val])(self)
}