From 26ba7d19837ce364f1b8ca94e4166a7f884a0704 Mon Sep 17 00:00:00 2001
From: Tw <wei.tan@intel.com>
Date: Fri, 12 Nov 2021 09:57:02 +0800
Subject: [PATCH] Also support single location in
 Implementation/TypeDefinition/Definition response

Signed-off-by: Tw <wei.tan@intel.com>
---
 internal/lsp/protocol/compat.go   | 28 ++++++++++++++++++++++++++++
 internal/lsp/protocol/tsserver.go |  6 +++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/internal/lsp/protocol/compat.go b/internal/lsp/protocol/compat.go
index e9fe721..8e3e6e8 100644
--- a/internal/lsp/protocol/compat.go
+++ b/internal/lsp/protocol/compat.go
@@ -67,3 +67,31 @@ func ToCodeActionOptions(v map[string]interface{}) (*CodeActionOptions, error) {
 	}
 	return &opt, nil
 }
+
+// Locations is a type which represents the union of Location and []Location
+type Locations []Location
+
+func (ls *Locations) UnmarshalJSON(data []byte) error {
+	d := strings.TrimSpace(string(data))
+	if len(d) == 0 && strings.EqualFold(d, "null") {
+		return nil
+	}
+
+	if d[0] == '[' {
+		var locations []Location
+		err := json.Unmarshal(data, &locations)
+		if err != nil {
+			return err
+		}
+		*ls = locations
+	} else {
+		var location Location
+		err := json.Unmarshal(data, &location)
+		if err != nil {
+			return err
+		}
+		*ls = append(*ls, location)
+	}
+
+	return nil
+}
diff --git a/internal/lsp/protocol/tsserver.go b/internal/lsp/protocol/tsserver.go
index 451278d..e5b4893 100644
--- a/internal/lsp/protocol/tsserver.go
+++ b/internal/lsp/protocol/tsserver.go
@@ -583,7 +583,7 @@ func (s *serverDispatcher) LogTraceNotification(ctx context.Context, params *Log
 	return s.Conn.Notify(ctx, "$/logTraceNotification", params)
 }
 func (s *serverDispatcher) Implementation(ctx context.Context, params *ImplementationParams) ([]Location, error) {
-	var result []Location
+	var result Locations
 	if err := s.Conn.Call(ctx, "textDocument/implementation", params, &result); err != nil {
 		return nil, err
 	}
@@ -591,7 +591,7 @@ func (s *serverDispatcher) Implementation(ctx context.Context, params *Implement
 }
 
 func (s *serverDispatcher) TypeDefinition(ctx context.Context, params *TypeDefinitionParams) ([]Location, error) {
-	var result []Location
+	var result Locations
 	if err := s.Conn.Call(ctx, "textDocument/typeDefinition", params, &result); err != nil {
 		return nil, err
 	}
@@ -691,7 +691,7 @@ func (s *serverDispatcher) SignatureHelp(ctx context.Context, params *SignatureH
 }
 
 func (s *serverDispatcher) Definition(ctx context.Context, params *DefinitionParams) ([]Location, error) {
-	var result []Location
+	var result Locations
 	if err := s.Conn.Call(ctx, "textDocument/definition", params, &result); err != nil {
 		return nil, err
 	}