From 65dddcaae644f5739ec3d64662ba245689637ac7 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Sat, 19 Oct 2024 09:21:32 -0600
Subject: [PATCH] set import ast more precisely

---
 d2ir/compile.go     | 2 +-
 d2lsp/d2lsp_test.go | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/d2ir/compile.go b/d2ir/compile.go
index f5d5fa86b..a85f9ddb2 100644
--- a/d2ir/compile.go
+++ b/d2ir/compile.go
@@ -874,7 +874,7 @@ func (c *compiler) _compileField(f *Field, refctx *RefContext) {
 		if !ok {
 			return
 		}
-		n.(Importable).SetImportAST(refctx.Key)
+		n.(Importable).SetImportAST(refctx.Key.Value.Import)
 		switch n := n.(type) {
 		case *Field:
 			if n.Primary_ != nil {
diff --git a/d2lsp/d2lsp_test.go b/d2lsp/d2lsp_test.go
index bae41ed95..63c6b8654 100644
--- a/d2lsp/d2lsp_test.go
+++ b/d2lsp/d2lsp_test.go
@@ -99,6 +99,7 @@ okay
 	assert.Equal(t, "ok.d2", ranges[0].Path)
 	assert.Equal(t, 1, len(importRanges))
 	assert.Equal(t, 3, importRanges[0].Start.Line)
+	assert.Equal(t, 5, importRanges[0].Start.Column)
 
 	ranges, _, err = d2lsp.GetRefRanges("ok.d2", fs, nil, "hi")
 	assert.Success(t, err)