ast/topdown: add strings.split_n builtin

ast/builtins.go

@@ -206,6 +206,8 @@ var Upper = v1.Upper
 
 var Split = v1.Split
 
+var SplitN = v1.SplitN
+
 var Replace = v1.Replace
 
 var ReplaceN = v1.ReplaceN

builtin_metadata.json

@@ -185,6 +185,7 @@
       "strings.render_template",
       "strings.replace_n",
       "strings.reverse",
+      "strings.split_n",
       "substring",
       "trim",
       "trim_left",
@@ -22620,6 +22621,36 @@
     },
     "wasm": true
   },
+  "strings.split_n": {
+    "args": [
+      {
+        "description": "string that is split",
+        "name": "x",
+        "type": "string"
+      },
+      {
+        "description": "delimiter used for splitting",
+        "name": "delimiter",
+        "type": "string"
+      },
+      {
+        "description": "positive: at most n substrings from left; negative: last |n| substrings from right; zero: nil",
+        "name": "n",
+        "type": "number"
+      }
+    ],
+    "available": [
+      "edge"
+    ],
+    "description": "Split returns an array containing elements of the input string split on a delimiter. A positive n limits the number of substrings returned from the left. A negative n returns the last |n| substrings from the right. Zero returns nil.",
+    "introduced": "edge",
+    "result": {
+      "description": "split parts",
+      "name": "ys",
+      "type": "array[string]"
+    },
+    "wasm": false
+  },
   "substring": {
     "args": [
       {

capabilities.json

@@ -4129,6 +4129,29 @@
         "type": "function"
       }
     },
+    {
+      "name": "strings.split_n",
+      "decl": {
+        "args": [
+          {
+            "type": "string"
+          },
+          {
+            "type": "string"
+          },
+          {
+            "type": "number"
+          }
+        ],
+        "result": {
+          "dynamic": {
+            "type": "string"
+          },
+          "type": "array"
+        },
+        "type": "function"
+      }
+    },
     {
       "name": "substring",
       "decl": {

v1/ast/builtins.go

@@ -141,6 +141,7 @@ var DefaultBuiltins = [...]*Builtin{
 	StartsWith,
 	EndsWith,
 	Split,
+	SplitN,
 	Replace,
 	ReplaceN,
 	Trim,
@@ -1279,6 +1280,21 @@ var Split = &Builtin{
 	CanSkipBctx: true,
 }
 
+var SplitN = &Builtin{
+	Name:        "strings.split_n",
+	Description: "Split returns an array containing elements of the input string split on a delimiter. A positive n limits the number of substrings returned from the left. A negative n returns the last |n| substrings from the right. Zero returns nil.",
+	Decl: types.NewFunction(
+		types.Args(
+			types.Named("x", types.S).Description("string that is split"),
+			types.Named("delimiter", types.S).Description("delimiter used for splitting"),
+			types.Named("n", types.N).Description("positive: at most n substrings from left; negative: last |n| substrings from right; zero: nil"),
+		),
+		types.Named("ys", types.NewArray(nil, types.S)).Description("split parts"),
+	),
+	Categories:  stringsCat,
+	CanSkipBctx: true,
+}
+
 var Replace = &Builtin{
 	Name:        "replace",
 	Description: "Replace replaces all instances of a sub-string.",

v1/test/cases/testdata/v1/strings/test-splitn.yaml

@@ -0,0 +1,101 @@
+---
+cases:
+  - note: "strings/split_n: basic split with limit"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("a.b.c.d", ".", 2)
+    want_result:
+      - x: ["a", "b.c.d"]
+  - note: "strings/split_n: limit greater than splits"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("a.b", ".", 5)
+    want_result:
+      - x: ["a", "b"]
+  - note: "strings/split_n: negative limit takes last n from right"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("a.b.c", ".", -1)
+    want_result:
+      - x: ["c"]
+  - note: "strings/split_n: negative limit takes last 2 from right"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("a;b;c;d", ";", -2)
+    want_result:
+      - x: ["c", "d"]
+  - note: "strings/split_n: negative limit exceeds splits returns all"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("a.b", ".", -5)
+    want_result:
+      - x: ["a", "b"]
+  - note: "strings/split_n: zero limit returns nil"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("a.b.c", ".", 0)
+    want_result:
+      - x: []
+  - note: "strings/split_n: limit of 1 returns whole string"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("a.b.c", ".", 1)
+    want_result:
+      - x: ["a.b.c"]
+  - note: "strings/split_n: no delimiter found"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("abc", ".", 2)
+    want_result:
+      - x: ["abc"]
+  - note: "strings/split_n: empty string"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("", ".", 2)
+    want_result:
+      - x: [""]
+  - note: "strings/split_n: empty delimiter"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("abc", "", 2)
+    want_result:
+      - x: ["a", "bc"]
+  - note: "strings/split_n: multi-char delimiter"
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p := strings.split_n("a::b::c::d", "::", 3)
+    want_result:
+      - x: ["a", "b", "c::d"]

v1/topdown/strings.go

@@ -538,6 +538,44 @@ func builtinSplit(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) e
 	return iter(ast.ArrayTerm(util.SplitMap(text, delim, ast.InternedTerm)...))
 }
 
+func builtinSplitN(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
+	s, err := builtins.StringOperand(operands[0].Value, 1)
+	if err != nil {
+		return err
+	}
+
+	d, err := builtins.StringOperand(operands[1].Value, 2)
+	if err != nil {
+		return err
+	}
+
+	n, err := builtins.IntOperand(operands[2].Value, 3)
+	if err != nil {
+		return err
+	}
+
+	text, delim := string(s), string(d)
+
+	var elems []string
+	if n < 0 {
+		// Negative n: split all, then take the last |n| elements.
+		all := strings.Split(text, delim)
+		if -n >= len(all) {
+			elems = all
+		} else {
+			elems = all[len(all)+n:]
+		}
+	} else {
+		elems = strings.SplitN(text, delim, n)
+	}
+
+	arr := make([]*ast.Term, len(elems))
+	for i := range elems {
+		arr[i] = ast.InternedTerm(elems[i])
+	}
+	return iter(ast.ArrayTerm(arr...))
+}
+
 func builtinReplace(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
 	s, err := builtins.StringOperand(operands[0].Value, 1)
 	if err != nil {
@@ -801,6 +839,7 @@ func init() {
 	RegisterBuiltinFunc(ast.Upper.Name, builtinUpper)
 	RegisterBuiltinFunc(ast.Lower.Name, builtinLower)
 	RegisterBuiltinFunc(ast.Split.Name, builtinSplit)
+	RegisterBuiltinFunc(ast.SplitN.Name, builtinSplitN)
 	RegisterBuiltinFunc(ast.Replace.Name, builtinReplace)
 	RegisterBuiltinFunc(ast.ReplaceN.Name, builtinReplaceN)
 	RegisterBuiltinFunc(ast.Trim.Name, builtinTrim)