From 54d51dce5c4322871f55d2751a0cf6d6c51b04a8 Mon Sep 17 00:00:00 2001
From: Adam Snyder <armsnyder@gmail.com>
Date: Sun, 23 Jun 2024 19:17:39 -0700
Subject: [PATCH] fix: fix a bug where the internal document model was not
 updating after file changes

---
 internal/analysis/handler.go      |  2 +
 internal/analysis/handler_test.go | 96 +++++++++++++++++++++++++------
 2 files changed, 81 insertions(+), 17 deletions(-)

diff --git a/internal/analysis/handler.go b/internal/analysis/handler.go
index 078928b..5bf5668 100644
--- a/internal/analysis/handler.go
+++ b/internal/analysis/handler.go
@@ -82,6 +82,8 @@ func (h *Handler) HandleChange(params types.DidChangeTextDocumentParams) error {
 		}
 	}
 
+	f.document = yaml.Document{}
+
 	return nil
 }
 
diff --git a/internal/analysis/handler_test.go b/internal/analysis/handler_test.go
index 32e9f2b..723f388 100644
--- a/internal/analysis/handler_test.go
+++ b/internal/analysis/handler_test.go
@@ -147,6 +147,60 @@ bar:
 	}
 }
 
+func TestHandler_HandleChangeThenHandleDefinition(t *testing.T) {
+	var h Handler
+
+	// The file starts with only a definition.
+
+	if err := h.HandleOpen(types.DidOpenTextDocumentParams{
+		TextDocument: types.TextDocumentItem{
+			URI: "file:///foo.yaml",
+			Text: `bar:
+  baz:
+    type: object`,
+		},
+	}); err != nil {
+		t.Fatalf("HandleOpen: %v", err)
+	}
+
+	// Trigger a HandleDefinition call so tha the yaml is parsed once.
+
+	if _, err := h.HandleDefinition(types.DefinitionParams{
+		TextDocumentPositionParams: positionParams("file:///foo.yaml", "0:0"),
+	}); err != nil {
+		t.Fatalf("HandleDefinition: %v", err)
+	}
+
+	// Add the reference to the file.
+
+	if err := h.HandleChange(types.DidChangeTextDocumentParams{
+		TextDocument: types.TextDocumentIdentifier{URI: "file:///foo.yaml"},
+		ContentChanges: []types.TextDocumentContentChangeEvent{{
+			Text: `foo:
+  $ref: "#/bar/baz"
+`,
+			Range: toPtr(newRange("0:0-0:0")),
+		}},
+	}); err != nil {
+		t.Fatalf("HandleChange: %v", err)
+	}
+
+	// Now that the reference has been added, we should be able to find the
+	// definition.
+
+	got, err := h.HandleDefinition(types.DefinitionParams{
+		TextDocumentPositionParams: positionParams("file:///foo.yaml", "1:8"),
+	})
+	if err != nil {
+		t.Fatalf("HandleDefinition: %v", err)
+	}
+
+	want := locations("file:///foo.yaml", "3:2-3:5")
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("HandleDefinition() = %v, want %v", got, want)
+	}
+}
+
 func loadFile(uri, text string) HandlerSetupFunc {
 	return func(t *testing.T, h *Handler) {
 		if err := h.HandleOpen(types.DidOpenTextDocumentParams{
@@ -196,27 +250,35 @@ func positionParams(uri, position string) types.TextDocumentPositionParams {
 	}
 }
 
+func newRange(r string) types.Range {
+	match := regexp.MustCompile(`^(\d+):(\d+)-(\d+):(\d+)$`).FindStringSubmatch(r)
+	if match == nil {
+		panic("invalid range")
+	}
+
+	return types.Range{
+		Start: types.Position{
+			Line:      mustAtoi(match[1]),
+			Character: mustAtoi(match[2]),
+		},
+		End: types.Position{
+			Line:      mustAtoi(match[3]),
+			Character: mustAtoi(match[4]),
+		},
+	}
+}
+
+func toPtr[T any](v T) *T {
+	return &v
+}
+
 func locations(uri string, ranges ...string) []types.Location {
-	pat := regexp.MustCompile(`^(\d+):(\d+)-(\d+):(\d+)$`)
 	locs := make([]types.Location, len(ranges))
-	for i, rng := range ranges {
-		match := pat.FindStringSubmatch(rng)
-		if match == nil {
-			panic("invalid range")
-		}
 
+	for i, rng := range ranges {
 		locs[i] = types.Location{
-			URI: uri,
-			Range: types.Range{
-				Start: types.Position{
-					Line:      mustAtoi(match[1]),
-					Character: mustAtoi(match[2]),
-				},
-				End: types.Position{
-					Line:      mustAtoi(match[3]),
-					Character: mustAtoi(match[4]),
-				},
-			},
+			URI:   uri,
+			Range: newRange(rng),
 		}
 	}