Skip to content

Commit 299d3fc

Browse files
SeanLedfordSeanLedford
authored
Semantic Token Support for Comprehensions and Constructs (#1883)
* Fix bug causing lsp to error when new empty rego files were created Convert QueryEval return type to struct in router handler for easier field access Signed-off-by: Sean Ledford <s_ledford@apple.com> * Add functionality for comprehensions and ever/some constructs (Pending testing, working out edge cases/robustness) Signed-off-by: Sean Ledford <s_ledford@apple.com> * Update rego rules to properly differentiate constructs and correctly assign modifiers Add test cases for new functionality Signed-off-by: Sean Ledford <s_ledford@apple.com> * Refactored rules to use functions from the AST package Reorganized different token types into their own individual files Signed-off-by: Sean Ledford <s_ledford@apple.com> * Further refactored file structure and rule format Added tests for each individual token type and updated existing tests Signed-off-by: Sean Ledford <s_ledford@apple.com> * Address PR feedback regarding: - Test organization - Rule optimization - Code quality tweaks - Custom unmarshaling for location type Signed-off-by: Sean Ledford <s_ledford@apple.com> * Address PR feedback regarding: - Test organization - Rule optimization - Code quality tweaks - Custom unmarshaling for location type Signed-off-by: Sean Ledford <s_ledford@apple.com> * Addressing final PR feedback: - Address var refs not being parsed from rule head - Touch up syntax/metadate for comprehensions, import and packages rules - Update token field types to uint32 Signed-off-by: Sean Ledford <s_ledford@apple.com> --------- Signed-off-by: Sean Ledford <s_ledford@apple.com>
1 parent fd85e2d commit 299d3fc

20 files changed

Lines changed: 831 additions & 297 deletions

‎bundle/regal/ast/search.rego‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,15 @@ found.symbols[rule_index] contains value.symbols if {
209209
walk(_rules[i], [_, value])
210210
}
211211

212+
# METADATA
213+
# description: all every constructs found in module
214+
found.every[rule_index] contains terms if {
215+
some rule_index, i
216+
found.expressions[rule_index][i].terms.domain
217+
218+
terms := found.expressions[rule_index][i].terms
219+
}
220+
212221
# METADATA
213222
# description: all comprehensions found in module
214223
found.comprehensions[rule_index] contains value if {

‎bundle/regal/lsp/semantictokens/semantictokens.rego‎

Lines changed: 21 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -4,77 +4,38 @@
44
# - declarations of function args in text documents
55
# - variable references that are used in function calls
66
# - variable references that are used in expressions
7+
# - variables declarations and references in comprehensions
8+
# - variables declarations and references in every/some keyword domains
79
# related_resources:
810
# - https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_semanticTokens
911
# schemas:
1012
# - input: schema.regal.lsp.common
1113
# - input.params: schema.regal.lsp.semantictokens
1214
package regal.lsp.semantictokens
1315

16+
import data.regal.lsp.semantictokens.vars.comprehensions
17+
import data.regal.lsp.semantictokens.vars.every_expr
18+
import data.regal.lsp.semantictokens.vars.function_args
19+
import data.regal.lsp.semantictokens.vars.imports
20+
import data.regal.lsp.semantictokens.vars.packages
21+
import data.regal.lsp.semantictokens.vars.some_expr
22+
1423
# METADATA
1524
# description: Get the module from workspace
1625
module := data.workspace.parsed[input.params.textDocument.uri]
1726

18-
# METADATA
19-
# entrypoint: true
20-
result.response := {
21-
"arg_tokens": arg_tokens,
22-
"package_tokens": package_tokens,
23-
"import_tokens": import_tokens,
24-
}
25-
26-
# METADATA
27-
# description: Extract import tokens - return only last term of the path
28-
import_tokens contains last_term if {
29-
some import_statement in module.imports
30-
import_path := import_statement.path.value
31-
32-
last_term := import_path[count(import_path) - 1]
33-
}
34-
35-
# METADATA
36-
# description: Extract function argument declarations
37-
arg_tokens.declaration contains arg if {
38-
some rule in module.rules
39-
some arg in rule.head.args
40-
arg.type == "var"
41-
}
42-
43-
# METADATA
44-
# description: Extract variable references in function calls
45-
arg_tokens.reference contains arg if {
46-
some rule in module.rules
47-
48-
rule.head.args
49-
50-
arg_names := {term.value | some term in rule.head.args}
51-
52-
walk(rule.body, [_, expr])
53-
54-
expr.terms[0].type == "ref"
55-
56-
some arg in array.slice(expr.terms, 1, count(expr.terms))
57-
58-
arg.type == "var"
59-
arg.value in arg_names
60-
}
27+
# This is handling the case where the module from the parsed workspace is empty
28+
default result["response"] := {}
6129

6230
# METADATA
63-
# description: Extract variable references in call expressions
64-
arg_tokens.reference contains arg if {
65-
some rule in module.rules
66-
arg_names := {term.value | some term in rule.head.args}
67-
walk(rule.body, [_, expr])
68-
69-
some term in expr.terms
70-
term.type == "call"
71-
72-
some arg in term.value
73-
arg.type == "var"
74-
75-
arg.value in arg_names
31+
# entrypoint: true
32+
result["response"] := {
33+
"packages": packages.result,
34+
"imports": imports.result,
35+
"vars": {
36+
"function_args": function_args.result,
37+
"comprehensions": comprehensions.result,
38+
"every_expr": every_expr.result,
39+
"some_expr": some_expr.result,
40+
},
7641
}
77-
78-
# METADATA
79-
# description: Extract package tokens - return full package path
80-
package_tokens := module.package.path

‎bundle/regal/lsp/semantictokens/semantictokens_test.rego‎

Lines changed: 24 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -2,71 +2,37 @@ package regal.lsp.semantictokens_test
22

33
import data.regal.lsp.semantictokens
44

5-
test_arg_tokens_with_declarations_and_references if {
6-
policy := `package regal.woo
7-
8-
test_function(param1, param2) := result if {
9-
calc1 := param1 * 2
10-
calc2 := param2 + 10
11-
result := calc1 + calc2
12-
13-
calc3 := 1
14-
calc3 == param1
15-
}`
16-
tokens := semantictokens.arg_tokens with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
17-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy)
18-
19-
# Check declarations
20-
{"location": "3:15:3:21", "value": "param1", "type": "var"} in tokens.declaration
21-
{"location": "3:23:3:29", "value": "param2", "type": "var"} in tokens.declaration
22-
23-
# Check references
24-
{"location": "4:11:4:17", "value": "param1", "type": "var"} in tokens.reference
25-
{"location": "5:11:5:17", "value": "param2", "type": "var"} in tokens.reference
26-
{"location": "9:11:9:17", "value": "param1", "type": "var"} in tokens.reference
27-
}
28-
29-
test_arg_tokens_declarations_only if {
30-
policy := `package regal.woo
31-
32-
test_function(param1, param2) := result if {
33-
true
34-
}`
35-
tokens := semantictokens.arg_tokens with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
36-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy)
37-
38-
# Check declarations
39-
{"location": "3:15:3:21", "value": "param1", "type": "var"} in tokens.declaration
40-
{"location": "3:23:3:29", "value": "param2", "type": "var"} in tokens.declaration
41-
42-
# Should have no references since variables aren't used
43-
count(tokens.reference) == 0
44-
}
45-
46-
test_import_tokens if {
5+
test_result if {
476
policy := `package regal.woo
487
498
import data.regal.ast
509
5110
test_function(param1, param2) := result if {
5211
ast.is_constant
5312
}`
54-
tokens := semantictokens.import_tokens with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
55-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy)
56-
57-
# Check imports
58-
{"location": "3:19:3:22", "type": "string", "value": "ast"} in tokens
59-
}
60-
61-
test_arg_tokens_no_variables if {
62-
policy := `package regal.woo
63-
64-
allow := true`
65-
66-
tokens := semantictokens.arg_tokens with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
13+
result := semantictokens.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
6714
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy)
6815

69-
# Should find no tokens for a policy with no function arguments
70-
count(tokens.declaration) == 0
71-
count(tokens.reference) == 0
16+
result == {"response": {
17+
"imports": {{"location": "3:19:3:22", "type": "string", "value": "ast"}},
18+
"packages": {{"location": "1:15:1:18", "type": "string", "value": "woo"}},
19+
"vars": {
20+
"comprehensions": {
21+
"declaration": set(),
22+
"reference": set(),
23+
},
24+
"every_expr": {
25+
"declaration": set(),
26+
"reference": set(),
27+
},
28+
"function_args": {
29+
"declaration": {
30+
{"location": "5:15:5:21", "type": "var", "value": "param1"},
31+
{"location": "5:23:5:29", "type": "var", "value": "param2"},
32+
},
33+
"reference": set(),
34+
},
35+
"some_expr": {"declaration": set(), "reference": set()},
36+
},
37+
}}
7238
}

