d2ir: References wip

d2ast/d2ast.go

@@ -651,7 +651,7 @@ type KeyPath struct {
 }
 
 func MakeKeyPath(a []string) *KeyPath {
-	var kp *KeyPath
+	kp := &KeyPath{}
 	for _, el := range a {
 		kp.Path = append(kp.Path, MakeValueBox(RawString(el, true)).StringBox())
 	}

d2ir/compile.go

@@ -50,7 +50,7 @@ func (c *compiler) compileKey(dst *Map, k *d2ast.Key) {
 }
 
 func (c *compiler) compileField(dst *Map, k *d2ast.Key) {
-	f, err := dst.Ensure(d2format.KeyPath(k.Key))
+	f, err := dst.EnsureField(d2format.KeyPath(k.Key))
 	if err != nil {
 		c.errorf(k, err.Error())
 		return
@@ -88,7 +88,7 @@ func (c *compiler) compileField(dst *Map, k *d2ast.Key) {
 
 func (c *compiler) compileEdges(dst *Map, k *d2ast.Key) {
 	if k.Key != nil && len(k.Key.Path) > 0 {
-		f, err := dst.Ensure(d2format.KeyPath(k.Key))
+		f, err := dst.EnsureField(d2format.KeyPath(k.Key))
 		if err != nil {
 			c.errorf(k, err.Error())
 			return
@@ -115,7 +115,17 @@ func (c *compiler) compileEdges(dst *Map, k *d2ast.Key) {
 			}
 			e = ea[0]
 		} else {
-			var err error
+			_, err := dst.EnsureField(eid.SrcPath)
+			if err != nil {
+				c.errorf(k.Edges[i].Src, err.Error())
+				continue
+			}
+			_, err = dst.EnsureField(eid.DstPath)
+			if err != nil {
+				c.errorf(k.Edges[i].Dst, err.Error())
+				continue
+			}
+
 			e, err = dst.EnsureEdge(eid)
 			if err != nil {
 				c.errorf(k.Edges[i], err.Error())
@@ -123,17 +133,6 @@ func (c *compiler) compileEdges(dst *Map, k *d2ast.Key) {
 			}
 		}
 
-		_, err := dst.Ensure(eid.SrcPath)
-		if err != nil {
-			c.errorf(k.Edges[i].Src, err.Error())
-			continue
-		}
-		_, err = dst.Ensure(eid.DstPath)
-		if err != nil {
-			c.errorf(k.Edges[i].Dst, err.Error())
-			continue
-		}
-
 		if k.EdgeKey != nil {
 			if e.Map == nil {
 				e.Map = &Map{

d2ir/compile_test.go

@@ -79,7 +79,7 @@ func assertField(t testing.TB, n d2ir.Node, nfields, nedges int, primary interfa
 
 	var f *d2ir.Field
 	if len(ida) > 0 {
-		f = m.Get(ida)
+		f = m.GetField(ida)
 		if f == nil {
 			t.Fatalf("expected field %v in map %s", ida, m)
 		}
@@ -255,7 +255,7 @@ func testCompileEdge(t *testing.T) {
 	t.Parallel()
 	tca := []testCase{
 		{
-			name: "edge",
+			name: "root",
 			run: func(t testing.TB, m *d2ir.Map) {
 				err := parse(t, m, `x -> y`)
 				assert.Success(t, err)
@@ -285,14 +285,14 @@ func testCompileEdge(t *testing.T) {
 		{
 			name: "underscore",
 			run: func(t testing.TB, m *d2ir.Map) {
-				err := parse(t, m, `x._ -> z`)
+				err := parse(t, m, `p: { _.x -> z }`)
 				assert.Success(t, err)
 				assertField(t, m, 3, 1, nil)
 
 				assertField(t, m, 0, 0, nil, "x")
-				assertField(t, m, 0, 0, nil, "z")
+				assertField(t, m, 1, 0, nil, "p")
 
-				assertEdge(t, m, 0, nil, "(x -> z)[0]")
+				assertEdge(t, m, 0, nil, "(x -> p.z)[0]")
 			},
 		},
 	}

d2ir/d2ir.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"strings"
 
+	"oss.terrastruct.com/util-go/go2"
+
 	"oss.terrastruct.com/d2/d2ast"
 	"oss.terrastruct.com/d2/d2format"
 )
@@ -12,7 +14,9 @@ import (
 type Node interface {
 	node()
 	ast() d2ast.Node
-	Copy(newp Parent) Node
+	Parent() Node
+	Copy(newp Node) Node
+
 	fmt.Stringer
 }
 
@@ -22,16 +26,6 @@ var _ Node = &Edge{}
 var _ Node = &Array{}
 var _ Node = &Map{}
 
-type Parent interface {
-	Node
-	Parent() Parent
-}
-
-var _ Parent = &Field{}
-var _ Parent = &Edge{}
-var _ Parent = &Array{}
-var _ Parent = &Map{}
-
 type Value interface {
 	Node
 	value()
@@ -56,11 +50,11 @@ func (n *Edge) node()   {}
 func (n *Array) node()  {}
 func (n *Map) node()    {}
 
-func (n *Scalar) Parent() Parent { return n.parent }
+func (n *Scalar) Parent() Node { return n.parent }
-func (n *Field) Parent() Parent  { return n.parent }
+func (n *Field) Parent() Node  { return n.parent }
-func (n *Edge) Parent() Parent   { return n.parent }
+func (n *Edge) Parent() Node   { return n.parent }
-func (n *Array) Parent() Parent  { return n.parent }
+func (n *Array) Parent() Node  { return n.parent }
-func (n *Map) Parent() Parent    { return n.parent }
+func (n *Map) Parent() Node    { return n.parent }
 
 func (n *Scalar) value() {}
 func (n *Array) value()  {}
@@ -76,11 +70,11 @@ func (n *Array) String() string  { return d2format.Format(n.ast()) }
 func (n *Map) String() string    { return d2format.Format(n.ast()) }
 
 type Scalar struct {
-	parent Parent
+	parent Node
 	Value  d2ast.Scalar `json:"value"`
 }
 
-func (s *Scalar) Copy(newp Parent) Node {
+func (s *Scalar) Copy(newp Node) Node {
 	tmp := *s
 	s = &tmp
 
@@ -99,12 +93,12 @@ func (s *Scalar) Equal(s2 *Scalar) bool {
 }
 
 type Map struct {
-	parent Parent
+	parent Node
 	Fields []*Field `json:"fields"`
 	Edges  []*Edge  `json:"edges"`
 }
 
-func (m *Map) Copy(newp Parent) Node {
+func (m *Map) Copy(newp Node) Node {
 	tmp := *m
 	m = &tmp
 
@@ -134,15 +128,15 @@ type Field struct {
 	Primary   *Scalar   `json:"primary,omitempty"`
 	Composite Composite `json:"composite,omitempty"`
 
-	References []KeyReference `json:"references,omitempty"`
+	References []FieldReference `json:"references,omitempty"`
 }
 
-func (f *Field) Copy(newp Parent) Node {
+func (f *Field) Copy(newp Node) Node {
 	tmp := *f
 	f = &tmp
 
 	f.parent = newp.(*Map)
-	f.References = append([]KeyReference(nil), f.References...)
+	f.References = append([]FieldReference(nil), f.References...)
 	if f.Primary != nil {
 		f.Primary = f.Primary.Copy(f).(*Scalar)
 	}
@@ -221,6 +215,30 @@ func (eid *EdgeID) Match(eid2 *EdgeID) bool {
 	return true
 }
 
+func (eid *EdgeID) resolveUnderscores(m *Map) (*EdgeID, *Map, error) {
+	eid = eid.Copy()
+	maxUnderscores := go2.Max(countUnderscores(eid.SrcPath), countUnderscores(eid.DstPath))
+	for i := 0; i < maxUnderscores; i++ {
+		if eid.SrcPath[0] == "_" {
+			eid.SrcPath = eid.SrcPath[1:]
+		} else {
+			mf := parentField(m)
+			eid.SrcPath = append([]string{mf.Name}, eid.SrcPath...)
+		}
+		if eid.DstPath[0] == "_" {
+			eid.DstPath = eid.DstPath[1:]
+		} else {
+			mf := parentField(m)
+			eid.DstPath = append([]string{mf.Name}, eid.DstPath...)
+		}
+		m = parentMap(m)
+		if m == nil {
+			return nil, nil, errors.New("invalid underscore")
+		}
+	}
+	return eid, m, nil
+}
+
 func (eid *EdgeID) trimCommon() (common []string, _ *EdgeID) {
 	eid = eid.Copy()
 	for len(eid.SrcPath) > 1 && len(eid.DstPath) > 1 {
@@ -245,7 +263,7 @@ type Edge struct {
 	References []EdgeReference `json:"references,omitempty"`
 }
 
-func (e *Edge) Copy(newp Parent) Node {
+func (e *Edge) Copy(newp Node) Node {
 	tmp := *e
 	e = &tmp
 
@@ -261,11 +279,11 @@ func (e *Edge) Copy(newp Parent) Node {
 }
 
 type Array struct {
-	parent Parent
+	parent Node
 	Values []Value `json:"values"`
 }
 
-func (a *Array) Copy(newp Parent) Node {
+func (a *Array) Copy(newp Node) Node {
 	tmp := *a
 	a = &tmp
 
@@ -277,14 +295,14 @@ func (a *Array) Copy(newp Parent) Node {
 	return a
 }
 
-type KeyReference struct {
+type FieldReference struct {
 	String  *d2ast.StringBox `json:"string"`
 	KeyPath *d2ast.KeyPath   `json:"key_path"`
 
 	Context *RefContext `json:"-"`
 }
 
-func (kr KeyReference) KeyPathIndex() int {
+func (kr FieldReference) KeyPathIndex() int {
 	for i, sb := range kr.KeyPath.Path {
 		if sb == kr.String {
 			return i
@@ -293,11 +311,11 @@ func (kr KeyReference) KeyPathIndex() int {
 	panic("d2ir.KeyReference.KeyPathIndex: String not in KeyPath?")
 }
 
-func (kr KeyReference) EdgeDest() bool {
+func (kr FieldReference) EdgeDest() bool {
 	return kr.KeyPath == kr.Context.Edge.Dst
 }
 
-func (kr KeyReference) InEdge() bool {
+func (kr FieldReference) InEdge() bool {
 	return kr.KeyPath != kr.Context.Key.Key
 }
 
@@ -311,7 +329,6 @@ type RefContext struct {
 	Scope *d2ast.Map
 
 	// UnresolvedScopeMap is prior to interpreting _
-	ScopeMap *Map
 	UnresolvedScopeMap *Map
 }
 
@@ -347,6 +364,11 @@ func (m *Map) EdgeCountRecursive() int {
 		return 0
 	}
 	acc := len(m.Edges)
+	for _, f := range m.Fields {
+		if f_m, ok := f.Composite.(*Map); ok {
+			acc += f_m.EdgeCountRecursive()
+		}
+	}
 	for _, e := range m.Edges {
 		if e.Map != nil {
 			acc += e.Map.EdgeCountRecursive()
@@ -355,7 +377,17 @@ func (m *Map) EdgeCountRecursive() int {
 	return acc
 }
 
-func (m *Map) Get(ida []string) *Field {
+func (m *Map) GetField(ida []string) *Field {
+	for len(ida) > 0 && ida[0] == "_" {
+		m = parentMap(m)
+		if m == nil {
+			return nil
+		}
+	}
+	return m.getField(ida)
+}
+
+func (m *Map) getField(ida []string) *Field {
 	if len(ida) == 0 {
 		return nil
 	}
@@ -363,6 +395,10 @@ func (m *Map) Get(ida []string) *Field {
 	s := ida[0]
 	rest := ida[1:]
 
+	if s == "_" {
+		return nil
+	}
+
 	for _, f := range m.Fields {
 		if !strings.EqualFold(f.Name, s) {
 			continue
@@ -371,20 +407,35 @@ func (m *Map) Get(ida []string) *Field {
 			return f
 		}
 		if f_m, ok := f.Composite.(*Map); ok {
-			return f_m.Get(rest)
+			return f_m.getField(rest)
 		}
 	}
 	return nil
 }
 
-func (m *Map) Ensure(ida []string) (*Field, error) {
+func (m *Map) EnsureField(ida []string) (*Field, error) {
+	for len(ida) > 0 && ida[0] == "_" {
+		m = parentMap(m)
+		if m == nil {
+			return nil, errors.New("invalid underscore")
+		}
+		ida = ida[1:]
+	}
+	return m.ensureField(ida)
+}
+
+func (m *Map) ensureField(ida []string) (*Field, error) {
 	if len(ida) == 0 {
-		return nil, errors.New("empty ida")
+		return nil, errors.New("invalid underscore")
 	}
 
 	s := ida[0]
 	rest := ida[1:]
 
+	if s == "_" {
+		return nil, errors.New(`parent "_" can only be used in the beginning of paths, e.g. "_.x"`)
+	}
+
 	for _, f := range m.Fields {
 		if !strings.EqualFold(f.Name, s) {
 			continue
@@ -394,14 +445,14 @@ func (m *Map) Ensure(ida []string) (*Field, error) {
 		}
 		switch fc := f.Composite.(type) {
 		case *Map:
-			return fc.Ensure(rest)
+			return fc.ensureField(rest)
 		case *Array:
 			return nil, errors.New("cannot index into array")
 		}
 		f.Composite = &Map{
 			parent: f,
 		}
-		return f.Composite.(*Map).Ensure(rest)
+		return f.Composite.(*Map).ensureField(rest)
 	}
 
 	f := &Field{
@@ -415,7 +466,7 @@ func (m *Map) Ensure(ida []string) (*Field, error) {
 	f.Composite = &Map{
 		parent: f,
 	}
-	return f.Composite.(*Map).Ensure(rest)
+	return f.Composite.(*Map).ensureField(rest)
 }
 
 func (m *Map) Delete(ida []string) bool {
@@ -442,9 +493,13 @@ func (m *Map) Delete(ida []string) bool {
 }
 
 func (m *Map) GetEdges(eid *EdgeID) []*Edge {
+	eid, m, err := eid.resolveUnderscores(m)
+	if err != nil {
+		return nil
+	}
 	common, eid := eid.trimCommon()
 	if len(common) > 0 {
-		f := m.Get(common)
+		f := m.GetField(common)
 		if f == nil {
 			return nil
 		}
@@ -464,9 +519,13 @@ func (m *Map) GetEdges(eid *EdgeID) []*Edge {
 }
 
 func (m *Map) EnsureEdge(eid *EdgeID) (*Edge, error) {
+	eid, m, err := eid.resolveUnderscores(m)
+	if err != nil {
+		return nil, err
+	}
 	common, eid := eid.trimCommon()
 	if len(common) > 0 {
-		f, err := m.Ensure(common)
+		f, err := m.EnsureField(common)
 		if err != nil {
 			return nil, err
 		}
@@ -560,8 +619,10 @@ func (m *Map) ast() d2ast.Node {
 		return nil
 	}
 	astMap := &d2ast.Map{}
-	if m.parent != nil {
+	if m.parent == nil {
-		astMap.Range = d2ast.MakeRange(",1:0:0-1:0:0")
+		astMap.Range = d2ast.MakeRange(",0:0:0-1:0:0")
+	} else {
+		astMap.Range = d2ast.MakeRange(",1:0:0-2:0:0")
 	}
 	for _, f := range m.Fields {
 		astMap.Nodes = append(astMap.Nodes, d2ast.MakeMapNodeBox(f.ast().(d2ast.MapNode)))
@@ -572,14 +633,14 @@ func (m *Map) ast() d2ast.Node {
 	return astMap
 }
 
-func (m *Map) appendKeyReferences(i int, kp *d2ast.KeyPath, refctx *RefContext) {
+func (m *Map) appendFieldReferences(i int, kp *d2ast.KeyPath, refctx *RefContext) {
 	sb := kp.Path[i]
-	f := m.Get([]string{sb.Unbox().ScalarString()})
+	f := m.GetField([]string{sb.Unbox().ScalarString()})
 	if f == nil {
 		return
 	}
 
-	f.References = append(f.References, KeyReference{
+	f.References = append(f.References, FieldReference{
 		String:  sb,
 		KeyPath: kp,
 		Context: refctx,
@@ -588,7 +649,7 @@ func (m *Map) appendKeyReferences(i int, kp *d2ast.KeyPath, refctx *RefContext)
 		return
 	}
 	if f_m, ok := f.Composite.(*Map); ok {
-		f_m.appendKeyReferences(i+1, kp, refctx)
+		f_m.appendFieldReferences(i+1, kp, refctx)
 	}
 }
 
@@ -596,6 +657,37 @@ func (m *Map) appendEdgeReferences(e *Edge, refctx *RefContext) {
 	e.References = append(e.References, EdgeReference{
 		Context: refctx,
 	})
-	m.appendKeyReferences(0, refctx.Edge.Src, refctx)
+	m.appendFieldReferences(0, refctx.Edge.Src, refctx)
-	m.appendKeyReferences(0, refctx.Edge.Dst, refctx)
+	m.appendFieldReferences(0, refctx.Edge.Dst, refctx)
+}
+
+func parentMap(n Node) *Map {
+	for n.Parent() != nil {
+		n = n.Parent()
+		if n_m, ok := n.(*Map); ok {
+			return n_m
+		}
+	}
+	return nil
+}
+
+func parentField(n Node) *Field {
+	for n.Parent() != nil {
+		n = n.Parent()
+		if n_f, ok := n.(*Field); ok {
+			return n_f
+		}
+	}
+	return nil
+}
+
+func countUnderscores(p []string) int {
+	var count int
+	for _, el := range p {
+		if el != "_" {
+			break
+		}
+		count++
+	}
+	return count
 }

testdata/d2ir/TestCompile/edge/root.exp.json

@@ -0,0 +1,25 @@
+{
+  "fields": [
+    {
+      "name": "x"
+    },
+    {
+      "name": "y"
+    }
+  ],
+  "edges": [
+    {
+      "edge_id": {
+        "src_path": [
+          "x"
+        ],
+        "src_arrow": false,
+        "dst_path": [
+          "y"
+        ],
+        "dst_arrow": true,
+        "index": 0
+      }
+    }
+  ]
+}

testdata/d2ir/TestCompile/edge/underscore.exp.json

@@ -1,29 +1,29 @@
 {
   "fields": [
     {
-      "name": "x",
+      "name": "p",
       "composite": {
         "fields": [
           {
-            "name": "_"
+            "name": "z"
           }
         ],
         "edges": null
       }
     },
     {
-      "name": "z"
+      "name": "x"
     }
   ],
   "edges": [
     {
       "edge_id": {
         "src_path": [
-          "x",
+          "x"
-          "_"
         ],
         "src_arrow": false,
         "dst_path": [
+          "p",
           "z"
         ],
         "dst_arrow": true,