Walk: skip path creation in wildcard assignment (#6267)

We do a lot of `walk`-ing in [Regal](https://docs.styra.com/regal).
So much that it's by far the single most expensive operation. That means
any optimization of the `walk` built-in function will be a win for us.

Seeing as we rarely make use of the `path` component when `walk`-ing
through AST inputs, I was curious to see if there was any optimization
we could take when the path is a wildcard assignment, and as such clearly
marked as unimportant. Turns out there is. Since the return value is
provided in the operators list, we can check the value provided for the
`path` part of the assigned array, and if it's a wildcard (`_`) skip
path construction entirely. Example:

```rego
walk(input, [_, value])
```

This greatly simplifies the walk, and the performance gains are
substantial. Traversing a ~7MB AST:

**main**
```shell
$ opa bench -d p.rego -i objects.json data.p.w
+-------------------------------------------+---------------+
| samples                                   |             6 |
| ns/op                                     |     168806625 |
| B/op                                      |     197364318 |
| allocs/op                                 |       3855327 |
| histogram_timer_rego_query_eval_ns_75%    |     169968114 |
| histogram_timer_rego_query_eval_ns_90%    |     170513459 |
| histogram_timer_rego_query_eval_ns_95%    |     170513459 |
| histogram_timer_rego_query_eval_ns_99%    |     170513459 |
| histogram_timer_rego_query_eval_ns_99.9%  |     170513459 |
| histogram_timer_rego_query_eval_ns_99.99% |     170513459 |
| histogram_timer_rego_query_eval_ns_count  |          6.00 |
| histogram_timer_rego_query_eval_ns_max    |     170513459 |
| histogram_timer_rego_query_eval_ns_mean   |     168789611 |
| histogram_timer_rego_query_eval_ns_median |     168924020 |
| histogram_timer_rego_query_eval_ns_min    |     166685000 |
| histogram_timer_rego_query_eval_ns_stddev |       1239390 |
+-------------------------------------------+---------------+
```

**no-path-walk**
```shell
$ opa bench -d p.rego -i objects.json data.p.w
+-------------------------------------------+--------------+
| samples                                   |           21 |
| ns/op                                     |     50629984 |
| B/op                                      |     38018790 |
| allocs/op                                 |      1025211 |
| histogram_timer_rego_query_eval_ns_75%    |     51239562 |
| histogram_timer_rego_query_eval_ns_90%    |     51540933 |
| histogram_timer_rego_query_eval_ns_95%    |     51674420 |
| histogram_timer_rego_query_eval_ns_99%    |     51688208 |
| histogram_timer_rego_query_eval_ns_99.9%  |     51688208 |
| histogram_timer_rego_query_eval_ns_99.99% |     51688208 |
| histogram_timer_rego_query_eval_ns_count  |         21.0 |
| histogram_timer_rego_query_eval_ns_max    |     51688208 |
| histogram_timer_rego_query_eval_ns_mean   |     50611103 |
| histogram_timer_rego_query_eval_ns_median |     50871459 |
| histogram_timer_rego_query_eval_ns_min    |     49518833 |
| histogram_timer_rego_query_eval_ns_stddev |       748688 |
+-------------------------------------------+--------------+
```

The real-world impact is not as dramatic, since we aren't *just*
walking, but normally need to actually **do** something with the
values returned, but consistently shaving off about 13% eval time
when linting one of the largest policy libraries isn't bad at all:

**Regal main**
```shell
go run main.go lint ~/tmp/kics/assets  162.16s user 6.04s system 593% cpu 28.362 total
```

**Regal walk-no-path**
```shell
go run main.go lint ~/tmp/kics/assets  145.51s user 5.01s system 597% cpu 25.176 total
```

Signed-off-by: Anders Eknert <anders@eknert.com>

ast/builtins.go

@@ -2470,7 +2470,7 @@ var WalkBuiltin = &Builtin{
 				types.A,
 			},
 			nil,
-		)).Description("pairs of `path` and `value`: `path` is an array representing the pointer to `value` in `x`"),
+		)).Description("pairs of `path` and `value`: `path` is an array representing the pointer to `value` in `x`. If `path` is assigned a wildcard (`_`), the `walk` function will skip path creation entirely for faster evaluation."),
 	),
 	Categories: graphs,
 }

