From 14703dd991306e8e7a8f04b2a9f371eea1683ddc Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Fri, 16 Jun 2023 16:23:08 -0700
Subject: [PATCH] save d2oracle

---
 d2oracle/edit.go      | 14 ++++++++++----
 d2oracle/edit_test.go | 32 ++++++++++++++++++++++++++++----
 d2oracle/get.go       | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 8 deletions(-)

diff --git a/d2oracle/edit.go b/d2oracle/edit.go
index 165d6886c..89fbcd143 100644
--- a/d2oracle/edit.go
+++ b/d2oracle/edit.go
@@ -22,19 +22,25 @@ import (
 	"oss.terrastruct.com/d2/d2target"
 )
 
-func Create(g *d2graph.Graph, key string) (_ *d2graph.Graph, newKey string, err error) {
+func Create(g *d2graph.Graph, boardPath []string, key string) (_ *d2graph.Graph, newKey string, err error) {
 	defer xdefer.Errorf(&err, "failed to create %#v", key)
 
-	newKey, edge, err := generateUniqueKey(g, key, nil, nil)
+	boardG := GetBoardGraph(g, boardPath)
+	if boardG == nil {
+		return nil, "", fmt.Errorf("board %v not found", boardPath)
+	}
+
+	newKey, edge, err := generateUniqueKey(boardG, key, nil, nil)
 	if err != nil {
 		return nil, "", err
 	}
 
 	if edge {
-		err = _set(g, key, nil, nil)
+		err = _set(boardG, key, nil, nil)
 	} else {
-		err = _set(g, newKey, nil, nil)
+		err = _set(boardG, newKey, nil, nil)
 	}
+	println(d2format.Format(boardG.AST))
 	if err != nil {
 		return nil, "", err
 	}
diff --git a/d2oracle/edit_test.go b/d2oracle/edit_test.go
index a48bcc0d1..f582aa306 100644
--- a/d2oracle/edit_test.go
+++ b/d2oracle/edit_test.go
@@ -26,9 +26,10 @@ func TestCreate(t *testing.T) {
 	t.Parallel()
 
 	testCases := []struct {
-		name string
-		text string
-		key  string
+		boardPath []string
+		name      string
+		text      string
+		key       string
 
 		expKey     string
 		expErr     string
@@ -455,6 +456,29 @@ rawr: {
 }
 
 after
+`,
+		},
+		{
+			name: "layers-basic",
+
+			text: `a
+layers: {
+  x: {
+    a
+  }
+}
+`,
+			key:       `b`,
+			boardPath: []string{"root", "layers", "x"},
+
+			expKey: `b`,
+			exp: `a
+layers: {
+  x: {
+    a
+    b
+  }
+}
 `,
 		},
 	}
@@ -469,7 +493,7 @@ after
 				text: tc.text,
 				testFunc: func(g *d2graph.Graph) (*d2graph.Graph, error) {
 					var err error
-					g, newKey, err = d2oracle.Create(g, tc.key)
+					g, newKey, err = d2oracle.Create(g, tc.boardPath, tc.key)
 					return g, err
 				},
 
diff --git a/d2oracle/get.go b/d2oracle/get.go
index f24178efd..c2bce9596 100644
--- a/d2oracle/get.go
+++ b/d2oracle/get.go
@@ -9,6 +9,48 @@ import (
 	"oss.terrastruct.com/d2/d2parser"
 )
 
+func GetBoardGraph(g *d2graph.Graph, boardPath []string) *d2graph.Graph {
+	if len(boardPath) == 0 {
+		return g
+	}
+	switch boardPath[0] {
+	case "root":
+		if g.Parent == nil {
+			return GetBoardGraph(g, boardPath[1:])
+		}
+		return nil
+	case "layers":
+		if len(boardPath) < 2 {
+			return nil
+		}
+		for i, b := range g.Layers {
+			if b.Name == boardPath[1] {
+				return GetBoardGraph(g.Layers[i], boardPath[2:])
+			}
+		}
+	case "scenarios":
+		if len(boardPath) < 2 {
+			return nil
+		}
+		for i, b := range g.Scenarios {
+			if b.Name == boardPath[1] {
+				return GetBoardGraph(g.Scenarios[i], boardPath[2:])
+			}
+		}
+	case "steps":
+		if len(boardPath) < 2 {
+			return nil
+		}
+		for i, b := range g.Steps {
+			if b.Name == boardPath[1] {
+				return GetBoardGraph(g.Steps[i], boardPath[2:])
+			}
+		}
+	}
+
+	return nil
+}
+
 func GetChildrenIDs(g *d2graph.Graph, absID string) (ids []string, _ error) {
 	mk, err := d2parser.ParseMapKey(absID)
 	if err != nil {