Regal Language Server (#532)

Signed-off-by: Charlie Egan <charlie@styra.com>

cmd/languageserver.go

@@ -0,0 +1,57 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/sourcegraph/jsonrpc2"
+	"github.com/spf13/cobra"
+
+	"github.com/styrainc/regal/internal/lsp"
+)
+
+func init() {
+	languageServerCommand := &cobra.Command{
+		Use:   "language-server",
+		Short: "Run the Regal Language Server",
+		Long:  `Start the Regal Language Server and listen on stdin/stdout for client editor messages.`,
+
+		RunE: wrapProfiling(func(args []string) error {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			opts := &lsp.LanguageServerOptions{
+				ErrorLog:       os.Stderr,
+				VerboseLogging: true,
+			}
+
+			ls := lsp.NewLanguageServer(opts)
+
+			conn := jsonrpc2.NewConn(
+				ctx,
+				jsonrpc2.NewBufferedStream(lsp.StdOutReadWriteCloser{}, jsonrpc2.VSCodeObjectCodec{}),
+				jsonrpc2.HandlerWithError(ls.Handle),
+			)
+
+			ls.SetConn(conn)
+			go ls.StartDiagnosticsWorker(ctx)
+
+			sigChan := make(chan os.Signal, 1)
+			signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+
+			select {
+			case <-conn.DisconnectNotify():
+				fmt.Fprint(os.Stderr, "Connection closed\n")
+			case sig := <-sigChan:
+				fmt.Fprint(os.Stderr, "signal: ", sig.String(), "\n")
+			}
+
+			return nil
+		}),
+	}
+
+	RootCommand.AddCommand(languageServerCommand)
+}

go.mod

@@ -21,6 +21,7 @@ require (
 	github.com/gdamore/encoding v1.0.0 // indirect
 	github.com/gdamore/tcell v1.4.0 // indirect
 	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
+	github.com/sourcegraph/jsonrpc2 v0.2.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 )
 

go.sum

@@ -71,6 +71,7 @@ github.com/google/uuid v1.5.0 h1:1p67kYwdtXjb0gL0BPiP1Av9wiZPo5A8z2cWkTZ+eyU=
 github.com/google/uuid v1.5.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
 github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
+github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.18.0 h1:RtRsiaGvWxcwd8y3BiRZxsylPT8hLWZ5SPcfI+3IDNk=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.18.0/go.mod h1:TzP6duP4Py2pHLVPPQp42aoYI92+PCrVotyR5e8Vqlk=
 github.com/ianlancetaylor/demangle v0.0.0-20210905161508-09a460cdf81d/go.mod h1:aYm2/VgdVmcIU8iMfdMvDMsRAQjcfZSKFby6HOFvi/w=
@@ -133,6 +134,8 @@ github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjR
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sirupsen/logrus v1.9.4-0.20230606125235-dd1b4c2e81af h1:Sp5TG9f7K39yfB+If0vjp97vuT74F72r8hfRpP8jLU0=
 github.com/sirupsen/logrus v1.9.4-0.20230606125235-dd1b4c2e81af/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/sourcegraph/jsonrpc2 v0.2.0 h1:KjN/dC4fP6aN9030MZCJs9WQbTOjWHhrtKVpzzSrr/U=
+github.com/sourcegraph/jsonrpc2 v0.2.0/go.mod h1:ZafdZgk/axhT1cvZAPOhw+95nz2I/Ra5qMlU4gTRwIo=
 github.com/spf13/cobra v1.8.0 h1:7aJaZx1B85qltLMc546zn58BxxfZdR/W22ej9CFoEf0=
 github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=

internal/lsp/cache.go