‎bundle/regal/lsp/semantictokens/vars/comprehensions/comprehensions.rego‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
# METADATA
2+
# description: |
3+
# Helper package for semantictokens that returns variable references and declarations in comprehensions
4+
# schemas:
5+
# - input: schema.regal.lsp.common
6+
# - input.params: schema.regal.lsp.semantictokens
7+
package regal.lsp.semantictokens.vars.comprehensions
8+
9+
import data.regal.ast
10+
11+
# METADATA
12+
# description: Extract comprehension variable declarations from array/set/object comprehensions
13+
result.declaration contains var if {
14+
comprehension := ast.found.comprehensions[_][_]
15+
16+
some expr in comprehension.value.body
17+
some symbol in expr.terms.symbols
18+
some var in array.slice(symbol.value, 1, count(symbol.value) - 1)
19+
var.type == "var"
20+
}
21+
22+
# METADATA
23+
# description: Extract comprehension variable references in the output
24+
result.reference contains var if {
25+
comprehension := ast.found.comprehensions[_][_]
26+
27+
output_vars := array.flatten([
28+
_comprehension_key(comprehension),
29+
_comprehension_value(comprehension),
30+
_comprehension_body_vars(comprehension.value.body),
31+
])
32+
some var in output_vars
33+
var.type == "var"
34+
var.value in {v.value | some v in result.declaration}
35+
}
36+
37+
default _comprehension_value(_) := set()
38+
39+
# Helper to get variables from differing comprehensions
40+
_comprehension_value(comprehension) := value if {
41+
comprehension.type in ["arraycomprehension", "setcomprehension"]
42+
value := comprehension.value.term
43+
} else := value if {
44+
comprehension.type == "objectcomprehension"
45+
value := comprehension.value.value
46+
}
47+
48+
default _comprehension_key(_) := set()
49+
50+
# Helper to get variables from differing comprehensions
51+
_comprehension_key(comprehension) := value if {
52+
comprehension.type == "objectcomprehension"
53+
value := comprehension.value.key
54+
}
55+
56+
# Helper to get variables from differing comprehensions
57+
_comprehension_body_vars(body) := [value |
58+
some expr in body
59+
not expr.terms.symbols
60+
some value in expr.terms
61+
value.type == "var"
62+
]

