From 26a72e3e0c7c812ea89aae5b8392356d0d603774 Mon Sep 17 00:00:00 2001
From: Anmol Sethi <anmol@terrastruct.com>
Date: Wed, 18 Jan 2023 07:44:34 -0800
Subject: [PATCH] d2ir: Add recursive scenario/step test

---
 d2ir/compile.go      |  4 ++++
 d2ir/compile_test.go | 51 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/d2ir/compile.go b/d2ir/compile.go
index e5ca7baf1..69660b3a9 100644
--- a/d2ir/compile.go
+++ b/d2ir/compile.go
@@ -43,6 +43,8 @@ func (c *compiler) compileScenarios(m *Map) {
 		}
 		base := m.Copy(sf).(*Map)
 		sf.Composite = Overlay(base, sf.Map())
+		c.compileScenarios(sf.Map()
+		c.compileSteps(sf.Map()
 	}
 }
 
@@ -69,6 +71,8 @@ func (c *compiler) compileSteps(m *Map) {
 			base = steps.Fields[i-1].Map().Copy(sf).(*Map)
 		}
 		sf.Composite = Overlay(base, sf.Map())
+		c.compileScenarios(sf.Map()
+		c.compileSteps(sf.Map()
 	}
 }
 
diff --git a/d2ir/compile_test.go b/d2ir/compile_test.go
index 9154c6280..a3b5ec53e 100644
--- a/d2ir/compile_test.go
+++ b/d2ir/compile_test.go
@@ -413,6 +413,57 @@ steps: {
 				assertQuery(t, m, 0, 0, nil, "steps.nuclear.quiche")
 			},
 		},
+		{
+			name: "recursive",
+			run: func(t testing.TB) {
+				m, err := compile(t, `x -> y
+steps: {
+	bingo: { p.q.z }
+	nuclear: {
+		quiche
+		scenarios: {
+			bavarian: {
+				perseverance
+			}
+		}
+	}
+}`)
+				assert.Success(t, err)
+
+				assertQuery(t, m, 16, 3, nil, "")
+
+				assertQuery(t, m, 0, 0, nil, "x")
+				assertQuery(t, m, 0, 0, nil, "y")
+				assertQuery(t, m, 0, 0, nil, `(x -> y)[0]`)
+
+				assertQuery(t, m, 5, 1, nil, "steps.bingo")
+				assertQuery(t, m, 0, 0, nil, "steps.bingo.x")
+				assertQuery(t, m, 0, 0, nil, "steps.bingo.y")
+				assertQuery(t, m, 0, 0, nil, `steps.bingo.(x -> y)[0]`)
+				assertQuery(t, m, 2, 0, nil, "steps.bingo.p")
+				assertQuery(t, m, 1, 0, nil, "steps.bingo.p.q")
+				assertQuery(t, m, 0, 0, nil, "steps.bingo.p.q.z")
+
+				assertQuery(t, m, 6, 1, nil, "steps.nuclear")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.x")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.y")
+				assertQuery(t, m, 0, 0, nil, `steps.nuclear.(x -> y)[0]`)
+				assertQuery(t, m, 2, 0, nil, "steps.nuclear.p")
+				assertQuery(t, m, 1, 0, nil, "steps.nuclear.p.q")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.p.q.z")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.quiche")
+
+				assertQuery(t, m, 6, 1, nil, "steps.nuclear.scenarios.bavarian")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.scenarios.bavarian.x")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.scenarios.bavarian.y")
+				assertQuery(t, m, 0, 0, nil, `steps.nuclear.scenarios.bavarian.(x -> y)[0]`)
+				assertQuery(t, m, 2, 0, nil, "steps.nuclear.scenarios.bavarian.p")
+				assertQuery(t, m, 1, 0, nil, "steps.nuclear.scenarios.bavarian.p.q")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.scenarios.bavarian.p.q.z")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.scenarios.bavarian.quiche")
+				assertQuery(t, m, 0, 0, nil, "steps.nuclear.scenarios.bavarian.perseverance")
+			},
+		},
 	}
 	runa(t, tca)
 }