From 9c0f49829a14ba0a3b70b8769e4ea4d35c8d2534 Mon Sep 17 00:00:00 2001
From: Bernard Xie <bernard@terrastruct.com>
Date: Mon, 19 Jun 2023 17:33:41 -0700
Subject: [PATCH] add test

---
 d2ast/d2ast.go          |  12 ++++
 d2format/format.go      |  40 +++++++++++-
 d2format/format_test.go | 133 +++++++++++++++++++++++++++++++++++++++-
 3 files changed, 182 insertions(+), 3 deletions(-)

diff --git a/d2ast/d2ast.go b/d2ast/d2ast.go
index 63a7d07f9..43348804d 100644
--- a/d2ast/d2ast.go
+++ b/d2ast/d2ast.go
@@ -760,6 +760,18 @@ func (mb MapNodeBox) Unbox() MapNode {
 	}
 }
 
+func (mb MapNodeBox) IsSpecialBoard() bool {
+	if mb.MapKey == nil || mb.MapKey.Key == nil {
+		return false
+	}
+	switch mb.MapKey.Key.Path[0].Unbox().ScalarString() {
+	case "layers", "scenarios", "steps":
+		return true
+	default:
+		return false
+	}
+}
+
 // ArrayNodeBox is used to box ArrayNode for JSON persistence.
 type ArrayNodeBox struct {
 	Comment            *Comment            `json:"comment,omitempty"`
diff --git a/d2format/format.go b/d2format/format.go
index e979a6fea..82fd9daac 100644
--- a/d2format/format.go
+++ b/d2format/format.go
@@ -253,9 +253,38 @@ func (p *printer) _map(m *d2ast.Map) {
 		}
 	}
 
-	prev := d2ast.Node(m)
+	nodes := []d2ast.MapNodeBox{}
+	// extract out layer, scenario, and step nodes
+	layerNodes := []d2ast.MapNodeBox{}
+	scenarioNodes := []d2ast.MapNodeBox{}
+	stepNodes := []d2ast.MapNodeBox{}
 	for i := 0; i < len(m.Nodes); i++ {
-		nb := m.Nodes[i]
+		node := m.Nodes[i]
+		if node.MapKey != nil && node.MapKey.Key != nil {
+			key := node.MapKey.Key
+			switch key.Path[0].Unbox().ScalarString() {
+			case "layers":
+				layerNodes = append(layerNodes, node)
+			case "scenarios":
+				scenarioNodes = append(scenarioNodes, node)
+			case "steps":
+				stepNodes = append(stepNodes, node)
+			default:
+				nodes = append(nodes, node)
+			}
+		} else {
+			nodes = append(nodes, node)
+		}
+	}
+
+	// append layers, scenarios, and steps at the end
+	nodes = append(nodes, layerNodes...)
+	nodes = append(nodes, scenarioNodes...)
+	nodes = append(nodes, stepNodes...)
+
+	prev := d2ast.Node(m)
+	for i := 0; i < len(nodes); i++ {
+		nb := nodes[i]
 		n := nb.Unbox()
 
 		// Handle inline comments.
@@ -280,6 +309,13 @@ func (p *printer) _map(m *d2ast.Map) {
 			p.sb.WriteString("; ")
 		}
 
+		if m.IsFileMap() && nb.IsSpecialBoard() {
+			currString := p.sb.String()
+			// if the the character before the special board is not a newline, we add one
+			if currString[len(currString)-2:] != "\n\n" {
+				p.newline()
+			}
+		}
 		p.node(n)
 		prev = n
 	}
diff --git a/d2format/format_test.go b/d2format/format_test.go
index d25a45241..6f0616db0 100644
--- a/d2format/format_test.go
+++ b/d2format/format_test.go
@@ -153,13 +153,13 @@ meow
     diagram: int {constraint: foreign_key}
   }
 
+  meow <- diagrams.id
   steps: {
     shape: sql_table
     id: {type: int; constraint: primary_key}
     representation: {type: jsonb}
     diagram: int {constraint: foreign_key}
   }
-  meow <- diagrams.id
 }
 
 D2 AST Parser: {
@@ -617,6 +617,137 @@ y
 x <= y
 `,
 			exp: `x <- = y
+`,
+		},
+		{
+			name: "layers_scenarios_steps_bottom_simple",
+			in: `a
+
+layers: {
+  b: {
+    scenarios: {
+      p: { 
+        x 
+      }
+    }
+    e
+  }
+}
+
+g
+`,
+			exp: `a
+
+g
+
+layers: {
+  b: {
+    e
+    scenarios: {
+      p: {
+        x
+      }
+    }
+  }
+}
+`,
+		},
+		{
+			name: "layers_scenarios_steps_bottom_complex",
+			in: `a
+
+scenarios: {
+
+
+
+	scenario-1: {
+		steps: {
+			step-1: {
+        Test
+			}
+			step-2
+		}
+		non-step
+	}
+}
+
+layers: {
+  Test super nested: {
+		base-layer
+		layers: {
+			layers: {
+				grand-child-layer: {
+					grand-child-board
+				}
+			}
+			layer-board
+		}
+		last-layer
+	}
+}
+b
+steps: {
+	1: {
+		step-1-content
+	}
+}
+
+
+scenarios: {
+	scenario-2: {
+		scenario-2-content
+	}
+}
+
+c
+d
+`,
+			exp: `a
+
+b
+
+c
+d
+
+layers: {
+  Test super nested: {
+    base-layer
+
+    last-layer
+    layers: {
+      layer-board
+      layers: {
+        grand-child-layer: {
+          grand-child-board
+        }
+      }
+    }
+  }
+}
+
+scenarios: {
+  scenario-1: {
+    non-step
+    steps: {
+      step-1: {
+        Test
+      }
+      step-2
+    }
+  }
+}
+
+scenarios: {
+  scenario-2: {
+    scenario-2-content
+  }
+}
+
+steps: {
+  1: {
+    step-1-content
+  }
+}
 `,
 		},
 	}