Skip to content

Commit 0a00617

Browse files
SeanLedfordSeanLedford
authored
Refactor: Sending Variable Semantic Tokens Data from Rego (#1961)
* Refactor variable semantic tokens implementation to send token data straight from Rego! Signed-off-by: Sean Ledford <s_ledford@apple.com> * Update Tests Signed-off-by: Sean Ledford <s_ledford@apple.com> * Update modifiers to differentiate between definitions and declarations Signed-off-by: Sean Ledford <s_ledford@apple.com> * Remove custom unmarshaling logic Signed-off-by: Sean Ledford <s_ledford@apple.com> * Formatting/cleanup Signed-off-by: Sean Ledford <s_ledford@apple.com> --------- Signed-off-by: Sean Ledford <s_ledford@apple.com>
1 parent 784fa86 commit 0a00617

12 files changed

Lines changed: 279 additions & 338 deletions

File tree

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

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,5 @@ default result["response"] := {}
3232
result["response"] := {
3333
"packages": packages.result,
3434
"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-
},
35+
"vars": union({function_args.result, comprehensions.result, every_expr.result, some_expr.result}),
4136
}

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

Lines changed: 3 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@ test_function(param1, param2) := result if {
1212
}`
1313
module := regal.parse_module("p.rego", policy)
1414
result := semantictokens.result with data.workspace.parsed["file:///p.rego"] as module
15-
with input.regal as module.regal
1615
with input.params.textDocument.uri as "file:///p.rego"
16+
with input.regal.file.lines as split(policy, "\n")
1717

1818
result == {"response": {
1919
"imports": {
@@ -25,22 +25,8 @@ test_function(param1, param2) := result if {
2525
{"col": 14, "length": 3, "line": 0, "modifiers": 0, "type": 0},
2626
},
2727
"vars": {
28-
"comprehensions": {
29-
"declaration": set(),
30-
"reference": set(),
31-
},
32-
"every_expr": {
33-
"declaration": set(),
34-
"reference": set(),
35-
},
36-
"function_args": {
37-
"declaration": {
38-
{"location": "5:15:5:21", "type": "var", "value": "param1"},
39-
{"location": "5:23:5:29", "type": "var", "value": "param2"},
40-
},
41-
"reference": set(),
42-
},
43-
"some_expr": {"declaration": set(), "reference": set()},
28+
{"col": 14, "length": 6, "line": 4, "modifiers": 1, "type": 1},
29+
{"col": 22, "length": 6, "line": 4, "modifiers": 1, "type": 1},
4430
},
4531
}}
4632
}

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

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,21 +7,43 @@
77
package regal.lsp.semantictokens.vars.comprehensions
88

99
import data.regal.ast
10+
import data.regal.util
11+
12+
_declared_var_names contains name if {
13+
comprehension := ast.found.comprehensions[_][_]
14+
15+
some expr in comprehension.value.body
16+
some symbol in expr.terms.symbols
17+
some var in array.slice(symbol.value, 1, 100)
18+
var.type == "var"
19+
not startswith(var.value, "$")
20+
name := var.value
21+
}
1022

1123
# METADATA
1224
# description: Extract comprehension variable declarations from array/set/object comprehensions
13-
result.declaration contains var if {
25+
result contains token if {
1426
comprehension := ast.found.comprehensions[_][_]
1527

1628
some expr in comprehension.value.body
1729
some symbol in expr.terms.symbols
18-
some var in array.slice(symbol.value, 1, count(symbol.value) - 1)
30+
some var in array.slice(symbol.value, 1, 100)
1931
var.type == "var"
32+
33+
tloc := util.to_location_object(var.location)
34+
35+
token := {
36+
"line": tloc.row - 1,
37+
"col": tloc.col - 1,
38+
"length": tloc.end.col - tloc.col,
39+
"type": 1,
40+
"modifiers": bits.lsh(1, 1),
41+
}
2042
}
2143

2244
# METADATA
2345
# description: Extract comprehension variable references in the output
24-
result.reference contains var if {
46+
result contains token if {
2547
comprehension := ast.found.comprehensions[_][_]
2648

2749
output_vars := array.flatten([
@@ -31,7 +53,17 @@ result.reference contains var if {
3153
])
3254
some var in output_vars
3355
var.type == "var"
34-
var.value in {v.value | some v in result.declaration}
56+
var.value in _declared_var_names
57+
58+
tloc := util.to_location_object(var.location)
59+
60+
token := {
61+
"line": tloc.row - 1,
62+
"col": tloc.col - 1,
63+
"length": tloc.end.col - tloc.col,
64+
"type": 1,
65+
"modifiers": bits.lsh(1, 2),
66+
}
3567
}
3668

3769
default _comprehension_value(_) := set()

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

Lines changed: 34 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -2,80 +2,63 @@ package regal.lsp.semantictokens.vars.comprehensions_test
22

33
import data.regal.lsp.semantictokens.vars.comprehensions
44

5-
test_array_comprehension[note] if {
6-
policy_one := `package regal.woo
5+
test_array_comprehension if {
6+
policy := `package regal.woo
77
88
array_comprehensions := [x |
99
some i, x in [1, 2, 3]
1010
i == 2
1111
]`
1212

13-
array_comp_tokens := comprehensions.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
14-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy_one)
13+
tokens := comprehensions.result with data.workspace.parsed["file:///p.rego"] as regal.parse_module("p.rego", policy)
14+
with input.params.textDocument.uri as "file:///p.rego"
15+
with input.regal.file.lines as split(policy, "\n")
1516