‎bundle/regal/lsp/semantictokens/vars/comprehensions/comprehensions_test.rego‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package regal.lsp.semantictokens.vars.comprehensions_test
2+
3+
import data.regal.lsp.semantictokens.vars.comprehensions
4+
5+
test_array_comprehension[note] if {
6+
policy_one := `package regal.woo
7+
8+
array_comprehensions := [x |
9+
some i, x in [1, 2, 3]
10+
i == 2
11+
]`
12+
some note, tc in {"array comprehensions": {
13+
"declarations": {
14+
{"location": "4:10:4:11", "type": "var", "value": "x"},
15+
{"location": "4:7:4:8", "type": "var", "value": "i"},
16+
},
17+
"references": {
18+
{"location": "3:26:3:27", "type": "var", "value": "x"},
19+
{"location": "5:2:5:3", "type": "var", "value": "i"},
20+
},
21+
}}
22+
array_comp_tokens := comprehensions.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
23+
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy_one)
24+
25+
tc.declarations == array_comp_tokens.declaration
26+
tc.references == array_comp_tokens.reference
27+
}
28+
29+
test_set_comprehensions[note] if {
30+
policy_one := `package regal.woo
31+
32+
set_comprehensions := {x |
33+
some i, x in [1, 2, 3]
34+
i == 2
35+
}`
36+
some note, tc in {"set comprehensions": {
37+
"declarations": {
38+
{"location": "4:10:4:11", "type": "var", "value": "x"},
39+
{"location": "4:7:4:8", "type": "var", "value": "i"},
40+
},
41+
"references": {
42+
{"location": "3:24:3:25", "type": "var", "value": "x"},
43+
{"location": "5:2:5:3", "type": "var", "value": "i"},
44+
},
45+
}}
46+
set_comp_tokens := comprehensions.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
47+
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy_one)
48+
49+
tc.declarations == set_comp_tokens.declaration
50+
tc.references == set_comp_tokens.reference
51+
}
52+
53+
test_object_comprehension[note] if {
54+
policy_one := `package regal.woo
55+
56+
object_comprehensions := {k: v |
57+
some k, v in [1, 2, 3]
58+
v == 2
59+
}`
60+
some note, tc in {"object comprehensions": {
61+
"declarations": {
62+
{"location": "4:10:4:11", "type": "var", "value": "v"},
63+
{"location": "4:7:4:8", "type": "var", "value": "k"},
64+
},
65+
"references": {
66+
{"location": "3:27:3:28", "type": "var", "value": "k"},
67+
{"location": "3:30:3:31", "type": "var", "value": "v"},
68+
{"location": "5:2:5:3", "type": "var", "value": "v"},
69+
},
70+
}}
71+
object_comp_tokens := comprehensions.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
72+
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy_one)
73+
74+
tc.declarations == object_comp_tokens.declaration
75+
tc.references == object_comp_tokens.reference
76+
}

