From c9a6e65c6be6930562f06860c511b5251051829e Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Wed, 13 Sep 2023 15:02:26 -0700
Subject: [PATCH] shift edges with grid

---
 d2graph/d2graph.go               |  7 +++++++
 d2graph/layout.go                |  7 ++++++-
 d2layouts/d2grid/grid_diagram.go |  3 +++
 d2layouts/d2grid/layout.go       | 12 +++++++++---
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/d2graph/d2graph.go b/d2graph/d2graph.go
index df9d12a9a..1452cf4bc 100644
--- a/d2graph/d2graph.go
+++ b/d2graph/d2graph.go
@@ -1169,6 +1169,13 @@ func (e *Edge) Text() *d2target.MText {
 	}
 }
 
+func (e *Edge) Move(dx, dy float64) {
+	for _, p := range e.Route {
+		p.X += dx
+		p.Y += dy
+	}
+}
+
 func (e *Edge) AbsID() string {
 	srcIDA := e.Src.AbsIDArray()
 	dstIDA := e.Dst.AbsIDArray()
diff --git a/d2graph/layout.go b/d2graph/layout.go
index c6899517e..9f2702df7 100644
--- a/d2graph/layout.go
+++ b/d2graph/layout.go
@@ -81,7 +81,12 @@ func pluckObjAndEdges(g *Graph, obj *Object) (descendantsObjects []*Object, edge
 
 func (g *Graph) InjectNestedGraph(tempGraph *Graph, parent *Object) {
 	obj := tempGraph.Root.ChildrenArray[0]
-	obj.MoveWithDescendantsTo(0, 0)
+	dx := 0 - obj.TopLeft.X
+	dy := 0 - obj.TopLeft.Y
+	obj.MoveWithDescendants(dx, dy)
+	for _, e := range tempGraph.Edges {
+		e.Move(dx, dy)
+	}
 	obj.Parent = parent
 	for _, obj := range tempGraph.Objects {
 		obj.Graph = g
diff --git a/d2layouts/d2grid/grid_diagram.go b/d2layouts/d2grid/grid_diagram.go
index 0aa3ce463..cc56e82de 100644
--- a/d2layouts/d2grid/grid_diagram.go
+++ b/d2layouts/d2grid/grid_diagram.go
@@ -108,6 +108,9 @@ func (gd *gridDiagram) shift(dx, dy float64) {
 	for _, obj := range gd.objects {
 		obj.MoveWithDescendants(dx, dy)
 	}
+	for _, e := range gd.edges {
+		e.Move(dx, dy)
+	}
 }
 
 func (gd *gridDiagram) cleanup(obj *d2graph.Object, graph *d2graph.Graph) {
diff --git a/d2layouts/d2grid/layout.go b/d2layouts/d2grid/layout.go
index 709bb4848..5e8d0125f 100644
--- a/d2layouts/d2grid/layout.go
+++ b/d2layouts/d2grid/layout.go
@@ -212,16 +212,22 @@ func withoutGridDiagrams(ctx context.Context, g *d2graph.Graph, layout d2graph.L
 		// simple straight line edge routing between grid objects
 		for i, e := range g.Edges {
 			edgeOrder[e.AbsID()] = i
-			if e.Dst.Parent.ClosestGridDiagram() != obj {
+			if !e.Src.Parent.IsDescendantOf(obj) && !e.Dst.Parent.IsDescendantOf(obj) {
 				continue
 			}
+			// if edge is within grid, remove it from outer layout
+			gd.edges = append(gd.edges, e)
+			edgeToRemove[e] = struct{}{}
+
+			if e.Src.Parent != obj || e.Dst.Parent != obj {
+				continue
+			}
+			// if edge is grid child, use simple routing
 			e.Route = []*geo.Point{e.Src.Center(), e.Dst.Center()}
 			e.TraceToShape(e.Route, 0, 1)
 			if e.Label.Value != "" {
 				e.LabelPosition = go2.Pointer(string(label.InsideMiddleCenter))
 			}
-			gd.edges = append(gd.edges, e)
-			edgeToRemove[e] = struct{}{}
 		}
 
 		return nil