From f538a8dae4cacbfb3f125dd09bc038ce2ae43a51 Mon Sep 17 00:00:00 2001
From: Dan Kortschak <dan.kortschak@elastic.co>
Date: Fri, 19 Jul 2024 06:12:04 +0930
Subject: [PATCH] lib: relax type hints on flatten function

The documented type signature of the flatten function is <list<dyn>> ->
<list<dyn>>, but the implementation was <list<typeV>> -> <list<typeV>>.
This has an impact on uses of flatten where the input is a nested set of
lists, the exact intended use-case. In this situation, the type checker
is told that the returned value is a list of lists..., resulting in
rejection of valid programs. There is a work-around that can be used,
but instead, just make the implementation match the documentation and
intention.
---
 lib/collections.go         |  4 +-
 testdata/flatten_types.txt | 87 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 testdata/flatten_types.txt

diff --git a/lib/collections.go b/lib/collections.go
index 10073a7..c15d8f8 100644
--- a/lib/collections.go
+++ b/lib/collections.go
@@ -358,8 +358,8 @@ func (collectionsLib) CompileOptions() []cel.EnvOption {
 		cel.Function("flatten",
 			cel.MemberOverload(
 				"list_flatten",
-				[]*cel.Type{listV},
-				listV,
+				[]*cel.Type{listDyn},
+				listDyn,
 				cel.UnaryBinding(catch(flatten)),
 			),
 		),
diff --git a/testdata/flatten_types.txt b/testdata/flatten_types.txt
new file mode 100644
index 0000000..e2de303
--- /dev/null
+++ b/testdata/flatten_types.txt
@@ -0,0 +1,87 @@
+mito -use collections -data state.json src.cel
+! stderr .
+cmp stdout want.txt
+
+-- state.json --
+{
+	"aye": {
+		"value": 1
+	},
+	"bee": {
+		"value": "two"
+	},
+	"one": [
+		{
+			"value": 1
+		},
+		{
+			"value": "two"
+		}
+	],
+	"two": [
+		{
+			"value": 3
+		},
+		{
+			"value": "four"
+		}
+	]
+}
+
+-- src.cel --
+[
+	"one",
+	"two",
+].map(k, state[k].map(e, {
+	k: e,
+})).flatten().map(d, d.with({
+	"aye": state.aye,
+	"bee": state.bee,
+}))
+-- want.txt --
+[
+	{
+		"aye": {
+			"value": 1
+		},
+		"bee": {
+			"value": "two"
+		},
+		"one": {
+			"value": 1
+		}
+	},
+	{
+		"aye": {
+			"value": 1
+		},
+		"bee": {
+			"value": "two"
+		},
+		"one": {
+			"value": "two"
+		}
+	},
+	{
+		"aye": {
+			"value": 1
+		},
+		"bee": {
+			"value": "two"
+		},
+		"two": {
+			"value": 3
+		}
+	},
+	{
+		"aye": {
+			"value": 1
+		},
+		"bee": {
+			"value": "two"
+		},
+		"two": {
+			"value": "four"
+		}
+	}
+]