16-
some note, tc in {"array comprehensions": {
17-
"declarations": {
18-
{"location": "4:10:4:11", "type": "var", "value": "x"},
19-
{"location": "4:7:4:8", "type": "var", "value": "i"},
20-
},
21-
"references": {
22-
{"location": "3:26:3:27", "type": "var", "value": "x"},
23-
{"location": "5:2:5:3", "type": "var", "value": "i"},
24-
},
25-
}}
26-
27-
tc.declarations == array_comp_tokens.declaration
28-
tc.references == array_comp_tokens.reference
17+
tokens == {
18+
{"col": 6, "length": 1, "line": 3, "modifiers": 2, "type": 1},
19+
{"col": 9, "length": 1, "line": 3, "modifiers": 2, "type": 1},
20+
{"col": 25, "length": 1, "line": 2, "modifiers": 4, "type": 1},
21+
{"col": 1, "length": 1, "line": 4, "modifiers": 4, "type": 1},
22+
}
2923
}
3024

31-
test_set_comprehensions[note] if {
32-
policy_one := `package regal.woo
25+
test_set_comprehension if {
26+
policy := `package regal.woo
3327
3428
set_comprehensions := {x |
3529
some i, x in [1, 2, 3]
3630
i == 2
3731
}`
3832

39-
set_comp_tokens := comprehensions.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
40-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy_one)
41-
42-
some note, tc in {"set comprehensions": {
43-
"declarations": {
44-
{"location": "4:10:4:11", "type": "var", "value": "x"},
45-
{"location": "4:7:4:8", "type": "var", "value": "i"},
46-
},
47-
"references": {
48-
{"location": "3:24:3:25", "type": "var", "value": "x"},
49-
{"location": "5:2:5:3", "type": "var", "value": "i"},
50-
},
51-
}}
33+
tokens := comprehensions.result with data.workspace.parsed["file:///p.rego"] as regal.parse_module("p.rego", policy)
34+
with input.params.textDocument.uri as "file:///p.rego"
35+
with input.regal.file.lines as split(policy, "\n")
5236

53-
tc.declarations == set_comp_tokens.declaration
54-
tc.references == set_comp_tokens.reference
37+
tokens == {
38+
{"col": 6, "length": 1, "line": 3, "modifiers": 2, "type": 1},
39+
{"col": 9, "length": 1, "line": 3, "modifiers": 2, "type": 1},
40+
{"col": 23, "length": 1, "line": 2, "modifiers": 4, "type": 1},
41+
{"col": 1, "length": 1, "line": 4, "modifiers": 4, "type": 1},
42+
}
5543
}
5644

57-
test_object_comprehension[note] if {
58-
policy_one := `package regal.woo
45+
test_object_comprehension if {
46+
policy := `package regal.woo
5947
6048
object_comprehensions := {k: v |
6149
some k, v in [1, 2, 3]
6250
v == 2
6351
}`
64-
object_comp_tokens := comprehensions.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
65-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy_one)
6652

67-
some note, tc in {"object comprehensions": {
68-
"declarations": {
69-
{"location": "4:10:4:11", "type": "var", "value": "v"},
70-
{"location": "4:7:4:8", "type": "var", "value": "k"},
71-
},
72-
"references": {
73-
{"location": "3:27:3:28", "type": "var", "value": "k"},
74-
{"location": "3:30:3:31", "type": "var", "value": "v"},
75-
{"location": "5:2:5:3", "type": "var", "value": "v"},
76-
},
77-
}}
53+
tokens := comprehensions.result with data.workspace.parsed["file:///p.rego"] as regal.parse_module("p.rego", policy)
54+
with input.params.textDocument.uri as "file:///p.rego"
55+
with input.regal.file.lines as split(policy, "\n")
7856

79-
tc.declarations == object_comp_tokens.declaration
80-
tc.references == object_comp_tokens.reference
57+
tokens == {
58+
{"col": 6, "length": 1, "line": 3, "modifiers": 2, "type": 1},
59+
{"col": 9, "length": 1, "line": 3, "modifiers": 2, "type": 1},
60+
{"col": 26, "length": 1, "line": 2, "modifiers": 4, "type": 1},
61+
{"col": 29, "length": 1, "line": 2, "modifiers": 4, "type": 1},
62+
{"col": 1, "length": 1, "line": 4, "modifiers": 4, "type": 1},
63+
}
8164
}

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

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,22 +7,33 @@
77
package regal.lsp.semantictokens.vars.every_expr
88

