Skip to content

Commit

Permalink
Precompute refs in file at update (#859)
Browse files Browse the repository at this point in the history
Signed-off-by: Charlie Egan <[email protected]>
  • Loading branch information
charlieegan3 committed Jun 20, 2024
1 parent 82fd171 commit 72f63aa
Show file tree
Hide file tree
Showing 7 changed files with 176 additions and 65 deletions.
21 changes: 6 additions & 15 deletions bundle/regal/lsp/completion/providers/rulerefs/rulerefs.rego
Original file line number Diff line number Diff line change
Expand Up @@ -9,17 +9,8 @@ import data.regal.lsp.completion.location
ref_is_internal(ref) if contains(ref, "._")

workspace_rule_refs contains ref if {
some file, parsed in data.workspace.parsed

package_name := concat(".", [path.value |
some i, path in parsed["package"].path
])

some rule in parsed.rules

rule_ref := ast.ref_to_string(rule.head.ref)

ref := concat(".", [package_name, rule_ref])
some refs in data.workspace.defined_refs
some ref in refs
}

parsed_current_file := data.workspace.parsed[input.regal.file.uri]
Expand Down Expand Up @@ -98,10 +89,10 @@ grouped_refs[size] contains ref if {
size := count(indexof_n(ref, "."))
}

default defermine_ref_prefix(_) := ""
default determine_ref_prefix(_) := ""

defermine_ref_prefix(word) := word if {
not word in {":="}
determine_ref_prefix(word) := word if {
word != ":="
}

items := [item |
Expand All @@ -116,7 +107,7 @@ items := [item |

last_word := regal.last(regex.split(`\s+`, trim_space(line)))

prefix := defermine_ref_prefix(last_word)
prefix := determine_ref_prefix(last_word)

sorted_counts := sort(object.keys(grouped_refs))

Expand Down
35 changes: 30 additions & 5 deletions bundle/regal/lsp/completion/providers/rulerefs/rulerefs_test.rego
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package regal.lsp.completion.providers.rulerefs_test

import data.regal.ast
import data.regal.lsp.completion.providers.rulerefs
import rego.v1

Expand Down Expand Up @@ -38,6 +39,20 @@ parsed_modules[file_uri] := parsed_module if {
parsed_module := regal.parse_module(file_uri, contents)
}

defined_refs[file_uri] contains ref if {
some file_uri, parsed_module in parsed_modules

package_name := concat(".", [path.value |
some i, path in parsed_module["package"].path
])

some rule in parsed_module.rules

rule_ref := ast.ref_to_string(rule.head.ref)

ref := concat(".", [package_name, rule_ref])
}

test_rule_refs_no_word if {
current_file_contents := concat("", [workspace["current_file.rego"], `
another_local_rule := `])
Expand All @@ -54,7 +69,9 @@ another_local_rule := `])
}},
}}

items := rulerefs.items with input as regal_module with data.workspace.parsed as parsed_modules
items := rulerefs.items with input as regal_module
with data.workspace.parsed as parsed_modules
with data.workspace.defined_refs as defined_refs

labels := [item.label | some item in items]

Expand Down Expand Up @@ -85,7 +102,9 @@ another_local_rule := imp`])
}},
}}

items := rulerefs.items with input as regal_module with data.workspace.parsed as parsed_modules
items := rulerefs.items with input as regal_module
with data.workspace.parsed as parsed_modules
with data.workspace.defined_refs as defined_refs

labels := [item.label | some item in items]

Expand Down Expand Up @@ -116,7 +135,9 @@ a`])
}},
}}

items := rulerefs.items with input as regal_module with data.workspace.parsed as parsed_modules
items := rulerefs.items with input as regal_module
with data.workspace.parsed as parsed_modules
with data.workspace.defined_refs as defined_refs

count(items) == 0
}
Expand All @@ -140,7 +161,9 @@ local_rule if local`])
}},
}}

items := rulerefs.items with input as regal_module with data.workspace.parsed as parsed_modules
items := rulerefs.items with input as regal_module
with data.workspace.parsed as parsed_modules
with data.workspace.defined_refs as defined_refs

