From 456799430d3e827a60dfce3e3d75bf46f182c0a8 Mon Sep 17 00:00:00 2001
From: Bernard Xie <bernard@terrastruct.com>
Date: Tue, 20 Jun 2023 11:40:24 -0700
Subject: [PATCH] refactor

---
 d2format/format.go | 33 +++++++++++++--------------------
 1 file changed, 13 insertions(+), 20 deletions(-)

diff --git a/d2format/format.go b/d2format/format.go
index 34b00410e..d0e455d83 100644
--- a/d2format/format.go
+++ b/d2format/format.go
@@ -275,12 +275,24 @@ func (p *printer) _map(m *d2ast.Map) {
 		}
 	}
 
+	layerNodes := []d2ast.MapNodeBox{}
+	scenarioNodes := []d2ast.MapNodeBox{}
+	stepNodes := []d2ast.MapNodeBox{}
+
 	prev := d2ast.Node(m)
 	for i := 0; i < len(m.Nodes); i++ {
 		nb := m.Nodes[i]
 		n := nb.Unbox()
-		// skip board nodes as we'll write them at the end
+		// extract out layer, scenario, and step nodes and skip
 		if nb.IsBoardNode() {
+			switch nb.MapKey.Key.Path[0].Unbox().ScalarString() {
+			case "layers":
+				layerNodes = append(layerNodes, nb)
+			case "scenarios":
+				scenarioNodes = append(scenarioNodes, nb)
+			case "steps":
+				stepNodes = append(stepNodes, nb)
+			}
 			prev = n
 			continue
 		}
@@ -311,25 +323,6 @@ func (p *printer) _map(m *d2ast.Map) {
 		prev = n
 	}
 
-	// extract out layer, scenario, and step nodes
-	layerNodes := []d2ast.MapNodeBox{}
-	scenarioNodes := []d2ast.MapNodeBox{}
-	stepNodes := []d2ast.MapNodeBox{}
-	for i := 0; i < len(m.Nodes); i++ {
-		node := m.Nodes[i]
-		if !node.IsBoardNode() {
-			continue
-		}
-		switch node.MapKey.Key.Path[0].Unbox().ScalarString() {
-		case "layers":
-			layerNodes = append(layerNodes, node)
-		case "scenarios":
-			scenarioNodes = append(scenarioNodes, node)
-		case "steps":
-			stepNodes = append(stepNodes, node)
-		}
-	}
-
 	boards := []d2ast.MapNodeBox{}
 	boards = append(boards, layerNodes...)
 	boards = append(boards, scenarioNodes...)