From 8d9b9d3e7ceb3b0530a73950093f2a7eb2e99b79 Mon Sep 17 00:00:00 2001
From: Fabrizio Sestito <fabrizio.sestito@suse.com>
Date: Fri, 16 Aug 2024 19:27:00 +0200
Subject: [PATCH] Add list flatten() (#980)

* Add list flatten()

Signed-off-by: Fabrizio Sestito <fabrizio.sestito@suse.com>
---
 ext/lists.go      | 103 +++++++++++++++++++++++++++++++++++++++++++---
 ext/lists_test.go |   6 +++
 2 files changed, 104 insertions(+), 5 deletions(-)

diff --git a/ext/lists.go b/ext/lists.go
index 08751d08..de2fd709 100644
--- a/ext/lists.go
+++ b/ext/lists.go
@@ -16,6 +16,7 @@ package ext
 
 import (
 	"fmt"
+	"math"
 
 	"github.com/google/cel-go/cel"
 	"github.com/google/cel-go/common/types"
@@ -35,21 +36,67 @@ import (
 //
 //	[1,2,3,4].slice(1, 3) // return [2, 3]
 //	[1,2,3,4].slice(2, 4) // return [3 ,4]
-func Lists() cel.EnvOption {
-	return cel.Lib(listsLib{})
+//
+// # Flatten
+//
+// Flattens a list recursively.
+// If an optional depth is provided, the list is flattened to a the specificied level.
+// A negative depth value flattens the list recursively to its deepest level.
+//
+//	<list>.flatten(<list>) -> <list>
+//	<list>.flatten(<list>, <int>) -> <list>
+//
+// Examples:
+//
+// [1,[2,3],[4]].flatten() // return [1, 2, 3, 4]
+// [1,[2,[3,4]]].flatten() // return [1, 2, [3, 4]]
+// [1,2,[],[],[3,4]].flatten() // return [1, 2, 3, 4]
+// [1,[2,[3,[4]]]].flatten(2) // return [1, 2, 3, [4]]
+// [1,[2,[3,[4]]]].flatten(-1) // return [1, 2, 3, 4]
+func Lists(options ...ListsOption) cel.EnvOption {
+	l := &listsLib{
+		version: math.MaxUint32,
+	}
+	for _, o := range options {
+		l = o(l)
+	}
+
+	return cel.Lib(l)
 }
 
-type listsLib struct{}
+type listsLib struct {
+	version uint32
+}
 
 // LibraryName implements the SingletonLibrary interface method.
 func (listsLib) LibraryName() string {
 	return "cel.lib.ext.lists"
 }
 
+// ListsOption is a functional interface for configuring the strings library.
+type ListsOption func(*listsLib) *listsLib
+
+// ListsVersion configures the version of the string library.
+//
+// The version limits which functions are available. Only functions introduced
+// below or equal to the given version included in the library. If this option
+// is not set, all functions are available.
+//
+// See the library documentation to determine which version a function was introduced.
+// If the documentation does not state which version a function was introduced, it can
+// be assumed to be introduced at version 0, when the library was first created.
+func ListsVersion(version uint32) ListsOption {
+	return func(lib *listsLib) *listsLib {
+		lib.version = version
+		return lib
+	}
+}
+
 // CompileOptions implements the Library interface method.
-func (listsLib) CompileOptions() []cel.EnvOption {
+func (lib listsLib) CompileOptions() []cel.EnvOption {
 	listType := cel.ListType(cel.TypeParamType("T"))
-	return []cel.EnvOption{
+	listDyn := cel.ListType(cel.DynType)
+	opts := []cel.EnvOption{
 		cel.Function("slice",
 			cel.MemberOverload("list_slice",
 				[]*cel.Type{listType, cel.IntType, cel.IntType}, listType,
@@ -66,6 +113,33 @@ func (listsLib) CompileOptions() []cel.EnvOption {
 			),
 		),
 	}
+	if lib.version >= 1 {
+		opts = append(opts,
+			cel.Function("flatten",
+				cel.MemberOverload("list_flatten",
+					[]*cel.Type{listDyn}, listDyn,
+					cel.UnaryBinding(func(arg ref.Val) ref.Val {
+						list := arg.(traits.Lister)
+						flatList := flatten(list, 1)
+						return types.DefaultTypeAdapter.NativeToValue(flatList)
+					}),
+				),
+			),
+			cel.Function("flatten",
+				cel.MemberOverload("list_flatten_int",
+					[]*cel.Type{listDyn, types.IntType}, listDyn,
+					cel.BinaryBinding(func(arg1, arg2 ref.Val) ref.Val {
+						list := arg1.(traits.Lister)
+						depth := arg2.(types.Int)
+						flatList := flatten(list, int64(depth))
+						return types.DefaultTypeAdapter.NativeToValue(flatList)
+					}),
+				),
+			),
+		)
+	}
+
+	return opts
 }
 
 // ProgramOptions implements the Library interface method.
@@ -92,3 +166,22 @@ func slice(list traits.Lister, start, end types.Int) (ref.Val, error) {
 	}
 	return types.DefaultTypeAdapter.NativeToValue(newList), nil
 }
+
+func flatten(list traits.Lister, depth int64) []ref.Val {
+	var newList []ref.Val
+	iter := list.Iterator()
+
+	for iter.HasNext() == types.True {
+		val := iter.Next()
+		nestedList, isList := val.(traits.Lister)
+
+		if !isList || depth == 0 {
+			newList = append(newList, val)
+			continue
+		} else {
+			newList = append(newList, flatten(nestedList, depth-1)...)
+		}
+	}
+
+	return newList
+}
diff --git a/ext/lists_test.go b/ext/lists_test.go
index 74715e25..1851f838 100644
--- a/ext/lists_test.go
+++ b/ext/lists_test.go
@@ -36,6 +36,12 @@ func TestLists(t *testing.T) {
 		{expr: `[1,2,3,4].slice(0, 10)`, err: "cannot slice(0, 10), list is length 4"},
 		{expr: `[1,2,3,4].slice(-5, 10)`, err: "cannot slice(-5, 10), negative indexes not supported"},
 		{expr: `[1,2,3,4].slice(-5, -3)`, err: "cannot slice(-5, -3), negative indexes not supported"},
+		{expr: `[].flatten() == []`},
+		{expr: `[1,2,3,4].flatten() == [1,2,3,4]`},
+		{expr: `[1,[2,[3,4]]].flatten() == [1,2,[3,4]]`},
+		{expr: `[1,2,[],[],[3,4]].flatten() == [1,2,3,4]`},
+		{expr: `[1,[2,[3,4]]].flatten(2) == [1,2,3,4]`},
+		{expr: `[1,[2,[3,[4]]]].flatten(-1) == [1,2,3,4]`},
 	}
 
 	env := testListsEnv(t)