diff --git a/d2ir/compile.go b/d2ir/compile.go index 55f5aeee83..f5d5fa86b8 100644 --- a/d2ir/compile.go +++ b/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 { diff --git a/d2ir/d2ir.go b/d2ir/d2ir.go index 654c9cdfd5..33dbda9976 100644 --- a/d2ir/d2ir.go +++ b/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 diff --git a/d2lsp/d2lsp.go b/d2lsp/d2lsp.go index b80fa57700..21106ea746 100644 --- a/d2lsp/d2lsp.go +++ b/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()) } - } - 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) { diff --git a/d2lsp/d2lsp_test.go b/d2lsp/d2lsp_test.go index 4c25c8dae6..bae41ed95e 100644 --- a/d2lsp/d2lsp_test.go +++ b/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, 0, refs[0].AST().GetRange().Start.Line) - assert.Equal(t, 1, refs[1].AST().GetRange().Start.Line) - assert.Equal(t, 3, refs[2].AST().GetRange().Start.Line) + assert.Equal(t, 3, len(ranges)) + assert.Equal(t, 0, ranges[0].Start.Line) + assert.Equal(t, 1, ranges[1].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, 2, refs[0].AST().GetRange().Start.Line) + assert.Equal(t, 1, len(ranges)) + 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, 3, refs[0].AST().GetRange().Start.Line) + assert.Equal(t, 1, len(ranges)) + 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, 4, refs[0].AST().GetRange().Start.Line) + assert.Equal(t, 1, len(ranges)) + 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, 5, refs[0].AST().GetRange().Start.Line) + assert.Equal(t, 1, len(ranges)) + 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, 7, refs[0].AST().GetRange().Start.Line) + assert.Equal(t, 1, len(ranges)) + 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, 2, refs[0].AST().GetRange().Start.Line) + assert.Equal(t, 1, len(ranges)) + 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, "ok.d2", refs[0].AST().GetRange().Path) + assert.Equal(t, 1, len(ranges)) + 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) + + ranges, _, err = d2lsp.GetRefRanges("ok.d2", fs, nil, "hi") + assert.Success(t, err) + assert.Equal(t, 0, len(ranges)) - refs, err = d2lsp.GetRefs("ok.d2", fs, nil, "okay") + ranges, _, err = d2lsp.GetRefRanges("ok.d2", fs, nil, "okay") assert.Success(t, err) - assert.Equal(t, 1, len(refs)) + 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) + assert.Equal(t, 1, len(ranges)) + assert.Equal(t, 4, ranges[0].Start.Line) - refs, err = d2lsp.GetRefs("index.d2", fs, []string{"x"}, "hi") + ranges, _, err = d2lsp.GetRefRanges("index.d2", fs, []string{"x"}, "hi") assert.Success(t, err) - assert.Equal(t, 0, len(refs)) + assert.Equal(t, 0, len(ranges)) - _, err = d2lsp.GetRefs("index.d2", fs, []string{"y"}, "hello") + _, _, err = d2lsp.GetRefRanges("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) -}