builtin_metadata.json

@@ -18889,7 +18889,7 @@
     "introduced": "v0.17.0",
     "relation": true,
     "result": {
-      "description": "pairs of `path` and `value`: `path` is an array representing the pointer to `value` in `x`",
+      "description": "pairs of `path` and `value`: `path` is an array representing the pointer to `value` in `x`. If `path` is assigned a wildcard (`_`), the `walk` function will skip path creation entirely for faster evaluation.",
       "name": "output",
       "type": "array\u003carray[any], any\u003e"
     },

test/cases/testdata/walkbuiltin/test-walkbuiltin-wildcard-path.yaml

@@ -0,0 +1,28 @@
+cases:
+  - modules:
+      - |
+        package testing
+
+        obj := {
+            "bar": "baz",
+            "qux": [
+                1,
+                {"p": "rego", "q": "rules"},
+                {1, 2, 3, {"a": "b", "c": {"d", "e", 1}}}
+            ]
+        }
+
+        with_path[value] {
+            walk(obj, [path, value])
+        }
+
+        without_path[value] {
+            walk(obj, [_, value])
+        }
+        
+        same_values := with_path == without_path
+
+    note: "walkbuiltin/wildcard-path same values as when path provided"
+    query: x = data.testing.same_values
+    want_result:
+      - x: true

topdown/walk.go

@@ -10,6 +10,15 @@ import (
 
 func evalWalk(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
 	input := operands[0]
+
+	if pathIsWildcard(operands) {
+		// When the path assignment is a wildcard: walk(input, [_, value])
+		// we may skip the path construction entirely, and simply return
+		// same pointer in each iteration. This is a much more efficient
+		// path when only the values are needed.
+		return walkNoPath(input, iter)
+	}
+
 	filter := getOutputPath(operands)
 	return walk(filter, nil, input, iter)
 }
@@ -70,6 +79,33 @@ func walk(filter, path *ast.Array, input *ast.Term, iter func(*ast.Term) error)
 	return nil
 }
 
+var emptyArr = ast.ArrayTerm()
+
+func walkNoPath(input *ast.Term, iter func(*ast.Term) error) error {
+	if err := iter(ast.ArrayTerm(emptyArr, input)); err != nil {
+		return err
+	}
+
+	switch v := input.Value.(type) {
+	case ast.Object:
+		return v.Iter(func(_, v *ast.Term) error {
+			return walkNoPath(v, iter)
+		})
+	case *ast.Array:
+		for i := 0; i < v.Len(); i++ {
+			if err := walkNoPath(v.Elem(i), iter); err != nil {
+				return err
+			}
+		}
+	case ast.Set:
+		return v.Iter(func(elem *ast.Term) error {
+			return walkNoPath(elem, iter)
+		})
+	}
+
+	return nil
+}
+
 func pathAppend(path *ast.Array, key *ast.Term) *ast.Array {
 	if path == nil {
 		return ast.NewArray(key)
@@ -80,17 +116,26 @@ func pathAppend(path *ast.Array, key *ast.Term) *ast.Array {
 
 func getOutputPath(operands []*ast.Term) *ast.Array {
 	if len(operands) == 2 {
-		if arr, ok := operands[1].Value.(*ast.Array); ok {
-			if arr.Len() == 2 {
-				if path, ok := arr.Elem(0).Value.(*ast.Array); ok {
-					return path
-				}
+		if arr, ok := operands[1].Value.(*ast.Array); ok && arr.Len() == 2 {
+			if path, ok := arr.Elem(0).Value.(*ast.Array); ok {
+				return path
 			}
 		}
 	}
 	return nil
 }
 
+func pathIsWildcard(operands []*ast.Term) bool {
+	if len(operands) == 2 {
+		if arr, ok := operands[1].Value.(*ast.Array); ok && arr.Len() == 2 {
+			if v, ok := arr.Elem(0).Value.(ast.Var); ok {
+				return v.IsWildcard()
+			}
+		}
+	}
+	return false
+}
+
 func init() {
 	RegisterBuiltinFunc(ast.WalkBuiltin.Name, evalWalk)
 }