Merge pull request lightningnetwork#8729 from Chinwendu20/fn

Add new  functions to the fn package

fn/map.go

@@ -0,0 +1,44 @@
+package fn
+
+import (
+	"fmt"
+
+	"golang.org/x/exp/maps"
+)
+
+// KeySet converts a map into a Set containing the keys of the map.
+func KeySet[K comparable, V any](m map[K]V) Set[K] {
+	return NewSet(maps.Keys(m)...)
+}
+
+// NewSubMapIntersect returns a sub-map of `m` containing only the keys found in
+// both `m` and the `keys` slice.
+func NewSubMapIntersect[K comparable, V any](m map[K]V, keys []K) map[K]V {
+	result := make(map[K]V)
+	for _, k := range keys {
+		v, ok := m[k]
+		if !ok {
+			continue
+		}
+
+		result[k] = v
+	}
+
+	return result
+}
+
+// NewSubMap creates a sub-map from a given map using specified keys. It errors
+// if any of the keys is not found in the map.
+func NewSubMap[K comparable, V any](m map[K]V, keys []K) (map[K]V, error) {
+	result := make(map[K]V, len(keys))
+	for _, k := range keys {
+		v, ok := m[k]
+		if !ok {
+			return nil, fmt.Errorf("NewSubMap: missing key %v", k)
+		}
+
+		result[k] = v
+	}
+
+	return result, nil
+}

fn/map_test.go

