keeptrack / pkg / enum / enum.go

package enum

import (
	"fmt"
	"sort"
)

// A functional function that transforms items using a mapping function.
func Map[T any, U any](items []T, fn func(T) U) []U {
	var result []U

	for _, item := range items {
		result = append(result, fn(item))
	}

	return result
}

// A functional function that reduces items to a single value using a reduction function.
func Reduce[T any, U any](items []T, fn func(U, T) U, initial U) U {
	result := initial

	for _, item := range items {
		result = fn(result, item)
	}

	return result
}

// A functional function that removes duplicate items based on a key function. The returned slice
// will contain only the items that are unique according to the returned value of the key function.
func UniqueBy[T any, K comparable](items []T, fn func(T) K) []T {
	keys := make(map[K]struct{})
	var result []T

	for _, item := range items {
		key := fn(item)
		if _, exists := keys[key]; exists {
			continue
		}
		keys[key] = struct{}{}
		result = append(result, item)
	}

	return result
}

// GroupBy groups elements of a slice by a key derived from each element.
// The keyer function extracts the key from each element.
func GroupBy[T any, K comparable](items []T, keyer func(T) K) map[K][]T {
	groups := make(map[K][]T)
	for _, item := range items {
		key := keyer(item)
		groups[key] = append(groups[key], item)
	}
	return groups
}

// GroupByWithValue groups elements of a slice by a key derived from each element,
// and maps each element to a value derived from the element.
func GroupByWithValue[T any, K comparable, V any](items []T, keyer func(T) K, valuer func(T) V) map[K][]V {
	groups := make(map[K][]V)
	for _, item := range items {
		key := keyer(item)
		value := valuer(item)
		groups[key] = append(groups[key], value)
	}
	return groups
}

// Combine two slices into one map, where the first slice should have unique values that can be
// used as keys. If the slices are not the same length, the extra values will be ignored
func ZipMap[K comparable, T any](keys []K, values []T) map[K]T {
	result := make(map[K]T)
	for i, key := range keys {
		if i < len(values) {
			result[key] = values[i]
		}
	}
	return result
}

// Combine two slices into one map, where the first slice should have unique values that can be
// used as keys
func ZipMapSafe[K comparable, T any](keys []K, values []T) (map[K]T, error) {
	if len(keys) != len(values) {
		return nil, fmt.Errorf("keys and values must have the same length")
	}

	result := make(map[K]T)
	for i, key := range keys {
		if i < len(values) {
			result[key] = values[i]
		}
	}
	return result, nil
}

// Create a slice of values from values inside the given slice that match the predicate
func Filter[T any](items []T, predicate func(T) bool) []T {
	var result []T
	for _, item := range items {
		if predicate(item) {
			result = append(result, item)
		}
	}
	return result
}

// Return the first value from the values inside the given slice that match the predicate, if no
// values can be found, it will return an empty value with a false boolean
func Find[T any](items []T, predicate func(T) bool) (T, bool) {
	for _, item := range items {
		if predicate(item) {
			return item, true
		}
	}
	var zero T
	return zero, false
}

func Flatten[T any](items [][]T) []T {
	var result []T
	for _, item := range items {
		result = append(result, item...)
	}
	return result
}

// OrderByKeys reorders items so that their keys appear in the same order as inputKeys. Items whose
// key is not present in inputKeys are placed at the end (while preserving their relative order).
func OrderByKeys[T any, K comparable](items []T, keys []K, keyFn func(T) K) []T {
	if len(items) == 0 || len(keys) == 0 {
		return items
	}

	pos := make(map[K]int, len(keys))
	for i, k := range keys {
		pos[k] = i
	}

	const notFound = int(^uint(0) >> 1) // max int

	sort.SliceStable(items, func(i, j int) bool {
		pi, ok := pos[keyFn(items[i])]
		if !ok {
			pi = notFound
		}

		pj, ok := pos[keyFn(items[j])]
		if !ok {
			pj = notFound
		}

		return pi < pj
	})

	return items
}