arrays, maps: add indexed variant of collection function and minor cleanup (#15948)

vlib/arrays/arrays.v

@@ -241,6 +241,42 @@ pub fn reduce<T>(list []T, reduce_op fn (t1 T, t2 T) T) ?T {
 	}
 }
 
+// reduce_indexed sets `acc = list[0]`, then successively calls `acc = reduce_op(idx, acc, elem)` for each remaining element in `list`.
+// returns the accumulated value in `acc`.
+// returns an error if the array is empty.
+// See also: [fold_indexed](#fold_indexed).
+pub fn reduce_indexed<T>(list []T, reduce_op fn (int, T, T) T) ?T {
+	if list.len == 0 {
+		return error('Cannot reduce array of nothing.')
+	} else {
+		mut value := list[0]
+
+		for i, e in list {
+			if i == 0 {
+				continue
+			} else {
+				value = reduce_op(i, value, e)
+			}
+		}
+
+		return value
+	}
+}
+
+// filter_indexed filters elements based on predicate function
+// being invoked on each element with its index in the original array.
+pub fn filter_indexed<T>(list []T, predicate fn (idx int, e T) bool) []T {
+	mut result := []T{cap: list.len}
+
+	for i, e in list {
+		if predicate(i, e) {
+			result << e
+		}
+	}
+
+	return result
+}
+
 // fold sets `acc = init`, then successively calls `acc = fold_op(acc, elem)` for each element in `list`.
 // returns `acc`.
 // Example:
@@ -261,7 +297,19 @@ pub fn fold<T, R>(list []T, init R, fold_op fn (r R, t T) R) R {
 	return value
 }
 
-// flattens n + 1 dimensional array into n dimensional array
+// fold_indexed sets `acc = init`, then successively calls `acc = fold_op(idx, acc, elem)` for each element in `list`.
+// returns `acc`.
+pub fn fold_indexed<T, R>(list []T, init R, fold_op fn (idx int, r R, t T) R) R {
+	mut value := init
+
+	for i, e in list {
+		value = fold_op(i, value, e)
+	}
+
+	return value
+}
+
+// flatten flattens n + 1 dimensional array into n dimensional array
 // Example: arrays.flatten<int>([[1, 2, 3], [4, 5]]) // => [1, 2, 3, 4, 5]
 pub fn flatten<T>(list [][]T) []T {
 	// calculate required capacity
@@ -285,6 +333,42 @@ pub fn flatten<T>(list [][]T) []T {
 	return result
 }
 
+// flat_map creates a new array populated with the flattened result of calling transform function
+// being invoked on each element of `list`.
+pub fn flat_map<T, R>(list []T, transform fn (T) []R) []R {
+	mut result := [][]R{cap: list.len}
+
+	for v in list {
+		result << transform(v)
+	}
+
+	return flatten(result)
+}
+
+// flat_map_indexed creates a new array populated with the flattened result of calling the `transform` function
+// being invoked on each element with its index in the original array.
+pub fn flat_map_indexed<T, R>(list []T, transform fn (int, T) []R) []R {
+	mut result := [][]R{cap: list.len}
+
+	for i, v in list {
+		result << transform(i, v)
+	}
+
+	return flatten(result)
+}
+
+// map_indexed creates a new array populated with the result of calling the `transform` function
+// being invoked on each element with its index in the original array.
+pub fn map_indexed<T, R>(list []T, transform fn (int, T) R) []R {
+	mut result := []R{cap: list.len}
+
+	for i, v in list {
+		result << transform(i, v)
+	}
+
+	return result
+}
+
 // group_by groups together elements, for which the `grouping_op` callback produced the same result.
 // Example: arrays.group_by<int, string>(['H', 'el', 'lo'], fn (v string) int { return v.len }) // => {1: ['H'], 2: ['el', 'lo']}
 pub fn group_by<K, V>(list []V, grouping_op fn (v V) K) map[K][]V {

vlib/arrays/arrays_test.v

@@ -92,6 +92,29 @@ fn test_fixed_array_assignment() {
 	}
 }
 
+fn test_array_flat_map() {
+	a := ['Hello V', 'Hello World', 'V Lang']
+	assert flat_map<string, string>(a, fn (e string) []string {
+		return e.split(' ')
+	}) == ['Hello', 'V', 'Hello', 'World', 'V', 'Lang']
+}
+
+fn test_array_flat_map_indexed() {
+	a := ['AB', 'CD', 'EF']
+	assert flat_map_indexed<string, string>(a, fn (i int, e string) []string {
+		mut arr := [i.str()]
+		arr << e.split('')
+		return arr
+	}) == ['0', 'A', 'B', '1', 'C', 'D', '2', 'E', 'F']
+}
+
+fn test_map_indexed() {
+	a := [1, 2, 3]
+	assert map_indexed<int, int>(a, fn (i int, e int) int {
+		return i + e * e
+	}) == [1, 5, 11]
+}
+
 fn test_group() {
 	x := [4, 5, 6]
 	y := [2, 1, 3]
@@ -146,6 +169,21 @@ fn test_reduce() {
 	}) or { -1 } == -1
 }
 
+fn test_reduce_indexed() {
+	x := [1, 2, 3, 4, 5]
+	assert reduce_indexed<int>(x, fn (idx int, t1 int, t2 int) int {
+		return idx + t1 + t2
+	}) or { 0 } == 25
+}
+
+fn test_filter_indexed() {
+	x := [0, 1, 2, 3, 4, 5]
+
+	assert filter_indexed<int>(x, fn (idx int, e int) bool {
+		return idx % 2 == 0
+	}) == [0, 2, 4]
+}
+
 fn test_fold() {
 	x := [1, 2, 3, 4, 5]
 
@@ -160,6 +198,14 @@ fn test_fold() {
 	}) == -1
 }
 
+fn test_fold_indexed() {
+	x := [1, 2, 3, 4, 5]
+
+	assert fold_indexed<int, int>(x, 5, fn (idx int, r int, t int) int {
+		return idx + r + t
+	}) == 30
+}
+
 fn test_flatten() {
 	x := [[1, 2, 3], [4, 5, 6]]
 

vlib/maps/maps.v

@@ -24,8 +24,8 @@ pub fn to_array<K, V, I>(m map[K]V, f fn (K, V) I) []I {
 	return a
 }
 
-// flatten maps map entries into arrays and flattens into a one-dimensional array
-pub fn flatten<K, V, I>(m map[K]V, f fn (K, V) []I) []I {
+// flat_map maps map entries into arrays and flattens into a one-dimensional array
+pub fn flat_map<K, V, I>(m map[K]V, f fn (K, V) []I) []I {
 	mut a := []I{cap: m.len}
 
 	for k, v in m {

vlib/maps/maps_test.v

@@ -36,13 +36,13 @@ fn test_to_array() {
 	}) == ['abc', 'def', 'ghi']
 }
 
-fn test_flatten() {
+fn test_flat_map() {
 	m1 := {
 		1: [2, 3]
 		4: [5, 6]
 		7: [8, 9]
 	}
-	assert flatten<int, []int, int>(m1, fn (k int, v []int) []int {
+	assert flat_map<int, []int, int>(m1, fn (k int, v []int) []int {
 		mut a := [k]
 		a << v
 		return a