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
+	parent    Node
-	Fields []*Field `json:"fields"`
+	importAST d2ast.Node
-	Edges  []*Edge  `json:"edges"`
+	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,18 +11,18 @@ 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
@@ -31,7 +30,7 @@ func GetRefs(path string, fs map[string]string, boardPath []string, key string)
 		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())
+		}
+		if f.ImportAST() != nil {
+			importRanges = append(importRanges, f.ImportAST().GetRange())
 		}
 	}
-	return refs, nil
+	return ranges, importRanges, 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)
-			}
-		}
-		return true
-	})
-
-	return ranges, nil
 }
 
 func getBoardMap(path string, fs map[string]string, boardPath []string) (*d2ir.Map, error) {

d2lsp/d2lsp_test.go

@@ -7,7 +7,7 @@ import (
 	"oss.terrastruct.com/util-go/assert"
 )
 
-func TestGetFieldRefs(t *testing.T) {
+func TestGetFieldRanges(t *testing.T) {
 	script := `x
 x.a
 a.x
@@ -15,20 +15,20 @@ x -> y`
 	fs := map[string]string{
 		"index.d2": script,
 	}
-	refs, err := d2lsp.GetRefs("index.d2", fs, nil, "x")
+	ranges, _, err := d2lsp.GetRefRanges("index.d2", fs, nil, "x")
 	assert.Success(t, err)
-	assert.Equal(t, 3, len(refs))
+	assert.Equal(t, 3, len(ranges))
-	assert.Equal(t, 0, refs[0].AST().GetRange().Start.Line)
+	assert.Equal(t, 0, ranges[0].Start.Line)
-	assert.Equal(t, 1, refs[1].AST().GetRange().Start.Line)
+	assert.Equal(t, 1, ranges[1].Start.Line)
-	assert.Equal(t, 3, refs[2].AST().GetRange().Start.Line)
+	assert.Equal(t, 3, ranges[2].Start.Line)
 
-	refs, err = d2lsp.GetRefs("index.d2", fs, nil, "a.x")
+	ranges, _, err = d2lsp.GetRefRanges("index.d2", fs, nil, "a.x")
 	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 1, len(ranges))
-	assert.Equal(t, 2, refs[0].AST().GetRange().Start.Line)
+	assert.Equal(t, 2, ranges[0].Start.Line)
 }
 
-func TestGetEdgeRefs(t *testing.T) {
+func TestGetEdgeRanges(t *testing.T) {
 	script := `x
 x.a
 a.x
@@ -42,61 +42,74 @@ b: {
 	fs := map[string]string{
 		"index.d2": script,
 	}
-	refs, err := d2lsp.GetRefs("index.d2", fs, nil, "x -> y")
+	ranges, _, err := d2lsp.GetRefRanges("index.d2", fs, nil, "x -> y")
 	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 1, len(ranges))
-	assert.Equal(t, 3, refs[0].AST().GetRange().Start.Line)
+	assert.Equal(t, 3, ranges[0].Start.Line)
 
-	refs, err = d2lsp.GetRefs("index.d2", fs, nil, "y -> z")
+	ranges, _, err = d2lsp.GetRefRanges("index.d2", fs, nil, "y -> z")
 	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 1, len(ranges))
-	assert.Equal(t, 4, refs[0].AST().GetRange().Start.Line)
+	assert.Equal(t, 4, ranges[0].Start.Line)
 
-	refs, err = d2lsp.GetRefs("index.d2", fs, nil, "x -> z")
+	ranges, _, err = d2lsp.GetRefRanges("index.d2", fs, nil, "x -> z")
 	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 1, len(ranges))
-	assert.Equal(t, 5, refs[0].AST().GetRange().Start.Line)
+	assert.Equal(t, 5, ranges[0].Start.Line)
 
-	refs, err = d2lsp.GetRefs("index.d2", fs, nil, "a -> b")
+	ranges, _, err = d2lsp.GetRefRanges("index.d2", fs, nil, "a -> b")
 	assert.Success(t, err)
-	assert.Equal(t, 0, len(refs))
+	assert.Equal(t, 0, len(ranges))
 
-	refs, err = d2lsp.GetRefs("index.d2", fs, nil, "b.(x -> y)")
+	ranges, _, err = d2lsp.GetRefRanges("index.d2", fs, nil, "b.(x -> y)")
 	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 1, len(ranges))
-	assert.Equal(t, 7, refs[0].AST().GetRange().Start.Line)
+	assert.Equal(t, 7, ranges[0].Start.Line)
 }
 
-func TestGetRefsImported(t *testing.T) {
+func TestGetRangesImported(t *testing.T) {
 	fs := map[string]string{
 		"index.d2": `
 ...@ok
 hi
+hey: @ok
 `,
 		"ok.d2": `
+what
+lala
 okay
 `,
 	}
-	refs, err := d2lsp.GetRefs("index.d2", fs, nil, "hi")
+	ranges, importRanges, err := d2lsp.GetRefRanges("index.d2", fs, nil, "hi")
 	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 1, len(ranges))
-	assert.Equal(t, 2, refs[0].AST().GetRange().Start.Line)
+	assert.Equal(t, 2, ranges[0].Start.Line)
+	assert.Equal(t, 0, len(importRanges))
 
-	refs, err = d2lsp.GetRefs("index.d2", fs, nil, "okay")
+	ranges, importRanges, err = d2lsp.GetRefRanges("index.d2", fs, nil, "okay")
 	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 1, len(ranges))
-	assert.Equal(t, "ok.d2", refs[0].AST().GetRange().Path)
+	assert.Equal(t, "ok.d2", ranges[0].Path)
+	assert.Equal(t, 1, len(importRanges))
+	assert.Equal(t, 1, importRanges[0].Start.Line)
 
-	refs, err = d2lsp.GetRefs("ok.d2", fs, nil, "hi")
+	ranges, importRanges, err = d2lsp.GetRefRanges("index.d2", fs, nil, "hey.okay")
 	assert.Success(t, err)
-	assert.Equal(t, 0, len(refs))
+	assert.Equal(t, 1, len(ranges))
+	assert.Equal(t, "ok.d2", ranges[0].Path)
+	assert.Equal(t, 1, len(importRanges))
+	assert.Equal(t, 3, importRanges[0].Start.Line)
 
-	refs, err = d2lsp.GetRefs("ok.d2", fs, nil, "okay")
+	ranges, _, err = d2lsp.GetRefRanges("ok.d2", fs, nil, "hi")
 	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 0, len(ranges))
+
+	ranges, _, err = d2lsp.GetRefRanges("ok.d2", fs, nil, "okay")
+	assert.Success(t, err)
+	assert.Equal(t, 1, len(ranges))
 }
 
-func TestGetRefsBoards(t *testing.T) {
+func TestGetRangesBoards(t *testing.T) {
 	fs := map[string]string{
 		"index.d2": `
 hi
@@ -107,42 +120,15 @@ layers: {
 }
 `,
 	}
-	refs, err := d2lsp.GetRefs("index.d2", fs, []string{"x"}, "hello")
+	ranges, _, err := d2lsp.GetRefRanges("index.d2", fs, []string{"x"}, "hello")
-	assert.Success(t, err)
-	assert.Equal(t, 1, len(refs))
-	assert.Equal(t, 4, refs[0].AST().GetRange().Start.Line)
-
-	refs, err = d2lsp.GetRefs("index.d2", fs, []string{"x"}, "hi")
-	assert.Success(t, err)
-	assert.Equal(t, 0, len(refs))
-
-	_, err = d2lsp.GetRefs("index.d2", fs, []string{"y"}, "hello")
-	assert.Equal(t, `board "[y]" not found`, err.Error())
-}
-
-func TestGetImportRanges(t *testing.T) {
-	fs := map[string]string{
-		"yes/index.d2": `
-...@../fast/ok
-hi
-hey: {
-  ...@pok
-}
-`,
-		"fast/ok.d2": `
-okay
-`,
-		"yes/pok.d2": `
-des
-`,
-	}
-	ranges, err := d2lsp.GetImportRanges("yes/index.d2", fs["yes/index.d2"], "fast/ok.d2")
-	assert.Success(t, err)
-	assert.Equal(t, 1, len(ranges))
-	assert.Equal(t, 1, ranges[0].Start.Line)
-
-	ranges, err = d2lsp.GetImportRanges("yes/index.d2", fs["yes/index.d2"], "yes/pok.d2")
 	assert.Success(t, err)
 	assert.Equal(t, 1, len(ranges))
 	assert.Equal(t, 4, ranges[0].Start.Line)
+
+	ranges, _, err = d2lsp.GetRefRanges("index.d2", fs, []string{"x"}, "hi")
+	assert.Success(t, err)
+	assert.Equal(t, 0, len(ranges))
+
+	_, _, err = d2lsp.GetRefRanges("index.d2", fs, []string{"y"}, "hello")
+	assert.Equal(t, `board "[y]" not found`, err.Error())
 }