From 939eb500daa1a4c98a22433d952a5ae739fbe15c Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Mon, 8 May 2023 13:22:48 -0700
Subject: [PATCH] layout nested grids

---
 d2graph/d2graph.go               |   8 +++
 d2layouts/d2grid/grid_diagram.go |   3 +-
 d2layouts/d2grid/layout.go       | 105 ++++++++++++++++++-------------
 3 files changed, 70 insertions(+), 46 deletions(-)

diff --git a/d2graph/d2graph.go b/d2graph/d2graph.go
index aa211bb7e..788d1f640 100644
--- a/d2graph/d2graph.go
+++ b/d2graph/d2graph.go
@@ -1815,3 +1815,11 @@ func (g *Graph) ApplyTheme(themeID int64) error {
 	g.Theme = &theme
 	return nil
 }
+
+func (obj *Object) MoveWithDescendants(dx, dy float64) {
+	obj.TopLeft.X += dx
+	obj.TopLeft.Y += dy
+	for _, child := range obj.ChildrenArray {
+		child.MoveWithDescendants(dx, dy)
+	}
+}
diff --git a/d2layouts/d2grid/grid_diagram.go b/d2layouts/d2grid/grid_diagram.go
index c8b0c28a3..0d51a44ff 100644
--- a/d2layouts/d2grid/grid_diagram.go
+++ b/d2layouts/d2grid/grid_diagram.go
@@ -100,8 +100,7 @@ func newGridDiagram(root *d2graph.Object) *gridDiagram {
 
 func (gd *gridDiagram) shift(dx, dy float64) {
 	for _, obj := range gd.objects {
-		obj.TopLeft.X += dx
-		obj.TopLeft.Y += dy
+		obj.MoveWithDescendants(dx, dy)
 	}
 }
 
diff --git a/d2layouts/d2grid/layout.go b/d2layouts/d2grid/layout.go
index 5566bff49..7e26c521f 100644
--- a/d2layouts/d2grid/layout.go
+++ b/d2layouts/d2grid/layout.go
@@ -53,6 +53,65 @@ func withoutGridDiagrams(ctx context.Context, g *d2graph.Graph) (gridDiagrams ma
 	toRemove := make(map[*d2graph.Object]struct{})
 	gridDiagrams = make(map[string]*gridDiagram)
 
+	var processGrid func(obj *d2graph.Object) error
+	processGrid = func(obj *d2graph.Object) error {
+		// layout any nested grids first
+		for _, child := range obj.ChildrenArray {
+			if child.IsGridDiagram() {
+				if err := processGrid(child); err != nil {
+					return err
+				}
+			}
+		}
+
+		gd, err := layoutGrid(g, obj)
+		if err != nil {
+			return err
+		}
+		obj.Children = make(map[string]*d2graph.Object)
+		obj.ChildrenArray = nil
+
+		if obj.Box != nil {
+			// size shape according to grid
+			obj.SizeToContent(float64(gd.width), float64(gd.height), 2*CONTAINER_PADDING, 2*CONTAINER_PADDING)
+
+			// compute where the grid should be placed inside shape
+			dslShape := strings.ToLower(obj.Shape.Value)
+			shapeType := d2target.DSL_SHAPE_TO_SHAPE_TYPE[dslShape]
+			s := shape.NewShape(shapeType, geo.NewBox(geo.NewPoint(0, 0), obj.Width, obj.Height))
+			innerBox := s.GetInnerBox()
+			if innerBox.TopLeft.X != 0 || innerBox.TopLeft.Y != 0 {
+				gd.shift(innerBox.TopLeft.X, innerBox.TopLeft.Y)
+			}
+
+			var dx, dy float64
+			labelWidth := float64(obj.LabelDimensions.Width) + 2*label.PADDING
+			if labelWidth > obj.Width {
+				dx = (labelWidth - obj.Width) / 2
+				obj.Width = labelWidth
+			}
+			labelHeight := float64(obj.LabelDimensions.Height) + 2*label.PADDING
+			if labelHeight > CONTAINER_PADDING {
+				// if the label doesn't fit within the padding, we need to add more
+				grow := labelHeight - CONTAINER_PADDING
+				dy = grow / 2
+				obj.Height += grow
+			}
+			// we need to center children if we have to expand to fit the container label
+			if dx != 0 || dy != 0 {
+				gd.shift(dx, dy)
+			}
+		}
+
+		obj.LabelPosition = go2.Pointer(string(label.InsideTopCenter))
+		gridDiagrams[obj.AbsID()] = gd
+
+		for _, o := range gd.objects {
+			toRemove[o] = struct{}{}
+		}
+		return nil
+	}
+
 	if len(g.Objects) > 0 {
 		queue := make([]*d2graph.Object, 1, len(g.Objects))
 		queue[0] = g.Root
@@ -67,51 +126,9 @@ func withoutGridDiagrams(ctx context.Context, g *d2graph.Graph) (gridDiagrams ma
 				continue
 			}
 
-			gd, err := layoutGrid(g, obj)
-			if err != nil {
+			if err := processGrid(obj); err != nil {
 				return nil, nil, err
 			}
-			obj.Children = make(map[string]*d2graph.Object)
-			obj.ChildrenArray = nil
-
-			if obj.Box != nil {
-				// size shape according to grid
-				obj.SizeToContent(float64(gd.width), float64(gd.height), 2*CONTAINER_PADDING, 2*CONTAINER_PADDING)
-
-				// compute where the grid should be placed inside shape
-				dslShape := strings.ToLower(obj.Shape.Value)
-				shapeType := d2target.DSL_SHAPE_TO_SHAPE_TYPE[dslShape]
-				s := shape.NewShape(shapeType, geo.NewBox(geo.NewPoint(0, 0), obj.Width, obj.Height))
-				innerBox := s.GetInnerBox()
-				if innerBox.TopLeft.X != 0 || innerBox.TopLeft.Y != 0 {
-					gd.shift(innerBox.TopLeft.X, innerBox.TopLeft.Y)
-				}
-
-				var dx, dy float64
-				labelWidth := float64(obj.LabelDimensions.Width) + 2*label.PADDING
-				if labelWidth > obj.Width {
-					dx = (labelWidth - obj.Width) / 2
-					obj.Width = labelWidth
-				}
-				labelHeight := float64(obj.LabelDimensions.Height) + 2*label.PADDING
-				if labelHeight > CONTAINER_PADDING {
-					// if the label doesn't fit within the padding, we need to add more
-					grow := labelHeight - CONTAINER_PADDING
-					dy = grow / 2
-					obj.Height += grow
-				}
-				// we need to center children if we have to expand to fit the container label
-				if dx != 0 || dy != 0 {
-					gd.shift(dx, dy)
-				}
-			}
-
-			obj.LabelPosition = go2.Pointer(string(label.InsideTopCenter))
-			gridDiagrams[obj.AbsID()] = gd
-
-			for _, o := range gd.objects {
-				toRemove[o] = struct{}{}
-			}
 		}
 	}
 
@@ -821,7 +838,7 @@ func cleanup(graph *d2graph.Graph, gridDiagrams map[string]*gridDiagram, objects
 		gd, exists := gridDiagrams[graph.Root.AbsID()]
 		if exists {
 			gd.cleanup(graph.Root, graph)
-			return
+			// return
 		}
 	}