From acd0638325fa77c8460c41c46c429092a4454808 Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Mon, 3 Apr 2023 12:59:54 -0700
Subject: [PATCH] evenly size grid nodes

---
 d2layouts/d2grid/grid.go   | 45 +++++++++++++++++++++++++++++++++++---
 d2layouts/d2grid/layout.go | 15 +++++--------
 2 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/d2layouts/d2grid/grid.go b/d2layouts/d2grid/grid.go
index cd5e4ec8f..4a743c60f 100644
--- a/d2layouts/d2grid/grid.go
+++ b/d2layouts/d2grid/grid.go
@@ -1,6 +1,7 @@
 package d2grid
 
 import (
+	"math"
 	"strconv"
 
 	"oss.terrastruct.com/d2/d2graph"
@@ -12,8 +13,10 @@ type grid struct {
 	rows    int
 	columns int
 
-	width  float64
-	height float64
+	cellWidth  float64
+	cellHeight float64
+	width      float64
+	height     float64
 }
 
 func newGrid(root *d2graph.Object) *grid {
@@ -26,7 +29,6 @@ func newGrid(root *d2graph.Object) *grid {
 	}
 
 	// compute exact row/column count based on values entered
-	// TODO consider making this based on node dimensions
 	if g.rows == 0 {
 		// set rows based on number of columns
 		g.rows = len(g.nodes) / g.columns
@@ -48,6 +50,43 @@ func newGrid(root *d2graph.Object) *grid {
 		}
 	}
 
+	// if we have the following nodes for a 2 row, 3 column grid
+	// . ┌A──────────────────┐  ┌B─────┐  ┌C────────────┐  ┌D───────────┐  ┌E───────────────────┐
+	// . │                   │  │      │  │             │  │            │  │                    │
+	// . │                   │  │      │  │             │  │            │  │                    │
+	// . └───────────────────┘  │      │  │             │  │            │  │                    │
+	// .                        │      │  └─────────────┘  │            │  │                    │
+	// .                        │      │                   │            │  └────────────────────┘
+	// .                        └──────┘                   │            │
+	// .                                                   └────────────┘
+	// Then we must get the max width and max height to determine the grid cell size
+	// .                                          maxWidth├────────────────────┤
+	// .  ┌A───────────────────┐  ┌B───────────────────┐  ┌C───────────────────┐ ┬maxHeight
+	// .  │                    │  │                    │  │                    │ │
+	// .  │                    │  │                    │  │                    │ │
+	// .  │                    │  │                    │  │                    │ │
+	// .  │                    │  │                    │  │                    │ │
+	// .  │                    │  │                    │  │                    │ │
+	// .  │                    │  │                    │  │                    │ │
+	// .  └────────────────────┘  └────────────────────┘  └────────────────────┘ ┴
+	// .  ┌D───────────────────┐  ┌E───────────────────┐
+	// .  │                    │  │                    │
+	// .  │                    │  │                    │
+	// .  │                    │  │                    │
+	// .  │                    │  │                    │
+	// .  │                    │  │                    │
+	// .  │                    │  │                    │
+	// .  └────────────────────┘  └────────────────────┘
+	var maxWidth, maxHeight float64
+	for _, n := range g.nodes {
+		maxWidth = math.Max(maxWidth, n.Width)
+		maxHeight = math.Max(maxHeight, n.Height)
+	}
+	g.cellWidth = maxWidth
+	g.cellHeight = maxHeight
+	g.width = maxWidth + (float64(g.columns)-1)*(maxWidth+HORIZONTAL_PAD)
+	g.height = maxHeight + (float64(g.rows)-1)*(maxHeight+VERTICAL_PAD)
+
 	return &g
 }
 
diff --git a/d2layouts/d2grid/layout.go b/d2layouts/d2grid/layout.go
index 45966e3a2..e9d68bccd 100644
--- a/d2layouts/d2grid/layout.go
+++ b/d2layouts/d2grid/layout.go
@@ -2,7 +2,6 @@ package d2grid
 
 import (
 	"context"
-	"math"
 	"sort"
 
 	"oss.terrastruct.com/d2/d2graph"
@@ -69,7 +68,7 @@ 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+CONTAINER_PADDING*2, grid.height+CONTAINER_PADDING*2)
+			obj.Box = geo.NewBox(nil, grid.width+2*CONTAINER_PADDING, grid.height+2*CONTAINER_PADDING)
 			obj.LabelPosition = go2.Pointer(string(label.InsideTopCenter))
 			grids[obj.AbsID()] = grid
 
@@ -97,21 +96,17 @@ func layoutGrid(g *d2graph.Graph, obj *d2graph.Object) (*grid, error) {
 
 	// position nodes
 	cursor := geo.NewPoint(0, 0)
-	maxWidth := 0.
 	for i := 0; i < grid.rows; i++ {
-		maxHeight := 0.
 		for j := 0; j < grid.columns; j++ {
 			n := grid.nodes[i*grid.columns+j]
+			n.Width = grid.cellWidth
+			n.Height = grid.cellHeight
 			n.TopLeft = cursor.Copy()
-			cursor.X += n.Width + HORIZONTAL_PAD
-			maxHeight = math.Max(maxHeight, n.Height)
+			cursor.X += grid.cellWidth + HORIZONTAL_PAD
 		}
-		maxWidth = math.Max(maxWidth, cursor.X-HORIZONTAL_PAD)
 		cursor.X = 0
-		cursor.Y += float64(maxHeight) + VERTICAL_PAD
+		cursor.Y += float64(grid.cellHeight) + VERTICAL_PAD
 	}
-	grid.width = maxWidth
-	grid.height = cursor.Y - VERTICAL_PAD
 
 	// position labels and icons
 	for _, n := range grid.nodes {