From b629683c8bcd75c96b587e59e087a6ed3a3e9fc7 Mon Sep 17 00:00:00 2001
From: Hariom Verma <hariom18599@gmail.com>
Date: Sat, 6 Jan 2024 20:25:43 +0530
Subject: [PATCH] Definition support for symbols

Implement initial definition support, enabling users to navigate
to symbol definitions.
---
 internal/lsp/completion.go | 29 +++++++++----
 internal/lsp/definition.go | 84 ++++++++++++++++++++++++++++++++++++++
 internal/lsp/server.go     |  3 ++
 3 files changed, 107 insertions(+), 9 deletions(-)
 create mode 100644 internal/lsp/definition.go

diff --git a/internal/lsp/completion.go b/internal/lsp/completion.go
index 927f0cd..c6ff4a2 100644
--- a/internal/lsp/completion.go
+++ b/internal/lsp/completion.go
@@ -16,6 +16,7 @@ import (
 
 	"go.lsp.dev/jsonrpc2"
 	"go.lsp.dev/protocol"
+	"go.lsp.dev/uri"
 )
 
 type CompletionStore struct {
@@ -69,6 +70,8 @@ type Package struct {
 }
 
 type Symbol struct {
+	Position  token.Position
+	FileURI   uri.URI
 	Name      string
 	Doc       string
 	Signature string
@@ -172,11 +175,13 @@ func InitCompletionStore(dirs []string) *CompletionStore {
 
 func getSymbols(fname string) []*Symbol {
 	var symbols []*Symbol
-
-	bsrc, err := os.ReadFile(fname)
+	absPath, err := filepath.Abs(fname)
 	if err != nil {
-		// Ignore error and return empty symbol list
-		return symbols
+		return symbols // Ignore error and return empty symbol list
+	}
+	bsrc, err := os.ReadFile(absPath)
+	if err != nil {
+		return symbols // Ignore error and return empty symbol list
 	}
 	text := string(bsrc)
 
@@ -192,17 +197,19 @@ func getSymbols(fname string) []*Symbol {
 	ast.FileExports(file)
 
 	ast.Inspect(file, func(n ast.Node) bool {
-		var found *Symbol
+		var symbol *Symbol
 
 		switch n.(type) {
 		case *ast.FuncDecl:
-			found = function(n, text)
+			symbol = function(n, text)
 		case *ast.GenDecl:
-			found = declaration(n, text)
+			symbol = declaration(n, text)
 		}
 
-		if found != nil {
-			symbols = append(symbols, found)
+		if symbol != nil {
+			symbol.FileURI = getURI(absPath)
+			symbol.Position = fset.Position(n.Pos())
+			symbols = append(symbols, symbol)
 		}
 
 		return true
@@ -255,3 +262,7 @@ func typeName(t ast.TypeSpec) string {
 		return "type"
 	}
 }
+
+func getURI(absFilePath string) uri.URI {
+	return uri.URI("file://" + absFilePath)
+}
diff --git a/internal/lsp/definition.go b/internal/lsp/definition.go
new file mode 100644
index 0000000..4f97728
--- /dev/null
+++ b/internal/lsp/definition.go
@@ -0,0 +1,84 @@
+package lsp
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"log/slog"
+	"strings"
+
+	"go.lsp.dev/jsonrpc2"
+	"go.lsp.dev/protocol"
+)
+
+func (s *server) Definition(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
+	var params protocol.DefinitionParams
+	if err := json.Unmarshal(req.Params(), &params); err != nil {
+		sendParseError(ctx, reply, err)
+	}
+
+	uri := params.TextDocument.URI
+	file, ok := s.snapshot.Get(uri.Filename())
+	if !ok {
+		return reply(ctx, nil, errors.New("snapshot not found"))
+	}
+
+	offset := file.PositionToOffset(params.Position)
+	slog.Info("definition", "offset", offset)
+
+	pgf, err := file.ParseGno(ctx)
+	if err != nil {
+		return reply(ctx, nil, errors.New("cannot parse gno file"))
+	}
+
+	for _, spec := range pgf.File.Imports {
+		slog.Info("definition", "spec", spec.Path.Value, "pos", spec.Path.Pos(), "end", spec.Path.End())
+		if int(spec.Path.Pos()) <= offset && offset <= int(spec.Path.End()) {
+			// TODO: handle definition for imports
+			slog.Info("definition", "import", spec.Path.Value)
+			return reply(ctx, nil, nil)
+		}
+	}
+
+	token, err := file.TokenAt(params.Position)
+	if err != nil {
+		return reply(ctx, protocol.Hover{}, err)
+	}
+	text := strings.TrimSpace(token.Text)
+
+	// FIXME: Use the AST package to do this + get type of token.
+	//
+	// This is just a quick PoC to get something working.
+
+	// strings.Split(p.Body,
+	text = strings.Split(text, "(")[0]
+
+	text = strings.TrimSuffix(text, ",")
+	text = strings.TrimSuffix(text, ")")
+
+	// *mux.Request
+	text = strings.TrimPrefix(text, "*")
+
+	slog.Info("definition", "pkg", len(s.completionStore.pkgs))
+
+	parts := strings.Split(text, ".")
+	if len(parts) == 2 {
+		pkg := parts[0]
+		sym := parts[1]
+
+		slog.Info("definition", "pkg", pkg, "sym", sym)
+		symbol := s.completionStore.lookupSymbol(pkg, sym)
+		if symbol != nil {
+			slog.Info("definition", "URI", symbol.FileURI)
+			return reply(ctx, protocol.Location{
+				URI: symbol.FileURI,
+				Range: *posToRange(
+					symbol.Position.Line,
+					[]int{0, 0},
+				),
+			}, nil)
+		}
+	}
+
+	return reply(ctx, nil, nil)
+}
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 068409b..acbeafb 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -69,6 +69,8 @@ func (s *server) ServerHandler(ctx context.Context, reply jsonrpc2.Replier, req
 		return s.Hover(ctx, reply, req)
 	case "textDocument/completion":
 		return s.Completion(ctx, reply, req)
+	case "textDocument/definition":
+		return s.Definition(ctx, reply, req)
 	default:
 		return jsonrpc2.MethodNotFoundHandler(ctx, reply, req)
 	}
@@ -103,6 +105,7 @@ func (s *server) Initialize(ctx context.Context, reply jsonrpc2.Replier, req jso
 					"gnopls.version",
 				},
 			},
+			DefinitionProvider:         true,
 			DocumentFormattingProvider: true,
 		},
 	}, nil)