WIP

Author: charlieegan3

go.mod

@@ -22,6 +22,7 @@ require (
 	github.com/sourcegraph/jsonrpc2 v0.2.0
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/pflag v1.0.5
+	gopkg.in/yaml.v2 v2.4.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -89,6 +90,5 @@ require (
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240820151423-278611b39280 // indirect
 	google.golang.org/protobuf v1.34.2 // indirect
 	gopkg.in/warnings.v0 v0.1.2 // indirect
-	gopkg.in/yaml.v2 v2.4.0 // indirect
 	sigs.k8s.io/yaml v1.4.0 // indirect
 )

internal/lsp/cache/cache.go

@@ -197,33 +197,6 @@ func (c *Cache) SetAggregates(data map[string][]report.Aggregate) {
 	}
 }
 
-// GetFileComplimentAggregates returns all aggregate data other than for the
-// provided fileURIs. This is used when running file diagnostics while also
-// requiring the previous aggregate state to provide aggregate rule linting.
-func (c *Cache) GetFileComplimentAggregates(fileURIs ...string) map[string][]report.Aggregate {
-	c.aggregateDataMu.Lock()
-	defer c.aggregateDataMu.Unlock()
-
-	excludedFiles := make(map[string]struct{}, len(fileURIs))
-	for _, fileURI := range fileURIs {
-		excludedFiles[fileURI] = struct{}{}
-	}
-
-	allAggregates := make(map[string][]report.Aggregate)
-
-	for sourceFile, aggregates := range c.aggregateData {
-		if _, excluded := excludedFiles[sourceFile]; excluded {
-			continue
-		}
-
-		for _, aggregate := range aggregates {
-			allAggregates[aggregate.IndexKey()] = append(allAggregates[aggregate.IndexKey()], aggregate)
-		}
-	}
-
-	return allAggregates
-}
-
 // GetFileAggregates is used to get aggregate data for a given list of files.
 // This is only used in tests to validate the cache state.
 func (c *Cache) GetFileAggregates(fileURIs ...string) map[string][]report.Aggregate {
@@ -268,6 +241,28 @@ func (c *Cache) SetFileDiagnostics(fileURI string, diags []types.Diagnostic) {
 	c.diagnosticsFile[fileURI] = diags
 }
 
+// SetFileDiagnosticsForRules will perform a partial update of the diagnostics
+// for a file given a list of evaluated rules.
+func (c *Cache) SetFileDiagnosticsForRules(fileURI string, rules []string, diags []types.Diagnostic) {
+	c.diagnosticsFileMu.Lock()
+	defer c.diagnosticsFileMu.Unlock()
+
+	ruleKeys := make(map[string]struct{}, len(rules))
+	for _, rule := range rules {
+		ruleKeys[rule] = struct{}{}
+	}
+
+	preservedDiagnostics := make([]types.Diagnostic, 0)
+
+	for _, diag := range c.diagnosticsFile[fileURI] {
+		if _, ok := ruleKeys[diag.Code]; !ok {
+			preservedDiagnostics = append(preservedDiagnostics, diag)
+		}
+	}
+
+	c.diagnosticsFile[fileURI] = append(preservedDiagnostics, diags...)
+}
+
 func (c *Cache) ClearFileDiagnostics() {
 	c.diagnosticsFileMu.Lock()
 	defer c.diagnosticsFileMu.Unlock()

internal/lsp/cache/cache_test.go

@@ -4,6 +4,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/styrainc/regal/internal/lsp/types"
 	"github.com/styrainc/regal/pkg/report"
 )
 
@@ -89,12 +90,6 @@ func TestManageAggregates(t *testing.T) {
 		t.Fatalf("unexpected number of aggregates for file2.rego: %d", len(aggs2))
 	}
 
-	file1ComplimentAggs := c.GetFileComplimentAggregates("file1.rego")
-
-	if !reflect.DeepEqual(file1ComplimentAggs, aggs2) {
-		t.Fatalf("unexpected compliment aggregates for file1.rego, exp\n%v\ngot\n%v", aggs2, file1ComplimentAggs)
-	}
-
 	allAggs := c.GetFileAggregates()
 
 	if len(allAggs) != 2 {
@@ -126,3 +121,41 @@ func TestManageAggregates(t *testing.T) {
 		t.Fatalf("unexpected number of aggregates: %d", len(allAggs))
 	}
 }
+
+func TestPartialDiagnosticsUpdate(t *testing.T) {
+	t.Parallel()
+
+	c := NewCache()
+
+	diag1 := types.Diagnostic{Code: "code1"}
+	diag2 := types.Diagnostic{Code: "code2"}
+	diag3 := types.Diagnostic{Code: "code3"}
+
+	c.SetFileDiagnostics("foo.rego", []types.Diagnostic{
+		diag1, diag2,
+	})
+
+	foundDiags, ok := c.GetFileDiagnostics("foo.rego")
+	if !ok {
+		t.Fatalf("expected to get diags for foo.rego")
+	}
+
+	if !reflect.DeepEqual(foundDiags, []types.Diagnostic{diag1, diag2}) {
+		t.Fatalf("unexpected diagnostics: %v", foundDiags)
+	}
+
+	c.SetFileDiagnosticsForRules(
+		"foo.rego",
+		[]string{"code2", "code3"},
+		[]types.Diagnostic{diag3},
+	)
+
+	foundDiags, ok = c.GetFileDiagnostics("foo.rego")
+	if !ok {
+		t.Fatalf("expected to get diags for foo.rego")
+	}
+
+	if !reflect.DeepEqual(foundDiags, []types.Diagnostic{diag1, diag3}) {
+		t.Fatalf("unexpected diagnostics: %v", foundDiags)
+	}
+}

internal/lsp/lint.go

@@ -149,7 +149,8 @@ func updateFileDiagnostics(
 	cache *cache.Cache,
 	regalConfig *config.Config,
 	fileURI string,
-	workspaceRootURI string,
+	workspaceRootDir string,
+	updateDiagnosticsForRules []string,
 ) error {
 	module, ok := cache.GetModule(fileURI)
 	if !ok {
@@ -165,11 +166,12 @@ func updateFileDiagnostics(
 	input := rules.NewInput(map[string]string{fileURI: contents}, map[string]*ast.Module{fileURI: module})
 
 	regalInstance := linter.NewLinter().
-		WithAggregates(cache.GetFileComplimentAggregates(fileURI)).
-		WithAlwaysAggregate(true).
+		// needed to get the aggregateData for this file
+		WithCollectQuery(true).
+		// needed to get the aggregateData out so we can update the cache
 		WithExportAggregates(true).
 		WithInputModules(&input).
-		WithRootDir(workspaceRootURI)
+		WithRootDir(workspaceRootDir)
 
 	if regalConfig != nil {
 		regalInstance = regalInstance.WithUserConfig(*regalConfig)
@@ -180,7 +182,7 @@ func updateFileDiagnostics(
 		return fmt.Errorf("failed to lint: %w", err)
 	}
 
-	fileDiags := convertReportToDiagnostics(&rpt, workspaceRootURI)
+	fileDiags := convertReportToDiagnostics(&rpt, workspaceRootDir)
 
 	files := cache.GetAllFiles()
 
@@ -198,7 +200,7 @@ func updateFileDiagnostics(
 				fd = []types.Diagnostic{}
 			}
 
-			cache.SetFileDiagnostics(uri, fd)
+			cache.SetFileDiagnosticsForRules(uri, updateDiagnosticsForRules, fd)
 		}
 	}
 
@@ -211,32 +213,40 @@ func updateAllDiagnostics(
 	ctx context.Context,
 	cache *cache.Cache,
 	regalConfig *config.Config,
-	workspaceRootURI string,
+	workspaceRootDir string,
 	overwriteAggregates bool,
+	aggregatesReportOnly bool,
+	updateDiagnosticsForRules []string,
 ) error {
+	var err error
+
 	modules := cache.GetAllModules()
 	files := cache.GetAllFiles()
 
-	input := rules.NewInput(files, modules)
-
 	regalInstance := linter.NewLinter().
-		WithInputModules(&input).
-		WithRootDir(workspaceRootURI).
+		WithRootDir(workspaceRootDir).
 		// aggregates need only be exported if they're to be used to overwrite.
 		WithExportAggregates(overwriteAggregates)
 
 	if regalConfig != nil {
 		regalInstance = regalInstance.WithUserConfig(*regalConfig)
 	}
 
+	if aggregatesReportOnly {
+		regalInstance = regalInstance.
+			WithAggregates(cache.GetFileAggregates())
+	} else {
+		input := rules.NewInput(files, modules)
+		regalInstance = regalInstance.WithInputModules(&input)
+	}
+
 	rpt, err := regalInstance.Lint(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to lint: %w", err)
 	}
 
-	fileDiags := convertReportToDiagnostics(&rpt, workspaceRootURI)
+	fileDiags := convertReportToDiagnostics(&rpt, workspaceRootDir)
 
-	// Update diagnostics for all files
 	for uri := range files {
 		parseErrs, ok := cache.GetParseErrors(uri)
 		if ok && len(parseErrs) > 0 {
@@ -248,7 +258,14 @@ func updateAllDiagnostics(
 			fd = []types.Diagnostic{}
 		}
 
-		cache.SetFileDiagnostics(uri, fd)
+		// when only an aggregate report was run, then we must make sure to
+		// only update diagnostics from these rules. So the report is
+		// authoratative, but for those rules only.
+		if aggregatesReportOnly {
+			cache.SetFileDiagnosticsForRules(uri, updateDiagnosticsForRules, fd)
+		} else {
+			cache.SetFileDiagnostics(uri, fd)
+		}
 	}
 
 	if overwriteAggregates {

internal/lsp/race_off.go

@@ -0,0 +1,8 @@
+//go:build !race
+// +build !race
+
+package lsp
+
+func isRaceEnabled() bool {
+	return false
+}

internal/lsp/race_on.go

@@ -0,0 +1,8 @@
+//go:build race
+// +build race
+
+package lsp
+
+func isRaceEnabled() bool {
+	return true
+}