From 9c1f79a935f0987ab4e627f3f35af7fc9a4aa72c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?K=C3=A9vin=20Gomez?= <kevin.gomez@grafana.com>
Date: Tue, 19 Sep 2023 00:04:56 +0200
Subject: [PATCH] Finish Types.HasOnlyScalarOrArray() implementation

---
 internal/ast/types.go      |  20 ++++-
 internal/ast/types_test.go | 162 +++++++++++++++++++++++++++++++++++++
 2 files changed, 179 insertions(+), 3 deletions(-)
 create mode 100644 internal/ast/types_test.go

diff --git a/internal/ast/types.go b/internal/ast/types.go
index 4e36216ef..541ceb9aa 100644
--- a/internal/ast/types.go
+++ b/internal/ast/types.go
@@ -224,11 +224,17 @@ type Types []Type
 
 func (types Types) HasOnlyScalarOrArray() bool {
 	for _, t := range types {
-		if t.Kind != KindScalar && t.Kind != KindArray {
-			return false
+		if t.Kind == KindArray {
+			if !t.AsArray().IsArrayOfScalars() {
+				return false
+			}
+
+			continue
 		}
 
-		// FIXME: for arrays, we should also inspect them recursively to make sure only scalar types are used
+		if t.Kind != KindScalar {
+			return false
+		}
 	}
 
 	return true
@@ -289,6 +295,14 @@ type ArrayType struct {
 	ValueType Type
 }
 
+func (t ArrayType) IsArrayOfScalars() bool {
+	if t.ValueType.Kind == KindArray {
+		return t.ValueType.AsArray().IsArrayOfScalars()
+	}
+
+	return t.ValueType.Kind == KindScalar
+}
+
 type EnumType struct {
 	Values []EnumValue // possible values. Value types might be different
 }
diff --git a/internal/ast/types_test.go b/internal/ast/types_test.go
new file mode 100644
index 000000000..eaff218ee
--- /dev/null
+++ b/internal/ast/types_test.go
@@ -0,0 +1,162 @@
+package ast
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTypes_HasOnlyScalarOrArray(t *testing.T) {
+	testCases := []struct {
+		description string
+		types       Types
+		expected    bool
+	}{
+		{
+			description: "only scalars",
+			types: []Type{
+				NewScalar(KindString),
+				NewScalar(KindBool),
+			},
+			expected: true,
+		},
+		{
+			description: "scalars and an array of scalars",
+			types: []Type{
+				NewScalar(KindString),
+				NewArray(NewScalar(KindInt8)),
+			},
+			expected: true,
+		},
+		{
+			description: "scalars and an array of refs",
+			types: []Type{
+				NewScalar(KindString),
+				NewArray(NewRef("SomeType")),
+			},
+			expected: false,
+		},
+		{
+			description: "ref",
+			types: []Type{
+				NewRef("SomeType"),
+			},
+			expected: false,
+		},
+		{
+			description: "scalars and ref",
+			types: []Type{
+				NewScalar(KindString),
+				NewRef("SomeType"),
+			},
+			expected: false,
+		},
+	}
+
+	for _, testCase := range testCases {
+		tc := testCase
+
+		t.Run(tc.description, func(t *testing.T) {
+			req := require.New(t)
+
+			req.Equal(tc.expected, tc.types.HasOnlyScalarOrArray())
+		})
+	}
+}
+
+func TestTypes_HasOnlyRefs(t *testing.T) {
+	testCases := []struct {
+		description string
+		types       Types
+		expected    bool
+	}{
+		{
+			description: "only scalars",
+			types: []Type{
+				NewScalar(KindString),
+				NewScalar(KindBool),
+			},
+			expected: false,
+		},
+		{
+			description: "scalars and an array of scalars",
+			types: []Type{
+				NewScalar(KindString),
+				NewArray(NewScalar(KindInt8)),
+			},
+			expected: false,
+		},
+		{
+			description: "refs",
+			types: []Type{
+				NewRef("SomeType"),
+				NewRef("SomeOtherType"),
+			},
+			expected: true,
+		},
+		{
+			description: "ref",
+			types: []Type{
+				NewRef("SomeType"),
+			},
+			expected: true,
+		},
+	}
+
+	for _, testCase := range testCases {
+		tc := testCase
+
+		t.Run(tc.description, func(t *testing.T) {
+			req := require.New(t)
+
+			req.Equal(tc.expected, tc.types.HasOnlyRefs())
+		})
+	}
+}
+
+func TestArrayType_IsArrayOfScalars(t *testing.T) {
+	testCases := []struct {
+		description string
+		array       ArrayType
+		Expected    bool
+	}{
+		{
+			description: "array of scalars",
+			array: ArrayType{
+				ValueType: NewScalar(KindString),
+			},
+			Expected: true,
+		},
+		{
+			description: "array of array of scalars",
+			array: ArrayType{
+				ValueType: NewArray(NewScalar(KindString)),
+			},
+			Expected: true,
+		},
+		{
+			description: "array of refs",
+			array: ArrayType{
+				ValueType: NewRef("SomeType"),
+			},
+			Expected: false,
+		},
+		{
+			description: "array of structs",
+			array: ArrayType{
+				ValueType: NewStruct(NewStructField("Foo", NewScalar(KindString))),
+			},
+			Expected: false,
+		},
+	}
+
+	for _, testCase := range testCases {
+		tc := testCase
+
+		t.Run(tc.description, func(t *testing.T) {
+			req := require.New(t)
+
+			req.Equal(tc.Expected, tc.array.IsArrayOfScalars())
+		})
+	}
+}