Add edition 2024 support to LSP

This adds edition 2024 support to the LSP, notably in completion,
semantic tokens, and hover.

Intended to be a fast-follow to #4407.

Author: stefanvanburen

private/buf/buflsp/builtin.go

@@ -109,4 +109,15 @@ var builtinDocs = map[string][]string{
 	"map": {
 		"A set of distinct keys, each of which is associated with a value.",
 	},
+
+	"export": {
+		"Marks a definition as exported, making it visible outside the current file.",
+		"",
+		"Available in edition 2024 and later.",
+	},
+	"local": {
+		"Marks a definition as local, restricting its visibility to the current file.",
+		"",
+		"Available in edition 2024 and later.",
+	},
 }

private/buf/buflsp/completion.go

@@ -282,6 +282,7 @@ func completionItemsForDef(ctx context.Context, file *file, declPath []ast.DeclA
 	hasStart := false // Start is a newline or open parenthesis for the start of a definition
 	hasTypeModifier := false
 	hasDeclaration := false
+	hasVisibilityModifier := false
 	typeSpan := extractAroundOffset(file, offset,
 		func(tok token.Token) bool {
 			if isTokenTypeDelimiter(tok) {
@@ -296,6 +297,8 @@ func completionItemsForDef(ctx context.Context, file *file, declPath []ast.DeclA
 					hasDeclaration = hasDeclaration || isDeclaration
 					_, isFieldModifier := typeModifierSet[tok.Keyword()]
 					hasTypeModifier = hasTypeModifier || isFieldModifier
+					_, isVisibilityMod := visibilityModifierSet[tok.Keyword()]
+					hasVisibilityModifier = hasVisibilityModifier || isVisibilityMod
 				}
 			}
 			if isTokenSpace(tok) {
@@ -340,6 +343,7 @@ func completionItemsForDef(ctx context.Context, file *file, declPath []ast.DeclA
 		slog.Bool("has_start", hasStart),
 		slog.Bool("has_field_modifier", hasTypeModifier),
 		slog.Bool("has_declaration", hasDeclaration),
+		slog.Bool("has_visibility_modifier", hasVisibilityModifier),
 		slog.Bool("inside_map_type", insideMapType),
 		slog.Bool("is_map_key_position", isMapKeyPosition),
 	)
@@ -400,13 +404,29 @@ func completionItemsForDef(ctx context.Context, file *file, declPath []ast.DeclA
 		return completionItemsForOptions(ctx, file, parentDef, def, offset)
 	}
 
-	// If at the top level, and on the first item, return top level keywords.
+	// If at the top level, return top level keywords.
 	if parentDef.IsZero() {
-		showKeywords := beforeCount == 0
-		if showKeywords {
+		editions := isEditions(file)
+		switch {
+		case beforeCount == 0:
+			// At the start of a definition: show all top-level keywords, plus
+			// visibility modifiers in edition 2024+ files.
 			file.lsp.logger.DebugContext(ctx, "completion: definition returning top-level keywords")
+			kws := topLevelKeywords()
+			if editions {
+				kws = joinSequences(kws, visibilityModifierKeywords())
+			}
+			return slices.Collect(keywordToCompletionItem(
+				kws,
+				protocol.CompletionItemKindKeyword,
+				tokenSpan,
+				offset,
+			))
+		case editions && beforeCount == 1 && hasVisibilityModifier:
+			// After export/local, only type declaration keywords are valid.
+			file.lsp.logger.DebugContext(ctx, "completion: definition returning top-level type declaration keywords after visibility modifier")
 			return slices.Collect(keywordToCompletionItem(
-				topLevelKeywords(),
+				topLevelTypeDeclarationKeywords(),
 				protocol.CompletionItemKindKeyword,
 				tokenSpan,
 				offset,
@@ -422,13 +442,13 @@ func completionItemsForDef(ctx context.Context, file *file, declPath []ast.DeclA
 		// - Show keywords for the first values (but not when inside map key position)
 		// - Show types if no type declaration and at first, or second position with field modifier.
 		// - Always show types if cursor is inside map<...> angle brackets
+		editions := isEditions(file)
 		showKeywords := beforeCount == 0 && !(insideMapType && isMapKeyPosition)
 		showTypes := insideMapType || (!hasDeclaration && (beforeCount == 0 || (hasTypeModifier && beforeCount == 1)))
 		if showKeywords {
-			isProto2 := isProto2(file)
 			iters = append(iters,
 				keywordToCompletionItem(
-					messageLevelKeywords(isProto2),
+					messageLevelKeywords(isProto2(file)),
 					protocol.CompletionItemKindKeyword,
 					tokenSpan,
 					offset,
@@ -440,6 +460,22 @@ func completionItemsForDef(ctx context.Context, file *file, declPath []ast.DeclA
 					offset,
 				),
 			)
+			if editions {
+				iters = append(iters, keywordToCompletionItem(
+					visibilityModifierKeywords(),
+					protocol.CompletionItemKindKeyword,
+					tokenSpan,
+					offset,
+				))
+			}
+		} else if editions && beforeCount == 1 && hasVisibilityModifier {
+			// After export/local inside a message, only nested type declarations are valid.
+			iters = append(iters, keywordToCompletionItem(
+				messageLevelTypeDeclarationKeywords(),
+				protocol.CompletionItemKindKeyword,
+				tokenSpan,
+				offset,
+			))
 		}
 		if showTypes {
 			// When inside map angle brackets, use only the prefix of the tokenSpan for filtering
@@ -818,6 +854,15 @@ var typeModifierSet = func() map[keyword.Keyword]struct{} {
 	return m
 }()
 
+// visibilityModifierSet is the set of edition 2024+ visibility modifier keywords.
+var visibilityModifierSet = func() map[keyword.Keyword]struct{} {
+	m := make(map[keyword.Keyword]struct{})
+	for kw := range visibilityModifierKeywords() {
+		m[kw] = struct{}{}
+	}
+	return m
+}()
+
 // topLevelKeywords returns keywords for the top-level.
 func topLevelKeywords() iter.Seq[keyword.Keyword] {
 	return func(yield func(keyword.Keyword) bool) {
@@ -833,6 +878,33 @@ func topLevelKeywords() iter.Seq[keyword.Keyword] {
 	}
 }
 
+// topLevelTypeDeclarationKeywords returns the type declaration keywords that can
+// follow a visibility modifier (export/local) at the top level in edition 2024+.
+func topLevelTypeDeclarationKeywords() iter.Seq[keyword.Keyword] {
+	return func(yield func(keyword.Keyword) bool) {
+		_ = yield(keyword.Message) &&
+			yield(keyword.Enum) &&
+			yield(keyword.Service)
+	}
+}
+
+// messageLevelTypeDeclarationKeywords returns the type declaration keywords that can
+// follow a visibility modifier (export/local) inside a message in edition 2024+.
+func messageLevelTypeDeclarationKeywords() iter.Seq[keyword.Keyword] {
+	return func(yield func(keyword.Keyword) bool) {
+		_ = yield(keyword.Message) &&
+			yield(keyword.Enum)
+	}
+}
+
+// visibilityModifierKeywords returns the visibility modifier keywords for edition 2024+.
+func visibilityModifierKeywords() iter.Seq[keyword.Keyword] {
+	return func(yield func(keyword.Keyword) bool) {
+		_ = yield(keyword.Export) &&
+			yield(keyword.Local)
+	}
+}
+
 // messageLevelKeywords returns keywords for messages.
 func messageLevelKeywords(isProto2 bool) iter.Seq[keyword.Keyword] {
 	return func(yield func(keyword.Keyword) bool) {
@@ -1782,6 +1854,11 @@ func isProto2(file *file) bool {
 	return file.ir.Syntax() == syntax.Proto2
 }
 
+// isEditions returns true if the file uses editions syntax.
+func isEditions(file *file) bool {
+	return file.ir.Syntax().IsEdition()
+}
+
 // findTypeBySpan returns the IR Type that corresponds to the given AST span.
 // Returns a zero Type if no matching type is found.
 func findTypeBySpan(file *file, span source.Span) ir.Type {

private/buf/buflsp/completion_test.go

@@ -514,3 +514,94 @@ func TestCompletionOptions(t *testing.T) {
 		})
 	}
 }
+
+func TestCompletionEdition2024(t *testing.T) {
+	t.Parallel()
+
+	ctx := t.Context()
+
+	testProtoPath, err := filepath.Abs("testdata/completion/edition2024_test.proto")
+	require.NoError(t, err)
+
+	clientJSONConn, testURI := setupLSPServer(t, testProtoPath)
+
+	tests := []struct {
+		name                string
+		line                uint32
+		character           uint32
+		expectedContains    []string
+		expectedNotContains []string
+	}{
+		{
+			// At top level, "ex" prefix matches "export": visibility modifier is offered.
+			name:             "toplevel_export_keyword",
+			line:             5, // "ex"
+			character:        2, // After the "ex"
+			expectedContains: []string{"export"},
+		},
+		{
+			// At top level, "lo" prefix matches "local": visibility modifier is offered.
+			name:             "toplevel_local_keyword",
+			line:             8, // "lo"
+			character:        2, // After the "lo"
+			expectedContains: []string{"local"},
+		},
+		{
+			// After "export " at top level, "mess" prefix matches "message".
+			// Only type declaration keywords (message, enum, service) should be offered.
+			name:             "after_export_toplevel",
+			line:             11, // "export mess"
+			character:        11, // After the "mess" in "export mess"
+			expectedContains: []string{"message"},
+			// Non-type-declaration top-level keywords should not be offered.
+			expectedNotContains: []string{"edition", "import", "package", "option", "syntax"},
+		},
+		{
+			// Inside a message, "lo" prefix matches "local": visibility modifier is offered.
+			name:             "message_local_keyword",
+			line:             15, // "  lo" (inside ExportedMessage)
+			character:        4,  // After the "lo"
+			expectedContains: []string{"local"},
+		},
+		{
+			// After "local " inside a message, "mess" prefix matches "message".
+			// Only nested type declaration keywords (message, enum) should be offered.
+			name:             "after_local_in_message",
+			line:             18, // "  local mess" (inside ExportedMessage)
+			character:        12, // After the "mess" in "  local mess"
+			expectedContains: []string{"message"},
+			// Services cannot be nested; other keywords should not appear.
+			expectedNotContains: []string{"service", "option", "oneof", "repeated", "optional"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			var completionList *protocol.CompletionList
+			_, completionErr := clientJSONConn.Call(ctx, protocol.MethodTextDocumentCompletion, protocol.CompletionParams{
+				TextDocumentPositionParams: protocol.TextDocumentPositionParams{
+					TextDocument: protocol.TextDocumentIdentifier{
+						URI: testURI,
+					},
+					Position: protocol.Position{
+						Line:      tt.line,
+						Character: tt.character,
+					},
+				},
+			}, &completionList)
+			require.NoError(t, completionErr)
+			require.NotNil(t, completionList, "expected completion list to be non-nil")
+			labels := make([]string, 0, len(completionList.Items))
+			for _, item := range completionList.Items {
+				labels = append(labels, item.Label)
+			}
+			for _, expected := range tt.expectedContains {
+				assert.Contains(t, labels, expected, "expected completion list to contain %q", expected)
+			}
+			for _, notExpected := range tt.expectedNotContains {
+				assert.NotContains(t, labels, notExpected, "expected completion list to not contain %q", notExpected)
+			}
+		})
+	}
+}

private/buf/buflsp/file.go

@@ -41,6 +41,7 @@ import (
 	"github.com/bufbuild/protocompile/experimental/report"
 	"github.com/bufbuild/protocompile/experimental/seq"
 	"github.com/bufbuild/protocompile/experimental/source"
+	"github.com/bufbuild/protocompile/experimental/token/keyword"
 	"go.lsp.dev/protocol"
 )
 
@@ -537,6 +538,7 @@ func (f *file) irToSymbols(irSymbol ir.Symbol) ([]*symbol, []*symbol) {
 		}
 		msg.def = msg
 		resolved = append(resolved, msg)
+		resolved = append(resolved, f.visibilityPrefixSymbols(irSymbol, irSymbol.AsType().AST())...)
 		unresolved = append(unresolved, f.messageToSymbols(irSymbol.AsType().Options())...)
 	case ir.SymbolKindEnum:
 		enum := &symbol{
@@ -549,6 +551,7 @@ func (f *file) irToSymbols(irSymbol ir.Symbol) ([]*symbol, []*symbol) {
 		}
 		enum.def = enum
 		resolved = append(resolved, enum)
+		resolved = append(resolved, f.visibilityPrefixSymbols(irSymbol, irSymbol.AsType().AST())...)
 		unresolved = append(unresolved, f.messageToSymbols(irSymbol.AsType().Options())...)
 	case ir.SymbolKindEnumValue:
 		name := &symbol{
@@ -723,6 +726,7 @@ func (f *file) irToSymbols(irSymbol ir.Symbol) ([]*symbol, []*symbol) {
 		}
 		service.def = service
 		resolved = append(resolved, service)
+		resolved = append(resolved, f.visibilityPrefixSymbols(irSymbol, irSymbol.AsService().AST())...)
 		unresolved = append(unresolved, f.messageToSymbols(irSymbol.AsService().Options())...)
 	case ir.SymbolKindMethod:
 		method := &symbol{
@@ -837,6 +841,30 @@ func getKindForMapType(typeAST ast.TypeAny, mapField ir.Member, isKey bool) (kin
 	return &static{ast: mapField.AST()}, false
 }
 
+// visibilityPrefixSymbols returns keywordBuiltin symbols for any export/local visibility
+// prefix tokens on the given decl, to support hover documentation on those keywords.
+func (f *file) visibilityPrefixSymbols(irSymbol ir.Symbol, decl ast.DeclDef) []*symbol {
+	var syms []*symbol
+	for prefix := range decl.Prefixes() {
+		kw := prefix.Prefix()
+		if kw != keyword.Export && kw != keyword.Local {
+			continue
+		}
+		prefixTok := prefix.PrefixToken()
+		if prefixTok.IsZero() {
+			continue
+		}
+		kwSym := &symbol{
+			ir:   irSymbol,
+			file: f,
+			span: prefixTok.Span(),
+			kind: &keywordBuiltin{name: kw.String(), anchor: "symbol-visibility"},
+		}
+		syms = append(syms, kwSym)
+	}
+	return syms
+}
+
 // importToSymbol takes an [ir.Import] and returns a symbol for it.
 func (f *file) importToSymbol(imp ir.Import) *symbol {
 	return &symbol{

private/buf/buflsp/hover_test.go

@@ -181,6 +181,27 @@ func TestHover(t *testing.T) {
 			character:     8, // On "TestTopLevel"
 			expectNoHover: true,
 		},
+		{
+			name:             "hover_on_export_keyword",
+			protoFile:        "testdata/hover/edition2024.proto",
+			line:             5, // Line with "export message ExportedMessage {"
+			character:        1, // On "export"
+			expectedContains: "symbol-visibility",
+		},
+		{
+			name:             "hover_on_local_keyword",
+			protoFile:        "testdata/hover/edition2024.proto",
+			line:             10, // Line with "local message LocalMessage {"
+			character:        1,  // On "local"
+			expectedContains: "symbol-visibility",
+		},
+		{
+			name:             "hover_on_exported_message_name",
+			protoFile:        "testdata/hover/edition2024.proto",
+			line:             5,  // Line with "export message ExportedMessage {"
+			character:        16, // On "ExportedMessage"
+			expectedContains: "ExportedMessage is visible outside this file",
+		},
 	}
 
 	for _, tt := range tests {