internal/lsp: dynamically register semantic tokens

This CL switches from automatically registering the semantic tokens
capability to dynamic registration. This allows us to turn it on and
off as the option is switched--otherwise it defaults on always.

To achieve this, we also have to set session options on
didChangeConfiguration. It turns out that the passed-in "changed"
parameter can be null, which is why we always refetch the workspace
configuration.

Fixes golang/go#41963

Change-Id: I58d742577ce7f3da67db32011ba21bd9813eb203
Reviewed-on: https://go-review.googlesource.com/c/tools/+/263525
Trust: Rebecca Stambler <rstambler@golang.org>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
Reviewed-by: Heschi Kreinick <heschi@google.com>
diff --git a/internal/lsp/cache/session.go b/internal/lsp/cache/session.go
index 00fc455..9affdb5 100644
--- a/internal/lsp/cache/session.go
+++ b/internal/lsp/cache/session.go
@@ -27,7 +27,8 @@
 	cache *Cache
 	id    string
 
-	options *source.Options
+	optionsMu sync.Mutex
+	options   *source.Options
 
 	viewMu  sync.Mutex
 	views   []*View
@@ -118,10 +119,14 @@
 func (s *Session) String() string { return s.id }
 
 func (s *Session) Options() *source.Options {
+	s.optionsMu.Lock()
+	defer s.optionsMu.Unlock()
 	return s.options
 }
 
 func (s *Session) SetOptions(options *source.Options) {
+	s.optionsMu.Lock()
+	defer s.optionsMu.Unlock()
 	s.options = options
 }
 
diff --git a/internal/lsp/cmd/semantictokens.go b/internal/lsp/cmd/semantictokens.go
index 93f2cdb..0bcb487 100644
--- a/internal/lsp/cmd/semantictokens.go
+++ b/internal/lsp/cmd/semantictokens.go
@@ -114,8 +114,7 @@
 	tok := fset.File(f.Pos())
 	if tok == nil {
 		// can't happen; just parsed this file
-		log.Printf("tok is nil!")
-		return fmt.Errorf("can't find %s in fset!", args[0])
+		return fmt.Errorf("can't find %s in fset", args[0])
 	}
 	tc := span.NewContentConverter(args[0], buf)
 	colmap = &protocol.ColumnMapper{
diff --git a/internal/lsp/general.go b/internal/lsp/general.go
index 11f2f3e..a6de385 100644
--- a/internal/lsp/general.go
+++ b/internal/lsp/general.go
@@ -82,10 +82,14 @@
 		}
 	}
 
+	if st := params.Capabilities.TextDocument.SemanticTokens; st != nil {
+		rememberToks(st.TokenTypes, st.TokenModifiers)
+	}
+
 	goplsVer := &bytes.Buffer{}
 	debug.PrintVersionInfo(ctx, goplsVer, true, debug.PlainText)
 
