Use new Roast library for custom AST serialization (#1001)

And the `jsoniter` library for any JSON handling.

This is quite a big win for linting performance:

```
➜ hyperfine --warmup 2 'regal_main lint bundle' 'regal_jsoniter lint bundle'
Benchmark 1: regal_main lint bundle
  Time (mean ± σ):      2.707 s ±  0.035 s    [User: 19.171 s, System: 0.636 s]
  Range (min … max):    2.662 s …  2.785 s    10 runs

Benchmark 2: regal_jsoniter lint bundle
  Time (mean ± σ):      2.185 s ±  0.019 s    [User: 14.932 s, System: 0.466 s]
  Range (min … max):    2.147 s …  2.207 s    10 runs

Summary
  regal_jsoniter lint bundle ran
    1.24 ± 0.02 times faster than regal_main lint bundle
```

But importantly, it doesn't make the AST harder to read or work with — quite
the opposite actually!

See the Roast README for all the details:
https://github.com/anderseknert/roast

Signed-off-by: Anders Eknert <[email protected]>

README.md

@@ -1,7 +1,7 @@
 # Regal
 
 [![Build Status](https://github.com/styrainc/regal/workflows/Build/badge.svg?branch=main)](https://github.com/styrainc/regal/actions)
-![OPA v0.67.0](https://openpolicyagent.org/badge/v0.67.0)
+![OPA v0.67.1](https://openpolicyagent.org/badge/v0.67.1)
 
 Regal is a linter and language server for [Rego](https://www.openpolicyagent.org/docs/latest/policy-language/), helping
 you write better policies and have fun while doing it!

bundle/regal/ast/ast.rego

@@ -3,6 +3,7 @@ package regal.ast
 import rego.v1
 
 import data.regal.config
+import data.regal.util
 
 scalar_types := {"boolean", "null", "number", "string"}
 
@@ -62,34 +63,6 @@ _is_name(ref, pos) if {
 	ref.type == "string"
 }
 
-# METADATA
-# description: |
-#   answers if the body was generated or not, i.e. not seen
-#   in the original Rego file — for example `x := 1`
-# scope: document
-
-# METADATA
-# description: covers case of allow := true, which expands to allow = true { true }
-generated_body(rule) if rule.body[0].location == rule.head.value.location
-
-# METADATA
-# description: covers case of default rules
-generated_body(rule) if rule["default"] == true
-
-# METADATA
-# description: covers case of rule["message"] or rule contains "message"
-generated_body(rule) if {
-	rule.body[0].location.row == rule.head.key.location.row
-
-	# this is a quirk in the AST — the generated body will have a location
-	# set before the key, i.e. "message"
-	rule.body[0].location.col < rule.head.key.location.col
-}
-
-# METADATA
-# description: covers case of f("x")
-generated_body(rule) if rule.body[0].location == rule.head.location
-
 # METADATA
 # description: all the rules (excluding functions) in the input AST
 rules := [rule |
@@ -168,7 +141,7 @@ function_calls[rule_index] contains call if {
 
 	call := {
 		"name": ref_to_string(ref[0].value),
-		"location": ref[0].location,
+		"location": util.to_location_object(ref[0].location),
 		"args": args,
 	}
 }
@@ -270,15 +243,18 @@ function_ret_in_args(fn_name, terms) if {
 }
 
 # METADATA
-# description: |
-#   answers if provided rule is implicitly assigned boolean true, i.e. allow { .. } or not
+# description: answers if provided rule is implicitly assigned boolean true, i.e. allow { .. } or not
 # scope: document
 implicit_boolean_assignment(rule) if {
 	# note the missing location attribute here, which is how we distinguish
 	# between implicit and explicit assignments
 	rule.head.value == {"type": "boolean", "value": true}
 }
 
+# or sometimes, like this...
+implicit_boolean_assignment(rule) if rule.head.value.location == rule.head.location
+
+# or like this...
 implicit_boolean_assignment(rule) if {
 	# This handles the *quite* special case of
 	# `a.b if true`, which is "rewritten" to `a.b = true` *and*  where a location is still added to the value
@@ -293,6 +269,7 @@ implicit_boolean_assignment(rule) if {
 	# If you write Rego like that — you're not going to use Regal anyway, are you? ¯\_(ツ)_/¯
 	rule.head.value.type == "boolean"
 	rule.head.value.value == true
+
 	rule.head.value.location.col == 1
 }
 
@@ -325,8 +302,10 @@ negated_expressions[rule] contains value if {
 #       input.baz
 #   }
 is_chained_rule_body(rule, lines) if {
-	row_text := lines[rule.head.location.row - 1]
-	col_text := substring(row_text, rule.head.location.col - 1, -1)
+	head_loc := util.to_location_object(rule.head.location)
+
+	row_text := lines[head_loc.row - 1]
+	col_text := substring(row_text, head_loc.col - 1, -1)
 
 	startswith(col_text, "{")
 }

bundle/regal/ast/ast_test.rego

@@ -4,6 +4,7 @@ import rego.v1
 
 import data.regal.ast
 import data.regal.capabilities
+import data.regal.util
 
 # regal ignore:rule-length
 test_find_vars if {
@@ -138,21 +139,14 @@ allow := true
 	blocks := ast.comment_blocks(module.comments)
 	blocks == [
 		[
-			{
-				"Location": {
-					"col": 1,
-					"row": 3,
-					"text": "IyBNRVRBREFUQQojIHRpdGxlOiBmb28KIyBiYXI6IGludmFsaWQ=",
-				},
-				"Text": "IE1FVEFEQVRB",
-			},
-			{"Location": {"col": 1, "file": "p.rego", "row": 4}, "Text": "IHRpdGxlOiBmb28="},
-			{"Location": {"col": 1, "file": "p.rego", "row": 5}, "Text": "IGJhcjogaW52YWxpZA=="},
+			{"location": "3:1:IyBNRVRBREFUQQojIHRpdGxlOiBmb28KIyBiYXI6IGludmFsaWQ=", "text": "IE1FVEFEQVRB"},
+			{"location": "4:1:IyB0aXRsZTogZm9v", "text": "IHRpdGxlOiBmb28="},
+			{"location": "5:1:IyBiYXI6IGludmFsaWQ=", "text": "IGJhcjogaW52YWxpZA=="},
 		],
-		[{"Location": {"col": 1, "file": "p.rego", "row": 8}, "Text": "IG5vdCBtZXRhZGF0YQ=="}],
+		[{"location": "8:1:IyBub3QgbWV0YWRhdGE=", "text": "IG5vdCBtZXRhZGF0YQ=="}],
 		[
-			{"Location": {"col": 1, "file": "p.rego", "row": 10}, "Text": "IGFub3RoZXI="},
-			{"Location": {"col": 1, "file": "p.rego", "row": 11}, "Text": "IGJsb2Nr"},
+			{"location": "10:1:IyBhbm90aGVy", "text": "IGFub3RoZXI="},
+			{"location": "11:1:IyBibG9jaw==", "text": "IGJsb2Nr"},
 		],
 	]
 }
@@ -269,16 +263,6 @@ test_find_some_decl_names_in_scope if {
 
 var_names(vars) := {var.value | some var in vars}
 
-test_generated_body_function if {
-	policy := `package p
-
-	f("x")`
-
-	module := regal.parse_module("p.rego", policy)
-
-	ast.generated_body(module.rules[0])
-}
-
 test_all_refs if {
 	policy := `package policy
 
@@ -295,7 +279,7 @@ test_all_refs if {
 
 	r := ast.all_refs with input as module
 
-	text_refs := {base64.decode(ref.location.text) | some ref in r}
+	text_refs := {base64.decode(util.to_location_object(ref.location).text) | some ref in r}
 
 	text_refs == {":=", "data.foo.bar", "data.foo.bax", "data.foo.baz"}
 }

bundle/regal/ast/comments.rego

@@ -2,11 +2,13 @@ package regal.ast
 
 import rego.v1
 
+import data.regal.util
+
 # METADATA
-# description: all comments in the input AST with their "Text" attribute base64 decoded
+# description: all comments in the input AST with their `text` attribute base64 decoded
 comments_decoded := [decoded |
 	some comment in input.comments
-	decoded := object.union(comment, {"Text": base64.decode(comment.Text)})
+	decoded := object.union(comment, {"text": base64.decode(comment.text)})
 ]
 
 # METADATA
@@ -39,14 +41,15 @@ comments["annotation_match"](str) if regex.match(
 #   found in input AST, indexed by the row they're at
 ignore_directives[row] := rules if {
 	some comment in comments_decoded
-	text := trim_space(comment.Text)
+	text := trim_space(comment.text)
 
 	i := indexof(text, "regal ignore:")
 	i != -1
 
 	list := regex.replace(substring(text, i + 13, -1), `\s`, "")
+	loc := util.to_location_object(comment.location)
 
-	row := comment.Location.row + 1
+	row := loc.row + 1
 	rules := split(list, ",")
 }
 
@@ -57,7 +60,7 @@ ignore_directives[row] := rules if {
 comment_blocks(comments) := [partition |
 	rows := [row |
 		some comment in comments
-		row := comment.Location.row
+		row := util.to_location_object(comment.location).row
 	]
 	breaks := _splits(rows)
 

bundle/regal/ast/keywords.rego

@@ -2,6 +2,8 @@ package regal.ast
 
 import rego.v1
 
+import data.regal.util
+
 keywords[row] contains keyword if {
 	some idx, line in input.regal.file.lines
 
@@ -21,35 +23,37 @@ keywords[row] contains keyword if {
 	}
 }
 
-keywords[pkg.location.row] contains keyword if {
+keywords[loc.row] contains keyword if {
 	pkg := input["package"]
+	loc := util.to_location_object(pkg.location)
 
 	keyword := {
 		"name": "package",
 		"location": {
-			"row": pkg.location.row,
-			"col": pkg.location.col,
+			"row": loc.row,
+			"col": loc.col,
 		},
 	}
 }
 
-keywords[imp.location.row] contains keyword if {
+keywords[loc.row] contains keyword if {
 	some imp in input.imports
 
+	loc := util.to_location_object(imp.location)
+
 	keyword := {
 		"name": "import",
 		"location": {
-			"row": imp.location.row,
-			"col": imp.location.col,
+			"row": loc.row,
+			"col": loc.col,
 		},
 	}
 }
 
 keywords[loc.row] contains keyword if {
 	some rule in input.rules
 
-	loc := rule.head.location
-
+	loc := util.to_location_object(rule.head.location)
 	col := indexof(base64.decode(loc.text), " contains ")
 
 	col > 0
@@ -63,37 +67,38 @@ keywords[loc.row] contains keyword if {
 	}
 }
 
-keywords[value.row] contains keyword if {
-	some expr in exprs[_]
+keywords[loc.row] contains keyword if {
+	expr := exprs[_][_]
 
 	walk(expr.terms, [path, value])
 
-	value.col
-	value.row
+	regal.last(path) == "location"
 
-	name := _keyword_b64s[value.text]
+	loc := util.to_location_object(value)
+	name := _keyword_b64s[loc.text]
 
 	parent_path := array.slice(path, 0, count(path) - 1)
 	context := object.get(expr.terms, parent_path, {})
 
-	some keyword in _determine_keywords(context, value, name)
+	some keyword in _determine_keywords(context, loc, name)
 }
 
-keywords[value.row] contains keyword if {
+keywords[loc.row] contains keyword if {
 	some rule in input.rules
 	rule.head.assign
 
 	walk(rule.head.value, [path, value])
 
-	value.col
-	value.row
+	regal.last(path) == "location"
 
-	name := _keyword_b64s[value.text]
+	loc := util.to_location_object(value)
+
+	name := _keyword_b64s[loc.text]
 
 	parent_path := array.slice(path, 0, count(path) - 1)
 	context := object.get(rule.head.value, parent_path, {})
 
-	some keyword in _determine_keywords(context, value, name)
+	some keyword in _determine_keywords(context, loc, name)
 }
 
 _determine_keywords(_, value, name) := {keyword} if {
@@ -109,7 +114,8 @@ _determine_keywords(_, value, name) := {keyword} if {
 }
 
 _determine_keywords(context, value, "every") := keywords if {
-	text := base64.decode(context.value.location.text)
+	ctx_loc := util.to_location_object(context.value.location)
+	text := base64.decode(ctx_loc.text)
 
 	keywords := {
 		{
@@ -123,15 +129,13 @@ _determine_keywords(context, value, "every") := keywords if {
 			"name": "in",
 			"location": {
 				"row": value.row,
-				"col": (context.value.location.col + count(text)) + 1,
+				"col": (ctx_loc.col + count(text)) + 1,
 			},
 		},
 	}
 }
 
-_comment_row_index contains comment.Location.row if {
-	some comment in object.get(input, "comments", [])
-}
+_comment_row_index contains util.to_location_object(comment.location).row if some comment in input.comments
 
 _keyword_b64s := {
 	"aW4=": "in",

bundle/regal/ast/keywords_test.rego

@@ -215,7 +215,7 @@ allow if {
 }
 
 _keyword_on_row(kwds, row, keyword) if {
-	some kwd in object.get(kwds, row, {})
+	some kwd in kwds[row]
 
 	kwd.name == keyword.name
 	kwd.location.row == keyword.location.row

bundle/regal/ast/rule_heads.rego → bundle/regal/ast/rule_head_locations.rego

@@ -2,11 +2,13 @@ package regal.ast
 
 import rego.v1
 
+import data.regal.util
+
 # METADATA
 # description: |
 #   For a given rule head name, this rule contains a list of locations where
 #   there is a rule head with that name.
-rule_head_locations[name] contains info if {
+rule_head_locations[name] contains {"row": loc.row, "col": loc.col} if {
 	some rule in input.rules
 
 	name := concat(".", [
@@ -15,8 +17,5 @@ rule_head_locations[name] contains info if {
 		ref_static_to_string(rule.head.ref),
 	])
 
-	info := {
-		"row": rule.head.location.row,
-		"col": rule.head.location.col,
-	}
+	loc := util.to_location_object(rule.head.location)
 }