@@ -0,0 +1,224 @@
+package lsp
+
+import (
+	"fmt"
+	"net/url"
+	"os"
+	"sync"
+
+	"github.com/open-policy-agent/opa/ast"
+)
+
+// Cache is used to store: current file contents (which includes unsaved changes), the latest parsed modules, and
+// diagnostics for each file (including diagnostics gathered from linting files alongside other files).
+type Cache struct {
+	// fileContents is a map of file URI to raw file contents received from the client
+	fileContents   map[string]string
+	fileContentsMu sync.Mutex
+
+	// modules is a map of file URI to parsed AST modules from the latest file contents value
+	modules  map[string]*ast.Module
+	moduleMu sync.Mutex
+
+	// diagnosticsFile is a map of file URI to diagnostics for that file
+	diagnosticsFile   map[string][]Diagnostic
+	diagnosticsFileMu sync.Mutex
+
+	// diagnosticsAggregate is a map of file URI to aggregate diagnostics for that file
+	diagnosticsAggregate   map[string][]Diagnostic
+	diagnosticsAggregateMu sync.Mutex
+
+	// diagnosticsParseErrors is a map of file URI to parse errors for that file
+	diagnosticsParseErrors map[string][]Diagnostic
+	diagnosticsParseMu     sync.Mutex
+}
+
+func NewCache() *Cache {
+	return &Cache{
+		fileContents: make(map[string]string),
+		modules:      make(map[string]*ast.Module),
+
+		diagnosticsFile:        make(map[string][]Diagnostic),
+		diagnosticsAggregate:   make(map[string][]Diagnostic),
+		diagnosticsParseErrors: make(map[string][]Diagnostic),
+	}
+}
+
+func (c *Cache) GetAllDiagnosticsForURI(uri string) []Diagnostic {
+	parseDiags, ok := c.GetParseErrors(uri)
+	if ok && len(parseDiags) > 0 {
+		return parseDiags
+	}
+
+	allDiags := make([]Diagnostic, 0)
+
+	aggDiags, ok := c.GetAggregateDiagnostics(uri)
+	if ok {
+		allDiags = append(allDiags, aggDiags...)
+	}
+
+	fileDiags, ok := c.GetFileDiagnostics(uri)
+	if ok {
+		allDiags = append(allDiags, fileDiags...)
+	}
+
+	return allDiags
+}
+
+func (c *Cache) GetAllFiles() map[string]string {
+	c.fileContentsMu.Lock()
+	defer c.fileContentsMu.Unlock()
+
+	return c.fileContents
+}
+
+func (c *Cache) GetFileContents(uri string) (string, bool) {
+	c.fileContentsMu.Lock()
+	defer c.fileContentsMu.Unlock()
+
+	val, ok := c.fileContents[uri]
+
+	return val, ok
+}
+
+func (c *Cache) SetFileContents(uri string, content string) {
+	c.fileContentsMu.Lock()
+	defer c.fileContentsMu.Unlock()
+
+	c.fileContents[uri] = content
+}
+
+func (c *Cache) GetAllModules() map[string]*ast.Module {
+	c.moduleMu.Lock()
+	defer c.moduleMu.Unlock()
+
+	return c.modules
+}
+
+func (c *Cache) GetModule(uri string) (*ast.Module, bool) {
+	c.moduleMu.Lock()
+	defer c.moduleMu.Unlock()
+
+	val, ok := c.modules[uri]
+
+	return val, ok
+}
+
+func (c *Cache) SetModule(uri string, module *ast.Module) {
+	c.moduleMu.Lock()
+	defer c.moduleMu.Unlock()
+
+	c.modules[uri] = module
+}
+
+func (c *Cache) GetFileDiagnostics(uri string) ([]Diagnostic, bool) {
+	c.diagnosticsFileMu.Lock()
+	defer c.diagnosticsFileMu.Unlock()
+
+	val, ok := c.diagnosticsFile[uri]
+
+	return val, ok
+}
+
+func (c *Cache) SetFileDiagnostics(uri string, diags []Diagnostic) {
+	c.diagnosticsFileMu.Lock()
+	defer c.diagnosticsFileMu.Unlock()
+
+	c.diagnosticsFile[uri] = diags
+}
+
+func (c *Cache) ClearFileDiagnostics() {
+	c.diagnosticsFileMu.Lock()
+	defer c.diagnosticsFileMu.Unlock()
+
+	c.diagnosticsFile = make(map[string][]Diagnostic)
+}
+
+func (c *Cache) GetAggregateDiagnostics(uri string) ([]Diagnostic, bool) {
+	c.diagnosticsAggregateMu.Lock()
+	defer c.diagnosticsAggregateMu.Unlock()
+
+	val, ok := c.diagnosticsAggregate[uri]
+
+	return val, ok
+}
+
+func (c *Cache) SetAggregateDiagnostics(uri string, diags []Diagnostic) {
+	c.diagnosticsAggregateMu.Lock()
+	defer c.diagnosticsAggregateMu.Unlock()
+
+	c.diagnosticsAggregate[uri] = diags
+}
+
+func (c *Cache) ClearAggregateDiagnostics() {
+	c.diagnosticsAggregateMu.Lock()
+	defer c.diagnosticsAggregateMu.Unlock()
+
+	c.diagnosticsAggregate = make(map[string][]Diagnostic)
+}
+
+func (c *Cache) GetParseErrors(uri string) ([]Diagnostic, bool) {
+	c.diagnosticsParseMu.Lock()
+	defer c.diagnosticsParseMu.Unlock()
+
+	val, ok := c.diagnosticsParseErrors[uri]
+
+	return val, ok
+}
+
+func (c *Cache) SetParseErrors(uri string, diags []Diagnostic) {
+	c.diagnosticsParseMu.Lock()
+	defer c.diagnosticsParseMu.Unlock()
+
+	c.diagnosticsParseErrors[uri] = diags
+}
+
+// Delete removes all cached data for a given URI.
+func (c *Cache) Delete(uri string) {
+	c.fileContentsMu.Lock()
+	delete(c.fileContents, uri)
+	c.fileContentsMu.Unlock()
+
+	c.moduleMu.Lock()
+	delete(c.modules, uri)
+	c.moduleMu.Unlock()
+
+	c.diagnosticsFileMu.Lock()
+	delete(c.diagnosticsFile, uri)
+	c.diagnosticsFileMu.Unlock()
+
+	c.diagnosticsAggregateMu.Lock()
+	delete(c.diagnosticsAggregate, uri)
+	c.diagnosticsAggregateMu.Unlock()
+
+	c.diagnosticsParseMu.Lock()
+	delete(c.diagnosticsParseErrors, uri)
+	c.diagnosticsParseMu.Unlock()
+}
+
+func updateCacheForURIFromDisk(cache *Cache, uri string) (string, error) {
+	parsedURI, err := url.Parse(uri)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse URI: %w", err)
+	}
+
+	if parsedURI.Scheme != "file" {
+		return "", fmt.Errorf("only file:// URIs are supported, got %q", parsedURI.String())
+	}
+
+	content, err := os.ReadFile(parsedURI.Path)
+	if err != nil {
+		return "", fmt.Errorf("failed to read file: %w", err)
+	}
+
+	currentContent := string(content)
+
+	cachedContent, ok := cache.GetFileContents(uri)
+	if ok && cachedContent == currentContent {
+		return cachedContent, nil
+	}
+
+	cache.SetFileContents(uri, currentContent)
+
+	return currentContent, nil
+}