New flatten/coll.Flatten function

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

2 files changed, 137 insertions, 1 deletions

diff --git a/coll/coll.go b/coll/coll.go
index 99911d0a..8e41ed91 100644
--- a/coll/coll.go
+++ b/coll/coll.go
@@ -292,3 +292,31 @@ func sameTypes(a []interface{}) bool {
 	}
 	return true
 }
+
+// Flatten a nested array or slice to at most 'depth' levels. Use depth of -1
+// to completely flatten the input.
+// Returns a new slice without modifying the input.
+func Flatten(list interface{}, depth int) (out []interface{}, err error) {
+	l, err := interfaceSlice(list)
+	if err != nil {
+		return nil, err
+	}
+	if depth == 0 {
+		return l, nil
+	}
+	for _, v := range l {
+		s := reflect.ValueOf(v)
+		kind := s.Kind()
+		switch kind {
+		case reflect.Slice, reflect.Array:
+			vl, err := Flatten(v, depth-1)
+			if err != nil {
+				return nil, err
+			}
+			out = append(out, vl...)
+		default:
+			out = append(out, v)
+		}
+	}
+	return out, nil
+}
diff --git a/coll/coll_test.go b/coll/coll_test.go
index a4c34a5c..6110a78d 100644
--- a/coll/coll_test.go
+++ b/coll/coll_test.go
@@ -15,7 +15,6 @@ func TestSlice(t *testing.T) {
 }
 
 func TestHas(t *testing.T) {
-
 	in := map[string]interface{}{
 		"foo": "bar",
 		"baz": map[string]interface{}{
@@ -457,3 +456,112 @@ func TestSort(t *testing.T) {
 		})
 	}
 }
+
+func TestFlatten(t *testing.T) {
+	data := []struct {
+		depth    int
+		in       interface{}
+		expected []interface{}
+	}{
+		{0, []int{1, 2, 3}, []interface{}{1, 2, 3}},
+		{0, [3]int{1, 2, 3}, []interface{}{1, 2, 3}},
+		{0,
+			[]interface{}{[]string{}, []int{1, 2}, 3},
+			[]interface{}{[]string{}, []int{1, 2}, 3},
+		},
+		{0,
+			[]interface{}{[]string{"one"}, [][]int{{1, 2}}, 3},
+			[]interface{}{[]string{"one"}, [][]int{{1, 2}}, 3},
+		},
+
+		{1, []int{1, 2, 3}, []interface{}{1, 2, 3}},
+		{1, [3]int{1, 2, 3}, []interface{}{1, 2, 3}},
+		{1, []interface{}{[]string{}, []int{1, 2}, 3}, []interface{}{1, 2, 3}},
+		{1,
+			[]interface{}{[]string{"one"}, [][]int{{1, 2}}, 3},
+			[]interface{}{"one", []int{1, 2}, 3},
+		},
+
+		{2, []int{1, 2, 3}, []interface{}{1, 2, 3}},
+		{2, [3]int{1, 2, 3}, []interface{}{1, 2, 3}},
+		{2, []interface{}{[]string{}, []int{1, 2}, 3}, []interface{}{1, 2, 3}},
+		{2,
+			[]interface{}{[]string{"one"}, [][]int{{1, 2}}, 3},
+			[]interface{}{"one", 1, 2, 3},
+		},
+		{2,
+			[]interface{}{
+				[]string{"one"},
+				[]interface{}{
+					[]interface{}{
+						[]int{1},
+						[]interface{}{2, []int{3}},
+					},
+					[]int{4, 5},
+				},
+				6,
+			},
+			[]interface{}{"one", []int{1}, []interface{}{2, []int{3}}, 4, 5, 6},
+		},
+
+		{-1, []int{1, 2, 3}, []interface{}{1, 2, 3}},
+		{-1, [3]int{1, 2, 3}, []interface{}{1, 2, 3}},
+		{-1, []interface{}{[]string{}, []int{1, 2}, 3}, []interface{}{1, 2, 3}},
+		{-1,
+			[]interface{}{[]string{"one"}, [][]int{{1, 2}}, 3},
+			[]interface{}{"one", 1, 2, 3},
+		},
+		{-1,
+			[]interface{}{
+				[]string{"one"},
+				[]interface{}{
+					[]interface{}{
+						[]int{1},
+						[]interface{}{2, []int{3}},
+					},
+					[]int{4, 5},
+				},
+				6,
+			},
+			[]interface{}{"one", 1, 2, 3, 4, 5, 6},
+		},
+	}
+
+	for _, d := range data {
+		out, err := Flatten(d.in, d.depth)
+		assert.NoError(t, err)
+		assert.EqualValues(t, d.expected, out)
+	}
+
+	_, err := Flatten(42, -1)
+	assert.Error(t, err)
+}
+
+func BenchmarkFlatten(b *testing.B) {
+	data := []interface{}{
+		[]int{1, 2, 3},
+		[3]int{1, 2, 3},
+		[]interface{}{[]string{}, []int{1, 2}, 3},
+		[]interface{}{[]string{"one"}, [][]int{{1, 2}}, 3},
+		[]interface{}{
+			[]string{"one"},
+			[]interface{}{
+				[]interface{}{
+					[]int{1},
+					[]interface{}{2, []int{3}},
+				},
+				[]int{4, 5},
+			},
+			6,
+		},
+	}
+	for depth := -1; depth <= 2; depth++ {
+		for _, d := range data {
+			b.Run(fmt.Sprintf("depth%d %T(%v)", depth, d, d), func(b *testing.B) {
+				for i := 0; i < b.N; i++ {
+					Flatten(d, depth)
+				}
+			})
+		}
+	}
+}