implement GetFieldRefs

d2lsp/d2lsp.go

@@ -0,0 +1,57 @@
+// d2lsp contains functions useful for IDE clients
+package d2lsp
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+
+	"oss.terrastruct.com/d2/d2ir"
+	"oss.terrastruct.com/d2/d2parser"
+	"oss.terrastruct.com/d2/lib/memfs"
+)
+
+func GetFieldRefs(path string, fs map[string]string, key string) (refs []d2ir.Reference, _ error) {
+	if _, ok := fs["index"]; !ok {
+		return nil, errors.New("index not found")
+	}
+	r := strings.NewReader(fs["index"])
+	ast, err := d2parser.Parse(path, r, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	mfs, err := memfs.New(fs)
+	if err != nil {
+		return nil, err
+	}
+
+	mk, err := d2parser.ParseMapKey(key)
+	if err != nil {
+		return nil, err
+	}
+	if mk.Key == nil {
+		return nil, fmt.Errorf(`"%s" is invalid`, key)
+	}
+
+	ir, _, err := d2ir.Compile(ast, &d2ir.CompileOptions{
+		FS: mfs,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	var f *d2ir.Field
+	curr := ir
+	for _, p := range mk.Key.Path {
+		f = curr.GetField(p.Unbox().ScalarString())
+		if f == nil {
+			return nil, fmt.Errorf(`"%s" not found`, key)
+		}
+		curr = f.Map()
+	}
+	for _, ref := range f.References {
+		refs = append(refs, ref)
+	}
+	return refs, nil
+}

d2lsp/d2lsp_test.go

@@ -0,0 +1,29 @@
+package d2lsp_test
+
+import (
+	"testing"
+
+	"oss.terrastruct.com/d2/d2lsp"
+	"oss.terrastruct.com/util-go/assert"
+)
+
+func TestGetRefs(t *testing.T) {
+	script := `x
+x.a
+a.x
+x -> y`
+	fs := map[string]string{
+		"index": script,
+	}
+	refs, err := d2lsp.GetFieldRefs("", fs, "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)
+
+	refs, err = d2lsp.GetFieldRefs("", fs, "a.x")
+	assert.Success(t, err)
+	assert.Equal(t, 1, len(refs))
+	assert.Equal(t, 2, refs[0].AST().GetRange().Start.Line)
+}

lib/memfs/memfs.go

@@ -0,0 +1,68 @@
+// memfs implements an in-memory fs.FS implementation
+// This is useful in for running d2 in javascript environments where native file calls are not available
+package memfs
+
+import (
+	"errors"
+	"io/fs"
+	"path"
+	"path/filepath"
+	"time"
+)
+
+type MemoryFile struct {
+	name    string
+	content []byte
+	modTime time.Time
+	isDir   bool
+}
+
+type MemoryFS struct {
+	files map[string]*MemoryFile
+}
+
+func New(m map[string]string) (*MemoryFS, error) {
+	memFS := &MemoryFS{files: make(map[string]*MemoryFile)}
+
+	for p, s := range m {
+		p = filepath.Clean(p)
+		dirPath := path.Dir(p)
+		memFS.addFile(dirPath, nil, true)
+		memFS.addFile(p, []byte(s), false)
+	}
+	return memFS, nil
+}
+
+func (mfs *MemoryFS) addFile(p string, content []byte, isDir bool) {
+	mfs.files[p] = &MemoryFile{
+		name:    filepath.Base(p),
+		content: content,
+		modTime: time.Now(),
+		isDir:   isDir,
+	}
+}
+
+func (mfs *MemoryFS) Open(name string) (fs.File, error) {
+	file, ok := mfs.files[filepath.Clean(name)]
+	if !ok {
+		return nil, fs.ErrNotExist
+	}
+	return file, nil
+}
+
+func (mf *MemoryFile) Stat() (fs.FileInfo, error) { return mf, nil }
+func (mf *MemoryFile) Read(b []byte) (int, error) {
+	if mf.isDir {
+		return 0, errors.New("cannot read a directory")
+	}
+	copy(b, mf.content)
+	return len(mf.content), nil
+}
+func (mf *MemoryFile) Close() error { return nil }
+
+func (mf *MemoryFile) Name() string       { return mf.name }
+func (mf *MemoryFile) Size() int64        { return int64(len(mf.content)) }
+func (mf *MemoryFile) Mode() fs.FileMode  { return 0644 }
+func (mf *MemoryFile) ModTime() time.Time { return mf.modTime }
+func (mf *MemoryFile) IsDir() bool        { return mf.isDir }
+func (mf *MemoryFile) Sys() interface{}   { return nil }