@@ -0,0 +1,135 @@
+package fn
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestKeySet(t *testing.T) {
+	testMap := map[string]int{"a": 1, "b": 2, "c": 3}
+	expected := NewSet([]string{"a", "b", "c"}...)
+
+	require.Equal(t, expected, KeySet(testMap))
+}
+
+// TestNewSubMap tests the NewSubMap function with various input cases.
+func TestNewSubMap(t *testing.T) {
+	tests := []struct {
+		name     string
+		original map[int]string
+		keys     []int
+		expected map[int]string
+		wantErr  bool
+	}{
+		{
+			name: "Successful submap creation",
+			original: map[int]string{
+				1: "apple",
+				2: "banana",
+				3: "cherry",
+			},
+			keys: []int{1, 3},
+			expected: map[int]string{
+				1: "apple",
+				3: "cherry",
+			},
+			wantErr: false,
+		},
+		{
+			name: "Key not found",
+			original: map[int]string{
+				1: "apple",
+				2: "banana",
+			},
+			keys:     []int{1, 4},
+			expected: nil,
+			wantErr:  true,
+		},
+		{
+			name: "Empty keys list",
+			original: map[int]string{
+				1: "apple",
+				2: "banana",
+			},
+			keys:     []int{},
+			expected: map[int]string{},
+			wantErr:  false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := NewSubMap(tt.original, tt.keys)
+			if tt.wantErr {
+				require.ErrorContains(
+					t, err, "NewSubMap: missing key",
+				)
+
+				require.Nil(t, result)
+
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+// TestNewSubMapIntersect tests the NewSubMapIntersect function for correctness.
+func TestNewSubMapIntersect(t *testing.T) {
+	tests := []struct {
+		name     string
+		original map[int]string
+		keys     []int
+		expected map[int]string
+	}{
+		{
+			name: "Successful intersection",
+			original: map[int]string{
+				1: "apple",
+				2: "banana",
+				3: "cherry",
+				4: "date",
+			},
+			keys: []int{2, 3, 5},
+			expected: map[int]string{
+				2: "banana",
+				3: "cherry",
+			},
+		},
+		{
+			name: "No intersection",
+			original: map[int]string{
+				1: "apple",
+				2: "banana",
+			},
+			keys:     []int{3, 4},
+			expected: map[int]string{},
+		},
+		{
+			name:     "Empty original map",
+			original: map[int]string{},
+			keys:     []int{1, 2},
+			expected: map[int]string{},
+		},
+		{
+			name: "Empty keys list",
+			original: map[int]string{
+				1: "apple",
+				2: "banana",
+			},
+			keys:     []int{},
+			expected: map[int]string{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(
+				t, tt.expected,
+				NewSubMapIntersect(tt.original, tt.keys))
+		})
+	}
+}

fn/slice.go

@@ -1,5 +1,13 @@
 package fn
 
+import "golang.org/x/exp/constraints"
+
+// Number is a type constraint for all numeric types in Go (integers,
+// float and complex numbers)
+type Number interface {
+	constraints.Integer | constraints.Float | constraints.Complex
+}
+
 // All returns true when the supplied predicate evaluates to true for all of
 // the values in the slice.
 func All[A any](pred func(A) bool, s []A) bool {
@@ -168,3 +176,25 @@ func ZipWith[A, B, C any](f func(A, B) C, a []A, b []B) []C {
 
 	return res
 }
+
+// SliceToMap converts a slice to a map using the provided key and value
+// functions.
+func SliceToMap[A any, K comparable, V any](s []A, keyFunc func(A) K,
+	valueFunc func(A) V) map[K]V {
+
+	res := make(map[K]V, len(s))
+	for _, val := range s {
+		key := keyFunc(val)
+		value := valueFunc(val)
+		res[key] = value
+	}
+
+	return res
+}
+
+// Sum calculates the sum of a slice of numbers, `items`.
+func Sum[B Number](items []B) B {
+	return Foldl(func(a, b B) B {
+		return a + b
+	}, 0, items)
+}

fn/slice_test.go

@@ -1,6 +1,7 @@
 package fn
 
 import (
+	"fmt"
 	"slices"
 	"testing"
 
@@ -136,3 +137,107 @@ func TestZipWith(t *testing.T) {
 		z, []bool{false, true, false, false, false},
 	))
 }
+
+// TestSum checks if the Sum function correctly calculates the sum of the
+// numbers in the slice.
+func TestSum(t *testing.T) {
+	tests := []struct {
+		name   string
+		items  interface{}
+		result interface{}
+	}{
+		{
+			name:   "Sum of positive integers",
+			items:  []int{1, 2, 3},
+			result: 6,
+		},
+		{
+			name:   "Sum of negative integers",
+			items:  []int{-1, -2, -3},
+			result: -6,
+		},
+		{
+			name:   "Sum of float numbers",
+			items:  []float64{1.1, 2.2, 3.3},
+			result: 6.6,
+		},
+		{
+			name: "Sum of complex numbers",
+			items: []complex128{
+				complex(1, 1),
+				complex(2, 2),
+				complex(3, 3),
+			},
+			result: complex(6, 6),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			switch v := tt.items.(type) {
+			case []int:
+				require.Equal(t, tt.result, Sum(v))
+			case []float64:
+				require.Equal(t, tt.result, Sum(v))
+			case []complex128:
+				require.Equal(t, tt.result, Sum(v))
+			}
+		})
+	}
+}
+
+// TestSliceToMap tests the SliceToMap function.
+func TestSliceToMap(t *testing.T) {
+	tests := []struct {
+		name      string
+		slice     []int
+		keyFunc   func(int) int
+		valueFunc func(int) string
+		expected  map[int]string
+	}{
+		{
+			name:    "Integers to string map",
+			slice:   []int{1, 2, 3},
+			keyFunc: func(a int) int { return a },
+			valueFunc: func(a int) string {
+				return fmt.Sprintf("Value%d", a)
+			},
+			expected: map[int]string{
+				1: "Value1",
+				2: "Value2",
+				3: "Value3",
+			},
+		},
+		{
+			name:    "Duplicates in slice",
+			slice:   []int{1, 2, 2, 3},
+			keyFunc: func(a int) int { return a },
+			valueFunc: func(a int) string {
+				return fmt.Sprintf("Value%d", a)
+			},
+			expected: map[int]string{
+				1: "Value1",
+				2: "Value2",
+				3: "Value3",
+			},
+		},
+		{
+			name:    "Empty slice",
+			slice:   []int{},
+			keyFunc: func(a int) int { return a },
+			valueFunc: func(a int) string {
+				return fmt.Sprintf("Value%d", a)
+			},
+			expected: map[int]string{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(
+				t, tt.expected,
+				SliceToMap(tt.slice, tt.keyFunc, tt.valueFunc),
+			)
+		})
+	}
+}