d2lsp: attach import ref

d2ir/compile.go

@@ -550,6 +550,7 @@ func (c *compiler) compileMap(dst *Map, ast, scopeAST *d2ast.Map) {
 				c.errorf(n.Import, "cannot spread import non map into map")
 				continue
 			}
+			impn.(Importable).SetImportAST(n.Import)
 
 			for _, gctx := range impn.Map().globs {
 				if !gctx.refctx.Key.HasTripleGlob() {
@@ -873,6 +874,7 @@ func (c *compiler) _compileField(f *Field, refctx *RefContext) {
 		if !ok {
 			return
 		}
+		n.(Importable).SetImportAST(refctx.Key)
 		switch n := n.(type) {
 		case *Field:
 			if n.Primary_ != nil {
@@ -1196,6 +1198,7 @@ func (c *compiler) compileArray(dst *Array, a *d2ast.Array, scopeAST *d2ast.Map)
 			if !ok {
 				continue
 			}
+			n.(Importable).SetImportAST(v)
 			switch n := n.(type) {
 			case *Field:
 				if v.Spread {

d2ir/d2ir.go

@@ -166,9 +166,10 @@ func (s *Scalar) Equal(n2 Node) bool {
 }
 
 type Map struct {
-	parent Node
-	Fields []*Field `json:"fields"`
-	Edges  []*Edge  `json:"edges"`
+	parent    Node
+	importAST d2ast.Node
+	Fields    []*Field `json:"fields"`
+	Edges     []*Edge  `json:"edges"`
 
 	globs []*globContext
 }
@@ -184,6 +185,20 @@ func (m *Map) initRoot() {
 	}
 }
 
+func (m *Map) ImportAST() d2ast.Node {
+	return m.importAST
+}
+
+func (m *Map) SetImportAST(node d2ast.Node) {
+	m.importAST = node
+	for _, f := range m.Fields {
+		f.SetImportAST(node)
+	}
+	for _, e := range m.Edges {
+		e.SetImportAST(node)
+	}
+}
+
 func (m *Map) Copy(newParent Node) Node {
 	tmp := *m
 	m = &tmp
@@ -290,9 +305,19 @@ func NodeBoardKind(n Node) BoardKind {
 	}
 }
 
+type Importable interface {
+	ImportAST() d2ast.Node
+	SetImportAST(d2ast.Node)
+}
+
+var _ Importable = &Edge{}
+var _ Importable = &Field{}
+var _ Importable = &Map{}
+
 type Field struct {
 	// *Map.
-	parent Node
+	parent    Node
+	importAST d2ast.Node
 
 	Name string `json:"name"`
 
@@ -304,6 +329,17 @@ type Field struct {
 	References []*FieldReference `json:"references,omitempty"`
 }
 
+func (f *Field) ImportAST() d2ast.Node {
+	return f.importAST
+}
+
+func (f *Field) SetImportAST(node d2ast.Node) {
+	f.importAST = node
+	if f.Map() != nil {
+		f.Map().SetImportAST(node)
+	}
+}
+
 func (f *Field) Copy(newParent Node) Node {
 	tmp := *f
 	f = &tmp
@@ -450,7 +486,8 @@ func (eid *EdgeID) resolve(m *Map) (_ *EdgeID, _ *Map, common []string, _ error)
 
 type Edge struct {
 	// *Map
-	parent Node
+	parent    Node
+	importAST d2ast.Node
 
 	ID *EdgeID `json:"edge_id"`
 
@@ -460,6 +497,17 @@ type Edge struct {
 	References []*EdgeReference `json:"references,omitempty"`
 }
 
+func (e *Edge) ImportAST() d2ast.Node {
+	return e.importAST
+}
+
+func (e *Edge) SetImportAST(node d2ast.Node) {
+	e.importAST = node
+	if e.Map() != nil {
+		e.Map().SetImportAST(node)
+	}
+}
+
 func (e *Edge) Copy(newParent Node) Node {
 	tmp := *e
 	e = &tmp

d2lsp/d2lsp.go

@@ -3,7 +3,6 @@ package d2lsp
 
 import (
 	"fmt"
-	"path/filepath"
 	"strings"
 
 	"oss.terrastruct.com/d2/d2ast"
@@ -12,26 +11,26 @@ import (
 	"oss.terrastruct.com/d2/lib/memfs"
 )
 
-func GetRefs(path string, fs map[string]string, boardPath []string, key string) (refs []d2ir.Reference, _ error) {
+func GetRefRanges(path string, fs map[string]string, boardPath []string, key string) (ranges []d2ast.Range, importRanges []d2ast.Range, _ error) {
 	m, err := getBoardMap(path, fs, boardPath)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	mk, err := d2parser.ParseMapKey(key)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	if mk.Key == nil && len(mk.Edges) == 0 {
-		return nil, fmt.Errorf(`"%s" is invalid`, key)
+		return nil, nil, fmt.Errorf(`"%s" is invalid`, key)
 	}
 
 	var f *d2ir.Field
 	if mk.Key != nil {
 		for _, p := range mk.Key.Path {
 			f = m.GetField(p.Unbox().ScalarString())
 			if f == nil {
-				return nil, nil
+				return nil, nil, nil
 			}
 			m = f.Map()
 		}
@@ -44,40 +43,25 @@ func GetRefs(path string, fs map[string]string, boardPath []string, key string)
 			edges = append(edges, m.GetEdges(eid, nil, nil)...)
 		}
 		if len(edges) == 0 {
-			return nil, nil
+			return nil, nil, nil
 		}
 		for _, edge := range edges {
 			for _, ref := range edge.References {
-				refs = append(refs, ref)
+				ranges = append(ranges, ref.AST().GetRange())
+			}
+			if edge.ImportAST() != nil {
+				importRanges = append(importRanges, edge.ImportAST().GetRange())
 			}
 		}
-		return refs, nil
 	} else {
 		for _, ref := range f.References {
-			refs = append(refs, ref)
+			ranges = append(ranges, ref.AST().GetRange())
 		}
-	}
-	return refs, nil
-}
-
-func GetImportRanges(path, file string, importPath string) (ranges []d2ast.Range, _ error) {
-	r := strings.NewReader(file)
-	ast, err := d2parser.Parse(path, r, nil)
-	if err != nil {
-		return nil, err
-	}
-
-	d2ast.Walk(ast, func(n d2ast.Node) bool {
-		switch t := n.(type) {
-		case *d2ast.Import:
-			if (filepath.Join(filepath.Dir(path), t.PathWithPre()) + ".d2") == importPath {
-				ranges = append(ranges, t.Range)
-			}
+		if f.ImportAST() != nil {
+			importRanges = append(importRanges, f.ImportAST().GetRange())
 		}
-		return true
-	})
-
-	return ranges, nil
+	}
+	return ranges, importRanges, nil
 }
 
 func getBoardMap(path string, fs map[string]string, boardPath []string) (*d2ir.Map, error) {