Merge pull request #2151 from alixander/lsp-3

support board path
terrastruct · Oct 12, 2024 · 8269c69 · 8269c69
2 parents 6a56f2b + 1e66afa
commit 8269c69
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 10 deletions.
diff --git a/d2ir/d2ir.go b/d2ir/d2ir.go
@@ -1771,3 +1771,51 @@ func (m *Map) IsClass() bool {
 	}
 	return false
 }
+
+func (m *Map) FindBoardRoot(path []string) *Map {
+	if m == nil {
+		return nil
+	}
+	if len(path) == 0 {
+		return m
+	}
+
+	layersf := m.GetField("layers")
+	scenariosf := m.GetField("scenarios")
+	stepsf := m.GetField("steps")
+
+	if layersf != nil && layersf.Map() != nil {
+		for _, f := range layersf.Map().Fields {
+			if f.Name == path[0] {
+				if len(path) == 1 {
+					return f.Map()
+				}
+				return layersf.Map().FindBoardRoot(path[1:])
+			}
+		}
+	}
+
+	if scenariosf != nil && scenariosf.Map() != nil {
+		for _, f := range scenariosf.Map().Fields {
+			if f.Name == path[0] {
+				if len(path) == 1 {
+					return f.Map()
+				}
+				return scenariosf.Map().FindBoardRoot(path[1:])
+			}
+		}
+	}
+
+	if stepsf != nil && stepsf.Map() != nil {
+		for _, f := range stepsf.Map().Fields {
+			if f.Name == path[0] {
+				if len(path) == 1 {
+					return f.Map()
+				}
+				return stepsf.Map().FindBoardRoot(path[1:])
+			}
+		}
+	}
+
+	return nil
+}
diff --git a/d2lsp/d2lsp.go b/d2lsp/d2lsp.go
@@ -10,11 +10,11 @@ import (
 	"oss.terrastruct.com/d2/lib/memfs"
 )
 
-func GetFieldRefs(path, index string, fs map[string]string, key string) (refs []d2ir.Reference, _ error) {
-	if _, ok := fs[index]; !ok {
-		return nil, fmt.Errorf(`"%s" not found`, index)
+func GetFieldRefs(path string, fs map[string]string, key string, boardPath []string) (refs []d2ir.Reference, _ error) {
+	if _, ok := fs[path]; !ok {
+		return nil, fmt.Errorf(`"%s" not found`, path)
 	}
-	r := strings.NewReader(fs[index])
+	r := strings.NewReader(fs[path])
 	ast, err := d2parser.Parse(path, r, nil)
 	if err != nil {
 		return nil, err
@@ -40,6 +40,11 @@ func GetFieldRefs(path, index string, fs map[string]string, key string) (refs []
 		return nil, err
 	}
 
+	ir = ir.FindBoardRoot(boardPath)
+	if ir == nil {
+		return nil, fmt.Errorf(`board "%v" not found`, boardPath)
+	}
+
 	var f *d2ir.Field
 	curr := ir
 	for _, p := range mk.Key.Path {

diff --git a/d2lsp/d2lsp_test.go b/d2lsp/d2lsp_test.go
@@ -15,14 +15,14 @@ x -> y`
 	fs := map[string]string{
 		"index.d2": script,
 	}
-	refs, err := d2lsp.GetFieldRefs("", "index.d2", fs, "x")
+	refs, err := d2lsp.GetFieldRefs("index.d2", fs, "x", nil)
 	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)
 
-	refs, err = d2lsp.GetFieldRefs("", "index.d2", fs, "a.x")
+	refs, err = d2lsp.GetFieldRefs("index.d2", fs, "a.x", nil)
 	assert.Success(t, err)
 	assert.Equal(t, 1, len(refs))
 	assert.Equal(t, 2, refs[0].AST().GetRange().Start.Line)
@@ -38,21 +38,45 @@ hi
 okay
 `,
 	}
-	refs, err := d2lsp.GetFieldRefs("", "index.d2", fs, "hi")
+	refs, err := d2lsp.GetFieldRefs("index.d2", fs, "hi", nil)
 	assert.Success(t, err)
 	assert.Equal(t, 1, len(refs))
 	assert.Equal(t, 2, refs[0].AST().GetRange().Start.Line)
 
-	refs, err = d2lsp.GetFieldRefs("", "index.d2", fs, "okay")
+	refs, err = d2lsp.GetFieldRefs("index.d2", fs, "okay", nil)
 	assert.Success(t, err)
 	assert.Equal(t, 1, len(refs))
 	assert.Equal(t, "ok.d2", refs[0].AST().GetRange().Path)
 
-	refs, err = d2lsp.GetFieldRefs("", "ok.d2", fs, "hi")
+	refs, err = d2lsp.GetFieldRefs("ok.d2", fs, "hi", nil)
 	assert.Success(t, err)
 	assert.Equal(t, 0, len(refs))
 
-	refs, err = d2lsp.GetFieldRefs("", "ok.d2", fs, "okay")
+	refs, err = d2lsp.GetFieldRefs("ok.d2", fs, "okay", nil)
 	assert.Success(t, err)
 	assert.Equal(t, 1, len(refs))
 }
+
+func TestGetRefsBoards(t *testing.T) {
+	fs := map[string]string{
+		"index.d2": `
+hi
+layers: {
+  x: {
+    hello
+  }
+}
+`,
+	}
+	refs, err := d2lsp.GetFieldRefs("index.d2", fs, "hello", []string{"x"})
+	assert.Success(t, err)
+	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 4, refs[0].AST().GetRange().Start.Line)
+
+	refs, err = d2lsp.GetFieldRefs("index.d2", fs, "hi", []string{"x"})
+	assert.Success(t, err)
+	assert.Equal(t, 0, len(refs))
+
+	_, err = d2lsp.GetFieldRefs("index.d2", fs, "hello", []string{"y"})
+	assert.Equal(t, `board "[y]" not found`, err.Error())
+}