-	ans := &protocol.InitializeResult{
+	return &protocol.InitializeResult{
 		Capabilities: protocol.ServerCapabilities{
 			CallHierarchyProvider: true,
 			CodeActionProvider:    codeActionProvider,
@@ -131,25 +135,7 @@
 			Name:    "gopls",
 			Version: goplsVer.String(),
 		},
-	}
-
-	st := params.Capabilities.TextDocument.SemanticTokens
-	if st != nil {
-		tokTypes, tokModifiers := rememberToks(st.TokenTypes, st.TokenModifiers)
-		// check that st.TokenFormat is "relative"
-		v := &protocol.SemanticTokensOptions{
-			Legend: protocol.SemanticTokensLegend{
-				// TODO(pjw): trim these to what we use (and an unused one
-				// at position 0 of TokTypes, to catch typos)
-				TokenTypes:     tokTypes,
-				TokenModifiers: tokModifiers,
-			},
-			Range: true,
-			Full:  true,
-		}
-		ans.Capabilities.SemanticTokensProvider = v
-	}
-	return ans, nil
+	}, nil
 }
 
 func (s *Server) initialized(ctx context.Context, params *protocol.InitializedParams) error {
@@ -175,17 +161,22 @@
 	s.pendingFolders = nil
 
 	if options.ConfigurationSupported && options.DynamicConfigurationSupported {
-		if err := s.client.RegisterCapability(ctx, &protocol.RegistrationParams{
-			Registrations: []protocol.Registration{
-				{
-					ID:     "workspace/didChangeConfiguration",
-					Method: "workspace/didChangeConfiguration",
-				},
-				{
-					ID:     "workspace/didChangeWorkspaceFolders",
-					Method: "workspace/didChangeWorkspaceFolders",
-				},
+		registrations := []protocol.Registration{
+			{
+				ID:     "workspace/didChangeConfiguration",
+				Method: "workspace/didChangeConfiguration",
 			},
+			{
+				ID:     "workspace/didChangeWorkspaceFolders",
+				Method: "workspace/didChangeWorkspaceFolders",
+			},
+		}
+		if options.SemanticTokens {
+			registrations = append(registrations, semanticTokenRegistrations()...)
+
+		}
+		if err := s.client.RegisterCapability(ctx, &protocol.RegistrationParams{
+			Registrations: registrations,
 		}); err != nil {
 			return err
 		}
diff --git a/internal/lsp/semantic.go b/internal/lsp/semantic.go
index f883a67..b0f42d6 100644
--- a/internal/lsp/semantic.go
+++ b/internal/lsp/semantic.go
@@ -595,7 +595,7 @@
 }
 
 // save what the client sent
-func rememberToks(toks []string, mods []string) ([]string, []string) {
+func rememberToks(toks []string, mods []string) {
 	SemanticMemo = &SemMemo{
 		tokTypes: toks,
 		tokMods:  mods,
@@ -610,7 +610,6 @@
 	}
 	// we could have pruned or rearranged them.
 	// But then change the list in cmd.go too
-	return SemanticMemo.tokTypes, SemanticMemo.tokMods
 }
 
 // SemanticTypes to use in case there is no client, as in the command line, or tests
diff --git a/internal/lsp/workspace.go b/internal/lsp/workspace.go
index 72123a2..beced5a 100644
--- a/internal/lsp/workspace.go
+++ b/internal/lsp/workspace.go
@@ -41,8 +41,16 @@
 	return snapshot, release, err
 }
 
-func (s *Server) didChangeConfiguration(ctx context.Context, changed interface{}) error {
-	// go through all the views getting the config
+func (s *Server) didChangeConfiguration(ctx context.Context, _ *protocol.DidChangeConfigurationParams) error {
+	// Apply any changes to the session-level settings.
+	options := s.session.Options().Clone()
+	semanticTokensRegistered := options.SemanticTokens
+	if err := s.fetchConfig(ctx, "", "", options); err != nil {
+		return err
+	}
+	s.session.SetOptions(options)
+
+	// Go through each view, getting and updating its configuration.
 	for _, view := range s.session.Views() {
 		options := s.session.Options().Clone()
 		if err := s.fetchConfig(ctx, view.Name(), view.Folder(), options); err != nil {
@@ -58,5 +66,50 @@
 			s.diagnoseDetached(snapshot)
 		}()
 	}
+
+	// Update any session-specific registrations or unregistrations.
+	if !semanticTokensRegistered && options.SemanticTokens {
+		if err := s.client.RegisterCapability(ctx, &protocol.RegistrationParams{
+			Registrations: semanticTokenRegistrations(),
+		}); err != nil {
+			return err
+		}
+	} else if semanticTokensRegistered && !options.SemanticTokens {
+		var unregistrations []protocol.Unregistration
+		for _, r := range semanticTokenRegistrations() {
+			unregistrations = append(unregistrations, protocol.Unregistration{
+				ID:     r.ID,
+				Method: r.Method,
+			})
+		}
+		if err := s.client.UnregisterCapability(ctx, &protocol.UnregistrationParams{
+			Unregisterations: unregistrations,
+		}); err != nil {
+			return err
+		}
+	}
 	return nil
 }
+
+func semanticTokenRegistrations() []protocol.Registration {
+	var registrations []protocol.Registration
+	for _, method := range []string{
+		"textDocument/semanticTokens/full",
+		"textDocument/semanticTokens/full/delta",
+		"textDocument/semanticTokens/range",
+	} {
+		registrations = append(registrations, protocol.Registration{
+			ID:     method,
+			Method: method,
+			RegisterOptions: &protocol.SemanticTokensOptions{
+				Legend: protocol.SemanticTokensLegend{
+					// TODO(pjw): trim these to what we use (and an unused one
+					// at position 0 of TokTypes, to catch typos)
+					TokenTypes:     SemanticMemo.tokTypes,
+					TokenModifiers: SemanticMemo.tokMods,
+				},
+			},
+		})
+	}
+	return registrations
+}