count(items) == 0
}
Expand Down Expand Up @@ -168,7 +191,9 @@ local_func("foo") := local_f`])
}},
}}

items := rulerefs.items with input as regal_module with data.workspace.parsed as parsed_modules
items := rulerefs.items with input as regal_module
with data.workspace.parsed as parsed_modules
with data.workspace.defined_refs as defined_refs

count(items) == 0
}
6 changes: 5 additions & 1 deletion internal/lsp/completions/providers/policy_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,10 @@ import data.example
"file:///file1.rego": file1,
"file:///file2.rego": file2,
},
"defined_refs": map[string][]string{
"file:///file1.rego": {"example.foo"},
"file:///file2.rego": {},
},
},
})

Expand Down Expand Up @@ -135,7 +139,7 @@ allow if {
labels = append(labels, item.Label)
}

expected := []string{"example.foo"}
expected := []string{"example", "example.foo"}
if !slices.Equal(expected, labels) {
t.Fatalf("expected %v, got %v", expected, labels)
}
Expand Down
33 changes: 15 additions & 18 deletions internal/lsp/lint.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"encoding/json"
"errors"
"fmt"
"os"
"strings"

"github.com/open-policy-agent/opa/ast"
Expand Down Expand Up @@ -38,37 +39,33 @@ func updateParse(ctx context.Context, cache *cache.Cache, store storage.Store, f

cache.SetModule(fileURI, module)

txn, err := store.NewTransaction(ctx, storage.WriteParams)
err := PutFileMod(ctx, store, fileURI, module)
if err != nil {
return false, fmt.Errorf("failed to create transaction: %w", err)
return false, fmt.Errorf("failed to update rego store with parsed module: %w", err)
}

var stErr *storage.Error
definedRefs := refs.DefinedInModule(module)

err = store.Write(ctx, txn, storage.ReplaceOp, storage.Path{"workspace", "parsed", fileURI}, module)
cache.SetFileRefs(fileURI, definedRefs)

if errors.As(err, &stErr) && stErr.Code == storage.NotFoundErr {
err = store.Write(ctx, txn, storage.AddOp, storage.Path{"workspace", "parsed", fileURI}, module)
if err != nil {
store.Abort(ctx, txn)
// TODO: consider how we use and generate these to avoid needing to have in the cache and the store
var ruleRefs []string

return false, fmt.Errorf("failed to init module in store: %w", err)
for _, ref := range definedRefs {
if ref.Kind == types.Package {
continue
}
}

if err != nil {
store.Abort(ctx, txn)

return false, fmt.Errorf("failed to replace module in store: %w", err)
ruleRefs = append(ruleRefs, ref.Label)
}

err = store.Commit(ctx, txn)
fmt.Fprintln(os.Stderr, ruleRefs)

err = PutFileRefs(ctx, store, fileURI, ruleRefs)
if err != nil {
return false, fmt.Errorf("failed to commit transaction: %w", err)
return false, fmt.Errorf("failed to update rego store with defined refs: %w", err)
}

cache.SetFileRefs(fileURI, refs.DefinedInModule(module))

refNames, err := refs.UsedInModule(ctx, module)
if err != nil {
return false, fmt.Errorf("failed to get used refs: %w", err)
Expand Down
24 changes: 2 additions & 22 deletions internal/lsp/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ import (
"github.com/open-policy-agent/opa/ast"
"github.com/open-policy-agent/opa/format"
"github.com/open-policy-agent/opa/storage"
"github.com/open-policy-agent/opa/storage/inmem"

"github.com/styrainc/regal/bundle"
rio "github.com/styrainc/regal/internal/io"
Expand Down Expand Up @@ -57,11 +56,7 @@ type LanguageServerOptions struct {

func NewLanguageServer(opts *LanguageServerOptions) *LanguageServer {
c := cache.NewCache()
store := inmem.NewFromObject(map[string]interface{}{
"workspace": map[string]interface{}{
"parsed": map[string]interface{}{},
},
})
store := NewRegalStore()

ls := &LanguageServer{
cache: c,
Expand Down Expand Up @@ -347,24 +342,9 @@ func (l *LanguageServer) StartConfigWorker(ctx context.Context) {
l.cache.SetIgnoredFileContents(k, contents)
}

txn, err := l.regoStore.NewTransaction(ctx, storage.WriteParams)
if err != nil {
l.logError(fmt.Errorf("failed to create transaction to remove mod from store: %w", err))

continue
}

err = l.regoStore.Write(ctx, txn, storage.RemoveOp, storage.Path{"workspace", "parsed", k}, nil)
err := RemoveFileMod(ctx, l.regoStore, k)
if err != nil {
l.regoStore.Abort(ctx, txn)
l.logError(fmt.Errorf("failed to remove mod from store: %w", err))

continue
}

err = l.regoStore.Commit(ctx, txn)
if err != nil {
l.logError(fmt.Errorf("failed to commit transaction to remove mod from store: %w", err))
}
}

Expand Down
114 changes: 114 additions & 0 deletions internal/lsp/store.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
package lsp

import (
"context"
"errors"
"fmt"

"github.com/open-policy-agent/opa/ast"
"github.com/open-policy-agent/opa/storage"
"github.com/open-policy-agent/opa/storage/inmem"
)

func NewRegalStore() storage.Store {
return inmem.NewFromObject(map[string]interface{}{
"workspace": map[string]interface{}{
"parsed": map[string]interface{}{},
"defined_refs": map[string][]string{},
},
})
}

func transact(ctx context.Context, store storage.Store, writeMode bool, op func(txn storage.Transaction) error) error {
var params storage.TransactionParams
if writeMode {
params = storage.WriteParams
}

txn, err := store.NewTransaction(ctx, params)
if err != nil {
return fmt.Errorf("failed to create transaction: %w", err)
}

success := false
defer func() {
if !success {
store.Abort(ctx, txn)
}
}()

if err := op(txn); err != nil {
return err
}

if err := store.Commit(ctx, txn); err != nil {
return fmt.Errorf("failed to commit transaction: %w", err)
}

success = true

return nil
}

func PutFileMod(ctx context.Context, store storage.Store, fileURI string, mod *ast.Module) error {
return transact(ctx, store, true, func(txn storage.Transaction) error {
var stErr *storage.Error

err := store.Write(ctx, txn, storage.ReplaceOp, storage.Path{"workspace", "parsed", fileURI}, mod)
if errors.As(err, &stErr) && stErr.Code == storage.NotFoundErr {
err = store.Write(ctx, txn, storage.AddOp, storage.Path{"workspace", "parsed", fileURI}, mod)
if err != nil {
return fmt.Errorf("failed to init module in store: %w", err)
}
}

if err != nil {
return fmt.Errorf("failed to replace module in store: %w", err)
}

return nil
})
}

func RemoveFileMod(ctx context.Context, store storage.Store, fileURI string) error {
return transact(ctx, store, true, func(txn storage.Transaction) error {
var stErr *storage.Error

_, err := store.Read(ctx, txn, storage.Path{"workspace", "parsed", fileURI})
if errors.As(err, &stErr) && stErr.Code == storage.NotFoundErr {
// nothing to do
return nil
}

if err != nil {
return fmt.Errorf("failed to read module from store: %w", err)
}

err = store.Write(ctx, txn, storage.RemoveOp, storage.Path{"workspace", "parsed", fileURI}, nil)
if err != nil {
return fmt.Errorf("failed to remove module from store: %w", err)
}

return nil
})
}

func PutFileRefs(ctx context.Context, store storage.Store, fileURI string, refs []string) error {
return transact(ctx, store, true, func(txn storage.Transaction) error {
var stErr *storage.Error

err := store.Write(ctx, txn, storage.ReplaceOp, storage.Path{"workspace", "defined_refs", fileURI}, refs)
if errors.As(err, &stErr) && stErr.Code == storage.NotFoundErr {
err = store.Write(ctx, txn, storage.AddOp, storage.Path{"workspace", "defined_refs", fileURI}, refs)
if err != nil {
return fmt.Errorf("failed to init refs in store: %w", err)
}
}

if err != nil {
return fmt.Errorf("failed to replace refs in store: %w", err)
}

return nil
})
}
8 changes: 4 additions & 4 deletions internal/lsp/types/internal.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,13 @@ import "github.com/open-policy-agent/opa/ast"
// Ref is designed to be used in completions and provides information
// relevant to the object with that operation in mind.
type Ref struct {
Kind RefKind
Kind RefKind `json:"kind"`
// Label is a identifier for the object. e.g. data.package.rule.
Label string
Label string `json:"label"`
// Detail is a small amount of additional information about the object.
Detail string
Detail string `json:"detail"`
// Description is a longer description of the object and uses Markdown formatting.
Description string
Description string `json:"description"`
}

// RefKind represents the kind of object that a Ref represents.
Expand Down

0 comments on commit 72f63aa

Please sign in to comment.