feat: Go implementation for manifestYamlDoc and escapeStringJson (#2)

* Builtins for escapeStringJson and manifestYamlDoc

* Benchmark and tests

Author: jgraeger

builtin-benchmarks/escapeStringJson.jsonnet

@@ -0,0 +1 @@
+std.escapeStringJson("Lorem ipsum dolor sit amet, consectetur \"adipiscing\" elit. Nullam \\nec sagittis \\u0065lit, sed do \\teiusmod tempor incididunt ut labore et dolore magna aliqua.\nUt enim ad minim veniam, quis nostrud exercitation ullamco \flaboris nisi ut aliquip ex ea commodo consequat.\rDuis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.\n\tExcepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit \\anim id est laborum.\nCurabitur pretium tincidunt lacus. Nulla gravida orci a odio. \\Nullam varius, turpis et commodo pharetra, est eros bibendum elit, nec luctus magna felis sollicitudin mauris.\nInteger in mauris eu nibh euismod gravida. Duis ac tellus et risus vulputate vehicula. Donec lobortis risus a elit. Etiam tempor. Ut ullamcorper, ligula eu tempor congue, eros est euismod turpis, id tincidunt sapien risus a quam. Maecenas fermentum consequat mi. Donec fermentum. Pellentesque malesuada nulla a mi. Duis sapien sem, aliquet nec, commodo eget, consequat quis, neque. Aliquam faucibus, elit ut dictum aliquet, felis nisl adipiscing sapien, sed malesuada diam lacus eget erat. Cras mollis scelerisque nunc. Nullam arcu. Aliquam consequat. Curabitur augue lorem, dapibus quis, laoreet et, pretium ac, nisi. Aenean magna nisl, mollis quis, molestie eu, feugiat in, orci. In hac habitasse platea dictumst.")

builtin-benchmarks/manifestYamlDoc.jsonnet

@@ -0,0 +1,51 @@
+{
+  bar: {
+    prometheusOperator+: {
+      service+: {
+        spec+: {
+          ports: [
+            {
+              name: 'https',
+              port: 8443,
+              targetPort: 'https',
+            },
+          ],
+        },
+      },
+      serviceMonitor+: {
+        spec+: {
+          endpoints: [
+            {
+              port: 'https',
+              scheme: 'https',
+              honorLabels: true,
+              bearerTokenFile: '/var/run/secrets/kubernetes.io/serviceaccount/token',
+              tlsConfig: {
+                insecureSkipVerify: true,
+              },
+            },
+          ],
+        },
+      },
+      clusterRole+: {
+        rules+: [
+          {
+            apiGroups: ['authentication.k8s.io'],
+            resources: ['tokenreviews'],
+            verbs: ['create'],
+          },
+          {
+            apiGroups: ['authorization.k8s.io'],
+            resources: ['subjectaccessreviews'],
+            verbs: ['create'],
+          },
+        ],
+      },
+    },
+    additional+: {
+      '$schema': "http://json-schema.org/draft-07/schema#",
+      '09': ['no', 'yes'],
+    },
+  },
+  nothing: std.manifestYamlDoc(self.bar, indent_array_in_object=true, quote_keys=true),
+}

builtins.go

@@ -29,6 +29,7 @@ import (
 	"io"
 	"math"
 	"reflect"
+	"regexp"
 	"sort"
 	"strconv"
 	"strings"
@@ -233,17 +234,25 @@ func builtinLength(i *interpreter, x value) (value, error) {
 	return makeValueNumber(float64(num)), nil
 }
 
-func builtinToString(i *interpreter, x value) (value, error) {
+func valueToString(i *interpreter, x value) (string, error) {
 	switch x := x.(type) {
 	case valueString:
-		return x, nil
+		return x.getGoString(), nil
 	}
+
 	var buf bytes.Buffer
-	err := i.manifestAndSerializeJSON(&buf, x, false, "")
+	if err := i.manifestAndSerializeJSON(&buf, x, false, ""); err != nil {
+		return "", err
+	}
+	return buf.String(), nil
+}
+
+func builtinToString(i *interpreter, x value) (value, error) {
+	s, err := valueToString(i, x)
 	if err != nil {
 		return nil, err
 	}
-	return makeValueString(buf.String()), nil
+	return makeValueString(s), nil
 }
 
 func builtinTrace(i *interpreter, x value, y value) (value, error) {
@@ -1556,8 +1565,16 @@ func tomlIsSection(i *interpreter, val value) (bool, error) {
 	}
 }
 
-// tomlEncodeString encodes a string as quoted TOML string
-func tomlEncodeString(s string) string {
+func builtinEscapeStringJson(i *interpreter, v value) (value, error) {
+	s, err := valueToString(i, v)
+	if err != nil {
+		return nil, err
+	}
+
+	return makeValueString(unparseString(s)), nil
+}
+
+func escapeStringJson(s string) string {
 	res := "\""
 
 	for _, c := range s {
@@ -1589,6 +1606,11 @@ func tomlEncodeString(s string) string {
 	return res
 }
 
+// tomlEncodeString encodes a string as quoted TOML string
+func tomlEncodeString(s string) string {
+	return unparseString(s)
+}
+
 // tomlEncodeKey encodes a key - returning same string if it does not need quoting,
 // otherwise return it quoted; returns empty key as ”
 func tomlEncodeKey(s string) string {
@@ -1980,6 +2002,209 @@ func builtinManifestJSONEx(i *interpreter, arguments []value) (value, error) {
 	return makeValueString(finalString), nil
 }
 
+const (
+	yamlIndent = "  "
+)
+
+var (
+	yamlReserved = []string{
+		// Boolean types taken from https://yaml.org/type/bool.html
+		"true", "false", "yes", "no", "on", "off", "y", "n",
+		// Numerical words taken from https://yaml.org/type/float.html
+		".nan", "-.inf", "+.inf", ".inf", "null",
+		// Invalid keys that contain no invalid characters
+		"-", "---", "''",
+	}
+	yamlTimestampPattern = regexp.MustCompile(`^(?:[0-9]*-){2}[0-9]*$`)
+	yamlBinaryPattern    = regexp.MustCompile(`^[-+]?0b[0-1_]+$`)
+	yamlHexPattern       = regexp.MustCompile(`[-+]?0x[0-9a-fA-F_]+`)
+)
+
+func yamlReservedString(s string) bool {
+	for _, r := range yamlReserved {
+		if strings.EqualFold(s, r) {
+			return true
+		}
+	}
+	return false
+}
+
+func yamlBareSafe(s string) bool {
+	if len(s) == 0 {
+		return false
+	}
+
+	// String contains unsafe char
+	for _, c := range s {
+		isAlpha := (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z')
+		isDigit := c >= '0' && c <= '9'
+
+		if !isAlpha && !isDigit && c != '_' && c != '-' && c != '/' && c != '.' {
+			return false
+		}
+	}
+
+	if yamlReservedString(s) {
+		return false
+	}
+
+	if yamlTimestampPattern.MatchString(s) {
+		return false
+	}
+
+	// Check binary /
+	if yamlBinaryPattern.MatchString(s) || yamlHexPattern.MatchString(s) {
+		return false
+	}
+
+	// Is integer
+	if _, err := strconv.Atoi(s); err == nil {
+		return false
+	}
+	// Is float
+	if _, err := strconv.ParseFloat(s, 64); err == nil {
+		return false
+	}
+
+	return true
+}
+
+func builtinManifestYamlDoc(i *interpreter, arguments []value) (value, error) {
+	val := arguments[0]
+	vindentArrInObj, err := i.getBoolean(arguments[1])
+	if err != nil {
+		return nil, err
+	}
+	vQuoteKeys, err := i.getBoolean(arguments[2])
+	if err != nil {
+		return nil, err
+	}
+
+	var buf bytes.Buffer
+
+	var aux func(ov value, buf *bytes.Buffer, cindent string) error
+	aux = func(ov value, buf *bytes.Buffer, cindent string) error {
+		switch v := ov.(type) {
+		case *valueNull:
+			buf.WriteString("null")
+		case *valueBoolean:
+			if v.value {
+				buf.WriteString("true")
+			} else {
+				buf.WriteString("false")
+			}
+		case valueString:
+			s := v.getGoString()
+			if s == "" {
+				buf.WriteString(`""`)
+			} else if strings.HasSuffix(s, "\n") {
+				s := strings.TrimSuffix(s, "\n")
+				buf.WriteString("|")
+				for _, line := range strings.Split(s, "\n") {
+					buf.WriteByte('\n')
+					buf.WriteString(cindent)
+					buf.WriteString(yamlIndent)
+					buf.WriteString(line)
+				}
+			} else {
+				buf.WriteString(unparseString(s))
+			}
+		case *valueNumber:
+			buf.WriteString(strconv.FormatFloat(v.value, 'f', -1, 64))
+		case *valueArray:
+			if v.length() == 0 {
+				buf.WriteString("[]")
+				return nil
+			}
+			for ix, elem := range v.elements {
+				if ix != 0 {
+					buf.WriteByte('\n')
+					buf.WriteString(cindent)
+				}
+				thunkValue, err := elem.getValue(i)
+				if err != nil {
+					return err
+				}
+				buf.WriteByte('-')
+
+				if v, isArr := thunkValue.(*valueArray); isArr && v.length() > 0 {
+					buf.WriteByte('\n')
+					buf.WriteString(cindent)
+					buf.WriteString(yamlIndent)
+				} else {
+					buf.WriteByte(' ')
+				}
+
+				prevIndent := cindent
+				switch thunkValue.(type) {
+				case *valueArray, *valueObject:
+					cindent = cindent + yamlIndent
+				}
+
+				if err := aux(thunkValue, buf, cindent); err != nil {
+					return err
+				}
+				cindent = prevIndent
+			}
+		case *valueObject:
+			fields := objectFields(v, withoutHidden)
+			if len(fields) == 0 {
+				buf.WriteString("{}")
+				return nil
+			}
+			sort.Strings(fields)
+			for ix, fieldName := range fields {
+				fieldValue, err := v.index(i, fieldName)
+				if err != nil {
+					return err
+				}
+
+				if ix != 0 {
+					buf.WriteByte('\n')
+					buf.WriteString(cindent)
+				}
+
+				keyStr := fieldName
+				if vQuoteKeys.value || !yamlBareSafe(fieldName) {
+					keyStr = escapeStringJson(fieldName)
+				}
+				buf.WriteString(keyStr)
+				buf.WriteByte(':')
+
+				prevIndent := cindent
+				if v, isArr := fieldValue.(*valueArray); isArr && v.length() > 0 {
+					buf.WriteByte('\n')
+					buf.WriteString(cindent)
+					if vindentArrInObj.value {
+						buf.WriteString(yamlIndent)
+						cindent = cindent + yamlIndent
+					}
+				} else if v, isObj := fieldValue.(*valueObject); isObj {
+					if len(objectFields(v, withoutHidden)) > 0 {
+						buf.WriteByte('\n')
+						buf.WriteString(cindent)
+						buf.WriteString(yamlIndent)
+						cindent = cindent + yamlIndent
+					} else {
+						buf.WriteByte(' ')
+					}
+				} else {
+					buf.WriteByte(' ')
+				}
+				aux(fieldValue, buf, cindent)
+				cindent = prevIndent
+			}
+		}
+		return nil
+	}
+
+	if err := aux(val, &buf, ""); err != nil {
+		return nil, err
+	}
+
+	return makeValueString(buf.String()), nil
+}
+
 func builtinExtVar(i *interpreter, name value) (value, error) {
 	str, err := i.getString(name)
 	if err != nil {
@@ -2097,12 +2322,12 @@ func builtinAvg(i *interpreter, arrv value) (value, error) {
 	if err != nil {
 		return nil, err
 	}
-	
+
 	len := float64(arr.length())
 	if len == 0 {
 		return nil, i.Error("Cannot calculate average of an empty array.")
 	}
-	
+
 	sumValue, err := builtinSum(i, arrv)
 	if err != nil {
 		return nil, err
@@ -2112,7 +2337,7 @@ func builtinAvg(i *interpreter, arrv value) (value, error) {
 		return nil, err
 	}
 
-	avg := sum.value/len
+	avg := sum.value / len
 	return makeValueNumber(avg), nil
 }
 
@@ -2459,6 +2684,7 @@ var funcBuiltins = buildBuiltinMap([]builtin{
 	&unaryBuiltin{name: "extVar", function: builtinExtVar, params: ast.Identifiers{"x"}},
 	&unaryBuiltin{name: "length", function: builtinLength, params: ast.Identifiers{"x"}},
 	&unaryBuiltin{name: "toString", function: builtinToString, params: ast.Identifiers{"a"}},
+	&unaryBuiltin{name: "escapeStringJson", function: builtinEscapeStringJson, params: ast.Identifiers{"str_"}},
 	&binaryBuiltin{name: "trace", function: builtinTrace, params: ast.Identifiers{"str", "rest"}},
 	&binaryBuiltin{name: "makeArray", function: builtinMakeArray, params: ast.Identifiers{"sz", "func"}},
 	&binaryBuiltin{name: "flatMap", function: builtinFlatMap, params: ast.Identifiers{"func", "arr"}},
@@ -2524,6 +2750,11 @@ var funcBuiltins = buildBuiltinMap([]builtin{
 		{name: "newline", defaultValue: &valueFlatString{value: []rune("\n")}},
 		{name: "key_val_sep", defaultValue: &valueFlatString{value: []rune(": ")}}}},
 	&generalBuiltin{name: "manifestTomlEx", function: builtinManifestTomlEx, params: []generalBuiltinParameter{{name: "value"}, {name: "indent"}}},
+	&generalBuiltin{name: "manifestYamlDoc", function: builtinManifestYamlDoc, params: []generalBuiltinParameter{
+		{name: "value"},
+		{name: "indent_array_in_object", defaultValue: &valueBoolean{value: false}},
+		{name: "quote_keys", defaultValue: &valueBoolean{value: true}},
+	}},
 	&unaryBuiltin{name: "base64", function: builtinBase64, params: ast.Identifiers{"input"}},
 	&unaryBuiltin{name: "encodeUTF8", function: builtinEncodeUTF8, params: ast.Identifiers{"str"}},
 	&unaryBuiltin{name: "decodeUTF8", function: builtinDecodeUTF8, params: ast.Identifiers{"arr"}},

builtins_benchmark_test.go

@@ -56,6 +56,10 @@ func Benchmark_Builtin_base64_byte_array(b *testing.B) {
 	RunBenchmark(b, "base64_byte_array")
 }
 
+func Benchmark_Builtin_escapeStringJson(b *testing.B) {
+	RunBenchmark(b, "escapeStringJson")
+}
+
 func Benchmark_Builtin_manifestJsonEx(b *testing.B) {
 	RunBenchmark(b, "manifestJsonEx")
 }
@@ -64,6 +68,10 @@ func Benchmark_Builtin_manifestTomlEx(b *testing.B) {
 	RunBenchmark(b, "manifestTomlEx")
 }
 
+func Benchmark_Builtin_manifestYamlDoc(b *testing.B) {
+	RunBenchmark(b, "manifestYamlDoc")
+}
+
 func Benchmark_Builtin_comparison(b *testing.B) {
 	RunBenchmark(b, "comparison")
 }