From 89c25cc6aaac38ffd7c27c076f911f08142bd4e1 Mon Sep 17 00:00:00 2001
From: wfee2000 <ilmingwinnie19@gmail.com>
Date: Mon, 23 Dec 2024 22:39:37 +0100
Subject: [PATCH] fix: make permutation types generic

---
 iter/perm.v      | 18 +++++++++---------
 iter/perm_test.v | 17 ++++++++++++++++-
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/iter/perm.v b/iter/perm.v
index a18d5f5ba..3fb399389 100644
--- a/iter/perm.v
+++ b/iter/perm.v
@@ -3,7 +3,7 @@ module iter
 import vsl.util
 import math
 
-pub struct PermutationsIter {
+pub struct PermutationsIter[T] {
 mut:
 	pos    u64
 	idxs   []int
@@ -11,15 +11,15 @@ mut:
 pub:
 	repeat int
 	size   u64
-	data   []f64
+	data   []T
 }
 
 // PermutationsIter.new will return an iterator that allows
 // lazy computation for all length `r` permutations of `data`
-pub fn PermutationsIter.new(data []f64, r int) PermutationsIter {
+pub fn PermutationsIter.new[T](data []T, r int) PermutationsIter[T] {
 	n := data.len
 	if r > n {
-		return PermutationsIter{
+		return PermutationsIter[T]{
 			data:   data
 			repeat: r
 		}
@@ -27,7 +27,7 @@ pub fn PermutationsIter.new(data []f64, r int) PermutationsIter {
 	size := u64(math.factorial(n) / math.factorial(n - r))
 	idxs := util.arange(n)
 	cycles := util.range(n, n - r, step: -1)
-	return PermutationsIter{
+	return PermutationsIter[T]{
 		data:   data
 		repeat: r
 		size:   size
@@ -37,7 +37,7 @@ pub fn PermutationsIter.new(data []f64, r int) PermutationsIter {
 }
 
 // next will return next permutation if possible
-pub fn (mut o PermutationsIter) next() ?[]f64 {
+pub fn (mut o PermutationsIter[T]) next[T]() ?[]T {
 	// base case for every iterator
 	if o.pos == o.size {
 		return none
@@ -69,9 +69,9 @@ pub fn (mut o PermutationsIter) next() ?[]f64 {
 }
 
 // permutations returns successive `r` length permutations of elements in `data`
-pub fn permutations(data []f64, r int) [][]f64 {
-	mut perms := PermutationsIter.new(data, r)
-	mut result := [][]f64{cap: int(perms.size)}
+pub fn permutations[T](data []T, r int) [][]T {
+	mut perms := PermutationsIter.new[T](data, r)
+	mut result := [][]T{cap: int(perms.size)}
 	for perm in perms {
 		result << perm
 	}
diff --git a/iter/perm_test.v b/iter/perm_test.v
index e64e80581..b29541801 100644
--- a/iter/perm_test.v
+++ b/iter/perm_test.v
@@ -64,7 +64,22 @@ fn test_permutations_simple_5() {
 	assert assert_permutation(expected, result)
 }
 
-fn assert_permutation(a [][]f64, b [][]f64) bool {
+fn test_permutations_generic_type_string() {
+	data := ['a', 'b', 'c']
+	expected := [['a', 'b'], ['a', 'c'], ['b', 'a'], ['b', 'c'],
+		['c', 'a'], ['c', 'b']]
+	result := permutations(data, 2)
+	assert assert_permutation(expected, result)
+}
+
+fn test_permutations_generic_type_int() {
+	data := [1, 2, 3]
+	expected := [[1, 2], [1, 3], [2, 1], [2, 3], [3, 1], [3, 2]]
+	result := permutations(data, 2)
+	assert assert_permutation(expected, result)
+}
+
+fn assert_permutation[T](a [][]T, b [][]T) bool {
 	if a.len != b.len {
 		return false
 	}