lsp: Fix Code Lenses jank

internal/lsp/cache/cache.go

@@ -47,6 +47,10 @@ type Cache struct {
 	// to be used for hover hints.
 	keywordLocationsFile *concurrent.Map[string, map[uint][]types.KeywordLocation]
 
+	// when a file is successfully parsed, the number of lines in the file is stored
+	// here. This is used to gracefully fail when exiting unparsable files.
+	successfulParseLineCounts *concurrent.Map[string, int]
+
 	// fileRefs is a map of file URI to refs that are defined in that file. These are
 	// intended to be used for completions in other files.
 	// fileRefs is expected to be updated when a file is successfully parsed.
@@ -55,15 +59,16 @@ type Cache struct {
 
 func NewCache() *Cache {
 	return &Cache{
-		fileContents:           concurrent.MapOf(make(map[string]string)),
-		ignoredFileContents:    concurrent.MapOf(make(map[string]string)),
-		modules:                concurrent.MapOf(make(map[string]*ast.Module)),
-		aggregateData:          concurrent.MapOf(make(map[string][]report.Aggregate)),
-		diagnosticsFile:        concurrent.MapOf(make(map[string][]types.Diagnostic)),
-		diagnosticsParseErrors: concurrent.MapOf(make(map[string][]types.Diagnostic)),
-		builtinPositionsFile:   concurrent.MapOf(make(map[string]map[uint][]types.BuiltinPosition)),
-		keywordLocationsFile:   concurrent.MapOf(make(map[string]map[uint][]types.KeywordLocation)),
-		fileRefs:               concurrent.MapOf(make(map[string]map[string]types.Ref)),
+		fileContents:              concurrent.MapOf(make(map[string]string)),
+		ignoredFileContents:       concurrent.MapOf(make(map[string]string)),
+		modules:                   concurrent.MapOf(make(map[string]*ast.Module)),
+		aggregateData:             concurrent.MapOf(make(map[string][]report.Aggregate)),
+		diagnosticsFile:           concurrent.MapOf(make(map[string][]types.Diagnostic)),
+		diagnosticsParseErrors:    concurrent.MapOf(make(map[string][]types.Diagnostic)),
+		builtinPositionsFile:      concurrent.MapOf(make(map[string]map[uint][]types.BuiltinPosition)),
+		keywordLocationsFile:      concurrent.MapOf(make(map[string]map[uint][]types.KeywordLocation)),
+		fileRefs:                  concurrent.MapOf(make(map[string]map[string]types.Ref)),
+		successfulParseLineCounts: concurrent.MapOf(make(map[string]int)),
 	}
 }
 
@@ -166,6 +171,11 @@ func (c *Cache) Rename(oldKey, newKey string) {
 		c.fileRefs.Set(newKey, refs)
 		c.fileRefs.Delete(oldKey)
 	}
+
+	if lineCount, ok := c.successfulParseLineCounts.Get(oldKey); ok {
+		c.successfulParseLineCounts.Set(newKey, lineCount)
+		c.successfulParseLineCounts.Delete(oldKey)
+	}
 }
 
 // SetFileAggregates will only set aggregate data for the provided URI. Even if
@@ -287,6 +297,14 @@ func (c *Cache) GetAllFileRefs() map[string]map[string]types.Ref {
 	return c.fileRefs.Clone()
 }
 
+func (c *Cache) GetSuccessfulParseLineCount(fileURI string) (int, bool) {
+	return c.successfulParseLineCounts.Get(fileURI)
+}
+
+func (c *Cache) SetSuccessfulParseLineCount(fileURI string, count int) {
+	c.successfulParseLineCounts.Set(fileURI, count)
+}
+
 // Delete removes all cached data for a given URI. Ignored file contents are
 // also removed if found for a matching URI.
 func (c *Cache) Delete(fileURI string) {
@@ -299,6 +317,7 @@ func (c *Cache) Delete(fileURI string) {
 	c.builtinPositionsFile.Delete(fileURI)
 	c.keywordLocationsFile.Delete(fileURI)
 	c.fileRefs.Delete(fileURI)
+	c.successfulParseLineCounts.Delete(fileURI)
 }
 
 func UpdateCacheForURIFromDisk(cache *Cache, fileURI, path string) (bool, string, error) {

internal/lsp/lint.go

@@ -45,6 +45,7 @@ func updateParse(
 		// if the parse was ok, clear the parse errors
 		cache.SetParseErrors(fileURI, []types.Diagnostic{})
 		cache.SetModule(fileURI, module)
+		cache.SetSuccessfulParseLineCount(fileURI, len(lines))
 
 		if err := PutFileMod(ctx, store, fileURI, module); err != nil {
 			return false, fmt.Errorf("failed to update rego store with parsed module: %w", err)

internal/lsp/server.go

@@ -1660,9 +1660,27 @@ func (l *LanguageServer) handleTextDocumentInlayHint(params types.TextDocumentIn
 }
 
 func (l *LanguageServer) handleTextDocumentCodeLens(ctx context.Context, params types.CodeLensParams) (any, error) {
+	lastSuccessfullyParsedLineCount, everParsed := l.cache.GetSuccessfulParseLineCount(params.TextDocument.URI)
+
+	// if the file has always been unparsable, we can return early
+	// as there is no value to be gained from showing code lenses.
+	if !everParsed {
+		return noCodeLenses, nil
+	}
+
 	parseErrors, ok := l.cache.GetParseErrors(params.TextDocument.URI)
 	if ok && len(parseErrors) > 0 {
-		return noCodeLenses, nil
+		// if there are parse errors, but the line count is the same, then we
+		// still show them based on the last parsed module.
+		contents, ok := l.cache.GetFileContents(params.TextDocument.URI)
+		if !ok {
+			// we have no contents for an unknown reason
+			return noCodeLenses, nil
+		}
+
+		if len(strings.Split(contents, "\n")) != lastSuccessfullyParsedLineCount {
+			return noCodeLenses, nil
+		}
 	}
 
 	contents, module, ok := l.cache.GetContentAndModule(params.TextDocument.URI)