From fc06b09df6bbd3b180fb7df828113cc036244f00 Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Tue, 12 Sep 2023 22:55:31 -0700
Subject: [PATCH] validate nested grid edges (not implemented yet)

---
 d2compiler/compile.go                         | 25 ++++++++---
 d2compiler/compile_test.go                    |  7 ++-
 d2graph/grid_diagram.go                       | 19 ++++++++
 e2etests/stable_test.go                       |  1 +
 .../files/grid_nested_simple_edges.d2         | 43 +++++++++++++++++++
 .../TestCompile/grid_deeper_edge.exp.json     |  4 ++
 6 files changed, 93 insertions(+), 6 deletions(-)
 create mode 100644 e2etests/testdata/files/grid_nested_simple_edges.d2

diff --git a/d2compiler/compile.go b/d2compiler/compile.go
index 94063206c..88db84d7e 100644
--- a/d2compiler/compile.go
+++ b/d2compiler/compile.go
@@ -1075,14 +1075,29 @@ func (c *compiler) validateEdges(g *d2graph.Graph) {
 		srcGrid := edge.Src.Parent.ClosestGridDiagram()
 		dstGrid := edge.Dst.Parent.ClosestGridDiagram()
 		if srcGrid != nil || dstGrid != nil {
+			if top := srcGrid.TopGridDiagram(); srcGrid != top {
+				// valid: grid.child1 -> grid.child2
+				// invalid: grid.childGrid.child1 -> grid.childGrid.child2
+				c.errorf(edge.GetAstEdge(), "edge must be on direct child of grid diagram %#v", top.AbsID())
+				continue
+			}
+			if top := dstGrid.TopGridDiagram(); dstGrid != top {
+				// valid: grid.child1 -> grid.child2
+				// invalid: grid.childGrid.child1 -> grid.childGrid.child2
+				c.errorf(edge.GetAstEdge(), "edge must be on direct child of grid diagram %#v", top.AbsID())
+				continue
+			}
 			if srcGrid != dstGrid {
+				// valid: a -> grid
+				// invalid: a -> grid.child
 				c.errorf(edge.GetAstEdge(), "edges into grid diagrams are not supported yet")
 				continue
-			} else {
-				if srcGrid != edge.Src.Parent || dstGrid != edge.Dst.Parent {
-					c.errorf(edge.GetAstEdge(), "grid diagrams can only have edges between children right now")
-					continue
-				}
+			}
+			if srcGrid != edge.Src.Parent || dstGrid != edge.Dst.Parent {
+				// valid: grid.child1 -> grid.child2
+				// invalid: grid.child1 -> grid.child2.child1
+				c.errorf(edge.GetAstEdge(), "grid diagrams can only have edges between children right now")
+				continue
 			}
 		}
 	}
diff --git a/d2compiler/compile_test.go b/d2compiler/compile_test.go
index d9af486c3..a8788e95e 100644
--- a/d2compiler/compile_test.go
+++ b/d2compiler/compile_test.go
@@ -2496,9 +2496,14 @@ d2/testdata/d2compiler/TestCompile/grid_edge.d2:7:2: edges into grid diagrams ar
 	b: {
 		c -> d: not yet
 	}
+	a: {
+		grid-columns: 1
+		e -> f: also not yet
+	}
 }
 `,
-			expErr: `d2/testdata/d2compiler/TestCompile/grid_deeper_edge.d2:5:3: grid diagrams can only have edges between children right now`,
+			expErr: `d2/testdata/d2compiler/TestCompile/grid_deeper_edge.d2:9:3: edge must be on direct child of grid diagram "hey"
+d2/testdata/d2compiler/TestCompile/grid_deeper_edge.d2:5:3: grid diagrams can only have edges between children right now`,
 		},
 		{
 			name: "grid_nested",
diff --git a/d2graph/grid_diagram.go b/d2graph/grid_diagram.go
index b6968e8d4..62659df59 100644
--- a/d2graph/grid_diagram.go
+++ b/d2graph/grid_diagram.go
@@ -14,3 +14,22 @@ func (obj *Object) ClosestGridDiagram() *Object {
 	}
 	return obj.Parent.ClosestGridDiagram()
 }
+
+// TopGridDiagram returns the least nested (outermost) grid diagram
+func (obj *Object) TopGridDiagram() *Object {
+	if obj == nil {
+		return nil
+	}
+	var gd *Object
+	if obj.IsGridDiagram() {
+		gd = obj
+	}
+	curr := obj.Parent
+	for curr != nil {
+		if curr.IsGridDiagram() {
+			gd = curr
+		}
+		curr = curr.Parent
+	}
+	return gd
+}
diff --git a/e2etests/stable_test.go b/e2etests/stable_test.go
index 6b01e1076..0e82cb2a6 100644
--- a/e2etests/stable_test.go
+++ b/e2etests/stable_test.go
@@ -2834,6 +2834,7 @@ y: profits {
 		loadFromFile(t, "dagre_spacing"),
 		loadFromFile(t, "dagre_spacing_right"),
 		loadFromFile(t, "simple_grid_edges"),
+		loadFromFile(t, "grid_nested_simple_edges"),
 	}
 
 	runa(t, tcs)
diff --git a/e2etests/testdata/files/grid_nested_simple_edges.d2 b/e2etests/testdata/files/grid_nested_simple_edges.d2
new file mode 100644
index 000000000..8148fffed
--- /dev/null
+++ b/e2etests/testdata/files/grid_nested_simple_edges.d2
@@ -0,0 +1,43 @@
+direction: right
+outer-grid -> outer-container
+
+outer-grid: {
+  grid-columns: 1
+
+  inner-grid -> container -> etc
+
+  container: {
+    label.near: top-left
+    # edges not yet supported here since they must be direct grid children
+    a
+    b
+    c
+  }
+
+  inner-grid: {
+    grid-rows: 1
+    1
+    2
+    3
+    # edges here are not supported yet since this is inside another grid
+  }
+}
+
+outer-container: {
+  grid -> container
+
+  grid: {
+    grid-rows: 1
+    # direct child edges ok in least nested grid
+    1 -> 2 -> 3
+  }
+
+  container: {
+    # non grid edges ok
+    4 -> 5 -> 6
+    nested container: {
+      # nested non grid edges ok
+      7 -> 8
+    }
+  }
+}
diff --git a/testdata/d2compiler/TestCompile/grid_deeper_edge.exp.json b/testdata/d2compiler/TestCompile/grid_deeper_edge.exp.json
index d3247b6fb..abe81dc64 100644
--- a/testdata/d2compiler/TestCompile/grid_deeper_edge.exp.json
+++ b/testdata/d2compiler/TestCompile/grid_deeper_edge.exp.json
@@ -2,6 +2,10 @@
   "graph": null,
   "err": {
     "errs": [
+      {
+        "range": "d2/testdata/d2compiler/TestCompile/grid_deeper_edge.d2,8:2:86-8:8:92",
+        "errmsg": "d2/testdata/d2compiler/TestCompile/grid_deeper_edge.d2:9:3: edge must be on direct child of grid diagram \"hey\""
+      },
       {
         "range": "d2/testdata/d2compiler/TestCompile/grid_deeper_edge.d2,4:2:41-4:8:47",
         "errmsg": "d2/testdata/d2compiler/TestCompile/grid_deeper_edge.d2:5:3: grid diagrams can only have edges between children right now"