‎bundle/regal/lsp/semantictokens/vars/every_expr/every_expr.rego‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# METADATA
2+
# description: |
3+
# Helper package for semantictokens that returns variable references and declarations 'every' keyword domains
4+
# schemas:
5+
# - input: schema.regal.lsp.common
6+
# - input.params: schema.regal.lsp.semantictokens
7+
package regal.lsp.semantictokens.vars.every_expr
8+
9+
import data.regal.ast
10+
11+
# METADATA
12+
# description: Get the module from workspace
13+
module := data.workspace.parsed[input.params.textDocument.uri]
14+
15+
# METADATA
16+
# description: Extract variable declarations from every keyword domains
17+
result.declaration contains var if {
18+
some rule_index
19+
declared_vars := ast.found.vars[rule_index]["every"]
20+
some var in declared_vars
21+
}
22+
23+
# METADATA
24+
# description: Extract variable references in every keyword domains
25+
result.reference contains var if {
26+
some rule_index
27+
28+
declared_vars := ast.found.vars[rule_index]["every"]
29+
30+
declared_var_names := {v.value | some v in declared_vars}
31+
32+
some every_terms in ast.found.every[rule_index]
33+
walk(every_terms.body, [_, expr])
34+
35+
some var in expr.terms
36+
var.type == "var"
37+
var.value in declared_var_names
38+
not var in declared_vars
39+
}

0 commit comments

Comments
 (0)