Skip to content

Commit 12cc776

Browse files
anderseknertanderseknert
authored
Semantic package tokens in Rego (#1956)
Finally had a moment to play some with @SeanLedford's great work on semantic tokens, and I could not resist seeing what moving a first part of the token creation over to Rego. Feels pretty straightforward, although a bit of a pain updating the token tests. But that's a good problem to have! Signed-off-by: Anders Eknert <anders.eknert@apple.com>
1 parent 5c62c72 commit 12cc776

6 files changed

Lines changed: 121 additions & 87 deletions

File tree

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,17 @@ import data.regal.ast
1010
test_function(param1, param2) := result if {
1111
ast.is_constant
1212
}`
13-
result := semantictokens.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)
13+
module := regal.parse_module("p.rego", policy)
14+
result := semantictokens.result with data.workspace.parsed["file:///p.rego"] as module
15+
with input.regal as module.regal
16+
with input.params.textDocument.uri as "file:///p.rego"
1517

1618
result == {"response": {
1719
"imports": {{"location": "3:19:3:22", "type": "string", "value": "ast"}},
18-
"packages": {{"location": "1:15:1:18", "type": "string", "value": "woo"}},
20+
"packages": {
21+
{"col": 0, "length": 7, "line": 0, "modifiers": 0, "type": 3},
22+
{"col": 14, "length": 3, "line": 0, "modifiers": 0, "type": 0},
23+
},
1924
"vars": {
2025
"comprehensions": {
2126
"declaration": set(),

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

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,37 @@
66
# - input.params: schema.regal.lsp.semantictokens
77
package regal.lsp.semantictokens.vars.packages
88

9+
import data.regal.util
10+
911
# METADATA
1012
# description: Get the module from workspace
1113
module := data.workspace.parsed[input.params.textDocument.uri]
1214

1315
# METADATA
14-
# description: Extract package tokens - return full package path
15-
result contains regal.last(module.package.path)
16+
# description: Package keyword
17+
result contains token if {
18+
tloc := util.to_location_object(module.package.location)
19+
20+
token := {
21+
"line": tloc.row - 1,
22+
"col": tloc.col - 1,
23+
"length": tloc.end.col - tloc.col,
24+
"type": 3,
25+
"modifiers": 0,
26+
}
27+
}
28+
29+
# METADATA
30+
# description: Last package path term as token
31+
result contains token if {
32+
term := regal.last(module.package.path)
33+
tloc := util.to_location_object(term.location)
34+
35+
token := {
36+
"line": tloc.row - 1,
37+
"col": tloc.col - 1,
38+
"length": tloc.end.col - tloc.col,
39+
"type": 0,
40+
"modifiers": 0,
41+
}
42+
}

‎bundle/regal/lsp/semantictokens/vars/packages/packages_test.rego‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,15 @@ package regal.lsp.semantictokens.vars.packages_test
33
import data.regal.lsp.semantictokens.vars.packages
44

55
test_packages if {
6-
policy := `package regal.woo`
7-
tokens := packages.result with input as {"params": {"textDocument": {"uri": "file://p.rego"}}}
8-
with data.workspace.parsed["file://p.rego"] as regal.parse_module("p.rego", policy)
6+
policy := "package regal.woo\n"
7+
module := regal.parse_module("policy.rego", policy)
98

10-
{"location": "1:15:1:18", "type": "string", "value": "woo"} in tokens
9+
tokens := packages.result with data.workspace.parsed["file:///p.rego"] as module
10+
with input.params.textDocument.uri as "file:///p.rego"
11+
with input.regal.file.lines as split(policy, "\n")
12+
13+
tokens == {
14+
{"col": 0, "length": 7, "line": 0, "modifiers": 0, "type": 3},
15+
{"col": 14, "length": 3, "line": 0, "modifiers": 0, "type": 0},
16+
}
1117
}

‎internal/lsp/semantictokens/semantictokens.go‎

Lines changed: 31 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,36 @@ import (
1212
"github.com/open-policy-agent/regal/internal/lsp/types"
1313
)
1414

15+
type (
16+
TokenType = uint32
17+
TokenModifier = uint32
18+
)
19+
1520
const (
16-
TokenTypePackage = 0
17-
TokenTypeVariable = 1
18-
TokenTypeImport = 2
21+
Namespace TokenType = iota
22+
Variable
23+
Import
24+
Keyword
1925
)
2026

2127
const (
22-
ModifierDeclaration = 1 << 0
23-
ModifierReference = 1 << 1
28+
ModifierDeclaration TokenModifier = 1 << iota
29+
ModifierReference
2430
)
2531

32+
var Legend = types.SemanticTokensLegend{
33+
TokenTypes: []string{
34+
Namespace: "namespace",
35+
Variable: "variable",
36+
Import: "import",
37+
Keyword: "keyword",
38+
},
39+
TokenModifiers: []string{
40+
ModifierDeclaration: "declaration",
41+
ModifierReference: "reference",
42+
},
43+
}
44+
2645
type Token struct {
2746
Line uint32
2847
Col uint32
@@ -94,20 +113,15 @@ type VarsSection struct {
94113

95114
// Represents the structured result from the Rego query
96115
type SemanticTokensResult struct {
97-
PackageTokens []ASTLocation `json:"packages"`
116+
PackageTokens []Token `json:"packages"`
98117
ImportTokens []ASTLocation `json:"imports"`
99118
Vars VarsSection `json:"vars"`
100-
DebugInfo interface{} `json:"debug_info"`
119+
DebugInfo any `json:"debug_info"`
101120
}
102121

103122
func Full(ctx context.Context, result SemanticTokensResult) (*types.SemanticTokens, error) {
104-
tokens := make([]Token, 0)
105-
106-
packageTokens, err := processPackageTokens(result.PackageTokens)
107-
if err != nil {
108-
return nil, err
109-
}
110-
tokens = append(tokens, packageTokens...)
123+
tokens := make([]Token, 0, len(result.PackageTokens))
124+
tokens = append(tokens, result.PackageTokens...)
111125

112126
varTokens, err := processVariableTokens(result.Vars)
113127
if err != nil {
@@ -124,20 +138,6 @@ func Full(ctx context.Context, result SemanticTokensResult) (*types.SemanticToke
124138
return encodeTokens(tokens), nil
125139
}
126140

127-
func processPackageTokens(packageTokens []ASTLocation) ([]Token, error) {
128-
tokens := make([]Token, 0)
129-
130-
for _, pkgToken := range packageTokens {
131-
token, err := extractTokens(pkgToken, TokenTypePackage, 0)
132-
if err != nil {
133-
return nil, fmt.Errorf("failed to create package token: %w", err)
134-
}
135-
tokens = append(tokens, token)
136-
}
137-
138-
return tokens, nil
139-
}
140-
141141
func processVariableTokens(vars VarsSection) ([]Token, error) {
142142
tokens := make([]Token, 0)
143143

@@ -172,15 +172,15 @@ func processTokenCategory(category ArgTokenCategory, categoryName string) ([]Tok
172172
tokens := make([]Token, 0)
173173

174174
for _, declToken := range category.Declaration {
175-
token, err := extractTokens(declToken, TokenTypeVariable, ModifierDeclaration)
175+
token, err := extractTokens(declToken, Variable, ModifierDeclaration)
176176
if err != nil {
177177
return nil, fmt.Errorf("failed to create %s declaration token: %w", categoryName, err)
178178
}
179179
tokens = append(tokens, token)
180180
}
181181

182182
for _, refToken := range category.Reference {
183-
token, err := extractTokens(refToken, TokenTypeVariable, ModifierReference)
183+
token, err := extractTokens(refToken, Variable, ModifierReference)
184184
if err != nil {
185185
return nil, fmt.Errorf("failed to create %s reference token: %w", categoryName, err)
186186
}
@@ -194,7 +194,7 @@ func processImportTokens(importTokens []ASTLocation) ([]Token, error) {
194194
tokens := make([]Token, 0)
195195

196196
for _, importToken := range importTokens {
197-
token, err := extractTokens(importToken, TokenTypeImport, 0)
197+
token, err := extractTokens(importToken, Import, 0)
198198
if err != nil {
199199
return nil, fmt.Errorf("failed to create import token: %w", err)
200200
}

‎internal/lsp/semantictokens_test.go‎

Lines changed: 40 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,8 @@ func TestFull(t *testing.T) {
2626
expectedTokens []semanticTokenInstance
2727
}{
2828
"package only": {
29-
policy: `package regal.woo`,
30-
expectedTokens: []semanticTokenInstance{
31-
{DeltaLine: 0, DeltaCol: 14, Length: 3, Type: 0, Modifier: 0},
32-
},
29+
policy: `package regal.woo`,
30+
expectedTokens: []semanticTokenInstance{{Length: 7, Type: 3}, {DeltaCol: 14, Length: 3}},
3331
},
3432
"variable declarations": {
3533
policy: `package regal.woo
@@ -39,21 +37,23 @@ test_function(param1, param2) := result if {
3937
}
4038
`,
4139
expectedTokens: []semanticTokenInstance{
42-
{DeltaLine: 0, DeltaCol: 14, Length: 3, Type: 0, Modifier: 0},
40+
{Length: 7, Type: 3},
41+
{DeltaLine: 0, DeltaCol: 14, Length: 3},
4342
{DeltaLine: 2, DeltaCol: 14, Length: 6, Type: 1, Modifier: 1},
4443
{DeltaLine: 0, DeltaCol: 8, Length: 6, Type: 1, Modifier: 1},
4544
},
4645
},
4746
"variable references": {
4847
policy: `package regal.woo
49-
48+
5049
test_function(param1) := result if {
5150
calc3 := 1
5251
calc3 == param1
5352
}
5453
`,
5554
expectedTokens: []semanticTokenInstance{
56-
{DeltaLine: 0, DeltaCol: 14, Length: 3, Type: 0, Modifier: 0},
55+
{Length: 7, Type: 3},
56+
{DeltaLine: 0, DeltaCol: 14, Length: 3},
5757
{DeltaLine: 2, DeltaCol: 14, Length: 6, Type: 1, Modifier: 1},
5858
{DeltaLine: 2, DeltaCol: 15, Length: 6, Type: 1, Modifier: 2},
5959
},
@@ -62,34 +62,36 @@ test_function(param1) := result if {
6262
policy: `package regal.woo
6363
6464
import data.regal.ast
65-
65+
6666
test_function(param1) := result if {
6767
calc3 := 1
6868
calc3 == param1
6969
ast.is_constant
7070
}
7171
`,
7272
expectedTokens: []semanticTokenInstance{
73-
{DeltaLine: 0, DeltaCol: 14, Length: 3, Type: 0, Modifier: 0},
74-
{DeltaLine: 2, DeltaCol: 18, Length: 3, Type: 2, Modifier: 0},
73+
{Length: 7, Type: 3},
74+
{DeltaLine: 0, DeltaCol: 14, Length: 3},
75+
{DeltaLine: 2, DeltaCol: 18, Length: 3, Type: 2},
7576
{DeltaLine: 2, DeltaCol: 14, Length: 6, Type: 1, Modifier: 1},
7677
{DeltaLine: 2, DeltaCol: 15, Length: 6, Type: 1, Modifier: 2},
7778
},
7879
},
7980
"full policy with package, declarations and references": {
8081
policy: `package regal.woo
81-
82+
8283
test_function(param1) := result if {
8384
calc1 := param1 * 2
8485
calc2 := param2 + 10
8586
result := calc1 + calc2
86-
87+
8788
calc3 := 1
8889
calc3 == param1
8990
}
9091
`,
9192
expectedTokens: []semanticTokenInstance{
92-
{DeltaLine: 0, DeltaCol: 14, Length: 3, Type: 0, Modifier: 0},
93+
{Length: 7, Type: 3},
94+
{DeltaLine: 0, DeltaCol: 14, Length: 3},
9395
{DeltaLine: 2, DeltaCol: 14, Length: 6, Type: 1, Modifier: 1},
9496
{DeltaLine: 1, DeltaCol: 12, Length: 6, Type: 1, Modifier: 2},
9597
{DeltaLine: 5, DeltaCol: 15, Length: 6, Type: 1, Modifier: 2},
@@ -98,22 +100,23 @@ test_function(param1) := result if {
98100
"comprehensions": {
99101
policy: `package regal.woo
100102
101-
array_comprehensions := [x |
102-
some i, x in [1, 2, 3]
103-
i == 2
103+
array_comprehensions := [x |
104+
some i, x in [1, 2, 3]
105+
i == 2
104106
]
105107
106-
set_comprehensions := {x |
107-
some i, x in [1, 2, 3]
108-
i == 2
108+
set_comprehensions := {x |
109+
some i, x in [1, 2, 3]
110+
i == 2
109111
}
110112
111-
object_comprehensions := {k: v |
112-
some k, v in [1, 2, 3]
113-
v == 2
113+
object_comprehensions := {k: v |
114+
some k, v in [1, 2, 3]
115+
v == 2
114116
}`,
115117
expectedTokens: []semanticTokenInstance{
116-
{DeltaLine: 0, DeltaCol: 14, Length: 3, Type: 0, Modifier: 0},
118+
{Length: 7, Type: 3},
119+
{DeltaLine: 0, DeltaCol: 14, Length: 3},
117120
{DeltaLine: 2, DeltaCol: 25, Length: 1, Type: 1, Modifier: 2},
118121
{DeltaLine: 1, DeltaCol: 9, Length: 1, Type: 1, Modifier: 1},
119122
{DeltaLine: 0, DeltaCol: 0, Length: 1, Type: 1, Modifier: 1},
@@ -139,19 +142,20 @@ object_comprehensions := {k: v |
139142
policy: `package regal.woo
140143
141144
every_two_vars_construct if {
142-
every k, v in input.object {
143-
is_string(k)
144-
v > 0
145+
every k, v in input.object {
146+
is_string(k)
147+
v > 0
145148
}
146149
}
147150
148151
every_one_var_construct if {
149-
every k in input.object {
150-
is_string(k)
152+
every k in input.object {
153+
is_string(k)
151154
}
152155
}`,
153156
expectedTokens: []semanticTokenInstance{
154-
{DeltaLine: 0, DeltaCol: 14, Length: 3, Type: 0, Modifier: 0},
157+
{Length: 7, Type: 3},
158+
{DeltaLine: 0, DeltaCol: 14, Length: 3},
155159
{DeltaLine: 3, DeltaCol: 10, Length: 1, Type: 1, Modifier: 1},
156160
{DeltaLine: 0, DeltaCol: 3, Length: 1, Type: 1, Modifier: 1},
157161
{DeltaLine: 1, DeltaCol: 18, Length: 1, Type: 1, Modifier: 2},
@@ -164,17 +168,18 @@ every_one_var_construct if {
164168
policy: `package regal.woo
165169
166170
some_two_vars_construct if {
167-
some i, item in input.array
168-
i < 10
169-
item > 0
171+
some i, item in input.array
172+
i < 10
173+
item > 0
170174
}
171175
172176
some_one_var_construct if {
173-
some i in input.array
174-
i < 10
177+
some i in input.array
178+
i < 10
175179
}`,
176180
expectedTokens: []semanticTokenInstance{
177-
{DeltaLine: 0, DeltaCol: 14, Length: 3, Type: 0, Modifier: 0},
181+
{Length: 7, Type: 3},
182+
{DeltaLine: 0, DeltaCol: 14, Length: 3},
178183
{DeltaLine: 3, DeltaCol: 9, Length: 1, Type: 1, Modifier: 1},
179184
{DeltaLine: 0, DeltaCol: 3, Length: 4, Type: 1, Modifier: 1},
180185
{DeltaLine: 1, DeltaCol: 4, Length: 1, Type: 1, Modifier: 2},

‎internal/lsp/server.go‎

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ import (
4040
"github.com/open-policy-agent/regal/internal/lsp/log"
4141
"github.com/open-policy-agent/regal/internal/lsp/rego"
4242
"github.com/open-policy-agent/regal/internal/lsp/rego/query"
43+
"github.com/open-policy-agent/regal/internal/lsp/semantictokens"
4344
"github.com/open-policy-agent/regal/internal/lsp/types"
4445
"github.com/open-policy-agent/regal/internal/lsp/uri"
4546
rparse "github.com/open-policy-agent/regal/internal/parse"
@@ -1841,18 +1842,8 @@ func (l *LanguageServer) handleInitialize(ctx context.Context, params types.Init
18411842
SelectionRangeProvider: true,
18421843
LinkedEditingRangeProvider: true,
18431844
SemanticTokensProvider: types.SemanticTokensOptions{
1844-
Legend: types.SemanticTokensLegend{
1845-
TokenTypes: []string{
1846-
"namespace",
1847-
"variable",
1848-
"namespace",
1849-
},
1850-
TokenModifiers: []string{
1851-
"declaration",
1852-
"reference",
1853-
},
1854-
},
1855-
Full: true,
1845+
Legend: semantictokens.Legend,
1846+
Full: true,
18561847
},
18571848
// 'Experimental' is LSP terminology, we are using these to be
18581849
// 'custom' additions that are ready for use, but not in the base

0 commit comments

Comments
 (0)