From aaaa995073045d6ba202046d056c0551eca5a1c9 Mon Sep 17 00:00:00 2001
From: Deep Goel <deepgoel@google.com>
Date: Sun, 13 Aug 2023 17:26:08 +0530
Subject: [PATCH] Feat: Implementing std.groupBy

---
 builtins.go                           | 42 +++++++++++++++++++++++++++
 interpreter.go                        |  3 ++
 linter/internal/types/stdlib.go       |  1 +
 testdata/builtinGroupBy.golden        | 12 ++++++++
 testdata/builtinGroupBy.jsonnet       |  1 +
 testdata/builtinGroupBy.linter.golden |  0
 6 files changed, 59 insertions(+)
 create mode 100644 testdata/builtinGroupBy.golden
 create mode 100644 testdata/builtinGroupBy.jsonnet
 create mode 100644 testdata/builtinGroupBy.linter.golden

diff --git a/builtins.go b/builtins.go
index e426b069c..862c8b252 100644
--- a/builtins.go
+++ b/builtins.go
@@ -1992,6 +1992,47 @@ func builtinExtVar(i *interpreter, name value) (value, error) {
 	return nil, i.Error("Undefined external variable: " + string(index))
 }
 
+func builtinGroupBy(i *interpreter, arguments []value) (value, error) {
+	arrv := arguments[0]
+	keyFv := arguments[1]
+
+	arr, err := i.getArray(arrv)
+	if err != nil {
+		return nil, err
+	}
+	keyF, err := i.getFunction(keyFv)
+	if err != nil {
+		return nil, err
+	}
+	json := make(map[string]interface{})
+	for _, thunk := range arr.elements {
+		currentValue, err := keyF.call(i, args(thunk))
+		if err != nil {
+			return nil, err
+		}
+		toStrVal, err := builtinToString(i, currentValue)
+		if err != nil {
+			return nil, err
+		}
+		strVal, err := i.getString(toStrVal)
+		if err != nil {
+			return nil, err
+		}
+		thunkValue, err := thunk.getValue(i)
+		if err != nil {
+			return nil, err
+		}
+
+		if json[strVal.getGoString()] == nil {
+			json[strVal.getGoString()] = []interface{}{thunkValue}
+		} else {
+			temp := json[strVal.getGoString()].([]interface{})
+			json[strVal.getGoString()] = append(temp, thunkValue)
+		}
+	}
+	return jsonToValue(i, json)
+}
+
 func builtinMinArray(i *interpreter, arguments []value) (value, error) {
 	arrv := arguments[0]
 	keyFv := arguments[1]
@@ -2530,6 +2571,7 @@ var funcBuiltins = buildBuiltinMap([]builtin{
 	&generalBuiltin{name: "sort", function: builtinSort, params: []generalBuiltinParameter{{name: "arr"}, {name: "keyF", defaultValue: functionID}}},
 	&generalBuiltin{name: "minArray", function: builtinMinArray, params: []generalBuiltinParameter{{name: "arr"}, {name: "keyF", defaultValue: functionID}}},
 	&generalBuiltin{name: "maxArray", function: builtinMaxArray, params: []generalBuiltinParameter{{name: "arr"}, {name: "keyF", defaultValue: functionID}}},
+	&generalBuiltin{name: "groupBy", function: builtinGroupBy, params: []generalBuiltinParameter{{name: "arr"}, {name: "keyF", defaultValue: functionID}}},
 	&unaryBuiltin{name: "native", function: builtinNative, params: ast.Identifiers{"x"}},
 	&unaryBuiltin{name: "sum", function: builtinSum, params: ast.Identifiers{"arr"}},
 	&unaryBuiltin{name: "avg", function: builtinAvg, params: ast.Identifiers{"arr"}},
diff --git a/interpreter.go b/interpreter.go
index 390d459b6..8ec440d23 100644
--- a/interpreter.go
+++ b/interpreter.go
@@ -968,6 +968,9 @@ func jsonToValue(i *interpreter, v interface{}) (value, error) {
 	case string:
 		return makeValueString(v), nil
 
+	case value:
+		return v, nil
+
 	default:
 		return nil, i.Error(fmt.Sprintf("Not a json type: %#+v", v))
 	}
diff --git a/linter/internal/types/stdlib.go b/linter/internal/types/stdlib.go
index 63c0eed34..ba5479d09 100644
--- a/linter/internal/types/stdlib.go
+++ b/linter/internal/types/stdlib.go
@@ -151,6 +151,7 @@ func prepareStdlib(g *typeGraph) {
 		"sum":           g.newSimpleFuncType(numberType, "arr"),
 		"minArray":      g.newFuncType(anyArrayType, []ast.Parameter{required("arr"), optional("keyF")}),
 		"maxArray":      g.newFuncType(anyArrayType, []ast.Parameter{required("arr"), optional("keyF")}),
+		"groupBy":       g.newFuncType(anyArrayType, []ast.Parameter{required("arr"), optional("keyF")}),
 		"contains":      g.newSimpleFuncType(boolType, "arr", "elem"),
 		"avg":		    	 g.newSimpleFuncType(numberType, "arr"),	
 		"all":           g.newSimpleFuncType(boolArrayType, "arr"),
diff --git a/testdata/builtinGroupBy.golden b/testdata/builtinGroupBy.golden
new file mode 100644
index 000000000..38d76efc3
--- /dev/null
+++ b/testdata/builtinGroupBy.golden
@@ -0,0 +1,12 @@
+{
+   "1": [
+      1,
+      1.5
+   ],
+   "2": [
+      2
+   ],
+   "3": [
+      3
+   ]
+}
diff --git a/testdata/builtinGroupBy.jsonnet b/testdata/builtinGroupBy.jsonnet
new file mode 100644
index 000000000..cc17da2f5
--- /dev/null
+++ b/testdata/builtinGroupBy.jsonnet
@@ -0,0 +1 @@
+std.groupBy([1, 1.5, 2, 3], std.floor)
\ No newline at end of file
diff --git a/testdata/builtinGroupBy.linter.golden b/testdata/builtinGroupBy.linter.golden
new file mode 100644
index 000000000..e69de29bb