From ac0845da1c35e50c639d6248f9fa65feddeb6bfa Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Tue, 4 Apr 2023 19:07:49 -0700
Subject: [PATCH] center container if growing to fit label

---
 d2layouts/d2grid/layout.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/d2layouts/d2grid/layout.go b/d2layouts/d2grid/layout.go
index c16529fe0..69050c15d 100644
--- a/d2layouts/d2grid/layout.go
+++ b/d2layouts/d2grid/layout.go
@@ -69,7 +69,28 @@ func withoutGrids(ctx context.Context, g *d2graph.Graph) (idToGrid map[string]*g
 			}
 			obj.Children = make(map[string]*d2graph.Object)
 			obj.ChildrenArray = nil
-			obj.Box = geo.NewBox(nil, grid.width+2*CONTAINER_PADDING, grid.height+2*CONTAINER_PADDING)
+
+			var dx, dy float64
+			width := grid.width + 2*CONTAINER_PADDING
+			labelWidth := float64(obj.LabelDimensions.Width) + 2*label.PADDING
+			if labelWidth > width {
+				dx = (labelWidth - width) / 2
+				width = labelWidth
+			}
+			height := grid.height + 2*CONTAINER_PADDING
+			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
+				height += grow
+			}
+			// we need to center children if we have to expand to fit the container label
+			if dx != 0 || dy != 0 {
+				grid.shift(dx, dy)
+			}
+			obj.Box = geo.NewBox(nil, width, height)
+
 			obj.LabelPosition = go2.Pointer(string(label.InsideTopCenter))
 			grids[obj.AbsID()] = grid