From 7ac72140c7178c723f8e68163d952acee70761ba Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Tue, 4 Apr 2023 13:42:25 -0700
Subject: [PATCH] fixing rows creation

---
 d2layouts/d2grid/layout.go             | 57 +++++++++++++++++++++-----
 e2etests/testdata/files/dagger_grid.d2 | 46 +++++++++++++++++----
 2 files changed, 85 insertions(+), 18 deletions(-)

diff --git a/d2layouts/d2grid/layout.go b/d2layouts/d2grid/layout.go
index 180ae0cf2..c16529fe0 100644
--- a/d2layouts/d2grid/layout.go
+++ b/d2layouts/d2grid/layout.go
@@ -123,27 +123,62 @@ func layoutGrid(g *d2graph.Graph, obj *d2graph.Object) (*grid, error) {
 		targetWidth := totalWidth / float64(grid.rows)
 		rowWidth := 0.
 		rowIndex := 0
-		for i, n := range grid.nodes {
+		addRow := func() {
+			layout = append(layout, []int{})
+			rowIndex++
+			rowWidth = 0
+		}
+		addNode := func(i int, n *d2graph.Object) {
 			layout[rowIndex] = append(layout[rowIndex], i)
 			rowWidth += n.Width + HORIZONTAL_PAD
-			// add a new row if we pass the target width and there are more nodes
-			if rowWidth > targetWidth && i < len(grid.nodes)-1 {
-				layout = append(layout, []int{})
-				rowIndex++
-				rowWidth = 0
+		}
+
+		for i, n := range grid.nodes {
+			// if the next node will be past the target, start a new row
+			if rowWidth+n.Width+HORIZONTAL_PAD > targetWidth {
+				// if the node is mostly past the target, put it on the next row
+				if rowWidth+n.Width/2 > targetWidth {
+					addRow()
+					addNode(i, n)
+				} else {
+					addNode(i, n)
+					if i < len(grid.nodes)-1 {
+						addRow()
+					}
+				}
+			} else {
+				addNode(i, n)
 			}
 		}
 	} else {
 		targetHeight := totalHeight / float64(grid.columns)
 		colHeight := 0.
 		colIndex := 0
-		for i, n := range grid.nodes {
+		addCol := func() {
+			layout = append(layout, []int{})
+			colIndex++
+			colHeight = 0
+		}
+		addNode := func(i int, n *d2graph.Object) {
 			layout[colIndex] = append(layout[colIndex], i)
 			colHeight += n.Height + VERTICAL_PAD
-			if colHeight > targetHeight && i < len(grid.nodes)-1 {
-				layout = append(layout, []int{})
-				colIndex++
-				colHeight = 0
+		}
+
+		for i, n := range grid.nodes {
+			// if the next node will be past the target, start a new row
+			if colHeight+n.Height+VERTICAL_PAD > targetHeight {
+				// if the node is mostly past the target, put it on the next row
+				if colHeight+n.Height/2 > targetHeight {
+					addCol()
+					addNode(i, n)
+				} else {
+					addNode(i, n)
+					if i < len(grid.nodes)-1 {
+						addCol()
+					}
+				}
+			} else {
+				addNode(i, n)
 			}
 		}
 	}
diff --git a/e2etests/testdata/files/dagger_grid.d2 b/e2etests/testdata/files/dagger_grid.d2
index b3db8ac9a..9bbe12df2 100644
--- a/e2etests/testdata/files/dagger_grid.d2
+++ b/e2etests/testdata/files/dagger_grid.d2
@@ -1,12 +1,44 @@
-rows: 4
+rows: 5
 style.fill: black
 
-flow: "" {
-  width: 800
-  height: 200
-  style: {
-    fill: cornflowerblue
-  }
+flow1: "" {
+  width: 120
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow2: "" {
+  width: 120
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow3: "" {
+  width: 120
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow4: "" {
+  width: 120
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow5: "" {
+  width: 120
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow6: "" {
+  width: 240
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow7: "" {
+  width: 100
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow8: "" {
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow9: "" {
+  width: 120
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
+}
+flow10: "" {
+  width: 120
+  style: {fill: white; stroke: cornflowerblue; stroke-width: 10}
 }
 
 DAGGER ENGINE: {