99
import data.regal.ast
10+
import data.regal.util
1011

1112
# METADATA
1213
# description: Get the module from workspace
1314
module := data.workspace.parsed[input.params.textDocument.uri]
1415

1516
# METADATA
16-
# description: Extract variable declarations from every keyword domains
17-
result.declaration contains var if {
17+
# description: Extract variable definitions from every keyword domains
18+
result contains token if {
1819
some rule_index
1920
declared_vars := ast.found.vars[rule_index]["every"]
2021
some var in declared_vars
22+
23+
tloc := util.to_location_object(var.location)
24+
25+
token := {
26+
"line": tloc.row - 1,
27+
"col": tloc.col - 1,
28+
"length": tloc.end.col - tloc.col,
29+
"type": 1,
30+
"modifiers": bits.lsh(1, 1),
31+
}
2132
}
2233

2334
# METADATA
2435
# description: Extract variable references in every keyword domains
25-
result.reference contains var if {
36+
result contains token if {
2637
some rule_index
2738

2839
declared_vars := ast.found.vars[rule_index]["every"]
@@ -36,4 +47,14 @@ result.reference contains var if {
3647
var.type == "var"
3748
var.value in declared_var_names
3849
not var in declared_vars
50+
51+
tloc := util.to_location_object(var.location)
52+
53+
token := {
54+
"line": tloc.row - 1,
55+
"col": tloc.col - 1,
56+
"length": tloc.end.col - tloc.col,
57+
"type": 1,
58+
"modifiers": bits.lsh(1, 2),
59+
}
3960
}

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

Lines changed: 19 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ package regal.lsp.semantictokens.vars.every_expr_test
22

33
import data.regal.lsp.semantictokens.vars.every_expr
44

5-
test_every_vars[note] if {
5+
test_every_two_vars if {
66
policy := `package regal.woo
77
88
every_two_vars if {
@@ -12,25 +12,19 @@ every_two_vars if {
1212
}
1313
}`
1414

15-
tokens := every_expr.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
16-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy)
17-
18-
some note, tc in {"every expression variables": {
19-
"declarations": {
20-
{"location": "4:11:4:12", "type": "var", "value": "v"},
21-
{"location": "4:8:4:9", "type": "var", "value": "k"},
22-
},
23-
"references": {
24-
{"location": "5:13:5:14", "type": "var", "value": "k"},
25-
{"location": "6:3:6:4", "type": "var", "value": "v"},
26-
},
27-
}}
28-
29-
tc.declarations == tokens.declaration
30-
tc.references == tokens.reference
15+
tokens := every_expr.result with data.workspace.parsed["file:///p.rego"] as regal.parse_module("p.rego", policy)
16+
with input.params.textDocument.uri as "file:///p.rego"
17+
with input.regal.file.lines as split(policy, "\n")
18+
19+
tokens == {
20+
{"col": 7, "length": 1, "line": 3, "modifiers": 2, "type": 1},
21+
{"col": 10, "length": 1, "line": 3, "modifiers": 2, "type": 1},
22+
{"col": 12, "length": 1, "line": 4, "modifiers": 4, "type": 1},
23+
{"col": 2, "length": 1, "line": 5, "modifiers": 4, "type": 1},
24+
}
3125
}
3226

33-
test_every_single_var_case if {
27+
test_every_single_var if {
3428
policy := `package regal.woo
3529
3630
every_one_var if {
@@ -39,14 +33,12 @@ every_one_var if {
3933
}
4034
}`
4135

42-
tokens := every_expr.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
43-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy)
36+
tokens := every_expr.result with data.workspace.parsed["file:///p.rego"] as regal.parse_module("p.rego", policy)
37+
with input.params.textDocument.uri as "file:///p.rego"
38+
with input.regal.file.lines as split(policy, "\n")
4439

45-
some note, tc in {"every expression variables": {
46-
"declarations": {{"location": "4:8:4:9", "type": "var", "value": "k"}},
47-
"references": {{"location": "5:13:5:14", "type": "var", "value": "k"}},
48-
}}
49-
50-
tc.declarations == tokens.declaration
51-
tc.references == tokens.reference
40+
tokens == {
41+
{"col": 7, "length": 1, "line": 3, "modifiers": 2, "type": 1},
42+
{"col": 12, "length": 1, "line": 4, "modifiers": 4, "type": 1},
43+
}
5244
}

0 commit comments

Comments
 (0)